Skip to content

Commit 3f63fc6

Browse files
shiyasmohdshiyasmohd
authored
store: ensure consistent id type in FindRangeQuery SQL output (#6080)
* store: ensure consistent id type in FindRangeQuery SQL output * tests: add integration tests * tests: add unit test for sql query
1 parent 0e192f8 commit 3f63fc6

File tree

7 files changed

+160
-7
lines changed

7 files changed

+160
-7
lines changed

‎store/postgres/src/relational/query_tests.rs

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,16 @@ use std::{collections::BTreeSet, sync::Arc};
22

33
use diesel::{debug_query, pg::Pg};
44
use graph::{
5+
data_source::CausalityRegion,
56
prelude::{r, serde_json as json, DeploymentHash, EntityFilter},
67
schema::InputSchema,
78
};
89

910
use crate::{
11+
block_range::BoundSide,
1012
layout_for_tests::{make_dummy_site, Namespace},
1113
relational::{Catalog, ColumnType, Layout},
12-
relational_queries::FromColumnValue,
14+
relational_queries::{FindRangeQuery, FromColumnValue},
1315
};
1416

1517
use crate::relational_queries::Filter;
@@ -86,3 +88,112 @@ fn prefix() {
8688
let filter = EntityFilter::In("address".to_string(), vec!["0xbeef".into()]);
8789
filter_contains(filter, r#"substring(c."address", 1, 64) in ($1)"#);
8890
}
91+
92+
#[test]
93+
fn find_range_query_id_type_casting() {
94+
let string_schema = "
95+
type StringEntity @entity {
96+
id: String!,
97+
name: String
98+
}";
99+
100+
let bytes_schema = "
101+
type BytesEntity @entity {
102+
id: Bytes!,
103+
address: Bytes
104+
}";
105+
106+
let int8_schema = "
107+
type Int8Entity @entity {
108+
id: Int8!,
109+
value: Int8
110+
}";
111+
112+
let string_layout = test_layout(string_schema);
113+
let bytes_layout = test_layout(bytes_schema);
114+
let int8_layout = test_layout(int8_schema);
115+
116+
let string_table = string_layout
117+
.table_for_entity(
118+
&string_layout
119+
.input_schema
120+
.entity_type("StringEntity")
121+
.unwrap(),
122+
)
123+
.unwrap();
124+
let bytes_table = bytes_layout
125+
.table_for_entity(
126+
&bytes_layout
127+
.input_schema
128+
.entity_type("BytesEntity")
129+
.unwrap(),
130+
)
131+
.unwrap();
132+
let int8_table = int8_layout
133+
.table_for_entity(&int8_layout.input_schema.entity_type("Int8Entity").unwrap())
134+
.unwrap();
135+
136+
let causality_region = CausalityRegion::ONCHAIN;
137+
let bound_side = BoundSide::Lower;
138+
let block_range = 100..200;
139+
140+
test_id_type_casting(
141+
string_table.as_ref(),
142+
"id::bytea",
143+
"String ID should be cast to bytea",
144+
);
145+
test_id_type_casting(bytes_table.as_ref(), "id", "Bytes ID should remain as id");
146+
test_id_type_casting(
147+
int8_table.as_ref(),
148+
"id::text::bytea",
149+
"Int8 ID should be cast to text then bytea",
150+
);
151+
152+
let tables = vec![
153+
string_table.as_ref(),
154+
bytes_table.as_ref(),
155+
int8_table.as_ref(),
156+
];
157+
let query = FindRangeQuery::new(&tables, causality_region, bound_side, block_range);
158+
let sql = debug_query::<Pg, _>(&query).to_string();
159+
160+
assert!(
161+
sql.contains("id::bytea"),
162+
"String entity ID casting should be present in UNION query"
163+
);
164+
assert!(
165+
sql.contains("id as id"),
166+
"Bytes entity ID should be present in UNION query"
167+
);
168+
assert!(
169+
sql.contains("id::text::bytea"),
170+
"Int8 entity ID casting should be present in UNION query"
171+
);
172+
173+
assert!(
174+
sql.contains("union all"),
175+
"Multiple tables should generate UNION ALL queries"
176+
);
177+
assert!(
178+
sql.contains("order by block_number, entity, id"),
179+
"Query should end with proper ordering"
180+
);
181+
}
182+
183+
fn test_id_type_casting(table: &crate::relational::Table, expected_cast: &str, test_name: &str) {
184+
let causality_region = CausalityRegion::ONCHAIN;
185+
let bound_side = BoundSide::Lower;
186+
let block_range = 100..200;
187+
188+
let tables = vec![table];
189+
let query = FindRangeQuery::new(&tables, causality_region, bound_side, block_range);
190+
let sql = debug_query::<Pg, _>(&query).to_string();
191+
192+
assert!(
193+
sql.contains(expected_cast),
194+
"{}: Expected '{}' in SQL, got: {}",
195+
test_name,
196+
expected_cast,
197+
sql
198+
);
199+
}

‎store/postgres/src/relational_queries.rs

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1884,7 +1884,19 @@ impl<'a> QueryFragment<Pg> for FindRangeQuery<'a> {
18841884
} else {
18851885
self.mut_range.compare_column(&mut out)
18861886
}
1887-
out.push_sql("as block_number, id, vid\n");
1887+
// Cast id to bytea to ensure consistent types across UNION
1888+
// The actual id type can be text, bytea, or numeric depending on the entity
1889+
out.push_sql("as block_number, ");
1890+
let pk_column = table.primary_key();
1891+
1892+
// We only support entity id types of string, bytes, and int8.
1893+
match pk_column.column_type {
1894+
ColumnType::String => out.push_sql("id::bytea"),
1895+
ColumnType::Bytes => out.push_sql("id"),
1896+
ColumnType::Int8 => out.push_sql("id::text::bytea"),
1897+
_ => out.push_sql("id::bytea"),
1898+
}
1899+
out.push_sql(" as id, vid\n");
18881900
out.push_sql(" from ");
18891901
out.push_sql(table.qualified_name.as_str());
18901902
out.push_sql(" e\n where");
@@ -1906,7 +1918,7 @@ impl<'a> QueryFragment<Pg> for FindRangeQuery<'a> {
19061918
// In case we have only immutable entities, the upper range will not create any
19071919
// select statement. So here we have to generate an SQL statement thet returns
19081920
// empty result.
1909-
out.push_sql("select 'dummy_entity' as entity, to_jsonb(1) as data, 1 as block_number, 1 as id, 1 as vid where false");
1921+
out.push_sql("select 'dummy_entity' as entity, to_jsonb(1) as data, 1 as block_number, '\\x'::bytea as id, 1 as vid where false");
19101922
} else {
19111923
out.push_sql("\norder by block_number, entity, id");
19121924
}

‎tests/integration-tests/source-subgraph/schema.graphql

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,9 @@ type Block2 @entity(immutable: true) {
1010
hash: Bytes!
1111
testMessage: String
1212
}
13+
14+
type Block3 @entity(immutable: true) {
15+
id: Bytes!
16+
number: BigInt!
17+
testMessage: String
18+
}

‎tests/integration-tests/source-subgraph/src/mapping.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { ethereum, log, store } from '@graphprotocol/graph-ts';
2-
import { Block, Block2 } from '../generated/schema';
2+
import { Block, Block2, Block3 } from '../generated/schema';
33

44
export function handleBlock(block: ethereum.Block): void {
55
log.info('handleBlock {}', [block.number.toString()]);
@@ -22,4 +22,10 @@ export function handleBlock(block: ethereum.Block): void {
2222
blockEntity3.hash = block.hash;
2323
blockEntity3.testMessage = block.number.toString().concat('-message');
2424
blockEntity3.save();
25+
26+
let id4 = block.hash;
27+
let blockEntity4 = new Block3(id4);
28+
blockEntity4.number = block.number;
29+
blockEntity4.testMessage = block.number.toString().concat('-message');
30+
blockEntity4.save();
2531
}

‎tests/integration-tests/subgraph-data-sources/schema.graphql

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,9 @@ type MirrorBlock @entity {
44
hash: Bytes!
55
testMessage: String
66
}
7+
8+
type MirrorBlockBytes @entity {
9+
id: Bytes!
10+
number: BigInt!
11+
testMessage: String
12+
}

‎tests/integration-tests/subgraph-data-sources/src/mapping.ts

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { log, store } from '@graphprotocol/graph-ts';
2-
import { Block, Block2 } from '../generated/subgraph-QmWi3H11QFE2PiWx6WcQkZYZdA5UasaBptUJqGn54MFux5';
3-
import { MirrorBlock } from '../generated/schema';
2+
import { Block, Block2, Block3 } from '../generated/subgraph-QmRWTEejPDDwALaquFGm6X2GBbbh5osYDXwCRRkoZ6KQhb';
3+
import { MirrorBlock, MirrorBlockBytes } from '../generated/schema';
44

55
export function handleEntity(block: Block): void {
66
let id = block.id;
@@ -23,6 +23,16 @@ export function handleEntity2(block: Block2): void {
2323
blockEntity.save();
2424
}
2525

26+
export function handleEntity3(block: Block3): void {
27+
let id = block.id;
28+
29+
let blockEntity = new MirrorBlockBytes(id);
30+
blockEntity.number = block.number;
31+
blockEntity.testMessage = block.testMessage;
32+
33+
blockEntity.save();
34+
}
35+
2636
export function loadOrCreateMirrorBlock(id: string): MirrorBlock {
2737
let block = MirrorBlock.load(id);
2838
if (!block) {

‎tests/integration-tests/subgraph-data-sources/subgraph.yaml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ dataSources:
66
name: Contract
77
network: test
88
source:
9-
address: 'QmWi3H11QFE2PiWx6WcQkZYZdA5UasaBptUJqGn54MFux5'
9+
address: 'QmRWTEejPDDwALaquFGm6X2GBbbh5osYDXwCRRkoZ6KQhb'
1010
startBlock: 0
1111
mapping:
1212
apiVersion: 0.0.7
@@ -18,4 +18,6 @@ dataSources:
1818
entity: Block
1919
- handler: handleEntity2
2020
entity: Block2
21+
- handler: handleEntity3
22+
entity: Block3
2123
file: ./src/mapping.ts

0 commit comments

Comments
 (0)