tests: add unit test for sql query

Author: shiyasmohd

store/postgres/src/relational/query_tests.rs

@@ -2,14 +2,16 @@ use std::{collections::BTreeSet, sync::Arc};
 
 use diesel::{debug_query, pg::Pg};
 use graph::{
+    data_source::CausalityRegion,
     prelude::{r, serde_json as json, DeploymentHash, EntityFilter},
     schema::InputSchema,
 };
 
 use crate::{
+    block_range::BoundSide,
     layout_for_tests::{make_dummy_site, Namespace},
     relational::{Catalog, ColumnType, Layout},
-    relational_queries::FromColumnValue,
+    relational_queries::{FindRangeQuery, FromColumnValue},
 };
 
 use crate::relational_queries::Filter;
@@ -86,3 +88,112 @@ fn prefix() {
     let filter = EntityFilter::In("address".to_string(), vec!["0xbeef".into()]);
     filter_contains(filter, r#"substring(c."address", 1, 64) in ($1)"#);
 }
+
+#[test]
+fn find_range_query_id_type_casting() {
+    let string_schema = "
+    type StringEntity @entity {
+        id: String!,
+        name: String
+    }";
+
+    let bytes_schema = "
+    type BytesEntity @entity {
+        id: Bytes!,
+        address: Bytes
+    }";
+
+    let int8_schema = "
+    type Int8Entity @entity {
+        id: Int8!,
+        value: Int8
+    }";
+
+    let string_layout = test_layout(string_schema);
+    let bytes_layout = test_layout(bytes_schema);
+    let int8_layout = test_layout(int8_schema);
+
+    let string_table = string_layout
+        .table_for_entity(
+            &string_layout
+                .input_schema
+                .entity_type("StringEntity")
+                .unwrap(),
+        )
+        .unwrap();
+    let bytes_table = bytes_layout
+        .table_for_entity(
+            &bytes_layout
+                .input_schema
+                .entity_type("BytesEntity")
+                .unwrap(),
+        )
+        .unwrap();
+    let int8_table = int8_layout
+        .table_for_entity(&int8_layout.input_schema.entity_type("Int8Entity").unwrap())
+        .unwrap();
+
+    let causality_region = CausalityRegion::ONCHAIN;
+    let bound_side = BoundSide::Lower;
+    let block_range = 100..200;
+
+    test_id_type_casting(
+        string_table.as_ref(),
+        "id::bytea",
+        "String ID should be cast to bytea",
+    );
+    test_id_type_casting(bytes_table.as_ref(), "id", "Bytes ID should remain as id");
+    test_id_type_casting(
+        int8_table.as_ref(),
+        "id::text::bytea",
+        "Int8 ID should be cast to text then bytea",
+    );
+
+    let tables = vec![
+        string_table.as_ref(),
+        bytes_table.as_ref(),
+        int8_table.as_ref(),
+    ];
+    let query = FindRangeQuery::new(&tables, causality_region, bound_side, block_range);
+    let sql = debug_query::<Pg, _>(&query).to_string();
+
+    assert!(
+        sql.contains("id::bytea"),
+        "String entity ID casting should be present in UNION query"
+    );
+    assert!(
+        sql.contains("id as id"),
+        "Bytes entity ID should be present in UNION query"
+    );
+    assert!(
+        sql.contains("id::text::bytea"),
+        "Int8 entity ID casting should be present in UNION query"
+    );
+
+    assert!(
+        sql.contains("union all"),
+        "Multiple tables should generate UNION ALL queries"
+    );
+    assert!(
+        sql.contains("order by block_number, entity, id"),
+        "Query should end with proper ordering"
+    );
+}
+
+fn test_id_type_casting(table: &crate::relational::Table, expected_cast: &str, test_name: &str) {
+    let causality_region = CausalityRegion::ONCHAIN;
+    let bound_side = BoundSide::Lower;
+    let block_range = 100..200;
+
+    let tables = vec![table];
+    let query = FindRangeQuery::new(&tables, causality_region, bound_side, block_range);
+    let sql = debug_query::<Pg, _>(&query).to_string();
+
+    assert!(
+        sql.contains(expected_cast),
+        "{}: Expected '{}' in SQL, got: {}",
+        test_name,
+        expected_cast,
+        sql
+    );
+}