fix(graph): allow more complex dataset and table names

Author: isum

graph/src/amp/manifest/data_source/raw.rs

@@ -144,7 +144,9 @@ impl RawSource {
             end_block,
         } = self;
 
-        validate_ident(&dataset).map_err(|e| e.source_context("invalid `dataset`"))?;
+        if dataset.is_empty() {
+            return Err(Error::InvalidValue(anyhow!("`dataset` cannot be empty")));
+        }
         Self::validate_tables(&tables)?;
 
         let address = address.unwrap_or(Address::ZERO);
@@ -180,8 +182,11 @@ impl RawSource {
         }
 
         for (i, table) in tables.iter().enumerate() {
-            validate_ident(table)
-                .map_err(|e| e.source_context(format!("invalid `table` at index {i}")))?;
+            if table.is_empty() {
+                return Err(Error::InvalidValue(anyhow!(
+                    "`table` at index {i} cannot be empty"
+                )));
+            }
         }
 
         Ok(())

graph/src/amp/sql/query_builder/block_range_query.rs

@@ -8,7 +8,7 @@ use ahash::RandomState;
 use alloy::primitives::BlockNumber;
 use sqlparser_latest::ast::{self, VisitMut, VisitorMut};
 
-use super::{extract_tables, normalize_table, parse_query};
+use super::{extract_tables, parse_query, TableReference};
 
 /// Limits the query execution to the specified block range.
 ///
@@ -56,7 +56,7 @@ pub(super) fn new_block_range_query<'a>(
 fn new_tables_to_ctes_mapping(
     query: &ast::Query,
     hasher: &mut impl Hasher,
-) -> BTreeMap<String, String> {
+) -> BTreeMap<TableReference, String> {
     extract_tables(query)
         .into_iter()
         .map(|table| {
@@ -69,12 +69,12 @@ fn new_tables_to_ctes_mapping(
 
 /// Visits the SQL query AST and replaces referenced table names with CTE names.
 struct TableReplacer {
-    tables_to_ctes_mapping: BTreeMap<String, String>,
+    tables_to_ctes_mapping: BTreeMap<TableReference, String>,
 }
 
 impl TableReplacer {
     /// Creates a new table replacer.
-    fn new(tables_to_ctes_mapping: BTreeMap<String, String>) -> Self {
+    fn new(tables_to_ctes_mapping: BTreeMap<TableReference, String>) -> Self {
         Self {
             tables_to_ctes_mapping,
         }
@@ -86,7 +86,10 @@ impl TableReplacer {
             return;
         };
 
-        let Some(cte_table) = self.tables_to_ctes_mapping.get(&normalize_table(name)) else {
+        let Some(cte_table) = self
+            .tables_to_ctes_mapping
+            .get(&TableReference::with_object_name(name))
+        else {
             return;
         };
 
@@ -133,20 +136,20 @@ mod tests {
         assert_eq!(
             block_range_query,
             parse_query(
-                "
-                WITH block_range_14621009630487609643 AS (
-                    SELECT * FROM d WHERE _block_num BETWEEN 0 AND 1000000
+                r#"
+                WITH block_range_1164572571450379730 AS (
+                    SELECT * FROM "d" WHERE _block_num BETWEEN 0 AND 1000000
                 ),
-                source_14621009630487609643 AS (
-                    SELECT a, b, c FROM block_range_14621009630487609643 AS d
+                source_1164572571450379730 AS (
+                    SELECT a, b, c FROM block_range_1164572571450379730 AS d
                 )
                 SELECT
-                    source_14621009630487609643.*
+                    source_1164572571450379730.*
                 FROM
-                    source_14621009630487609643
+                    source_1164572571450379730
                 ORDER BY
-                    source_14621009630487609643.b
-                "
+                    source_1164572571450379730.b
+                "#
             )
             .unwrap(),
         )
@@ -162,23 +165,23 @@ mod tests {
         assert_eq!(
             block_range_query,
             parse_query(
-                "
-                WITH block_range_14621009630487609643 AS (
-                    SELECT * FROM d WHERE _block_num BETWEEN 0 AND 1000000
+                r#"
+                WITH block_range_1164572571450379730 AS (
+                    SELECT * FROM "d" WHERE _block_num BETWEEN 0 AND 1000000
                 ),
-                block_range_12377422807768256314 AS (
-                    SELECT * FROM e WHERE _block_num BETWEEN 0 AND 1000000
+                block_range_13063992259633584610 AS (
+                    SELECT * FROM "e" WHERE _block_num BETWEEN 0 AND 1000000
                 ),
-                source_12377422807768256314 AS (
-                    SELECT a, b, c FROM block_range_14621009630487609643 AS d JOIN block_range_12377422807768256314 AS e ON e.e = d.d
+                source_13063992259633584610 AS (
+                    SELECT a, b, c FROM block_range_1164572571450379730 AS d JOIN block_range_13063992259633584610 AS e ON e.e = d.d
                 )
                 SELECT
-                    source_12377422807768256314.*
+                    source_13063992259633584610.*
                 FROM
-                    source_12377422807768256314
+                    source_13063992259633584610
                 ORDER BY
-                    source_12377422807768256314.b
-                "
+                    source_13063992259633584610.b
+                "#
             )
             .unwrap(),
         )

graph/src/amp/sql/query_builder/mod.rs

@@ -22,7 +22,7 @@ use self::{
     event_signature_resolver::resolve_event_signatures,
     parser::parse_query,
     source_address_resolver::resolve_source_address,
-    table_extractor::{extract_tables, normalize_table},
+    table_extractor::{extract_tables, TableReference},
     table_validator::validate_tables,
 };
 

graph/src/amp/sql/query_builder/table_extractor.rs

@@ -1,34 +1,84 @@
-use std::{collections::BTreeSet, ops::ControlFlow};
+use std::{collections::BTreeSet, fmt, ops::ControlFlow};
 
-use itertools::Itertools;
 use sqlparser_latest::ast::{self, Visit, Visitor};
 
 /// Returns all tables that are referenced by the SQL query.
 ///
 /// The table names are lowercased and quotes are ignored.
-pub(super) fn extract_tables(query: &ast::Query) -> BTreeSet<String> {
+pub(super) fn extract_tables(query: &ast::Query) -> BTreeSet<TableReference> {
     let mut table_extractor = TableExtractor::new();
     let _: ControlFlow<()> = Visit::visit(query, &mut table_extractor);
 
     table_extractor.tables
 }
 
-/// Returns the normalized table name.
+/// Contains a normalized table reference.
 ///
-/// The table name is lowercased and quotes are ignored.
-pub(super) fn normalize_table(object_name: &ast::ObjectName) -> String {
-    object_name
-        .0
-        .iter()
-        .map(|part| match part {
-            ast::ObjectNamePart::Identifier(ident) => ident.value.to_lowercase(),
-        })
-        .join(".")
+/// Used to compare physical table references with CTE names and custom tables.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub(super) struct TableReference(ast::ObjectName);
+
+impl TableReference {
+    const QUOTE_STYLE: char = '"';
+
+    /// Creates a new table reference from a custom dataset and table.
+    pub(super) fn new(dataset: &str, table: &str) -> Self {
+        Self(
+            vec![
+                ast::Ident::with_quote(Self::QUOTE_STYLE, dataset),
+                ast::Ident::with_quote(Self::QUOTE_STYLE, table),
+            ]
+            .into(),
+        )
+    }
+
+    /// Creates a new table reference from an object name.
+    pub(super) fn with_object_name(object_name: &ast::ObjectName) -> Self {
+        Self::with_idents(
+            object_name
+                .0
+                .iter()
+                .map(|object_name_part| match object_name_part {
+                    ast::ObjectNamePart::Identifier(ident) => ident,
+                }),
+        )
+    }
+
+    /// Creates a new table reference from a list of identifiers.
+    pub(super) fn with_idents<'a>(idents: impl IntoIterator<Item = &'a ast::Ident>) -> Self {
+        Self(
+            idents
+                .into_iter()
+                .map(|ident| {
+                    let ast::Ident {
+                        value,
+                        quote_style,
+                        span: _,
+                    } = ident;
+
+                    ast::Ident::with_quote(Self::QUOTE_STYLE, {
+                        if quote_style.is_none() {
+                            value.to_lowercase()
+                        } else {
+                            value.to_owned()
+                        }
+                    })
+                })
+                .collect::<Vec<_>>()
+                .into(),
+        )
+    }
+}
+
+impl fmt::Display for TableReference {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
 }
 
 /// Visits the SQL query AST and extracts referenced table names, ignoring CTEs.
 struct TableExtractor {
-    tables: BTreeSet<String>,
+    tables: BTreeSet<TableReference>,
     cte_stack: CteStack,
 }
 
@@ -47,13 +97,12 @@ impl TableExtractor {
             return;
         };
 
-        let table = normalize_table(name);
-
-        if self.cte_stack.contains(&table) {
+        let table_reference = TableReference::with_object_name(name);
+        if self.cte_stack.contains(&table_reference) {
             return;
         }
 
-        self.tables.insert(table);
+        self.tables.insert(table_reference);
     }
 }
 
@@ -81,7 +130,7 @@ impl Visitor for TableExtractor {
 
 /// Maintains a list of active CTEs for each subquery scope.
 struct CteStack {
-    stack: Vec<BTreeSet<String>>,
+    stack: Vec<BTreeSet<TableReference>>,
 }
 
 impl CteStack {
@@ -90,9 +139,11 @@ impl CteStack {
         Self { stack: Vec::new() }
     }
 
-    /// Returns `true` if the `table_name` is present in the CTE list at any scope.
-    fn contains(&self, table_name: &str) -> bool {
-        self.stack.iter().any(|scope| scope.contains(table_name))
+    /// Returns `true` if the `table_reference` is present in the CTE list at any scope.
+    fn contains(&self, table_reference: &TableReference) -> bool {
+        self.stack
+            .iter()
+            .any(|scope| scope.contains(table_reference))
     }
 
     /// Creates a new subquery scope with all the CTEs of the current `query`.
@@ -101,7 +152,7 @@ impl CteStack {
             Some(with) => with
                 .cte_tables
                 .iter()
-                .map(|cte_table| cte_table.alias.name.value.to_lowercase())
+                .map(|cte_table| TableReference::with_idents([&cte_table.alias.name]))
                 .collect(),
             None => BTreeSet::new(),
         };
@@ -126,28 +177,31 @@ mod tests {
                 #[test]
                 fn $name() {
                     let query = parse_query($input).unwrap();
-                    assert_eq!(extract_tables(&query), $expected.into_iter().map(Into::into).collect());
+                    assert_eq!(
+                        extract_tables(&query).into_iter().map(|table| table.to_string()).collect::<Vec<_>>(),
+                        $expected.into_iter().map(|table| table.to_string()).collect::<Vec<_>>()
+                    );
                 }
             )*
         };
     }
 
     test_extract_tables! {
-        one_table: "SELECT a FROM b" => ["b"],
-        multiple_tables_with_one_join: "SELECT a FROM b JOIN c ON c.c = b.b" => ["b", "c"],
-        multiple_tables_with_multiple_joins: "SELECT a FROM b JOIN c ON c.c = b.b JOIN d ON d.d = b.b" => ["b", "c", "d"],
-        one_table_with_one_cte: "WITH a AS (SELECT * FROM b) SELECT * FROM a" => ["b"],
-        one_table_with_multiple_ctes: "WITH a AS (SELECT * FROM b), c AS (SELECT * FROM a) SELECT * FROM c" => ["b"],
-        multiple_tables_with_multiple_ctes: "WITH a AS (SELECT * FROM b), c AS (SELECT * FROM d) SELECT * FROM a JOIN c ON c.c = a.a" => ["b", "d"],
-        multiple_tables_with_nested_ctes: "WITH a AS (WITH b AS (SELECT * FROM c) SELECT * FROM d JOIN b ON b.b = d.d) SELECT * FROM a" => ["c", "d"],
-        multiple_tables_with_union: "SELECT a FROM b UNION SELECT c FROM d" => ["b", "d"],
-        multiple_tables_with_union_all: "SELECT a FROM b UNION ALL SELECT c FROM d" => ["b", "d"],
-
-        namespace_is_preserved: "SELECT a FROM b.c" => ["b.c"],
-        catalog_is_preserved: "SELECT a FROM b.c.d" => ["b.c.d"],
-        tables_are_lowercased: "SELECT a FROM B.C" => ["b.c"],
-        single_quotes_in_tables_are_ignored: "SELECT a FROM 'B'.'C'" => ["b.c"],
-        double_quotes_in_tables_are_ignored: r#"SELECT a FROM "B"."C""# => ["b.c"],
-        backticks_in_tables_are_ignored: "SELECT a FROM `B`.`C`" => ["b.c"],
+        one_table: "SELECT a FROM b" => [r#""b""#],
+        multiple_tables_with_one_join: "SELECT a FROM b JOIN c ON c.c = b.b" => [r#""b""#, r#""c""#],
+        multiple_tables_with_multiple_joins: "SELECT a FROM b JOIN c ON c.c = b.b JOIN d ON d.d = b.b" => [r#""b""#, r#""c""#, r#""d""#],
+        one_table_with_one_cte: "WITH a AS (SELECT * FROM b) SELECT * FROM a" => [r#""b""#],
+        one_table_with_multiple_ctes: "WITH a AS (SELECT * FROM b), c AS (SELECT * FROM a) SELECT * FROM c" => [r#""b""#],
+        multiple_tables_with_multiple_ctes: "WITH a AS (SELECT * FROM b), c AS (SELECT * FROM d) SELECT * FROM a JOIN c ON c.c = a.a" => [r#""b""#, r#""d""#],
+        multiple_tables_with_nested_ctes: "WITH a AS (WITH b AS (SELECT * FROM c) SELECT * FROM d JOIN b ON b.b = d.d) SELECT * FROM a" => [r#""c""#, r#""d""#],
+        multiple_tables_with_union: "SELECT a FROM b UNION SELECT c FROM d" => [r#""b""#, r#""d""#],
+        multiple_tables_with_union_all: "SELECT a FROM b UNION ALL SELECT c FROM d" => [r#""b""#, r#""d""#],
+
+        namespace_is_preserved: "SELECT a FROM b.c" => [r#""b"."c""#],
+        catalog_is_preserved: "SELECT a FROM b.c.d" => [r#""b"."c"."d""#],
+        unquoted_tables_are_lowercased: "SELECT a FROM B.C" => [r#""b"."c""#],
+        single_quotes_in_tables_are_converted_to_double_quotes: "SELECT a FROM 'B'.'C'" => [r#""B"."C""#],
+        double_quotes_in_tables_are_preserved: r#"SELECT a FROM "B"."C""# => [r#""B"."C""#],
+        backticks_in_tables_are_converted_to_double_quotes: "SELECT a FROM `B`.`C`" => [r#""B"."C""#],
     }
 }