diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs index 3daa93527a..9602a08aab 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs +++ b/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 = 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") diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs index e49ac7c85c..76695224b6 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs +++ b/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, 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 = 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::) + 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) -> 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 = None; let cast: Option<&Map> = value diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs index 5da96a56e2..106d04d447 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs +++ b/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. diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs index 77aab0583f..b28cd18944 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs +++ b/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 + ); + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs index aa29f89c8d..1c338650e8 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs +++ b/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, +} diff --git a/crates/sqlbuild-analysis/src/sql_tree/main/deeper_than.rs b/crates/sqlbuild-analysis/src/sql_tree/main/deeper_than.rs new file mode 100644 index 0000000000..30d6a9998d --- /dev/null +++ b/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 +} diff --git a/crates/sqlbuild-analysis/src/sql_tree/main/mod.rs b/crates/sqlbuild-analysis/src/sql_tree/main/mod.rs index 924d435a79..26be427137 100644 --- a/crates/sqlbuild-analysis/src/sql_tree/main/mod.rs +++ b/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; diff --git a/crates/sqlbuild-rules/src/query_facts/_helpers/facts.rs b/crates/sqlbuild-rules/src/query_facts/_helpers/facts.rs index 3b468b18cb..a8d3b8e977 100644 --- a/crates/sqlbuild-rules/src/query_facts/_helpers/facts.rs +++ b/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, String> { let tokens = sqlquack_syntax::tokenize(sql, dialect).map_err(|error| error.to_string())?; Ok(QueryTokens::new(tokens, sql))