feat(contracts): prove dynamic pivots from SQLQuack parse trees

crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs

@@ -6,6 +6,7 @@ use polyglot_sql::{ComplexityGuardOptions, DataType, Dialect, Expression, ParseO
 use serde_json::{Map, Value, json};
 
 use crate::assembly::analysis_session::models::ProjectionFacts;
+use crate::relation_markers::constants::MARKER_NAME_PREFIX;
 
 use crate::assembly::analysis_session::constants::{
     ALIAS_AST_KIND, ANNOTATED_AST_KIND, ARRAY_FIXTURE_TYPE, ARRAY_PAYLOAD_KIND, CATALOG_TABLE_ARG,
@@ -95,7 +96,8 @@ pub(crate) fn relation_name(node: Option<&Value>) -> Option<&str> {
     let key: &str = node_key(node);
     let payload: &Map<String, Value> = payload(node)?;
     if key == TABLE_AST_KIND || key == COLUMN_AST_KIND {
-        return identifier_name(payload.get("name"));
+        return identifier_name(payload.get("name"))
+            .filter(|name| !name.starts_with(MARKER_NAME_PREFIX));
     }
     let called: &str = payload
         .get("name")

crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs

@@ -7,20 +7,28 @@ use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use serde_json::{Map, Value};
 use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic;
 use sqlbuild_core::panics::main::native_failure::native_failure;
+use sqlbuild_sqltext::sql_scan::main::dialect_lexical_syntax::dialect_lexical_syntax;
+use sqlquack_syntax::{CallMode, Dialect, Expression, ParseOptions};
 
 use crate::assembly::analysis_session::_helpers::dict_walk::{
     casefold, column_name, dict_list, identifier_name, is_single_wildcard, nested, node_key,
-    parse_one_or_error, payload, relation_name, render_type, to_dict, truthy, upper,
+    payload, relation_name, render_type, truthy, upper,
 };
 use crate::assembly::analysis_session::constants::{
-    ALIAS_AST_KIND, CAST_AST_KIND, DUCKDB_DIALECT, DYNAMIC_VALUE_SOURCE_KINDS, MOTHERDUCK_DIALECT,
-    PIVOT_AST_KIND, QUERY_AST_KINDS, SELECT_AST_KIND, SIMPLIFIED_PIVOT_DIALECTS,
-    SUPPORTED_PIVOT_DIALECTS, TYPE_PRESERVING_AGGREGATES, UNKNOWN_NULLABILITY, UNKNOWN_TYPE,
+    ALIAS_AST_KIND, CAST_AST_KIND, COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES, DUCKDB_DIALECT,
+    DYNAMIC_VALUE_SOURCE_KINDS, MOTHERDUCK_DIALECT, PIVOT_AST_KIND, PIVOT_MAX_AST_DEPTH,
+    QUERY_AST_KINDS, SELECT_AST_KIND, SIMPLIFIED_PIVOT_DIALECTS, SNOWFLAKE_DIALECT,
+    SUPPORTED_PIVOT_DIALECTS, TABLE_AST_KIND, TYPE_PRESERVING_AGGREGATES, UNKNOWN_NULLABILITY,
+    UNKNOWN_TYPE,
 };
 use crate::assembly::analysis_session::models::{
     ColumnFact, ContractProof, DynamicFamily, PivotModel, PivotOutcome, PivotTables,
 };
 use crate::assembly::analysis_session::types::Shapes;
+use crate::relation_markers::constants::MARKER_NAME_PREFIX;
+use crate::relation_markers::main::markers_as_names::markers_as_names;
+use crate::relation_markers::types::MarkerCall;
+use crate::sql_tree::main::deeper_than::deeper_than;
 use sqlbuild_core::constants::CTE_AST_KIND;
 
 /// What an internal native failure of a pivot proof names.
@@ -115,11 +123,81 @@ pub(crate) fn pivot_outcome(
     .map_err(|reason| native_failure(PIVOT_CONTEXT, &reason))
 }
 
-/// Python's `parse_one(...).to_dict()`, or the parse error Python reports.
+/// The one statement of `sql` as SQLQuack's tagged JSON tree, marker relations as their names.
 fn parsed(sql: &str, dialect: &str) -> Result<Result<Value, String>, String> {
-    parse_one_or_error(sql, dialect)?
-        .map(|expression| to_dict(&expression))
-        .map_or_else(|error| Ok(Err(error)), |value| value.map(Ok))
+    let (lexical_dialect, dialect): (&str, Dialect) = if dialect == SNOWFLAKE_DIALECT {
+        (SNOWFLAKE_DIALECT, Dialect::Snowflake)
+    } else {
+        (DUCKDB_DIALECT, Dialect::DuckDB)
+    };
+    let (named, calls) = markers_as_names(sql, &dialect_lexical_syntax(lexical_dialect))?;
+    let parsed = sqlquack_syntax::parse_statements_with(
+        &named,
+        dialect,
+        &ParseOptions::default(),
+        CallMode::Faithful,
+    );
+    let mut statements: Vec<Expression> = match parsed {
+        Ok(statements) => statements,
+        Err(error) => return Ok(Err(error.to_string())),
+    };
+    if statements.len() != 1 {
+        return Ok(Err(format!(
+            "Expected 1 statement, found {}",
+            statements.len()
+        )));
+    }
+    let statement: Expression = statements.pop().ok_or_else(|| "no statement".to_owned())?;
+    if deeper_than(&statement, PIVOT_MAX_AST_DEPTH) {
+        return Ok(Err(format!(
+            "E_GUARD_AST_DEPTH_EXCEEDED: value {} exceeds configured limit {PIVOT_MAX_AST_DEPTH}",
+            PIVOT_MAX_AST_DEPTH + 1
+        )));
+    }
+    let mut root: Value = serde_json::to_value(statement).map_err(|error| error.to_string())?;
+    name_marker_tables(&mut root, &calls);
+    Ok(Ok(root))
+}
+
+/// Each `__ref`, `__source` or `__seed` marker, a table or column node, renamed to the relation its call names.
+fn name_marker_tables(root: &mut Value, calls: &[MarkerCall]) {
+    let mut stack: Vec<&mut Value> = vec![root];
+    while let Some(node) = stack.pop() {
+        match node {
+            Value::Object(object) => {
+                for kind in [TABLE_AST_KIND, COLUMN_AST_KIND] {
+                    if let Some(Value::Object(relation)) = object.get_mut(kind) {
+                        rename_marker(relation.get_mut("name"), calls);
+                    }
+                }
+                stack.extend(object.values_mut());
+            }
+            Value::Array(items) => stack.extend(items.iter_mut()),
+            _ => {}
+        }
+    }
+}
+
+/// The marker name in `name`, a plain string or an identifier object, replaced by its call's relation.
+fn rename_marker(name: Option<&mut Value>, calls: &[MarkerCall]) {
+    let text: Option<&mut Value> = match name {
+        Some(Value::Object(identifier)) => identifier.get_mut("name"),
+        other => other,
+    };
+    let Some(Value::String(text)) = text else {
+        return;
+    };
+    let Some(Ok(index)) = text
+        .strip_prefix(MARKER_NAME_PREFIX)
+        .map(str::parse::<usize>)
+    else {
+        return;
+    };
+    if let Some((function, relation)) = calls.get(index)
+        && DEPENDENCY_FUNCTION_NAMES.contains(&function.as_str())
+    {
+        *text = relation.clone();
+    }
 }
 
 fn proof(
@@ -362,6 +440,18 @@ fn has_dynamic_values(field: &Value) -> bool {
         .any(|value| DYNAMIC_VALUE_SOURCE_KINDS.contains(&node_key(Some(value))))
 }
 
+/// The one argument of a call SQLQuack keeps by name, such as Snowflake's `SUM(amount)`.
+fn sole_argument(payload: &Map<String, Value>) -> Option<&Value> {
+    match payload
+        .get("args")
+        .and_then(Value::as_array)
+        .map(Vec::as_slice)
+    {
+        Some([argument]) => Some(argument),
+        _ => None,
+    }
+}
+
 /// Python's `_aggregate_fact`: aggregate name, value column and type-preserving inferred type.
 fn aggregate_fact<'a>(
     aggregate: &'a Value,
@@ -380,7 +470,7 @@ fn aggregate_fact<'a>(
         name = upper(authored);
     }
     let preserving: bool = TYPE_PRESERVING_AGGREGATES.contains(&name.as_str());
-    let value: Option<&Value> = payload.get("this");
+    let value: Option<&Value> = payload.get("this").or_else(|| sole_argument(payload));
     let mut value_name: Option<&str> = column_name(value);
     let mut inferred: Option<String> = None;
     let cast: Option<&Map<String, Value>> = value

crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs

@@ -63,6 +63,7 @@ pub(crate) const SUPPORTED_PIVOT_DIALECTS: [&str; 3] = ["duckdb", "motherduck",
 pub(crate) const SIMPLIFIED_PIVOT_DIALECTS: [&str; 2] = ["duckdb", "motherduck"];
 pub(crate) const DUCKDB_DIALECT: &str = "duckdb";
 pub(crate) const MOTHERDUCK_DIALECT: &str = "motherduck";
+pub(crate) const SNOWFLAKE_DIALECT: &str = "snowflake";
 /// The `to_dict()` node keys Python's dynamic pivot proof walks.
 pub(crate) const PIVOT_AST_KIND: &str = "pivot";
 pub(crate) const SELECT_AST_KIND: &str = "select";
@@ -93,6 +94,8 @@ pub(crate) const RENDERED_TYPE_NAMES: [(&str, &str); 12] = [
 ];
 /// The stack a standalone dynamic pivot proof parses on, as the analysis workers have.
 pub(crate) const PIVOT_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024;
+/// The deepest pivot statement the proof serialises; deeper SQL fails the proof instead of the stack.
+pub(crate) const PIVOT_MAX_AST_DEPTH: usize = 512;
 /// Python's `COMPACT_RELATION_STUB_PREFIX`: the relation name a shared query binds instead.
 pub(crate) const RELATION_STUB_PREFIX: &str = "__sqlbuild_project_input_";
 /// Python's `COMPACT_REFERENCE_MARKER_PATTERN`: a reference call and the name it reads.

crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs

@@ -9,9 +9,9 @@ use crate::assembly::analysis_session::tests::helpers::{
     session_lines, session_steps,
 };
 use crate::assembly::analysis_session::tests::test_types::{
-    CteRecoveryTestCase, KeptAnalysisTestCase, LegacyAnalysisTestCase, ModelSpec, PivotTestCase,
-    SessionFactsTestCase, SessionFailureTestCase, SessionTestCase, StarExpansionTestCase,
-    UnscheduledTestCase,
+    CteRecoveryTestCase, DeepPivotTestCase, KeptAnalysisTestCase, LegacyAnalysisTestCase,
+    ModelSpec, PivotTestCase, SessionFactsTestCase, SessionFailureTestCase, SessionTestCase,
+    StarExpansionTestCase, UnscheduledTestCase,
 };
 
 /// A model whose missing column stays untyped over a known input, so completion re-analyses it.
@@ -444,6 +444,32 @@ fn given_dynamic_pivots_when_proving_then_matches_python_or_defers() {
                 false,
             ),
         },
+        PivotTestCase {
+            description: "a Snowflake cast aggregate renders its target type",
+            dialect: "snowflake",
+            sql: "WITH base AS (SELECT customer_id, status, amount FROM __source(\"raw_orders\")) \
+                SELECT * FROM base \
+                PIVOT (MAX(CAST(amount AS DECIMAL(12, 2))) FOR status IN (ANY ORDER BY status))",
+            family: Some(("status", "amount", "MAX")),
+            expected_outcome: proven(
+                &[("customer_id", "INTEGER", "non_null")],
+                Some("DECIMAL(12, 2)"),
+                "raw_orders",
+                false,
+            ),
+        },
+        PivotTestCase {
+            description: "a Snowflake passthrough reads its relation marker as the relation",
+            dialect: "snowflake",
+            sql: "SELECT * FROM __ref(\"status_amounts\")",
+            family: Some(("status", "amount", "MAX")),
+            expected_outcome: proven(
+                &[("customer_id", "INTEGER", "unknown")],
+                Some("DOUBLE"),
+                "status_amounts",
+                false,
+            ),
+        },
         PivotTestCase {
             description: "a passthrough redeclaring the upstream families",
             dialect: "duckdb",
@@ -512,6 +538,16 @@ fn given_dynamic_pivots_when_proving_then_matches_python_or_defers() {
                  families exactly",
             ),
         },
+        PivotTestCase {
+            description: "a passthrough from a call that names no dependency",
+            dialect: "duckdb",
+            sql: "SELECT * FROM __dbt_ref(\"shop\", \"status_amounts\")",
+            family: Some(("status", "amount", "MAX")),
+            expected_outcome: failed(
+                "output must be a wildcard projection from one supported dynamic pivot or a \
+                 proven dynamic-family passthrough",
+            ),
+        },
         PivotTestCase {
             description: "no declared families",
             dialect: "duckdb",
@@ -742,3 +778,52 @@ fn given_a_finished_session_when_reading_model_analyses_then_keeps_typed_columns
         );
     }
 }
+
+#[test]
+fn given_deeply_nested_pivot_input_when_proving_then_the_proof_fails_without_aborting() {
+    let too_deep = failed(
+        "dynamic pivot SQL could not be parsed: \
+         E_GUARD_AST_DEPTH_EXCEEDED: value 513 exceeds configured limit 512",
+    );
+    let test_cases = [
+        DeepPivotTestCase {
+            description: "a 50,000-term DuckDB projection",
+            dialect: "duckdb",
+            terms: 50_000,
+            expected_outcome: too_deep.clone(),
+        },
+        DeepPivotTestCase {
+            description: "a 50,000-term Snowflake projection",
+            dialect: "snowflake",
+            terms: 50_000,
+            expected_outcome: too_deep,
+        },
+    ];
+    for test_case in test_cases {
+        let sql = format!(
+            "WITH base AS (SELECT customer_id, status, amount, 1{} AS k \
+             FROM __source(\"raw_orders\")) \
+             SELECT * FROM base PIVOT (MAX(amount) FOR status IN (ANY ORDER BY status))",
+            " + 1".repeat(test_case.terms)
+        );
+        let request = pivot_request(test_case.dialect, &sql, Some(("status", "amount", "MAX")));
+        let session = finished_session(request.tables.clone());
+
+        let standalone = prove_dynamic_contracts(&request).expect("the proofs run");
+        let in_session =
+            prove_finished_dynamic_contracts(&session, &request.models).expect("the proofs run");
+
+        assert_eq!(
+            standalone,
+            vec![test_case.expected_outcome.clone(), PivotOutcome::Absent],
+            "{}",
+            test_case.description
+        );
+        assert_eq!(
+            in_session,
+            vec![test_case.expected_outcome, PivotOutcome::Absent],
+            "{}",
+            test_case.description
+        );
+    }
+}

crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs

@@ -230,3 +230,11 @@ pub(crate) struct TypedMarkerCacheTestCase {
     pub(crate) sql: &'static str,
     pub(crate) expected_hits: &'static [bool],
 }
+
+/// A pivot over a CTE whose projection is a left-nested sum of `terms` ones.
+pub(crate) struct DeepPivotTestCase {
+    pub(crate) description: &'static str,
+    pub(crate) dialect: &'static str,
+    pub(crate) terms: usize,
+    pub(crate) expected_outcome: PivotOutcome,
+}

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

@@ -0,0 +1,15 @@
+//! `deeper_than`: whether a parsed statement nests past a depth limit.
+
+use sqlquack_syntax::{Expression, ExpressionWalk};
+
+/// Whether `statement` nests more than `limit` levels, walked iteratively so depth cannot overflow.
+pub fn deeper_than(statement: &Expression, limit: usize) -> bool {
+    let mut stack: Vec<(&Expression, usize)> = vec![(statement, 1)];
+    while let Some((node, depth)) = stack.pop() {
+        if depth > limit {
+            return true;
+        }
+        stack.extend(node.children().into_iter().map(|child| (child, depth + 1)));
+    }
+    false
+}

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

@@ -1,4 +1,5 @@
 //! Entry points.
 
+pub mod deeper_than;
 pub mod shared_dialect;
 pub mod sql_tree_nodes;

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

@@ -8,6 +8,7 @@ use crate::query_facts::models::{
     CallKind, ColumnFact, ComparisonFacts, ComparisonSpan, CteFacts, ExpressionItem, ItemFacts,
     ModelShape, ParsedFixture, QueryFacts, RelationFacts, SelectFacts, SetFacts, SourceText,
 };
+use sqlbuild_analysis::sql_tree::main::deeper_than::deeper_than;
 use sqlbuild_analysis::sql_tree::main::shared_dialect::shared_dialect;
 use sqlquack_syntax::expressions::names::models::{Identifier, Literal, TableRef};
 use sqlquack_syntax::expressions::queries::models::{Select, With};
@@ -46,18 +47,6 @@ fn parse_query(sql: &str, dialect_name: &str) -> Result<(Expression, Dialect), S
     Ok((statement, dialect))
 }
 
-/// Whether `statement` nests deeper than `limit` levels, measured without recursion.
-fn deeper_than(statement: &Expression, limit: usize) -> bool {
-    let mut stack: Vec<(&Expression, usize)> = vec![(statement, 1)];
-    while let Some((node, depth)) = stack.pop() {
-        if depth > limit {
-            return true;
-        }
-        stack.extend(node.children().into_iter().map(|child| (child, depth + 1)));
-    }
-    false
-}
-
 fn code_tokens(sql: &str, dialect: Dialect) -> Result<QueryTokens<'_>, String> {
     let tokens = sqlquack_syntax::tokenize(sql, dialect).map_err(|error| error.to_string())?;
     Ok(QueryTokens::new(tokens, sql))