fix(rules): read relation markers as tables in custom-rule syntax trees

crates/sqlbuild-analysis/src/lib.rs

@@ -37,6 +37,7 @@ pub mod python_plan;
 pub mod python_selection;
 pub mod python_values;
 pub mod query_analysis;
+pub mod relation_markers;
 pub mod relation_names;
 pub mod required_columns;
 pub mod resource_diagnostics;

crates/sqlbuild-analysis/src/relation_markers/_helpers/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod pattern;

crates/sqlbuild-analysis/src/relation_markers/_helpers/pattern.rs

@@ -0,0 +1,15 @@
+//! The compiled relation marker call pattern.
+
+use std::sync::LazyLock;
+
+use regex::Regex;
+
+use crate::relation_markers::constants::RELATION_MARKER_PATTERN;
+
+static RELATION_MARKER: LazyLock<Result<Regex, String>> =
+    LazyLock::new(|| Regex::new(RELATION_MARKER_PATTERN).map_err(|error| error.to_string()));
+
+/// The compiled relation marker call pattern, or its compile error.
+pub(crate) fn relation_marker_pattern() -> Result<&'static Regex, String> {
+    RELATION_MARKER.as_ref().map_err(Clone::clone)
+}

crates/sqlbuild-analysis/src/relation_markers/constants.rs

@@ -0,0 +1,9 @@
+//! The relation marker call pattern and the names marker calls stand as.
+
+/// The prefix of the name a marker call stands as while parsed; the call's index follows it.
+pub const MARKER_NAME_PREFIX: &str = "__sqb_";
+/// Any relation marker call: `__name("x")`, or `__name("a", "b")` as `__dbt_ref` writes it.
+pub(crate) const RELATION_MARKER_PATTERN: &str =
+    r#"(?i)\b(__[a-z_]+)\(\s*"([^"]+)"\s*(?:,\s*"([^"]+)"\s*)?\)"#;
+/// The marker whose two names read as one, joined by `__`.
+pub(crate) const DBT_REF_FUNCTION: &str = "__dbt_ref";

crates/sqlbuild-analysis/src/relation_markers/main/markers_as_names.rs

@@ -0,0 +1,60 @@
+//! `markers_as_names`: relation marker calls in code as plain names that keep every position.
+
+use regex::Regex;
+use sqlbuild_sqltext::sql_scan::main::dialect_non_code_ranges::dialect_non_code_ranges;
+use sqlbuild_sqltext::sql_scan::models::LexicalSyntax;
+
+use crate::relation_markers::_helpers::pattern::relation_marker_pattern;
+use crate::relation_markers::constants::{DBT_REF_FUNCTION, MARKER_NAME_PREFIX};
+use crate::relation_markers::types::MarkerCall;
+
+/// `sql` with each marker call in code as `__sqb_<index>`, padded so later text keeps its line and column, and the calls by index.
+pub fn markers_as_names(
+    sql: &str,
+    syntax: &LexicalSyntax,
+) -> Result<(String, Vec<MarkerCall>), String> {
+    let pattern: &Regex = relation_marker_pattern()?;
+    let mut protected: Option<Vec<(usize, usize)>> = None;
+    let mut named: String = String::with_capacity(sql.len());
+    let mut calls: Vec<MarkerCall> = Vec::new();
+    let mut cursor: usize = 0;
+    for captures in pattern.captures_iter(sql) {
+        let (Some(full), Some(function), Some(first)) =
+            (captures.get(0), captures.get(1), captures.get(2))
+        else {
+            continue;
+        };
+        let ranges = protected.get_or_insert_with(|| dialect_non_code_ranges(sql, syntax));
+        let following: usize = ranges.partition_point(|(start, _)| *start <= full.start());
+        if following > 0 && full.start() < ranges[following - 1].1 {
+            continue;
+        }
+        let function_name: String = function.as_str().to_ascii_lowercase();
+        let referenced_name: String = match (captures.get(3), function_name == DBT_REF_FUNCTION) {
+            (None, _) => first.as_str().to_owned(),
+            (Some(second), true) => format!("{}__{}", first.as_str(), second.as_str()),
+            (Some(_), false) => continue,
+        };
+        named.push_str(&sql[cursor..full.start()]);
+        named.push_str(&padded_name(
+            &format!("{MARKER_NAME_PREFIX}{}", calls.len()),
+            full.as_str(),
+        ));
+        cursor = full.end();
+        calls.push((function_name, referenced_name));
+    }
+    named.push_str(&sql[cursor..]);
+    Ok((named, calls))
+}
+
+/// `name`, then the call's line breaks, then spaces as wide as the call's last line.
+fn padded_name(name: &str, call: &str) -> String {
+    let breaks: usize = call.matches('\n').count();
+    let last_line: &str = call.rsplit('\n').next().unwrap_or(call);
+    let width: usize = if breaks == 0 {
+        call.chars().count().saturating_sub(name.chars().count())
+    } else {
+        last_line.chars().count()
+    };
+    format!("{name}{}{}", "\n".repeat(breaks), " ".repeat(width))
+}

crates/sqlbuild-analysis/src/relation_markers/main/mod.rs

@@ -0,0 +1 @@
+pub mod markers_as_names;

crates/sqlbuild-analysis/src/relation_markers/mod.rs

@@ -0,0 +1,9 @@
+//! Relation marker calls, such as `__ref("orders")`, read as plain relation names by strict parsers.
+
+mod _helpers;
+pub mod constants;
+pub mod main;
+pub mod types;
+
+#[cfg(test)]
+mod tests;

crates/sqlbuild-analysis/src/relation_markers/tests.rs

@@ -0,0 +1,4 @@
+#[path = "tests/test_markers_as_names.rs"]
+mod test_markers_as_names;
+#[path = "tests/test_types.rs"]
+mod test_types;

crates/sqlbuild-analysis/src/relation_markers/tests/test_markers_as_names.rs

@@ -0,0 +1,72 @@
+use sqlbuild_sqltext::sql_scan::main::dialect_lexical_syntax::dialect_lexical_syntax;
+
+use crate::relation_markers::main::markers_as_names::markers_as_names;
+use crate::relation_markers::tests::test_types::MarkersAsNamesTestCase;
+
+#[test]
+fn given_marker_calls_when_reading_them_as_names_then_code_calls_become_names_and_positions_hold() {
+    let test_cases = [
+        MarkersAsNamesTestCase {
+            description: "each call becomes its indexed name, padded to the call's width",
+            dialect: "snowflake",
+            sql: r#"SELECT * FROM __ref("orders") o JOIN __source("raw") r ON TRUE"#,
+            expected_sql: format!(
+                "SELECT * FROM __sqb_0{} o JOIN __sqb_1{} r ON TRUE",
+                " ".repeat(8),
+                " ".repeat(8)
+            ),
+            expected_calls: &[("__ref", "orders"), ("__source", "raw")],
+        },
+        MarkersAsNamesTestCase {
+            description: "calls in quoted text and comments stay; an upper-case call folds",
+            dialect: "snowflake",
+            sql: "SELECT '__ref(\"a\")' AS q -- __ref(\"b\")\nFROM __REF(\"Orders\") // __ref(\"c\")",
+            expected_sql: format!(
+                "SELECT '__ref(\"a\")' AS q -- __ref(\"b\")\nFROM __sqb_0{} // __ref(\"c\")",
+                " ".repeat(8)
+            ),
+            expected_calls: &[("__ref", "Orders")],
+        },
+        MarkersAsNamesTestCase {
+            description: "a call over three lines keeps its line breaks and its last line's width",
+            dialect: "duckdb",
+            sql: "FROM __ref(\n  \"orders\"\n) AS o",
+            expected_sql: "FROM __sqb_0\n\n  AS o".to_owned(),
+            expected_calls: &[("__ref", "orders")],
+        },
+        MarkersAsNamesTestCase {
+            description: "two dbt names join; a non-ASCII name pads by characters, not bytes",
+            dialect: "postgres",
+            sql: r#"FROM __dbt_ref("shop", "customers") d JOIN __ref("café") c ON TRUE"#,
+            expected_sql: format!(
+                "FROM __sqb_0{} d JOIN __sqb_1{} c ON TRUE",
+                " ".repeat(23),
+                " ".repeat(6)
+            ),
+            expected_calls: &[("__dbt_ref", "shop__customers"), ("__ref", "café")],
+        },
+        MarkersAsNamesTestCase {
+            description: "calls with a second name, or a non-literal argument, stay calls",
+            dialect: "duckdb",
+            sql: r#"FROM __ref("a", "b") JOIN __table_fn("orders", 2) ON TRUE"#,
+            expected_sql: r#"FROM __ref("a", "b") JOIN __table_fn("orders", 2) ON TRUE"#.to_owned(),
+            expected_calls: &[],
+        },
+    ];
+    for test_case in test_cases {
+        let (named, calls) =
+            markers_as_names(test_case.sql, &dialect_lexical_syntax(test_case.dialect))
+                .expect(test_case.description);
+        let expected_calls: Vec<(String, String)> = test_case
+            .expected_calls
+            .iter()
+            .map(|(function, name)| ((*function).to_owned(), (*name).to_owned()))
+            .collect();
+        assert_eq!(
+            (named, calls),
+            (test_case.expected_sql, expected_calls),
+            "{}",
+            test_case.description
+        );
+    }
+}

crates/sqlbuild-analysis/src/relation_markers/tests/test_types.rs

@@ -0,0 +1,10 @@
+//! Test case types for relation markers read as names.
+
+/// SQL in one dialect, the SQL its markers read as, and the calls they stand for.
+pub(crate) struct MarkersAsNamesTestCase {
+    pub(crate) description: &'static str,
+    pub(crate) dialect: &'static str,
+    pub(crate) sql: &'static str,
+    pub(crate) expected_sql: String,
+    pub(crate) expected_calls: &'static [(&'static str, &'static str)],
+}

crates/sqlbuild-analysis/src/relation_markers/types.rs

@@ -0,0 +1,4 @@
+//! Shapes of relation marker calls.
+
+/// A marker call's lower-case function, such as `__ref`, and the name it references.
+pub type MarkerCall = (String, String);

crates/sqlbuild-analysis/src/sql_tree/main/sql_tree_nodes.rs

@@ -1,7 +1,9 @@
 //! `sql_tree_nodes`: a rule's SQL syntax tree as nodes in pre-order.
 
+use sqlbuild_sqltext::sql_scan::main::dialect_lexical_syntax::dialect_lexical_syntax;
 use sqlquack_syntax::Expression;
 
+use crate::relation_markers::main::markers_as_names::markers_as_names;
 use crate::sql_tree::_helpers::parse::parse;
 use crate::sql_tree::_helpers::positions::line_starts;
 use crate::sql_tree::_helpers::rules_tree_parser::rules_tree_parser;
@@ -11,7 +13,7 @@ use crate::sql_tree::errors::SqlTreeError;
 use crate::sql_tree::models::SqlTreeNode;
 use crate::worker_stack::main::run_on_worker_stack::run_on_worker_stack;
 
-/// Every node of `sql` parsed in `dialect`, in pre-order, joins and CTEs owning their expressions.
+/// Every node of `sql` parsed in `dialect`, in pre-order, joins and CTEs owning their expressions; relation markers read as tables.
 pub fn sql_tree_nodes(sql: &str, dialect: &str) -> Result<Vec<SqlTreeNode>, SqlTreeError> {
     run_on_worker_stack(TREE_WORKER_STACK_BYTES, "the parser panicked", || {
         tree_nodes(sql, dialect)
@@ -22,6 +24,8 @@ pub fn sql_tree_nodes(sql: &str, dialect: &str) -> Result<Vec<SqlTreeNode>, SqlT
 fn tree_nodes(sql: &str, dialect: &str) -> Result<Vec<SqlTreeNode>, SqlTreeError> {
     let parser = rules_tree_parser(dialect)
         .ok_or_else(|| SqlTreeError::UnavailableDialect(dialect.to_owned()))?;
-    let statements: Vec<Expression> = parse(sql, parser)?;
-    walk(&statements, &line_starts(sql))
+    let (named, _) = markers_as_names(sql, &dialect_lexical_syntax(dialect))
+        .map_err(SqlTreeError::Unanalysable)?;
+    let statements: Vec<Expression> = parse(&named, parser)?;
+    walk(&statements, &line_starts(&named))
 }

crates/sqlbuild-analysis/src/sql_tree/tests/test_sql_tree.rs

@@ -6,6 +6,8 @@ use crate::sql_tree::tests::test_types::{TreeFailureTestCase, TreeShapeTestCase}
 const CTE_AND_JOIN_SQL: &str =
     "WITH x AS (SELECT 1 AS id)\nSELECT *\n  FROM x JOIN (SELECT 2 AS id) s ON x.id = s.id";
 
+const SNOWFLAKE_MARKER_SQL: &str = "WITH o AS (SELECT * FROM __ref(\"orders\"))\nSELECT o.*, c.name\nFROM o JOIN __ref(\"customers\") c ON o.id = c.id JOIN (SELECT 2 AS id) s ON s.id = o.id\nWHERE o.id IN (SELECT id FROM __source(\"raw_returns\"))";
+
 #[test]
 fn given_sql_when_building_a_rule_tree_then_clauses_own_their_queries_at_their_source_positions() {
     let test_cases = [
@@ -48,6 +50,23 @@ fn given_sql_when_building_a_rule_tree_then_clauses_own_their_queries_at_their_s
                 ("subquery", 2, 3, Some("alias")),
             ],
         },
+        TreeShapeTestCase {
+            description: "Snowflake reads relation markers as tables; nodes after them keep their columns",
+            dialect: "snowflake",
+            sql: SNOWFLAKE_MARKER_SQL,
+            expected_nodes: vec![
+                ("select", 1, 1, None),
+                ("cte", 1, 6, Some("select")),
+                ("select", 1, 12, Some("cte")),
+                ("star", 1, 12, Some("select")),
+                ("star", 2, 8, Some("select")),
+                ("join", 3, 8, Some("select")),
+                ("join", 3, 49, Some("select")),
+                ("select", 3, 54, Some("subquery")),
+                ("subquery", 3, 54, Some("join")),
+                ("select", 4, 7, Some("in")),
+            ],
+        },
         TreeShapeTestCase {
             description: "a comma-separated FROM item after an explicit join is not a join",
             dialect: "duckdb",

crates/sqlbuild-rules/src/query_facts/_helpers/shape.rs

@@ -5,6 +5,7 @@ use crate::query_facts::constants::{
 use crate::query_facts::models::{CallKind, ModelCall, ModelShape};
 use crate::rules::main::compact_sql::compact_sql;
 use sqlbuild_analysis::semantic_validation::main::normalize_mapped::normalize_mapped;
+use sqlbuild_sqltext::sql_scan::main::dialect_lexical_syntax::dialect_lexical_syntax;
 use sqlbuild_sqltext::sql_scan::main::dialect_matching_paren::dialect_matching_paren;
 use sqlbuild_sqltext::sql_scan::main::dialect_non_code_end::dialect_non_code_end;
 use sqlbuild_sqltext::sql_scan::main::quote_policy::quote_policy;
@@ -15,7 +16,7 @@ pub(crate) fn model_shape(sql: &str, dialect: &str) -> Result<ModelShape, String
         dollar_quotes: true,
         ..quote_policy(dialect)
     };
-    let (stubbed, stub_offsets, calls) = stub_calls(sql, policy, &lexical_syntax(dialect));
+    let (stubbed, stub_offsets, calls) = stub_calls(sql, policy, &dialect_lexical_syntax(dialect));
     let normalized = normalize_mapped(&stubbed, dialect)?;
     let end = stub_offsets.last().copied().unwrap_or_default();
     let offsets: Vec<usize> = normalized
@@ -77,46 +78,6 @@ fn is_word_byte(byte: u8) -> bool {
     byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'$'
 }
 
-/// The comment and quoted-text syntax the rules parser tokenized for `dialect`.
-fn lexical_syntax(dialect: &str) -> LexicalSyntax {
-    let strings = |quotes: &[&str]| quotes.iter().map(|quote| (*quote).to_owned()).collect();
-    let lines = |prefixes: &[&str]| prefixes.iter().map(|prefix| (*prefix).to_owned()).collect();
-    match dialect.to_ascii_lowercase().as_str() {
-        "bigquery" => LexicalSyntax {
-            backslash_escape_quotes: strings(&["'", "\""]),
-            raw_string_prefix: true,
-            triple_quoted_strings: true,
-            line_comment_prefixes: lines(&["--", "#"]),
-            ..LexicalSyntax::default()
-        },
-        "snowflake" => LexicalSyntax {
-            backslash_escape_quotes: strings(&["'"]),
-            line_comment_prefixes: lines(&["--", "//"]),
-            ..LexicalSyntax::default()
-        },
-        "clickhouse" => LexicalSyntax {
-            backslash_escape_quotes: strings(&["'"]),
-            line_comment_prefixes: lines(&["--", "#", "//"]),
-            ..LexicalSyntax::default()
-        },
-        "databricks" | "duckdb" => LexicalSyntax {
-            line_comment_prefixes: lines(&["--"]),
-            ..LexicalSyntax::default()
-        },
-        "mssql" | "sqlserver" | "tsql" => LexicalSyntax {
-            nested_block_comments: true,
-            line_comment_prefixes: lines(&["--"]),
-            ..LexicalSyntax::default()
-        },
-        _ => LexicalSyntax {
-            escape_string_prefix: true,
-            nested_block_comments: true,
-            line_comment_prefixes: lines(&["--"]),
-            ..LexicalSyntax::default()
-        },
-    }
-}
-
 fn is_quote_start(byte: u8) -> bool {
     matches!(byte, b'\'' | b'"' | b'`' | b'$')
 }

crates/sqlbuild-sqltext/src/sql_scan/main/dialect_lexical_syntax.rs

@@ -0,0 +1,43 @@
+//! `dialect_lexical_syntax`: the comment and quoted-text syntax of a dialect's SQL, by dialect name.
+
+use crate::sql_scan::models::LexicalSyntax;
+
+/// The comment and quoted-text syntax that SQLQuack tokenizes for `dialect`.
+pub fn dialect_lexical_syntax(dialect: &str) -> LexicalSyntax {
+    let strings = |quotes: &[&str]| quotes.iter().map(|quote| (*quote).to_owned()).collect();
+    let lines = |prefixes: &[&str]| prefixes.iter().map(|prefix| (*prefix).to_owned()).collect();
+    match dialect.to_ascii_lowercase().as_str() {
+        "bigquery" => LexicalSyntax {
+            backslash_escape_quotes: strings(&["'", "\""]),
+            raw_string_prefix: true,
+            triple_quoted_strings: true,
+            line_comment_prefixes: lines(&["--", "#"]),
+            ..LexicalSyntax::default()
+        },
+        "snowflake" => LexicalSyntax {
+            backslash_escape_quotes: strings(&["'"]),
+            line_comment_prefixes: lines(&["--", "//"]),
+            ..LexicalSyntax::default()
+        },
+        "clickhouse" => LexicalSyntax {
+            backslash_escape_quotes: strings(&["'"]),
+            line_comment_prefixes: lines(&["--", "#", "//"]),
+            ..LexicalSyntax::default()
+        },
+        "databricks" | "duckdb" => LexicalSyntax {
+            line_comment_prefixes: lines(&["--"]),
+            ..LexicalSyntax::default()
+        },
+        "mssql" | "sqlserver" | "tsql" => LexicalSyntax {
+            nested_block_comments: true,
+            line_comment_prefixes: lines(&["--"]),
+            ..LexicalSyntax::default()
+        },
+        _ => LexicalSyntax {
+            escape_string_prefix: true,
+            nested_block_comments: true,
+            line_comment_prefixes: lines(&["--"]),
+            ..LexicalSyntax::default()
+        },
+    }
+}

crates/sqlbuild-sqltext/src/sql_scan/main/mod.rs

@@ -1,4 +1,5 @@
 pub mod comment_end;
+pub mod dialect_lexical_syntax;
 pub mod dialect_matching_paren;
 pub mod dialect_non_code_end;
 pub mod dialect_non_code_ranges;

tests/e2e/src/sqlbuild/cli/commands/main/rules/_test_types.py

@@ -384,3 +384,5 @@ class SyntaxTreeRuleE2ETestCase:
     description: str
     model_sql: str
     expected_messages: tuple[str, ...]
+    adapter: str = "duckdb"
+    other_models: dict[str, str] = field(default_factory=dict)