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 new file mode 100644 index 0000000000..d11f30c0ae --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs @@ -0,0 +1,156 @@ +//! Python's dict walk over a parsed query's `to_dict()` payload, on the same serde values. + +use serde_json::{Map, Value}; + +use crate::assembly::analysis_session::constants::{ + COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES, FUNCTION_AST_KIND, RENDERED_TYPE_NAMES, + STAR_AST_KIND, TABLE_AST_KIND, +}; + +/// Python's `bool(value)`. +pub(crate) fn truthy(value: Option<&Value>) -> bool { + match value { + Some(Value::Bool(flag)) => *flag, + Some(Value::Null) | None => false, + Some(Value::String(text)) => !text.is_empty(), + Some(Value::Array(items)) => !items.is_empty(), + Some(Value::Object(items)) => !items.is_empty(), + Some(Value::Number(number)) => number.as_f64() != Some(0.0), + } +} + +/// The single key of a one-key object, or "" for anything else. +pub(crate) fn node_key(node: Option<&Value>) -> &str { + let Some(object) = node + .and_then(Value::as_object) + .filter(|object| object.len() == 1) + else { + return ""; + }; + object.keys().next().map_or("", String::as_str) +} + +/// The payload under a node's single key. +pub(crate) fn payload(node: Option<&Value>) -> Option<&Map> { + let key: &str = node_key(node); + node.and_then(|node| node.get(key)) + .and_then(Value::as_object) +} + +/// Python's `_dict_list`: the objects of a list, or nothing. +pub(crate) fn dict_list(value: Option<&Value>) -> Vec<&Value> { + match value.and_then(Value::as_array) { + Some(items) => objects(items), + None => Vec::new(), + } +} + +fn objects(items: &[Value]) -> Vec<&Value> { + items.iter().filter(|item| item.is_object()).collect() +} + +/// Python's `_nested`: follow object keys. +pub(crate) fn nested<'a>(node: Option<&'a Value>, keys: &[&str]) -> Option<&'a Value> { + let mut current: Option<&Value> = node; + for key in keys { + current = current + .and_then(Value::as_object) + .and_then(|object| object.get(*key)); + } + current +} + +/// Python's `_identifier_name`. +pub(crate) fn identifier_name(node: Option<&Value>) -> Option<&str> { + match node { + Some(Value::String(name)) => Some(name), + Some(Value::Object(object)) => object + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()), + _ => None, + } +} + +/// Python's `_column_name`. +pub(crate) fn column_name(node: Option<&Value>) -> Option<&str> { + if node_key(node) != COLUMN_AST_KIND { + return None; + } + identifier_name(payload(node)?.get("name")) +} + +/// Python's `_relation_name`: a table, a column or a dependency call's first argument. +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")); + } + let called: &str = payload + .get("name") + .and_then(Value::as_str) + .unwrap_or_default(); + if key == FUNCTION_AST_KIND && DEPENDENCY_FUNCTION_NAMES.contains(&casefold(called).as_str()) { + return column_name(dict_list(payload.get("args")).first().copied()); + } + None +} + +/// Python's `_is_single_wildcard`. +pub(crate) fn is_single_wildcard(value: Option<&Value>) -> bool { + let expressions: Vec<&Value> = dict_list(value); + let [only] = expressions.as_slice() else { + return false; + }; + if node_key(Some(only)) != STAR_AST_KIND { + return false; + } + let Some(star) = only.get(STAR_AST_KIND).and_then(Value::as_object) else { + return false; + }; + !["except", "replace", "rename", "table"] + .iter() + .any(|key| truthy(star.get(*key))) +} + +/// Python's `_render_type` of a cast target. +pub(crate) fn render_type(node: Option<&Value>) -> Option { + let node: &Map = node?.as_object()?; + let raw_type: &str = node + .get("data_type") + .and_then(Value::as_str) + .filter(|raw| !raw.is_empty())?; + let folded: String = casefold(raw_type); + let rendered: String = RENDERED_TYPE_NAMES + .iter() + .find(|(name, _)| *name == folded) + .map_or_else( + || raw_type.replace('_', " ").to_ascii_uppercase(), + |(_, rendered)| (*rendered).to_owned(), + ); + if let Some(length) = python_int(node.get("length")) { + return Some(format!("{rendered}({length})")); + } + let Some(precision) = python_int(node.get("precision")) else { + return Some(rendered); + }; + Some(match python_int(node.get("scale")) { + Some(scale) => format!("{rendered}({precision}, {scale})"), + None => format!("{rendered}({precision})"), + }) +} + +/// Python's `str(value)` for a value `isinstance(value, int)` accepts. +fn python_int(value: Option<&Value>) -> Option { + match value? { + Value::Bool(flag) => Some(if *flag { "True" } else { "False" }.to_owned()), + Value::Number(number) if number.is_i64() || number.is_u64() => Some(number.to_string()), + _ => None, + } +} + +/// Python's `str.casefold` for the ASCII text a pivot proof compares. +pub(crate) fn casefold(text: &str) -> String { + text.to_ascii_lowercase() +} 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 new file mode 100644 index 0000000000..78f0b3bdba --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs @@ -0,0 +1,675 @@ +//! Python's `analyze_dynamic_column_contract` over the crate's parse of the model SQL. + +use std::collections::{HashMap, HashSet}; + +use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions}; +use serde_json::{Map, Value, json}; +use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic; + +use crate::assembly::analysis_session::_helpers::dict_walk::{ + casefold, column_name, dict_list, identifier_name, is_single_wildcard, nested, node_key, + payload, relation_name, render_type, truthy, +}; +use crate::assembly::analysis_session::constants::{ + ALIAS_AST_KIND, CAST_AST_KIND, CTE_AST_KIND, DUCKDB_DIALECT, DYNAMIC_VALUE_SOURCE_KINDS, + MAX_FUNCTION_CALL_DEPTH, MOTHERDUCK_DIALECT, PIVOT_AST_KIND, QUERY_AST_KINDS, SELECT_AST_KIND, + SIMPLIFIED_PIVOT_DIALECTS, SUPPORTED_PIVOT_DIALECTS, TYPE_PRESERVING_AGGREGATES, + UNKNOWN_NULLABILITY, UNKNOWN_TYPE, +}; +use crate::assembly::analysis_session::models::{ + ColumnFact, ContractProof, DynamicFamily, PivotOutcome, +}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; + +/// Stands in for Python's `{}`: neither has a node key. +static MISSING_EXPRESSION: Value = Value::Null; + +/// The relation facts and declared families a pivot proof reads. +pub(crate) struct PivotFacts<'a> { + pub(crate) dialect: &'a str, + pub(crate) column_types: &'a Shapes, + pub(crate) authoritative_types: &'a Shapes, + pub(crate) column_nullability: &'a Shapes, + pub(crate) families_by_table: &'a [(String, Vec)], +} + +/// One CTE body by folded name; None for a CTE with column aliases. +type CteMap<'a> = HashMap>; + +/// The resolved output boundary: a pivot payload or a passthrough relation. +enum Boundary<'a> { + Pivot(&'a Map), + Passthrough(&'a str), +} + +/// The proof Python computes, or a deferral where its result is not reproduced exactly. +pub(crate) fn pivot_outcome( + sql: &str, + families: &[DynamicFamily], + facts: &PivotFacts<'_>, +) -> PivotOutcome { + if families.is_empty() { + return PivotOutcome::Absent; + } + let folded_dialect: String = casefold(facts.dialect); + if !SUPPORTED_PIVOT_DIALECTS.contains(&folded_dialect.as_str()) { + return failure(format!( + "adapter dialect '{}' does not support compiler-proven dynamic pivots", + facts.dialect + )); + } + if !ascii_inputs(sql, families, facts) { + return PivotOutcome::Deferred; + } + let parse_dialect: &str = if folded_dialect == MOTHERDUCK_DIALECT { + DUCKDB_DIALECT + } else { + folded_dialect.as_str() + }; + let outcome: Result = catch_compiler_panic(|| { + Ok(match parsed(sql, parse_dialect)? { + Some(root) => proof(&root, &folded_dialect, families, facts), + None => PivotOutcome::Deferred, + }) + }); + match outcome { + Ok(outcome) => outcome, + Err(_) => PivotOutcome::Deferred, + } +} + +/// Python's `parse_one(...).to_dict()`, or None where Python reports a parse error. +fn parsed(sql: &str, dialect: &str) -> Result, String> { + let guard: ComplexityGuardOptions = + serde_json::from_value(json!({"maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH})) + .map_err(|error| error.to_string())?; + let options: ParseOptions = ParseOptions { + complexity_guard: Some(guard), + }; + let parsed: Result, _> = Dialect::get_by_name(dialect) + .ok_or_else(|| format!("unknown dialect {dialect}"))? + .parse_with_options(sql, &options); + let expressions: Vec = match parsed { + Ok(expressions) => expressions, + Err(_) => return Ok(None), + }; + let [expression] = expressions.as_slice() else { + return Ok(None); + }; + serde_json::to_value(expression) + .map(Some) + .map_err(|error| error.to_string()) +} + +/// Casefolding equals Python's only for ASCII text. +fn ascii_inputs(sql: &str, families: &[DynamicFamily], facts: &PivotFacts<'_>) -> bool { + let shapes = [ + facts.column_types, + facts.authoritative_types, + facts.column_nullability, + ]; + sql.is_ascii() + && families.iter().all(ascii_family) + && facts.families_by_table.iter().all(ascii_families) + && shapes.into_iter().flatten().all(ascii_shape) +} + +fn ascii_families((name, families): &(String, Vec)) -> bool { + name.is_ascii() && families.iter().all(ascii_family) +} + +fn ascii_shape((name, shape): &(String, Pairs)) -> bool { + name.is_ascii() && shape.iter().all(|(column, _)| column.is_ascii()) +} + +fn proof( + root: &Value, + dialect: &str, + families: &[DynamicFamily], + facts: &PivotFacts<'_>, +) -> PivotOutcome { + if !root.is_object() { + return failure("dynamic pivot SQL did not produce a structured query".to_owned()); + } + let ctes: CteMap<'_> = collect_ctes(root); + match resolve_boundary(Some(root), &ctes, &HashSet::new()) { + None => failure( + "output must be a wildcard projection from one supported dynamic pivot or a proven \ + dynamic-family passthrough" + .to_owned(), + ), + Some(Boundary::Passthrough(name)) => passthrough(name, families, facts), + Some(Boundary::Pivot(pivot)) => pivot_proof( + pivot, + &PivotContext { + dialect, + ctes: &ctes, + facts, + bare: node_key(Some(root)) == PIVOT_AST_KIND, + }, + families, + ), + } +} + +struct PivotContext<'a, 'b> { + dialect: &'a str, + ctes: &'a CteMap<'b>, + facts: &'a PivotFacts<'a>, + bare: bool, +} + +fn passthrough(name: &str, families: &[DynamicFamily], facts: &PivotFacts<'_>) -> PivotOutcome { + let upstream: Option<&Vec> = casefold_lookup(facts.families_by_table, name); + if !upstream.is_some_and(|upstream| families_equal(families, upstream)) { + return failure(format!( + "dynamic-family passthrough from '{name}' must redeclare the upstream families exactly" + )); + } + let Some(columns) = external_schema(name, facts.column_types, facts.column_nullability) else { + return failure(format!( + "dynamic-family passthrough source '{name}' has no authoritative schema" + )); + }; + PivotOutcome::Proof(ContractProof { + output_proven: true, + fixed_columns: columns, + families: families + .iter() + .map(|family| (family.name.clone(), Some(family.data_type.clone()))) + .collect(), + input_relations: vec![name.to_owned()], + failure_reason: None, + bare_dynamic_pivot: false, + }) +} + +fn pivot_proof( + pivot: &Map, + context: &PivotContext<'_, '_>, + families: &[DynamicFamily], +) -> PivotOutcome { + if truthy(pivot.get("unpivot")) { + return failure("UNPIVOT does not produce a dynamic column family".to_owned()); + } + let (pivot_columns, aggregates, dynamic, columns_valid) = components(pivot, context.dialect); + if !dynamic { + return failure( + "static pivot values must use ordinary exact column declarations".to_owned(), + ); + } + let [pivot_column] = pivot_columns.as_slice() else { + return failure( + "dynamic column contracts currently require exactly one pivot column".to_owned(), + ); + }; + if !columns_valid { + return failure("dynamic pivot columns must be direct column references".to_owned()); + } + let source: Option<&Value> = pivot.get("this"); + let schema_facts = SchemaFacts { + ctes: context.ctes, + types: context.facts.authoritative_types, + nullability: context.facts.column_nullability, + }; + let Some(source_columns) = relation_schema(source, &schema_facts, &HashSet::new()) else { + return failure( + "dynamic pivot input does not have an authoritative, explicit schema".to_owned(), + ); + }; + let Some(input_relation) = external_relation_name(source, context.ctes, &HashSet::new()) else { + return failure("dynamic pivot input relation could not be resolved".to_owned()); + }; + let mut matched: HashSet = HashSet::new(); + let mut family_proofs: Vec<(String, Option)> = Vec::new(); + for family in families { + if casefold(&family.pivot_column) != casefold(pivot_column) { + return failure(format!( + "dynamic family '{}' declares pivot_column '{}' but the output pivot uses '{}'", + family.name, family.pivot_column, pivot_column + )); + } + let matches: Vec<(usize, Option)> = aggregates + .iter() + .enumerate() + .filter_map(|(index, aggregate)| { + let (name, value, inferred) = aggregate_fact(aggregate, &source_columns); + let same_value: bool = + value.is_some_and(|value| casefold(value) == casefold(&family.value_column)); + (casefold(&name) == casefold(&family.aggregate) && same_value) + .then_some((index, inferred)) + }) + .collect(); + let [(index, inferred)] = matches.as_slice() else { + return failure(format!( + "dynamic family '{}' must match exactly one {}({}) pivot aggregate", + family.name, family.aggregate, family.value_column + )); + }; + if !matched.insert(*index) { + return failure( + "one pivot aggregate cannot satisfy multiple dynamic families".to_owned(), + ); + } + family_proofs.push((family.name.clone(), inferred.clone())); + } + if matched.len() != aggregates.len() { + return failure( + "every dynamic pivot aggregate must have one declared column family".to_owned(), + ); + } + let mut excluded: HashSet = HashSet::from([casefold(pivot_column)]); + excluded.extend(families.iter().map(|family| casefold(&family.value_column))); + let fixed_names: Vec = fixed_names(pivot, context.dialect, &excluded, &source_columns); + let by_name: HashMap = source_columns + .iter() + .map(|column| (casefold(&column.name), column)) + .collect(); + let mut fixed_columns: Vec = Vec::with_capacity(fixed_names.len()); + for name in fixed_names { + let Some(source) = by_name.get(&casefold(&name)) else { + return failure( + "dynamic pivot grouping columns are not present in the pivot input".to_owned(), + ); + }; + fixed_columns.push(ColumnFact { + name, + data_type: source.data_type.clone(), + nullability: source.nullability.clone(), + }); + } + PivotOutcome::Proof(ContractProof { + output_proven: true, + fixed_columns, + families: family_proofs, + input_relations: vec![input_relation.to_owned()], + failure_reason: None, + bare_dynamic_pivot: context.bare, + }) +} + +/// The fixed output names: DuckDB's GROUP BY columns, otherwise the unpivoted input columns. +fn fixed_names( + pivot: &Map, + dialect: &str, + excluded: &HashSet, + source_columns: &[ColumnFact], +) -> Vec { + let group: Option<&Value> = pivot.get("group"); + if SIMPLIFIED_PIVOT_DIALECTS.contains(&dialect) && group.is_some_and(Value::is_object) { + return dict_list(nested(group, &["group", "expressions"])) + .into_iter() + .filter_map(|expression| column_name(Some(expression)).map(str::to_owned)) + .collect(); + } + source_columns + .iter() + .filter(|column| !excluded.contains(&casefold(&column.name))) + .map(|column| column.name.clone()) + .collect() +} + +/// Python's `_pivot_components`: pivot columns, aggregates, dynamic values, columns valid. +fn components<'a>( + pivot: &'a Map, + dialect: &str, +) -> (Vec<&'a str>, Vec<&'a Value>, bool, bool) { + if SIMPLIFIED_PIVOT_DIALECTS.contains(&dialect) && truthy(pivot.get("using")) { + let expressions: Vec<&Value> = dict_list(pivot.get("expressions")); + let columns: Vec<&str> = expressions + .iter() + .filter_map(|expression| column_name(Some(expression))) + .collect(); + let valid: bool = columns.len() == expressions.len(); + return ( + columns, + dict_list(pivot.get("using")), + !truthy(pivot.get("fields")), + valid, + ); + } + let fields: Vec<&Value> = dict_list(pivot.get("fields")); + let columns: Vec<&str> = fields + .iter() + .filter_map(|field| column_name(nested(Some(field), &["in", "this"]))) + .collect(); + let dynamic: bool = fields.iter().any(|field| has_dynamic_values(field)); + let valid: bool = columns.len() == fields.len(); + (columns, dict_list(pivot.get("expressions")), dynamic, valid) +} + +/// Python's `_has_dynamic_value_source` for one pivot field. +fn has_dynamic_values(field: &Value) -> bool { + dict_list(nested(Some(field), &["in", "expressions"])) + .into_iter() + .any(|value| DYNAMIC_VALUE_SOURCE_KINDS.contains(&node_key(Some(value)))) +} + +/// Python's `_aggregate_fact`: aggregate name, value column and type-preserving inferred type. +fn aggregate_fact<'a>( + aggregate: &'a Value, + source_columns: &[ColumnFact], +) -> (String, Option<&'a str>, Option) { + let key: &str = node_key(Some(aggregate)); + let mut name: String = key.to_ascii_uppercase(); + let Some(payload) = aggregate.get(key).and_then(Value::as_object) else { + return (name, None, None); + }; + if let Some(authored) = payload + .get("name") + .and_then(Value::as_str) + .filter(|authored| !authored.is_empty()) + { + name = authored.to_ascii_uppercase(); + } + let preserving: bool = TYPE_PRESERVING_AGGREGATES.contains(&name.as_str()); + let value: Option<&Value> = payload.get("this"); + let mut value_name: Option<&str> = column_name(value); + let mut inferred: Option = None; + let cast: Option<&Map> = value + .filter(|value| node_key(Some(value)) == CAST_AST_KIND) + .and_then(|value| value.get(CAST_AST_KIND)) + .and_then(Value::as_object); + if let Some(cast) = cast { + value_name = column_name(cast.get("this")); + if preserving { + inferred = render_type(cast.get("to")); + } + } + if inferred.is_none() && preserving { + let by_name: HashMap = source_columns + .iter() + .map(|column| (casefold(&column.name), column)) + .collect(); + inferred = value_name + .and_then(|value| by_name.get(&casefold(value))) + .and_then(|column| column.data_type.clone()); + } + (name, value_name, inferred) +} + +fn resolve_boundary<'a>( + node: Option<&'a Value>, + ctes: &CteMap<'a>, + seen: &HashSet, +) -> Option> { + let key: &str = node_key(node); + let payload: &Map = payload(node)?; + if key == PIVOT_AST_KIND { + return Some(Boundary::Pivot(payload)); + } + if QUERY_AST_KINDS.contains(&key) { + return resolve_boundary(payload.get("this"), ctes, seen); + } + if key != SELECT_AST_KIND + || truthy(payload.get("joins")) + || !is_single_wildcard(payload.get("expressions")) + { + return None; + } + let relations: Vec<&Value> = + dict_list(payload.get("from").and_then(|from| from.get("expressions"))); + let [relation] = relations.as_slice() else { + return None; + }; + if node_key(Some(relation)) == PIVOT_AST_KIND { + return relation + .get(PIVOT_AST_KIND) + .and_then(Value::as_object) + .map(Boundary::Pivot); + } + let name: &str = relation_name(Some(relation))?; + let folded: String = casefold(name); + let Some(cte) = ctes.get(&folded) else { + return Some(Boundary::Passthrough(name)); + }; + let cte: &Value = (*cte)?; + if seen.contains(&folded) { + return None; + } + let mut next_seen: HashSet = seen.clone(); + next_seen.insert(folded); + resolve_boundary(Some(cte), ctes, &next_seen) +} + +/// The relation facts schema lookups read. +struct SchemaFacts<'a, 'b> { + ctes: &'a CteMap<'b>, + types: &'a Shapes, + nullability: &'a Shapes, +} + +fn relation_schema( + node: Option<&Value>, + facts: &SchemaFacts<'_, '_>, + seen: &HashSet, +) -> Option> { + let name: &str = relation_name(node)?; + let folded: String = casefold(name); + let Some(cte) = facts.ctes.get(&folded) else { + return external_schema(name, facts.types, facts.nullability); + }; + let cte: &Value = (*cte)?; + if seen.contains(&folded) { + return None; + } + let mut next_seen: HashSet = seen.clone(); + next_seen.insert(folded); + select_schema(Some(cte), facts, &next_seen) +} + +fn external_relation_name<'a>( + node: Option<&'a Value>, + ctes: &CteMap<'a>, + seen: &HashSet, +) -> Option<&'a str> { + let key: &str = node_key(node); + if let Some(payload) = payload(node) { + if QUERY_AST_KINDS.contains(&key) { + return external_relation_name(payload.get("this"), ctes, seen); + } + if key == SELECT_AST_KIND { + if truthy(payload.get("joins")) { + return None; + } + let relations: Vec<&Value> = + dict_list(payload.get("from").and_then(|from| from.get("expressions"))); + let [relation] = relations.as_slice() else { + return None; + }; + return external_relation_name(Some(relation), ctes, seen); + } + } + let name: &str = relation_name(node)?; + let folded: String = casefold(name); + let Some(cte) = ctes.get(&folded) else { + return Some(name); + }; + let cte: &Value = (*cte)?; + if seen.contains(&folded) { + return None; + } + let mut next_seen: HashSet = seen.clone(); + next_seen.insert(folded); + external_relation_name(Some(cte), ctes, &next_seen) +} + +fn select_schema( + node: Option<&Value>, + facts: &SchemaFacts<'_, '_>, + seen: &HashSet, +) -> Option> { + let key: &str = node_key(node); + let payload: &Map = payload(node)?; + if QUERY_AST_KINDS.contains(&key) { + return select_schema(payload.get("this"), facts, seen); + } + if key != SELECT_AST_KIND || truthy(payload.get("joins")) { + return None; + } + let relations: Vec<&Value> = + dict_list(payload.get("from").and_then(|from| from.get("expressions"))); + let [relation] = relations.as_slice() else { + return None; + }; + let source: Vec = relation_schema(Some(relation), facts, seen)?; + let expressions: Option<&Value> = payload.get("expressions"); + if is_single_wildcard(expressions) { + return Some(source); + } + let by_name: HashMap = source + .iter() + .map(|column| (casefold(&column.name), column)) + .collect(); + let mut output: Vec = Vec::new(); + for expression in expressions?.as_array()? { + if !expression.is_object() { + return None; + } + output.push(projected_column(expression, &by_name)?); + } + Some(output) +} + +/// One projection of a CTE body over its single source relation. +fn projected_column( + projection: &Value, + by_name: &HashMap, +) -> Option { + let (output_name, expression) = projection_parts(projection); + let (source_name, cast_type): (Option<&str>, Option) = + if node_key(Some(expression)) == CAST_AST_KIND { + let cast: &Map = expression.get(CAST_AST_KIND)?.as_object()?; + (column_name(cast.get("this")), render_type(cast.get("to"))) + } else { + (column_name(Some(expression)), None) + }; + let source: &ColumnFact = by_name.get(&casefold(source_name?))?; + Some(ColumnFact { + name: output_name?.to_owned(), + data_type: cast_type.or_else(|| source.data_type.clone()), + nullability: source.nullability.clone(), + }) +} + +/// Python's `_projection`: the output name and the aliased expression, `{}` when absent. +fn projection_parts(node: &Value) -> (Option<&str>, &Value) { + if node_key(Some(node)) != ALIAS_AST_KIND { + return (column_name(Some(node)), node); + } + let Some(alias) = node.get(ALIAS_AST_KIND).and_then(Value::as_object) else { + return (None, node); + }; + let expression: &Value = alias + .get("this") + .filter(|expression| expression.is_object()) + .unwrap_or(&MISSING_EXPRESSION); + (identifier_name(alias.get("alias")), expression) +} + +fn collect_ctes(root: &Value) -> CteMap<'_> { + let mut ctes: CteMap<'_> = HashMap::new(); + let Some(with) = payload(Some(root)).and_then(|payload| payload.get("with")) else { + return ctes; + }; + if !with.is_object() { + return ctes; + } + for cte in dict_list(with.get("ctes")) { + let body: &Value = if node_key(Some(cte)) == CTE_AST_KIND { + match cte.get(CTE_AST_KIND) { + Some(body) => body, + None => continue, + } + } else { + cte + }; + let Some(body) = body.as_object() else { + continue; + }; + let Some(name) = identifier_name(body.get("alias")) else { + continue; + }; + let query: Option<&Value> = body.get("this"); + if truthy(body.get("columns")) { + ctes.insert(casefold(name), None); + } else if let Some(query) = query.filter(|query| query.is_object()) { + ctes.insert(casefold(name), Some(query)); + } + } + ctes +} + +/// Python's `_external_schema`: the relation's typed columns with their nullability. +fn external_schema(name: &str, types: &Shapes, nullability: &Shapes) -> Option> { + let types = casefold_lookup(types, name)?; + let mut by_name: HashMap = HashMap::new(); + for (column, value) in casefold_lookup(nullability, name).into_iter().flatten() { + by_name.insert(casefold(column), value); + } + Some( + types + .iter() + .map(|(column, data_type)| ColumnFact { + name: column.clone(), + data_type: (data_type != UNKNOWN_TYPE).then(|| data_type.clone()), + nullability: by_name + .get(&casefold(column)) + .map_or(UNKNOWN_NULLABILITY, |value| value) + .to_owned(), + }) + .collect(), + ) +} + +/// Python's `_casefold_lookup`: the first entry whose folded name matches. +fn casefold_lookup<'a, T>(entries: &'a [(String, T)], key: &str) -> Option<&'a T> { + let requested: String = casefold(key); + entries + .iter() + .find(|(candidate, _)| casefold(candidate) == requested) + .map(|(_, value)| value) +} + +fn families_equal(left: &[DynamicFamily], right: &[DynamicFamily]) -> bool { + left.len() == right.len() + && left + .iter() + .zip(right) + .all(|(left, right)| normalized(left) == normalized(right)) +} + +fn failure(reason: String) -> PivotOutcome { + PivotOutcome::Proof(ContractProof { + output_proven: false, + fixed_columns: Vec::new(), + families: Vec::new(), + input_relations: Vec::new(), + failure_reason: Some(reason), + bare_dynamic_pivot: false, + }) +} + +fn ascii_family(family: &DynamicFamily) -> bool { + [ + &family.name, + &family.pivot_column, + &family.value_column, + &family.aggregate, + &family.data_type, + ] + .iter() + .all(|text| text.is_ascii()) +} + +/// Python's `_families_equal` key for one family. +fn normalized(family: &DynamicFamily) -> (String, String, String, String, String, Option<&str>) { + ( + casefold(&family.name), + casefold(&family.pivot_column), + casefold(&family.value_column), + casefold(&family.aggregate), + casefold(&family.data_type), + family.name_pattern.as_deref(), + ) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs new file mode 100644 index 0000000000..2e176ff81e --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs @@ -0,0 +1,338 @@ +//! Python's re-analysis of a model with its known input shapes, over the crate's query facts. + +use std::collections::{HashMap, HashSet}; + +use polyglot_sql::{AnalyzeQueryOptions, analyze_query}; +use serde_json::{Map, Value, json}; +use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic; + +use crate::assembly::analysis_session::_helpers::dict_walk::truthy; +use crate::assembly::analysis_session::constants::{ + CAST_TRANSFORM, CONFIDENCE_HIGH, CONFIDENCE_MEDIUM, CONFIDENCE_UNKNOWN, FILTER_CONTEXT, + MAX_FUNCTION_CALL_DEPTH, NULL_KEYWORD, RESOLVED_SOURCE_CONFIDENCE, SELECT_SHAPE, + SET_OPERATION_SHAPE, TRANSFORM_AGGREGATION, TRANSFORM_CAST, TRANSFORM_CONSTANT, + TRANSFORM_DIRECT, TRANSFORM_EXPRESSION, TRANSFORM_STAR, UNKNOWN_NULLABILITY, UNKNOWN_TYPE, + WILDCARD, +}; +use crate::assembly::analysis_session::models::{ColumnFact, LineageRow, ModelRequest}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::main::normalize::normalize_analysis_sql; +use crate::semantic_validation::models::NormalizationInput; + +/// One model's enrichment inputs: the model, its known input shapes and the profile facts. +pub(crate) struct EnrichmentInput<'a> { + pub(crate) model: &'a ModelRequest, + pub(crate) input_schemas: &'a Shapes, + pub(crate) dialect: &'a str, + pub(crate) function_return_types: &'a Pairs, +} + +/// Python's successful re-analysis: columns, plain lineage facts and the star flag. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Enrichment { + pub(crate) columns: Vec, + pub(crate) lineage: Vec, + pub(crate) has_star: bool, +} + +/// One projection's upstream sources and confidence code. +type Upstream = (Vec<(String, String, String)>, u8); + +/// Python's compact re-analysis, or None where Python would take another path. +pub(crate) fn enrichment(input: &EnrichmentInput<'_>) -> Option { + let enrichment: Result, String> = catch_compiler_panic(|| analysed(input)); + enrichment.unwrap_or_default() +} + +/// The re-analysis; Ok(None) where Python takes another path, Err for an unexpected failure. +fn analysed(input: &EnrichmentInput<'_>) -> Result, String> { + let normalized: Result = normalize_analysis_sql(NormalizationInput { + sql: input.model.query_sql.clone(), + dialect: input.dialect.to_owned(), + stubs: HashMap::new(), + placeholders: input.model.placeholders.iter().cloned().collect(), + }); + let cleaned: String = match normalized { + Ok(cleaned) => cleaned, + Err(_) => return Ok(None), + }; + let mut options: Map = Map::new(); + options.insert( + "dialect".to_owned(), + Value::String(input.dialect.to_owned()), + ); + if let Some(schema) = analysis_schema(input) { + options.insert("schema".to_owned(), schema); + } + options.insert( + "complexityGuard".to_owned(), + json!({"maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH}), + ); + let options: AnalyzeQueryOptions = + serde_json::from_value(Value::Object(options)).map_err(|error| error.to_string())?; + let analysis = match analyze_query(&cleaned, options) { + Ok(analysis) => analysis, + Err(_) => return Ok(None), + }; + let analysis: Value = serde_json::to_value(analysis).map_err(|error| error.to_string())?; + Ok(projected(input, &analysis)) +} + +/// Python's `_compact_analysis_schema` over the input shapes, nullability unknown. +fn analysis_schema(input: &EnrichmentInput<'_>) -> Option { + let shapes: HashMap<&str, &Pairs> = input + .input_schemas + .iter() + .map(|(name, shape)| (name.as_str(), shape)) + .collect(); + let mut names: Vec<&str> = input + .model + .references + .iter() + .map(|reference| reference.analysis_name.as_str()) + .collect::>() + .into_iter() + .collect(); + names.sort_unstable(); + let mut tables: Vec = Vec::new(); + for name in names { + let Some(shape) = shapes.get(name).filter(|shape| !shape.is_empty()) else { + continue; + }; + let columns: Vec = shape + .iter() + .map(|(column, data_type)| json!({"name": column, "type": data_type})) + .collect(); + tables.push(json!({"name": name, "columns": columns})); + } + (!tables.is_empty()).then(|| json!({"tables": tables})) +} + +/// Python's projection of `analyze_query` facts into columns and lineage. +fn projected(input: &EnrichmentInput<'_>, analysis: &Value) -> Option { + let projections: &Vec = analysis.get("projections")?.as_array()?; + let shape: Option<&str> = analysis.get("shape").and_then(Value::as_str); + if !matches!(shape, Some(SELECT_SHAPE | SET_OPERATION_SHAPE)) || !eligible(projections) { + return None; + } + if has_null_filter(analysis) || recovers_cte_facts(input, analysis) { + return None; + } + let resources: HashMap<&str, (&str, &str)> = input + .model + .lineage_references + .iter() + .map(|(name, kind, resource)| (name.as_str(), (kind.as_str(), resource.as_str()))) + .collect(); + let infer_nullability: bool = shape != Some(SET_OPERATION_SHAPE); + let mut columns: Vec = Vec::new(); + let mut lineage: Vec = Vec::new(); + for projection in projections { + let projection: &Map = projection.as_object()?; + if truthy(projection.get("isStar")) { + continue; + } + let name: &str = text(projection.get("name")); + if name.is_empty() || name == WILDCARD { + continue; + } + columns.push(ColumnFact { + name: name.to_owned(), + data_type: projection_type(projection, input.function_return_types)?, + nullability: projection_nullability(projection, infer_nullability).to_owned(), + }); + let (sources, confidence) = upstream(projection, &resources); + let transform: u8 = transform_code(projection, !sources.is_empty()); + lineage.push(LineageRow { + output_column: name.to_owned(), + confidence_code: if sources.is_empty() && transform != TRANSFORM_CONSTANT { + CONFIDENCE_UNKNOWN + } else { + confidence + }, + transform_code: transform, + sources, + }); + } + let has_star: bool = analysis + .get("starProjections") + .and_then(Value::as_array) + .is_some_and(|stars| !stars.is_empty()); + if has_star && input.model.references.len() == 1 { + let order: HashMap<&str, usize> = declared_order(input); + let fallback: usize = order.len(); + columns.sort_by_key(|column| order.get(column.name.as_str()).copied().unwrap_or(fallback)); + lineage.sort_by_key(|row| { + order + .get(row.output_column.as_str()) + .copied() + .unwrap_or(fallback) + }); + } + Some(Enrichment { + columns, + lineage, + has_star, + }) +} + +/// Whether Python's CTE pass-through recovery reads the parsed query, which stays Python's. +fn recovers_cte_facts(input: &EnrichmentInput<'_>, analysis: &Value) -> bool { + input.model.recover_cte_facts + && truthy(analysis.get("cteFacts")) + && !input.input_schemas.is_empty() +} + +/// Python's `_compact_analysis_is_eligible` for the projections. +fn eligible(projections: &[Value]) -> bool { + for projection in projections { + let Some(projection) = projection.as_object() else { + return false; + }; + let cast: bool = text(projection.get("transformKind")) == CAST_TRANSFORM; + let upstream_empty: bool = projection + .get("upstream") + .and_then(Value::as_array) + .is_none_or(Vec::is_empty); + if cast && upstream_empty { + return false; + } + } + true +} + +/// Whether a filter reads NULL, where Python parses the query for non-null outputs. +fn has_null_filter(analysis: &Value) -> bool { + let Some(uses) = analysis.get("columnUses").and_then(Value::as_array) else { + return false; + }; + uses.iter().filter_map(Value::as_object).any(|value| { + text(value.get("context")) == FILTER_CONTEXT + && text(value.get("expressionSql")) + .to_uppercase() + .contains(NULL_KEYWORD) + }) +} + +/// Python's `_compact_projection_type`; None where the function name is not ASCII. +fn projection_type( + projection: &Map, + function_return_types: &Pairs, +) -> Option> { + let cast_type: &str = text(projection.get("castType")); + if !cast_type.is_empty() && cast_type != UNKNOWN_TYPE { + return Some(Some(cast_type.to_owned())); + } + let function_name: Option<&str> = projection + .get("transformFunction") + .and_then(Value::as_object) + .and_then(|function| function.get("name")) + .and_then(Value::as_str); + if let Some(function_name) = function_name { + if !function_name.is_ascii() { + return None; + } + let upper: String = function_name.to_ascii_uppercase(); + if let Some((_, declared)) = function_return_types + .iter() + .find(|(name, _)| *name == upper) + { + return Some(Some(declared.clone())); + } + } + let type_hint: &str = text(projection.get("typeHint")); + if !type_hint.is_empty() && type_hint != UNKNOWN_TYPE { + return Some(Some(type_hint.to_owned())); + } + Some(None) +} + +fn projection_nullability(projection: &Map, infer: bool) -> &'static str { + if !infer { + return UNKNOWN_NULLABILITY; + } + match text(projection.get("nullability")) { + "non_null" => "non_null", + "nullable" => "nullable", + _ => UNKNOWN_NULLABILITY, + } +} + +/// Python's `_compact_lineage_upstream_columns`. +fn upstream(projection: &Map, resources: &HashMap<&str, (&str, &str)>) -> Upstream { + let Some(values) = projection.get("upstream").and_then(Value::as_array) else { + return (Vec::new(), CONFIDENCE_UNKNOWN); + }; + let mut sources: Vec<(String, String, String)> = Vec::new(); + let mut seen: HashSet<(String, String, String)> = HashSet::new(); + let mut confidence: u8 = CONFIDENCE_HIGH; + for value in values.iter().filter_map(Value::as_object) { + let column: &str = text(value.get("column")); + if column.is_empty() || column == WILDCARD { + continue; + } + let mut source_name: &str = text(value.get("sourceName")); + if source_name.is_empty() { + source_name = text(value.get("table")); + } + if source_name.is_empty() { + confidence = CONFIDENCE_UNKNOWN; + continue; + } + let Some((kind, resource)) = resources.get(source_name) else { + continue; + }; + let resolved: bool = text(value.get("confidence")) == RESOLVED_SOURCE_CONFIDENCE; + let aliased: bool = value.get("sourceAlias").is_some_and(Value::is_string); + if !resolved && !aliased { + confidence = CONFIDENCE_MEDIUM; + } + let key = ( + (*kind).to_owned(), + (*resource).to_owned(), + column.to_owned(), + ); + if seen.insert(key.clone()) { + sources.push(key); + } + } + sources.sort(); + (sources, confidence) +} + +/// Python's `_compact_transform_kind` as a compact code. +fn transform_code(projection: &Map, has_upstream: bool) -> u8 { + match text(projection.get("transformKind")) { + "star" => TRANSFORM_STAR, + "cast" => TRANSFORM_CAST, + "aggregation" => TRANSFORM_AGGREGATION, + "constant" => TRANSFORM_CONSTANT, + _ if !has_upstream => TRANSFORM_CONSTANT, + "direct" => TRANSFORM_DIRECT, + _ => TRANSFORM_EXPRESSION, + } +} + +/// The declared column order of the model's single reference, from its input shape. +fn declared_order<'a>(input: &EnrichmentInput<'a>) -> HashMap<&'a str, usize> { + let Some(reference) = input.model.references.first() else { + return HashMap::new(); + }; + let Some((_, shape)) = input + .input_schemas + .iter() + .find(|(name, _)| *name == reference.analysis_name) + else { + return HashMap::new(); + }; + let mut order: HashMap<&str, usize> = HashMap::new(); + for (index, (column, _)) in shape.iter().enumerate() { + order.entry(column.as_str()).or_insert(index); + } + order +} + +/// Python's `str(value or "")` for the string fields this projection reads. +fn text(value: Option<&Value>) -> &str { + value.and_then(Value::as_str).unwrap_or_default() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs index f992aa5b95..1ff9c1dda6 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs @@ -1,5 +1,8 @@ pub(crate) mod catalog_state; pub(crate) mod compact_batch; +pub(crate) mod dict_walk; +pub(crate) mod dynamic_pivot; +pub(crate) mod enrichment; pub(crate) mod mappings; pub(crate) mod publication; pub(crate) mod schedule; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs index de099cbc93..cf945f7c71 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs @@ -2,10 +2,16 @@ use std::collections::HashMap; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; + use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; use crate::assembly::analysis_session::_helpers::compact_batch::{ Batch, BatchMember, MemberAnalysis, MemberResult, }; +use crate::assembly::analysis_session::_helpers::dynamic_pivot::{PivotFacts, pivot_outcome}; +use crate::assembly::analysis_session::_helpers::enrichment::{ + Enrichment, EnrichmentInput, enrichment, +}; use crate::assembly::analysis_session::_helpers::mappings::{ ShapeTable, dict_from_pairs, same_keys, same_relations, }; @@ -16,7 +22,7 @@ use crate::assembly::analysis_session::constants::{ }; use crate::assembly::analysis_session::models::{ AnalysisSession, Awaiting, ColumnFact, Deferral, DeferredAnalysis, LineageFacts, ModelAnalysis, - ModelOutcome, ModelRequest, Phase, SessionOutcome, SessionRequest, SessionStep, + ModelOutcome, ModelRequest, Phase, PivotOutcome, SessionOutcome, SessionRequest, SessionStep, }; use crate::assembly::analysis_session::types::{Pairs, Shapes}; use crate::semantic_validation::models::ProjectCatalog; @@ -128,7 +134,12 @@ impl AnalysisSession { .get(&result.model) .ok_or("an enrichment answers no deferral")?; let outcome: &mut ModelOutcome = self.outcome_mut(result.model)?; - outcome.analysis = enriched_analysis(&outcome.analysis, result, pending); + outcome.analysis = enriched_analysis( + &outcome.analysis, + result, + pending, + LineageFacts::PythonEnrichment, + ); } self.phase = Phase::Finish(wave); } @@ -141,6 +152,7 @@ impl AnalysisSession { if !matches!(self.phase, Phase::Done) { return Err("the session has not finished".to_owned()); } + let dynamic_contracts: Vec = self.dynamic_contracts()?; let models: Vec = self .outcomes .into_iter() @@ -151,9 +163,32 @@ impl AnalysisSession { models, schema_additions, analysis_names, + dynamic_contracts, }) } + /// Python's dynamic pivot proof of each model over the relation facts before analysis. + fn dynamic_contracts(&self) -> Result, String> { + let models: &[ModelRequest] = &self.request.models; + if models.iter().all(|model| model.dynamic_families.is_empty()) { + return Ok(vec![PivotOutcome::Absent; models.len()]); + } + let facts = PivotFacts { + dialect: &self.request.dialect, + column_types: &self.request.column_types, + authoritative_types: &self.request.complete_schemas, + column_nullability: &self.request.column_nullability, + families_by_table: &self.request.dynamic_families_by_table, + }; + let pool = self.catalog.native.analysis_pool()?; + Ok(pool.install(|| { + models + .par_iter() + .map(|model| pivot_outcome(&model.pivot_sql, &model.dynamic_families, &facts)) + .collect() + })) + } + fn step(&mut self, deferrals: Vec) -> SessionStep { SessionStep { publications: std::mem::take(&mut self.publications), @@ -251,6 +286,7 @@ impl AnalysisSession { fn complete_wave(&mut self, wave: usize) -> Result { let mut deferrals: Vec = Vec::new(); let mut star_pending_by_model: HashMap = HashMap::new(); + let mut candidates: Vec<(usize, Shapes, bool)> = Vec::new(); for model in self.waves[wave].clone() { let request: &ModelRequest = &self.request.models[model]; let required: Vec = request.required_names.clone(); @@ -275,16 +311,53 @@ impl AnalysisSession { .flatten() .any(|column| column.data_type.is_none()); if !required.is_empty() && (untyped || star_pending) && !has_set_operation { + candidates.push((model, input_schemas, star_pending)); + } + } + let enrichments: Vec> = self.native_enrichments(&candidates)?; + for ((model, input_schemas, star_pending), native) in + candidates.into_iter().zip(enrichments) + { + let Some(native) = native else { star_pending_by_model.insert(model, star_pending); deferrals.push(Deferral::Enrichment { model, input_schemas, }); - } + continue; + }; + let outcome: &mut ModelOutcome = self.outcome_mut(model)?; + outcome.analysis = enriched_analysis( + &outcome.analysis, + native_answer(model, native.columns, native.has_star), + star_pending, + LineageFacts::NativeEnrichment(native.lineage), + ); } Ok((deferrals, star_pending_by_model)) } + /// Python's re-analysis of each candidate natively, None where Python must answer it. + fn native_enrichments( + &self, + candidates: &[(usize, Shapes, bool)], + ) -> Result>, String> { + if candidates.is_empty() { + return Ok(Vec::new()); + } + let inputs: Vec> = candidates + .iter() + .map(|(model, input_schemas, _)| EnrichmentInput { + model: &self.request.models[*model], + input_schemas, + dialect: &self.request.dialect, + function_return_types: &self.request.function_return_types, + }) + .collect(); + let pool = self.catalog.native.analysis_pool()?; + Ok(pool.install(|| inputs.par_iter().map(enrichment).collect())) + } + fn input_schemas(&self, required: &[String]) -> Shapes { let mut schemas: Shapes = Vec::with_capacity(required.len()); for name in required { @@ -470,6 +543,7 @@ fn enriched_analysis( analysis: &ModelAnalysis, enriched: DeferredAnalysis, star_pending: bool, + lineage: LineageFacts, ) -> ModelAnalysis { let mut recovered: HashMap = HashMap::new(); for column in enriched @@ -489,7 +563,7 @@ fn enriched_analysis( ModelAnalysis { analysis_succeeded: enriched.analysis_succeeded, columns: enriched.columns, - lineage: LineageFacts::PythonEnrichment, + lineage, has_star: enriched.has_star, star_resolved: enriched.star_resolved, binding_diagnostics: analysis.binding_diagnostics.clone(), @@ -508,3 +582,16 @@ fn enriched_analysis( merged.star_resolved = merged.star_resolved || star_expanded; merged } + +/// A native re-analysis as Python's successful answer. +fn native_answer(model: usize, columns: Vec, has_star: bool) -> DeferredAnalysis { + DeferredAnalysis { + model, + analysis_succeeded: true, + columns: Some(columns), + has_star, + star_resolved: false, + binding_diagnostics: Vec::new(), + binding_validated: false, + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs index 735e39dbc8..d0d6fa5520 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs @@ -38,3 +38,60 @@ pub(crate) const QUOTED_IDENTIFIER_DELIMITER: char = '"'; /// Python's names for a legacy-analysis and an enrichment deferral. pub(crate) const DEFERRAL_ANALYSIS: &str = "analysis"; pub(crate) const DEFERRAL_ENRICHMENT: &str = "enrichment"; +/// Python's `ColumnTransformKind` compact codes. +pub(crate) const TRANSFORM_DIRECT: u8 = 0; +pub(crate) const TRANSFORM_CAST: u8 = 1; +pub(crate) const TRANSFORM_EXPRESSION: u8 = 2; +pub(crate) const TRANSFORM_AGGREGATION: u8 = 3; +pub(crate) const TRANSFORM_STAR: u8 = 4; +pub(crate) const TRANSFORM_CONSTANT: u8 = 5; +/// Python's `ColumnLineageConfidence` compact codes. +pub(crate) const CONFIDENCE_UNKNOWN: u8 = 0; +pub(crate) const CONFIDENCE_HIGH: u8 = 1; +pub(crate) const CONFIDENCE_MEDIUM: u8 = 2; +/// The `analyze_query` facts Python's re-analysis reads. +pub(crate) const SELECT_SHAPE: &str = "select"; +pub(crate) const SET_OPERATION_SHAPE: &str = "set_operation"; +pub(crate) const CAST_TRANSFORM: &str = "cast"; +pub(crate) const FILTER_CONTEXT: &str = "filter"; +pub(crate) const NULL_KEYWORD: &str = "NULL"; +pub(crate) const RESOLVED_SOURCE_CONFIDENCE: &str = "resolved"; +pub(crate) const WILDCARD: &str = "*"; +/// SQLBuild's `analyze_query` complexity guard. +pub(crate) const MAX_FUNCTION_CALL_DEPTH: usize = 512; +/// The dialects whose dynamic pivots the compiler proves, and those with DuckDB's PIVOT form. +pub(crate) const SUPPORTED_PIVOT_DIALECTS: [&str; 3] = ["duckdb", "motherduck", "snowflake"]; +pub(crate) const SIMPLIFIED_PIVOT_DIALECTS: [&str; 2] = ["duckdb", "motherduck"]; +pub(crate) const DUCKDB_DIALECT: &str = "duckdb"; +pub(crate) const MOTHERDUCK_DIALECT: &str = "motherduck"; +/// 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"; +pub(crate) const QUERY_AST_KINDS: [&str; 2] = ["query", "subquery"]; +pub(crate) const CAST_AST_KIND: &str = "cast"; +pub(crate) const COLUMN_AST_KIND: &str = "column"; +pub(crate) const TABLE_AST_KIND: &str = "table"; +pub(crate) const FUNCTION_AST_KIND: &str = "function"; +pub(crate) const ALIAS_AST_KIND: &str = "alias"; +pub(crate) const CTE_AST_KIND: &str = "cte"; +pub(crate) const STAR_AST_KIND: &str = "star"; +pub(crate) const DYNAMIC_VALUE_SOURCE_KINDS: [&str; 3] = ["pivot_any", "select", "query"]; +pub(crate) const DEPENDENCY_FUNCTION_NAMES: [&str; 3] = ["__ref", "__source", "__seed"]; +pub(crate) const TYPE_PRESERVING_AGGREGATES: [&str; 3] = ["ANY_VALUE", "MAX", "MIN"]; +/// Python's `_render_type` names for cast targets. +pub(crate) const RENDERED_TYPE_NAMES: [(&str, &str); 12] = [ + ("big_int", "BIGINT"), + ("bool", "BOOLEAN"), + ("boolean", "BOOLEAN"), + ("decimal", "DECIMAL"), + ("double", "DOUBLE"), + ("float", "FLOAT"), + ("int", "INT"), + ("integer", "INT"), + ("text", "TEXT"), + ("timestamp", "TIMESTAMP"), + ("var_char", "VARCHAR"), + ("varchar", "VARCHAR"), +]; +/// 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; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs index 5893bfd5e8..6c5394355c 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs @@ -1,5 +1,6 @@ pub mod expression_shapes; pub mod finish_analysis_session; +pub mod prove_dynamic_contract; pub mod provide_deferred_analyses; pub mod run_analysis_session; pub mod start_analysis_session; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs new file mode 100644 index 0000000000..96759642d4 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs @@ -0,0 +1,26 @@ +//! Python's `analyze_dynamic_column_contract` for one model outside the analysis session. + +use crate::assembly::analysis_session::_helpers::dynamic_pivot::{PivotFacts, pivot_outcome}; +use crate::assembly::analysis_session::constants::PIVOT_WORKER_STACK_BYTES; +use crate::assembly::analysis_session::models::{PivotOutcome, PivotRequest}; + +/// The model's proof, or a deferral where Python's result is not reproduced exactly. +pub fn prove_dynamic_contract(request: &PivotRequest) -> Result { + let facts = PivotFacts { + dialect: &request.dialect, + column_types: &request.column_types, + authoritative_types: &request.authoritative_types, + column_nullability: &request.column_nullability, + families_by_table: &request.families_by_table, + }; + std::thread::scope(|scope| { + std::thread::Builder::new() + .stack_size(PIVOT_WORKER_STACK_BYTES) + .spawn_scoped(scope, || { + pivot_outcome(&request.sql, &request.families, &facts) + }) + .map_err(|error| error.to_string())? + .join() + .map_err(|_| "the native dynamic pivot proof panicked".to_owned()) + }) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs index a52f3e38ce..d10d56c0ad 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs @@ -33,6 +33,52 @@ pub struct ModelRequest { pub has_set_operation: bool, /// The snapshot validity columns a published shape appends. pub snapshot_columns: Option<(String, String)>, + /// The SQL Python's dynamic pivot proof parses: placeholders replaced by their defaults. + pub pivot_sql: String, + pub dynamic_families: Vec, +} + +/// One declared family of columns a runtime-dynamic pivot generates. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DynamicFamily { + pub name: String, + pub pivot_column: String, + pub value_column: String, + pub aggregate: String, + pub data_type: String, + pub name_pattern: Option, +} + +/// Python's `DynamicColumnContractProof`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ContractProof { + pub output_proven: bool, + pub fixed_columns: Vec, + /// Each family's name and inferred type. + pub families: Vec<(String, Option)>, + pub input_relations: Vec, + pub failure_reason: Option, + pub bare_dynamic_pivot: bool, +} + +/// One model's dynamic pivot proof inputs, for a model the session does not analyse. +#[derive(Debug, Clone, Default)] +pub struct PivotRequest { + pub dialect: String, + pub column_types: Shapes, + pub authoritative_types: Shapes, + pub column_nullability: Shapes, + pub families_by_table: Vec<(String, Vec)>, + pub sql: String, + pub families: Vec, +} + +/// A model's dynamic pivot proof: none declared, left to Python, or proven natively. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PivotOutcome { + Absent, + Deferred, + Proof(ContractProof), } /// Everything one compile's model analysis reads. @@ -49,6 +95,8 @@ pub struct SessionRequest { pub complete_schemas: Shapes, /// The Python binding catalog's `schemas` before analysis. pub catalog_schemas: Shapes, + /// Python's `dynamic_families_by_table`, in dict order. + pub dynamic_families_by_table: Vec<(String, Vec)>, pub models: Vec, } @@ -78,6 +126,8 @@ pub enum LineageFacts { PythonAnalysis, /// The lineage of the analysis Python returned for this model's enrichment deferral. PythonEnrichment, + /// Plain lineage facts of the native re-analysis with known inputs. + NativeEnrichment(Vec), } /// Python's `PolyglotAnalysisResult` for one model. @@ -147,6 +197,8 @@ pub struct SessionOutcome { pub schema_additions: Shapes, /// Relations whose analysis shapes the Python binding catalog records. pub analysis_names: Vec, + /// Each model's dynamic pivot proof, in request order. + pub dynamic_contracts: Vec, } /// One expression source's shape: inferred, absent, or left to Python. diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs index fccee1d92b..1f274fdfc1 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs @@ -4,8 +4,9 @@ use crate::assembly::analysis_session::main::provide_deferred_analyses::provide_ use crate::assembly::analysis_session::main::run_analysis_session::run_analysis_session; use crate::assembly::analysis_session::main::start_analysis_session::start_analysis_session; use crate::assembly::analysis_session::models::{ - AnalysisSession, Deferral, DeferredAnalysis, ModelOutcome, ModelReference, ModelRequest, - SessionRequest, SessionStep, + AnalysisSession, ColumnFact, ContractProof, Deferral, DeferredAnalysis, DynamicFamily, + ModelOutcome, ModelReference, ModelRequest, PivotOutcome, PivotRequest, SessionRequest, + SessionStep, }; use crate::assembly::analysis_session::tests::test_types::ModelSpec; use crate::assembly::analysis_session::types::{Pairs, Shapes}; @@ -79,6 +80,8 @@ pub(crate) fn model(name: &str, sql: &str, sources: &[&str], refs: &[&str]) -> M recover_cte_facts: false, has_set_operation: sql.contains("UNION"), snapshot_columns: None, + pivot_sql: sql.to_owned(), + dynamic_families: Vec::new(), } } @@ -100,6 +103,7 @@ pub(crate) fn orders_request(models: Vec) -> SessionRequest { )]), complete_schemas: raw_orders.clone(), catalog_schemas: raw_orders, + dynamic_families_by_table: Vec::new(), models, } } @@ -210,3 +214,86 @@ pub(crate) fn session_lines(models: &[ModelSpec]) -> (Vec>, Vec, +) -> PivotRequest { + let raw_orders: Shapes = shapes(&[( + "raw_orders", + &[ + ("order_id", "INTEGER"), + ("customer_id", "INTEGER"), + ("status", "VARCHAR"), + ("amount", "DOUBLE"), + ], + )]); + let upstream: Shapes = shapes(&[("status_amounts", &[("customer_id", "INTEGER")])]); + let families: Vec = family + .map(|(pivot_column, value_column, aggregate)| { + amounts(pivot_column, value_column, aggregate) + }) + .into_iter() + .collect(); + PivotRequest { + dialect: dialect.to_owned(), + column_types: [raw_orders.clone(), upstream].concat(), + authoritative_types: raw_orders, + column_nullability: shapes(&[("raw_orders", &[("customer_id", "non_null")])]), + families_by_table: vec![( + "status_amounts".to_owned(), + vec![amounts("status", "amount", "MAX")], + )], + sql: sql.to_owned(), + families, + } +} + +fn amounts(pivot_column: &str, value_column: &str, aggregate: &str) -> DynamicFamily { + DynamicFamily { + name: "amounts".to_owned(), + pivot_column: pivot_column.to_owned(), + value_column: value_column.to_owned(), + aggregate: aggregate.to_owned(), + data_type: "DOUBLE".to_owned(), + name_pattern: None, + } +} + +/// A failed proof with Python's reason. +pub(crate) fn failed(reason: &str) -> PivotOutcome { + PivotOutcome::Proof(ContractProof { + output_proven: false, + fixed_columns: Vec::new(), + families: Vec::new(), + input_relations: Vec::new(), + failure_reason: Some(reason.to_owned()), + bare_dynamic_pivot: false, + }) +} + +/// A proven output: `(name, type, nullability)` fixed columns and the `amounts` family type. +pub(crate) fn proven( + fixed: &[(&str, &str, &str)], + family_type: Option<&str>, + input: &str, + bare: bool, +) -> PivotOutcome { + PivotOutcome::Proof(ContractProof { + output_proven: true, + fixed_columns: fixed + .iter() + .map(|(name, data_type, nullability)| ColumnFact { + name: (*name).to_owned(), + data_type: Some((*data_type).to_owned()), + nullability: (*nullability).to_owned(), + }) + .collect(), + families: vec![("amounts".to_owned(), family_type.map(str::to_owned))], + input_relations: vec![input.to_owned()], + failure_reason: None, + bare_dynamic_pivot: bare, + }) +} 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 4f53c3093f..c56625f012 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 @@ -1,8 +1,12 @@ +use crate::assembly::analysis_session::main::prove_dynamic_contract::prove_dynamic_contract; use crate::assembly::analysis_session::main::start_analysis_session::start_analysis_session; +use crate::assembly::analysis_session::models::PivotOutcome; use crate::assembly::analysis_session::tests::helpers::{ - catalog, model_requests, orders_request, session_lines, + catalog, failed, model_requests, orders_request, pivot_request, proven, session_lines, +}; +use crate::assembly::analysis_session::tests::test_types::{ + PivotTestCase, SessionTestCase, UnscheduledTestCase, }; -use crate::assembly::analysis_session::tests::test_types::{SessionTestCase, UnscheduledTestCase}; #[test] fn given_models_when_running_the_session_then_defers_and_publishes_as_python_does() { @@ -53,14 +57,14 @@ fn given_models_when_running_the_session_then_defers_and_publishes_as_python_doe ], }, SessionTestCase { - description: "an unknown column is a binding diagnostic and an untyped output defers", + description: "an unknown column is a binding diagnostic and enriches natively", models: &[( "orders_bad", "SELECT missing_column FROM __source(\"raw_orders\")", &["raw_orders"], &[], )], - expected_steps: &[&["defer enrichment 0"], &[]], + expected_steps: &[&[]], expected_outcomes: &[&[ "missing_column - unknown", "succeeded=true star=false/false diagnostics=[\"B002\"]", @@ -83,7 +87,6 @@ fn given_models_when_running_the_session_then_defers_and_publishes_as_python_doe ), ], expected_steps: &[ - &["defer enrichment 0"], &["publish events event_id:UNKNOWN", "defer analysis 1"], &[], ], @@ -164,3 +167,127 @@ fn given_unschedulable_models_when_starting_then_python_analyses() { ); } } + +#[test] +fn given_dynamic_pivots_when_proving_then_matches_python_or_defers() { + let test_cases = [ + PivotTestCase { + description: "a bare DuckDB pivot keeps its GROUP BY columns", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status USING MAX(amount) GROUP BY customer_id", + family: Some(("status", "amount", "MAX")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "non_null")], + Some("DOUBLE"), + "raw_orders", + true, + ), + }, + PivotTestCase { + description: "a cast aggregate renders its target type", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status \ + USING MIN(CAST(amount AS DECIMAL(12, 2))) GROUP BY customer_id", + family: Some(("status", "amount", "MIN")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "non_null")], + Some("DECIMAL(12, 2)"), + "raw_orders", + true, + ), + }, + PivotTestCase { + description: "a CTE input without GROUP BY keeps the remaining input columns", + dialect: "snowflake", + sql: "WITH base AS (SELECT customer_id, status, amount FROM __source(\"raw_orders\")) \ + SELECT * FROM base PIVOT (SUM(amount) FOR status IN (ANY ORDER BY status))", + family: Some(("status", "amount", "SUM")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "non_null")], + None, + "raw_orders", + false, + ), + }, + PivotTestCase { + description: "a passthrough redeclaring the upstream families", + dialect: "duckdb", + 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: "static pivot values", + dialect: "snowflake", + sql: "SELECT * FROM __source(\"raw_orders\") PIVOT (MAX(amount) FOR status IN ('placed'))", + family: Some(("status", "amount", "MAX")), + expected_outcome: failed( + "static pivot values must use ordinary exact column declarations", + ), + }, + PivotTestCase { + description: "DuckDB values on the ON clause", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status IN ('placed') USING MAX(amount)", + family: Some(("status", "amount", "MAX")), + expected_outcome: failed( + "dynamic column contracts currently require exactly one pivot column", + ), + }, + PivotTestCase { + description: "a dialect without dynamic pivots", + dialect: "postgres", + sql: "SELECT 1", + family: Some(("status", "amount", "MAX")), + expected_outcome: failed( + "adapter dialect 'postgres' does not support compiler-proven dynamic pivots", + ), + }, + PivotTestCase { + description: "a declared pivot column the pivot does not use", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status USING MAX(amount) GROUP BY customer_id", + family: Some(("order_id", "amount", "MAX")), + expected_outcome: failed( + "dynamic family 'amounts' declares pivot_column 'order_id' but the output pivot uses 'status'", + ), + }, + PivotTestCase { + description: "SQL Python reports as a parse error", + dialect: "duckdb", + sql: "PIVOT FROM WHERE", + family: Some(("status", "amount", "MAX")), + expected_outcome: PivotOutcome::Deferred, + }, + PivotTestCase { + description: "non-ASCII names Python casefolds differently", + dialect: "duckdb", + sql: "SELECT * FROM __ref(\"stra\u{df}e\")", + family: Some(("status", "amount", "MAX")), + expected_outcome: PivotOutcome::Deferred, + }, + PivotTestCase { + description: "no declared families", + dialect: "duckdb", + sql: "SELECT 1", + family: None, + expected_outcome: PivotOutcome::Absent, + }, + ]; + for test_case in test_cases { + let request = pivot_request(test_case.dialect, test_case.sql, test_case.family); + + let outcome = prove_dynamic_contract(&request).expect("the proof runs"); + + assert_eq!( + outcome, test_case.expected_outcome, + "{}", + 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 f94d532933..c09b783048 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 @@ -1,3 +1,5 @@ +use crate::assembly::analysis_session::models::PivotOutcome; + /// A model as `(name, sql, source names, model refs)`. pub(crate) type ModelSpec = ( &'static str, @@ -40,3 +42,12 @@ pub(crate) struct WavesTestCase { pub(crate) producers: &'static [&'static [usize]], pub(crate) expected_waves: Option<&'static [&'static [usize]]>, } + +pub(crate) struct PivotTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) sql: &'static str, + /// `(pivot column, value column, aggregate)` of one `amounts` family, or none declared. + pub(crate) family: Option<(&'static str, &'static str, &'static str)>, + pub(crate) expected_outcome: PivotOutcome, +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs index 2e3c48a906..87af8f9b33 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs @@ -4,13 +4,14 @@ use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyRef, PyResult, Python}; use pyo3::{pyclass, pyfunction, pymethods, wrap_pyfunction}; use sqlbuild_analysis::assembly::analysis_session::main::expression_shapes::expression_shapes; use sqlbuild_analysis::assembly::analysis_session::main::finish_analysis_session::finish_analysis_session; +use sqlbuild_analysis::assembly::analysis_session::main::prove_dynamic_contract::prove_dynamic_contract; use sqlbuild_analysis::assembly::analysis_session::main::provide_deferred_analyses::provide_deferred_analyses; use sqlbuild_analysis::assembly::analysis_session::main::run_analysis_session::run_analysis_session; use sqlbuild_analysis::assembly::analysis_session::main::start_analysis_session::start_analysis_session; use sqlbuild_analysis::assembly::analysis_session::models::{ - AnalysisSession, ColumnFact, Deferral, DeferredAnalysis, ExpressionShape, + AnalysisSession, ColumnFact, Deferral, DeferredAnalysis, DynamicFamily, ExpressionShape, ExpressionShapeRequest, LineageFacts, LineageRow, ModelOutcome, ModelReference, ModelRequest, - SessionOutcome, SessionRequest, SessionStep, + PivotOutcome, PivotRequest, SessionOutcome, SessionRequest, SessionStep, }; use sqlbuild_analysis::assembly::analysis_session::types::{Pairs, Shapes}; use sqlbuild_analysis::semantic_validation::types::DiagnosticRow; @@ -25,7 +26,9 @@ use crate::bindings::types::CompilerDetach; type ColumnRow = (String, Option, String); /// `(output column, transform code, confidence code, [(resource type, resource, column)])`. type LineageItem = (String, u8, u8, Vec<(String, String, String)>); -/// A model's name, SQL, placeholders, references, lineage, required names and flags. +/// `(name, pivot column, value column, aggregate, type, name pattern)`. +type FamilyRow = (String, String, String, String, String, Option); +/// A model's name, SQL, placeholders, references, lineage, required names, flags and pivot. type ModelRow = ( String, String, @@ -36,6 +39,8 @@ type ModelRow = ( bool, bool, Option<(String, String)>, + String, + Vec, ); /// The profile settings, relation shapes and models one session analyses. type RequestRow = ( @@ -47,6 +52,7 @@ type RequestRow = ( Shapes, Shapes, Shapes, + Vec<(String, Vec)>, Vec, ); /// `(kind, model, cleaned_sql, schemas, binding diagnostics, lineage)`. @@ -82,8 +88,29 @@ type OutcomeRow = ( bool, String, ); -/// `(models, catalog schema additions, analysis-shape names)`. -type FinishRow = (Vec, Shapes, Vec); +/// `(output proven, fixed columns, (family, inferred type), inputs, failure, bare pivot)`. +type ProofRow = ( + bool, + Vec, + Vec<(String, Option)>, + Vec, + Option, + bool, +); +/// The dialect, relation facts, families by table, SQL and families of one pivot proof. +type PivotRequestRow = ( + String, + Shapes, + Shapes, + Shapes, + Vec<(String, Vec)>, + String, + Vec, +); +/// `("absent" | "deferred" | "proof", proof)`. +type ContractRow = (&'static str, Option); +/// `(models, catalog schema additions, analysis-shape names, dynamic pivot proofs)`. +type FinishRow = (Vec, Shapes, Vec, Vec); /// One compile's native model analysis; any failure hands the whole analysis back to Python. #[pyclass(module = "sqlbuild._native")] @@ -204,6 +231,7 @@ fn session_request(request: RequestRow) -> SessionRequest { column_nullability, complete_schemas, catalog_schemas, + dynamic_families_by_table, models, ) = request; SessionRequest { @@ -215,6 +243,7 @@ fn session_request(request: RequestRow) -> SessionRequest { column_nullability, complete_schemas, catalog_schemas, + dynamic_families_by_table: families_by_table(dynamic_families_by_table), models: models.into_iter().map(model_request).collect(), } } @@ -230,6 +259,8 @@ fn model_request(row: ModelRow) -> ModelRequest { recover_cte_facts, has_set_operation, snapshot_columns, + pivot_sql, + dynamic_families, ) = row; ModelRequest { name, @@ -247,6 +278,8 @@ fn model_request(row: ModelRow) -> ModelRequest { recover_cte_facts, has_set_operation, snapshot_columns, + pivot_sql, + dynamic_families: dynamic_families.into_iter().map(dynamic_family).collect(), } } @@ -325,15 +358,57 @@ fn finish_row(outcome: SessionOutcome) -> FinishRow { outcome.models.into_iter().map(outcome_row).collect(), outcome.schema_additions, outcome.analysis_names, + outcome + .dynamic_contracts + .into_iter() + .map(contract_row) + .collect(), ) } +fn families_by_table(rows: Vec<(String, Vec)>) -> Vec<(String, Vec)> { + rows.into_iter() + .map(|(name, families)| (name, families.into_iter().map(dynamic_family).collect())) + .collect() +} + +fn dynamic_family(row: FamilyRow) -> DynamicFamily { + let (name, pivot_column, value_column, aggregate, data_type, name_pattern) = row; + DynamicFamily { + name, + pivot_column, + value_column, + aggregate, + data_type, + name_pattern, + } +} + +fn contract_row(outcome: PivotOutcome) -> ContractRow { + match outcome { + PivotOutcome::Absent => ("absent", None), + PivotOutcome::Deferred => ("deferred", None), + PivotOutcome::Proof(proof) => ( + "proof", + Some(( + proof.output_proven, + proof.fixed_columns.into_iter().map(column_row).collect(), + proof.families, + proof.input_relations, + proof.failure_reason, + proof.bare_dynamic_pivot, + )), + ), + } +} + fn outcome_row(outcome: ModelOutcome) -> OutcomeRow { let analysis = outcome.analysis; let (source, lineage) = match analysis.lineage { LineageFacts::Native(rows) => ("native", lineage_items(rows)), LineageFacts::PythonAnalysis => ("analysis", Vec::new()), LineageFacts::PythonEnrichment => ("enrichment", Vec::new()), + LineageFacts::NativeEnrichment(rows) => ("facts", lineage_items(rows)), }; ( analysis.analysis_succeeded, @@ -350,9 +425,28 @@ fn outcome_row(outcome: ModelOutcome) -> OutcomeRow { ) } +/// One model's dynamic pivot proof, or `("deferred", None)` for Python's proof. +#[pyfunction] +fn prove_dynamic_column_contract(py: Python<'_>, request: PivotRequestRow) -> ContractRow { + let (dialect, column_types, authoritative_types, column_nullability, by_table, sql, families) = + request; + let request = PivotRequest { + dialect, + column_types, + authoritative_types, + column_nullability, + families_by_table: families_by_table(by_table), + sql, + families: families.into_iter().map(dynamic_family).collect(), + }; + py.compiler_detach(|| prove_dynamic_contract(&request)) + .map_or(("deferred", None), contract_row) +} + pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_class::()?; module.add_function(wrap_pyfunction!(start_model_analysis_session, module)?)?; module.add_function(wrap_pyfunction!(infer_expression_source_shapes, module)?)?; + module.add_function(wrap_pyfunction!(prove_dynamic_column_contract, module)?)?; Ok(()) } diff --git a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py index 2c00e8a6d6..4a7c34a23e 100644 --- a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py +++ b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py @@ -13,6 +13,18 @@ f"SELECT s.*, s.amount * 2 AS doubled_amount\nFROM {_ORDERS} s\n" ) +_PIVOT_PATH: str = "models/marts/status_amounts.sql" + + +def _pivot_model(*, pivot_column: str, body: str) -> str: + return ( + 'MODEL (\n description "Order amounts pivoted by status",\n contract enforced,\n' + " materialized table,\n columns (customer_id (type INTEGER)),\n" + f" dynamic_columns (\n status_amounts (\n pivot_column {pivot_column},\n" + " value_column amount,\n aggregate MAX,\n type DOUBLE\n )\n ),\n);\n\n" + f"{body}" + ) + def analysis_session_failure_cases() -> tuple[FailureCase, ...]: """Return this lane's failure cases; the lane appends here without editing shared lists.""" @@ -47,4 +59,39 @@ def analysis_session_failure_cases() -> tuple[FailureCase, ...]: "SELECT customer_id, total_amount, order_count FROM totals\n", }, ), + failure_case( + name="analysis-session-enriched-subquery-type-mismatch", + expected_code="K002", + files={ + FAILURE_MART_PATH: 'MODEL (\n description "Order totals per customer",\n' + " contract enforced,\n" + " columns (\n customer_id (type INTEGER),\n amount (type VARCHAR),\n ),\n" + ");\n\n" + f"SELECT customer_id, amount\nFROM (SELECT * FROM {_ORDERS}) staged\n", + }, + ), + failure_case( + name="analysis-session-pivot-values-on-clause", + expected_code="K011", + expected_message="exactly one pivot column", + files={ + _PIVOT_PATH: _pivot_model( + pivot_column="status", + body="PIVOT __source(\"raw_orders\")\nON status IN ('placed')\n" + "USING MAX(amount)\nGROUP BY customer_id\n", + ), + }, + ), + failure_case( + name="analysis-session-pivot-family-column-mismatch", + expected_code="K011", + expected_message="declares pivot_column 'region'", + files={ + _PIVOT_PATH: _pivot_model( + pivot_column="region", + body='PIVOT __source("raw_orders")\nON status\nUSING MAX(amount)\n' + "GROUP BY customer_id\n", + ), + }, + ), ) diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 2fdc19708c..d9199a38ca 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -397,6 +397,20 @@ class NativeModelAnalysisSession: ], list[tuple[str, list[tuple[str, str]]]], list[str], + list[ + tuple[ + str, + tuple[ + bool, + list[tuple[str, str | None, str]], + list[tuple[str, str | None]], + list[str], + str | None, + bool, + ] + | None, + ] + ], ] | None ): ... @@ -406,6 +420,29 @@ class NativeModelAnalysisSession: def start_model_analysis_session( catalog: object, request: tuple[object, ...], / ) -> NativeModelAnalysisSession | None: ... +def prove_dynamic_column_contract( + request: tuple[ + str, + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, list[tuple[str, str, str, str, str, str | None]]]], + str, + list[tuple[str, str, str, str, str, str | None]], + ], + /, +) -> tuple[ + str, + tuple[ + bool, + list[tuple[str, str | None, str]], + list[tuple[str, str | None]], + list[str], + str | None, + bool, + ] + | None, +]: ... def infer_expression_source_shapes( catalog: object, request: tuple[str, bool, list[tuple[str, str]], list[str]], / ) -> tuple[list[tuple[bool, list[tuple[str, str]] | None]], str | None]: ... diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py b/src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py new file mode 100644 index 0000000000..803935326f --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py @@ -0,0 +1,52 @@ +"""Prove one model's dynamic column contract natively, outside the analysis session.""" + +from __future__ import annotations + +import sqlbuild._native as _native +from sqlbuild.compiler.analysis_session._helpers.deferral_records import record_analysis_deferral +from sqlbuild.compiler.analysis_session._helpers.session_rows import ( + contract_proof, + family_rows, + shape_rows, +) +from sqlbuild.compiler.analysis_session.constants import CONTRACT_DEFERRED, DEFERRAL_DYNAMIC_PIVOT +from sqlbuild.compiler.analysis_session.types import ContractRow, ShapeRows +from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis +from sqlbuild.compiler.compile.models import DynamicColumnContractProof +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily + + +def native_dynamic_contract( + *, + query_sql: str, + dialect: str | None, + families: tuple[SchemaDynamicColumnFamily, ...], + column_types_by_table: dict[str, dict[str, str]], + authoritative_column_types_by_table: dict[str, dict[str, str]], + column_nullability_by_table: dict[str, dict[str, InferredNullability]], + dynamic_families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], +) -> DynamicColumnContractProof | None: + """Python's proof for declared families, or None (recorded) where Python must prove it.""" + + shapes: tuple[ShapeRows | None, ...] = ( + shape_rows(column_types_by_table), + shape_rows(authoritative_column_types_by_table), + shape_rows(column_nullability_by_table), + ) + row: ContractRow = (CONTRACT_DEFERRED, None) + if all(rows is not None for rows in shapes): + row = _native.prove_dynamic_column_contract( + ( + dialect or "generic", + shapes[0] or [], + shapes[1] or [], + shapes[2] or [], + family_rows(dynamic_families_by_table), + query_sql, + [PythonModelAnalysis.family_row(family) for family in families], + ) + ) + if row[0] == CONTRACT_DEFERRED: + record_analysis_deferral(kind=DEFERRAL_DYNAMIC_PIVOT) + return contract_proof(row[1]) diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py index 88b2dab191..093c361e6e 100644 --- a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py +++ b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py @@ -10,18 +10,25 @@ ColumnRow, DeferredRow, DiagnosticRow, + FamilyRow, LineageItem, + ProofRow, ShapeRows, ) from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis from sqlbuild.compiler.compile.models import ( CompactLineageFacts, + CompiledLineageColumnFact, + CompiledLineageSourceFact, + DynamicColumnContractProof, + DynamicColumnFamilyProof, InferredColumn, PolyglotAnalysisResult, ) from sqlbuild.compiler.lineage.types import InferredNullability from sqlbuild.compiler.sql_analysis.constants import CASE_SENSITIVE_BINDING_DIALECTS from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily def shape_rows(shapes: Mapping[str, Mapping[str, object]]) -> ShapeRows | None: @@ -61,10 +68,30 @@ def session_request( list(profile.function_return_types.items()), request.rich_type_inference, *shapes, + family_rows(request.dynamic_families_by_table), python.model_rows(), ) +def contract_proof(row: ProofRow | None) -> DynamicColumnContractProof | None: + """A native dynamic pivot proof as Python's, None when absent or left to Python.""" + + if row is None: + return None + output_proven, columns, families, input_relations, failure_reason, bare = row + return DynamicColumnContractProof( + output_proven=output_proven, + fixed_columns=inferred_columns(columns) or (), + families=tuple( + DynamicColumnFamilyProof(name=name, inferred_type=inferred_type) + for name, inferred_type in families + ), + input_relations=tuple(input_relations), + failure_reason=failure_reason, + bare_dynamic_pivot=bare, + ) + + def case_sensitive_shapes(*, profile: ExpressionInferenceProfile, dialect: str | None) -> bool: """Python's `inferred_binding_shape` test for keeping authored identifier quoting.""" @@ -143,3 +170,38 @@ def compact_lineage(rows: list[LineageItem]) -> CompactLineageFacts: ) ) return CompactLineageFacts(string_pool=tuple(pool), rows=tuple(compact_rows)) + + +def lineage_facts(rows: list[LineageItem]) -> tuple[CompiledLineageColumnFact, ...]: + """Native lineage rows as the plain facts Python's re-analysis returns.""" + + return tuple( + CompiledLineageColumnFact( + output_column=output_column, + upstream_columns=_source_facts(sources), + transform_kind=CompactLineageFacts.transform_kind(transform_code), + confidence=CompactLineageFacts.confidence(confidence_code), + ) + for output_column, transform_code, confidence_code, sources in rows + ) + + +def _source_facts(sources: list[tuple[str, str, str]]) -> tuple[CompiledLineageSourceFact, ...]: + return tuple( + CompiledLineageSourceFact( + resource_type=resource_type, resource_name=resource_name, column_name=column_name + ) + for resource_type, resource_name, column_name in sources + ) + + +def family_rows( + families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], +) -> list[tuple[str, list[FamilyRow]]]: + """Declared dynamic families by relation, as native request rows.""" + + return [(name, _family_rows(families)) for name, families in families_by_table.items()] + + +def _family_rows(families: tuple[SchemaDynamicColumnFamily, ...]) -> list[FamilyRow]: + return [PythonModelAnalysis.family_row(family) for family in families] diff --git a/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py b/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py index 7674467755..8fa6049f45 100644 --- a/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py +++ b/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py @@ -12,14 +12,19 @@ from sqlbuild.compiler.analysis_session._helpers.session_rows import ( binding_diagnostics, compact_lineage, + contract_proof, deferred_row, inferred_columns, + lineage_facts, ) from sqlbuild.compiler.analysis_session.constants import ( + CONTRACT_DEFERRED, DEFERRAL_ANALYSIS, + DEFERRAL_DYNAMIC_PIVOT, DEFERRAL_ENRICHMENT, DEFERRAL_LEGACY_ANALYSIS, DEFERRAL_SESSION, + LINEAGE_FACTS, LINEAGE_NATIVE, NATIVE_ANALYSIS_FAILURE_MESSAGE, NATIVE_SESSION_FAILURE_MESSAGE, @@ -30,6 +35,7 @@ DeferredRow, FinishRow, OutcomeRow, + ProofRow, ShapeRows, StepRow, ) @@ -79,19 +85,22 @@ def analyses( sqlbuild_error=session.failure, ) return None - outcomes, schema_additions, analysis_names = finished + outcomes, schema_additions, analysis_names, contracts = finished self._record_catalog_changes( schema_additions=schema_additions, analysis_names=analysis_names ) + self._deferrals[DEFERRAL_DYNAMIC_PIVOT] += sum( + kind == CONTRACT_DEFERRED for kind, _ in contracts + ) for kind, count in self._deferrals.items(): record_analysis_deferral(kind=kind, count=count) self._python.record_uncached() return { request.model_input.model_file.file_path.stem: self._model_analysis( - index=index, request=request, outcome=outcome + index=index, request=request, outcome=outcome, contract=contract ) - for index, (request, outcome) in enumerate( - zip(self._python.requests, outcomes, strict=True) + for index, (request, outcome, (_, contract)) in enumerate( + zip(self._python.requests, outcomes, contracts, strict=True) ) } @@ -168,7 +177,12 @@ def _record_catalog_changes( ) def _model_analysis( - self, *, index: int, request: ModelSqlAnalysisRequest, outcome: OutcomeRow + self, + *, + index: int, + request: ModelSqlAnalysisRequest, + outcome: OutcomeRow, + contract: ProofRow | None, ) -> ModelSqlAnalysis: ( succeeded, @@ -184,6 +198,8 @@ def _model_analysis( lineage: Sequence[CompiledLineageColumnFact] = ( (compact_lineage(lineage_rows) if succeeded else ()) if lineage_source == LINEAGE_NATIVE + else lineage_facts(lineage_rows) + if lineage_source == LINEAGE_FACTS else self._kept[(index, lineage_source)].lineage_columns ) return ModelSqlAnalysis( @@ -200,4 +216,5 @@ def _model_analysis( fused_binding_validated=True, cleaned_sql=cleaned_sql, validated_schema=request.binding_schema, + dynamic_column_contract=contract_proof(contract), ) diff --git a/src/sqlbuild/compiler/analysis_session/constants.py b/src/sqlbuild/compiler/analysis_session/constants.py index 51a32d88cf..9d819d795d 100644 --- a/src/sqlbuild/compiler/analysis_session/constants.py +++ b/src/sqlbuild/compiler/analysis_session/constants.py @@ -6,10 +6,13 @@ DEFERRAL_SESSION: str = "session" DEFERRAL_LEGACY_ANALYSIS: str = "legacy_analysis" DEFERRAL_ENRICHMENT: str = "input_enrichment" +DEFERRAL_DYNAMIC_PIVOT: str = "dynamic_pivot" DEFERRAL_EXPRESSION_SHAPES: str = "expression_shapes" DEFERRAL_ANALYSIS: str = "analysis" DEFERRAL_ENRICH: str = "enrichment" LINEAGE_NATIVE: str = "native" +LINEAGE_FACTS: str = "facts" +CONTRACT_DEFERRED: str = "deferred" NATIVE_ANALYSIS_FAILURE_MESSAGE: str = "native compact query analysis failed" NATIVE_SESSION_FAILURE_MESSAGE: str = "native model analysis failed; falling back" NATIVE_SHAPES_FAILURE_MESSAGE: str = "native expression-source shape inference failed; falling back" diff --git a/src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py b/src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py new file mode 100644 index 0000000000..9d47358f42 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py @@ -0,0 +1,31 @@ +"""Prove a dynamic column contract natively for the preview compiler engine.""" + +from __future__ import annotations + +from sqlbuild.compiler.analysis_session._helpers.dynamic_contract import native_dynamic_contract +from sqlbuild.compiler.compile.models import DynamicColumnContractProof +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily + + +def prove_native_dynamic_contract( + *, + query_sql: str, + dialect: str | None, + families: tuple[SchemaDynamicColumnFamily, ...], + column_types_by_table: dict[str, dict[str, str]], + authoritative_column_types_by_table: dict[str, dict[str, str]], + column_nullability_by_table: dict[str, dict[str, InferredNullability]], + dynamic_families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], +) -> DynamicColumnContractProof | None: + """Return Python's proof for declared families, or None to prove it in Python.""" + + return native_dynamic_contract( + query_sql=query_sql, + dialect=dialect, + families=families, + column_types_by_table=column_types_by_table, + authoritative_column_types_by_table=authoritative_column_types_by_table, + column_nullability_by_table=column_nullability_by_table, + dynamic_families_by_table=dynamic_families_by_table, + ) diff --git a/src/sqlbuild/compiler/analysis_session/types.py b/src/sqlbuild/compiler/analysis_session/types.py index 746ce29cc0..8fcf150865 100644 --- a/src/sqlbuild/compiler/analysis_session/types.py +++ b/src/sqlbuild/compiler/analysis_session/types.py @@ -22,4 +22,9 @@ bool, str, ] -type FinishRow = tuple[list[OutcomeRow], ShapeRows, list[str]] +type FamilyRow = tuple[str, str, str, str, str, str | None] +type ProofRow = tuple[ + bool, list[ColumnRow], list[tuple[str, str | None]], list[str], str | None, bool +] +type ContractRow = tuple[str, ProofRow | None] +type FinishRow = tuple[list[OutcomeRow], ShapeRows, list[str], list[ContractRow]] diff --git a/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py b/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py index fc4fe25baa..145745939d 100644 --- a/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py +++ b/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py @@ -12,6 +12,9 @@ from sqlbuild.compiler.analysis_session.main._infer_native_expression_source_shapes import ( infer_native_expression_source_shapes, ) +from sqlbuild.compiler.analysis_session.main._prove_native_dynamic_contract import ( + prove_native_dynamic_contract, +) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( get_expression_source_shapes, @@ -91,8 +94,18 @@ def dynamic_column_contract_by_engine( sql_analysis: ModelSqlAnalysis | None, python_proof: partial[DynamicColumnContractProof | None], ) -> DynamicColumnContractProof | None: - """Return the proof the model's native analysis carries, or prove the contract in Python.""" + """Return the proof native analysis carries or proves alone, or prove it in Python.""" if sql_analysis is not None and sql_analysis.dynamic_column_contract is not None: return sql_analysis.dynamic_column_contract + if ( + sql_analysis is None + and python_proof.keywords["families"] + and native_stage_enabled(NativeStage.MODEL_ANALYSIS) + ): + native: DynamicColumnContractProof | None = prove_native_dynamic_contract( + **python_proof.keywords + ) + if native is not None: + return native return python_proof() diff --git a/src/sqlbuild/compiler/compile/classes/python_model_analysis.py b/src/sqlbuild/compiler/compile/classes/python_model_analysis.py index 20a7c1b100..025f8187d0 100644 --- a/src/sqlbuild/compiler/compile/classes/python_model_analysis.py +++ b/src/sqlbuild/compiler/compile/classes/python_model_analysis.py @@ -6,7 +6,10 @@ from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.compiler.compile._helpers.analysis.cache import record_analysis_cache_metrics -from sqlbuild.compiler.compile._helpers.analysis.columns import table_function_analysis_name +from sqlbuild.compiler.compile._helpers.analysis.columns import ( + substitute_placeholder_defaults, + table_function_analysis_name, +) from sqlbuild.compiler.compile._helpers.analysis.compact import ( analyze_columns_and_lineage_with_polyglot, ) @@ -33,6 +36,7 @@ from sqlbuild.compiler.planner.types import ContractPolicy, MaterializationType from sqlbuild.compiler.references.types import SqlReferenceKind from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily _SET_OPERATION_PATTERN: re.Pattern[str] = re.compile( r"\b(?:UNION|INTERSECT|EXCEPT)\b", re.IGNORECASE @@ -116,6 +120,19 @@ def enrich( recover_cte_facts=_recovers_cte_facts(request.model_input), ) + @staticmethod + def family_row(family: SchemaDynamicColumnFamily) -> tuple[str, str, str, str, str, str | None]: + """One declared dynamic column family as the native request row.""" + + return ( + family.name, + family.pivot_column, + family.value_column, + family.aggregate, + family.type, + family.name_pattern, + ) + @staticmethod def legacy_precomputed( *, @@ -189,6 +206,21 @@ def _model_row(request: ModelSqlAnalysisRequest) -> tuple[object, ...]: if values.get("materialized") == MaterializationType.SNAPSHOT else None ), + ( + substitute_placeholder_defaults( + query_sql=request.query_sql, placeholders=request.placeholders + ) + if request.placeholders + else request.query_sql + ), + [ + PythonModelAnalysis.family_row(family) + for family in ( + model_input.schema_entry.dynamic_columns + if model_input.schema_entry is not None + else () + ) + ], ) diff --git a/tests/e2e/scripts/compiler_differential/test_compiler_differential.py b/tests/e2e/scripts/compiler_differential/test_compiler_differential.py index 62806b8a70..edc38239ff 100644 --- a/tests/e2e/scripts/compiler_differential/test_compiler_differential.py +++ b/tests/e2e/scripts/compiler_differential/test_compiler_differential.py @@ -436,14 +436,13 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness "test_case", [ WheelSiteReportTestCase( - description="python_and_preview_wheel_calls_with_python_fallback_deferrals", + description="python_and_preview_wheel_calls_without_preview_deferrals", perturbation="", expected_lines=( "Polyglot wheel calls (python):", "Polyglot wheel calls (native-preview):", "Analysis deferrals (python): none recorded", - "Analysis deferrals (native-preview):", - " input_enrichment analysis_session (project ", + "Analysis deferrals (native-preview): none recorded", "Compiler differential passed: 1 projects identical", ), expected_sites=frozenset( @@ -452,7 +451,7 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness "_validate_sql_syntax_with_message parse_one", } ), - expected_deferrals=frozenset({"input_enrichment analysis_session"}), + expected_deferrals=frozenset(), ), WheelSiteReportTestCase( description="preview_deferrals_are_counted_per_kind_and_site", @@ -463,12 +462,7 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness " legacy_fallback orders.sql (project ", ), expected_sites=frozenset(), - expected_deferrals=frozenset( - { - "legacy_fallback orders.sql", - "input_enrichment analysis_session", - } - ), + expected_deferrals=frozenset({"legacy_fallback orders.sql"}), ), ], ids=lambda case: case.description, diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py index 8c14cdd29c..f3ea3a0aac 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py @@ -14,6 +14,10 @@ class GeneratedAnalysisParityTestCase: dialects: tuple[str | None, ...] expected_minimum_native: int expected_minimum_expression_shapes: int + expected_minimum_pivot_proofs: int + expected_minimum_proven_pivots: int + expected_maximum_enrichment_deferrals: int + expected_minimum_native_enrichments: int @dataclass(frozen=True) @@ -35,3 +39,13 @@ class SessionFailureTestCase: seed: int model_count: int expected_kinds: dict[str, int] + + +@dataclass(frozen=True) +class StandalonePivotProofTestCase: + """A pivot model left out of model analysis, whose proof assembly takes natively.""" + + description: str + analysed_models: frozenset[str] + expected_native_proofs: int + expected_proven_by_model: dict[str, bool | None] diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py index 080c39fe00..47095774c0 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py @@ -14,7 +14,9 @@ import pytest import sqlbuild._native as native_module +import sqlbuild.compiler.analysis_session.classes.native_model_analysis as native_model_analysis import sqlbuild.compiler.compile._helpers.assembly.project as project_assembly +import sqlbuild.compiler.compile._helpers.native_stages.assembly as native_stage_assembly from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.compiler.analysis_session.main._analyze_native_model_sql import ( @@ -23,6 +25,9 @@ from sqlbuild.compiler.analysis_session.main._infer_native_expression_source_shapes import ( infer_native_expression_source_shapes, ) +from sqlbuild.compiler.analysis_session.main._prove_native_dynamic_contract import ( + prove_native_dynamic_contract, +) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( get_expression_source_shapes, @@ -32,7 +37,9 @@ from sqlbuild.compiler.compile.models import ( CompactLineageFacts, CompileAdapterContext, + CompiledLineageColumnFact, CompileProjectInputs, + DynamicColumnContractProof, ModelSqlAnalysis, ) from sqlbuild.compiler.discovery.main.discover import discover_project_inputs @@ -60,6 +67,17 @@ description: Customers feed. expression: >- (SELECT 10 AS customer_id, 'Ada' AS Customer_Name, 'east' AS region) + - name: raw_payments + description: Typed payments feed. + expression: >- + (SELECT 10 AS customer_id, 'east' AS region, CAST(5 AS DECIMAL(10, 2)) AS amount) + columns: + - name: customer_id + type: INTEGER + - name: region + type: VARCHAR + - name: amount + type: DECIMAL(10, 2) - name: raw_events description: Untyped events feed. columns: @@ -78,26 +96,80 @@ python_functions_inherit_default_namespace=True, sql_lexical_syntax=DuckDbAdapter.sql_lexical_syntax, ) -_PIVOT_MODEL: str = """MODEL ( +_PIVOT_HEADER: str = """MODEL ( description "Order amounts pivoted by status", contract enforced, materialized table, columns (customer_id (type INTEGER)), dynamic_columns ( status_amounts ( - pivot_column status, - value_column amount, - aggregate MAX, + pivot_column {pivot}, + value_column {value}, + aggregate {aggregate}, type DOUBLE ) ), ); -PIVOT __source("raw_orders") -ON status -USING MAX(amount) -GROUP BY customer_id """ +_PIVOT_BODIES: tuple[str, ...] = ( + 'PIVOT __source("raw_orders")\nON status\nUSING {aggregate}(amount)\nGROUP BY customer_id', + 'PIVOT __source("raw_orders")\nON status\nUSING {aggregate}(CAST(amount AS DECIMAL(12, 2)))', + "WITH base AS (\n SELECT customer_id, status, CAST(amount AS DOUBLE) AS amount\n" + ' FROM __source("raw_orders")\n)\nSELECT * FROM base\n' + "PIVOT ({aggregate}(amount) FOR status IN (ANY ORDER BY status))", + 'WITH base AS (\n SELECT * FROM __source("raw_orders")\n),\npivoted AS (\n' + " PIVOT base ON status USING {aggregate}(amount) GROUP BY customer_id\n)\n" + "SELECT * FROM pivoted", + "PIVOT __source(\"raw_orders\")\nON status IN ('placed', 'shipped')\n" + "USING {aggregate}(amount)\nGROUP BY customer_id", + 'PIVOT __source("raw_events")\nON kind\nUSING {aggregate}(order_id)\nGROUP BY event_id', + 'SELECT * FROM __source("raw_orders") o\nJOIN __source("raw_customers") c USING (customer_id)', +) +_PIVOT_AGGREGATES: tuple[str, ...] = ("MAX", "MIN", "ANY_VALUE", "SUM") +_TYPED_PIVOT: str = ( + _PIVOT_HEADER.format(pivot="status", value="amount", aggregate="MAX") + + _PIVOT_BODIES[0].format(aggregate="MAX") + + "\n" +) +_TYPED_PIVOT_PASSTHROUGH: str = ( + _PIVOT_HEADER.format(pivot="status", value="amount", aggregate="MAX") + + 'SELECT * FROM __ref("status_amounts")\n' +) +_CONTRACT_CTE_MODELS: dict[str, str] = { + "models/marts/payments_by_customer.sql": ( + 'MODEL (\n description "Payments passed through a CTE",\n contract enforced,\n' + " columns (customer_id (type INTEGER), region (type VARCHAR)),\n);\n\n" + 'WITH base AS (SELECT customer_id, region FROM __source("raw_payments"))\n' + "SELECT customer_id, region FROM base\n" + ), + "models/marts/payment_totals.sql": ( + 'MODEL (\n description "Payment totals aggregated in a CTE",\n contract enforced,\n' + " columns (total_amount (type DECIMAL(38, 2)), region (type VARCHAR)),\n);\n\n" + "WITH base AS (\n SELECT SUM(amount) AS total_amount, region\n" + ' FROM __source("raw_payments")\n GROUP BY region\n)\n' + "SELECT total_amount, region FROM base\n" + ), +} + + +def _pivot_models(rng: random.Random) -> dict[str, str]: + """A typed pivot every passthrough reads, and seeded pivot shapes proven or refused.""" + + files: dict[str, str] = { + "models/marts/status_amounts.sql": _TYPED_PIVOT, + "models/marts/status_passthrough.sql": _TYPED_PIVOT_PASSTHROUGH, + } + for index, body in enumerate(rng.sample(_PIVOT_BODIES, k=4)): + aggregate: str = rng.choice(_PIVOT_AGGREGATES) + header: str = _PIVOT_HEADER.format( + pivot=rng.choice(("status", "status", "kind")), + value=rng.choice(("amount", "amount", "order_id")), + aggregate=aggregate, + ) + files[f"models/marts/pivot_{index}.sql"] = header + body.format(aggregate=aggregate) + "\n" + return files + type _Relation = tuple[str, tuple[str, ...]] type _Template = Callable[[random.Random, list[_Relation]], tuple[str, str, tuple[str, ...]]] @@ -215,13 +287,55 @@ def _subquery(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tu ) +def pivot_project_files() -> dict[str, str]: + """A typed pivot with its passthrough and one plain model, for selective assembly.""" + + return { + "sqlbuild_project.toml": _PROJECT_TOML, + "sources/raw.yml": _SOURCES, + "models/marts/status_amounts.sql": _TYPED_PIVOT, + "models/marts/status_passthrough.sql": _TYPED_PIVOT_PASSTHROUGH, + "models/staging/orders_list.sql": ( + 'MODEL (description "Order ids");\n\nSELECT order_id FROM __source("raw_orders")\n' + ), + } + + +def native_pivot_proofs( + *, monkeypatch: pytest.MonkeyPatch +) -> list[DynamicColumnContractProof | None]: + """Record every standalone native pivot proof assembly takes, failing any wheel pivot proof.""" + + proofs: list[DynamicColumnContractProof | None] = [] + native_proof: Callable[..., DynamicColumnContractProof | None] = ( + native_stage_assembly.prove_native_dynamic_contract + ) + + def recorded(**keywords: Any) -> DynamicColumnContractProof | None: + proofs.append(native_proof(**keywords)) + return proofs[-1] + + python_proof: Callable[..., DynamicColumnContractProof | None] = ( + project_assembly.analyze_dynamic_column_contract + ) + + def wheel_not_expected(**keywords: Any) -> DynamicColumnContractProof | None: + assert not keywords["families"], "the Python wheel proved a pivot native should prove" + return python_proof(**keywords) + + monkeypatch.setattr(native_stage_assembly, "prove_native_dynamic_contract", recorded) + monkeypatch.setattr(project_assembly, "analyze_dynamic_column_contract", wheel_not_expected) + return proofs + + def generated_analysis_files(*, rng: random.Random, model_count: int) -> dict[str, str]: """Sources (typed, inferred and untyped), a pivot and models drawing on every template.""" files: dict[str, str] = { "sqlbuild_project.toml": _PROJECT_TOML, "sources/raw.yml": _SOURCES, - "models/marts/status_amounts.sql": _PIVOT_MODEL, + **_pivot_models(rng), + **_CONTRACT_CTE_MODELS, } relations: list[_Relation] = list(_INITIAL_RELATIONS) for index in range(model_count): @@ -258,6 +372,10 @@ class AnalysisParity: native: list[object] = field(default_factory=list) analysed_models: int = 0 expression_shapes: int = 0 + pivot_proofs: int = 0 + standalone_proofs: int = 0 + proven_pivots: int = 0 + native_enrichments: int = 0 def compare_analyses( @@ -295,6 +413,25 @@ def models_by_both( _catalog_view(profile.binding_catalog), _catalog_view(native_catalog), ) + return native or python + + def proofs_by_both( + *, + sql_analysis: ModelSqlAnalysis | None, + python_proof: partial[DynamicColumnContractProof | None], + ) -> DynamicColumnContractProof | None: + native: DynamicColumnContractProof | None = getattr( + sql_analysis, "dynamic_column_contract", None + ) + python: DynamicColumnContractProof | None = python_proof() + standalone: DynamicColumnContractProof | None = prove_native_dynamic_contract( + **python_proof.keywords + ) + parity.pivot_proofs += native is not None + parity.proven_pivots += native is not None and native.output_proven + parity.standalone_proofs += standalone is not None + _append(parity, "dynamic pivot proof", python, native or python) + _append(parity, "standalone dynamic pivot proof", python, standalone or python) return python def shapes_by_both( @@ -317,9 +454,19 @@ def shapes_by_both( ) return python + native_lineage_facts: Callable[..., tuple[CompiledLineageColumnFact, ...]] = ( + native_model_analysis.lineage_facts + ) + + def counted_native_enrichment(rows: list[Any]) -> tuple[CompiledLineageColumnFact, ...]: + parity.native_enrichments += 1 + return native_lineage_facts(rows) + with monkeypatch.context() as patch: + patch.setattr(native_model_analysis, "lineage_facts", counted_native_enrichment) patch.setattr(project_assembly, "analyze_model_sql_by_engine", models_by_both) patch.setattr(project_assembly, "expression_source_shapes_by_engine", shapes_by_both) + patch.setattr(project_assembly, "dynamic_column_contract_by_engine", proofs_by_both) with suppress(CompileInputError): _ = project_assembly.assemble_compiled_project( inputs=inputs, @@ -387,7 +534,6 @@ def _analysis_views(analyses: dict[str, ModelSqlAnalysis] | None) -> object: analysis.polyglot_analysis.binding_validated, analysis.cleaned_sql, analysis.placeholders, - analysis.dynamic_column_contract, ) for name, analysis in analyses.items() } diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py index ee5c88edd7..4e9527ea64 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py @@ -14,12 +14,20 @@ analyze_native_model_sql, ) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest -from sqlbuild.compiler.compile.models import CompileProjectInputs +from sqlbuild.compiler.compile._helpers.assembly.project import assemble_compiled_project +from sqlbuild.compiler.compile.models import ( + CompiledProject, + CompileProjectInputs, + DynamicColumnContractProof, +) +from sqlbuild.compiler.frontier.constants import COMPILER_ENGINE_ENV_VAR +from sqlbuild.compiler.frontier.types import CompilerEngine from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR from tests.integration.src.sqlbuild.compiler.analysis_session._test_types import ( AnalysisFallbackTestCase, GeneratedAnalysisParityTestCase, SessionFailureTestCase, + StandalonePivotProofTestCase, ) from tests.integration.src.sqlbuild.compiler.analysis_session.helpers import ( AnalysisParity, @@ -30,6 +38,8 @@ deferral_kinds, failing_provide_sessions, generated_analysis_files, + native_pivot_proofs, + pivot_project_files, ) from tests.integration.src.sqlbuild.compiler.helpers import mismatches @@ -44,13 +54,17 @@ "test_case", [ GeneratedAnalysisParityTestCase( - description="stars, CTEs, set operations, untyped inputs, contracts and a pivot", + description="stars, CTEs, set operations, untyped inputs, contracts, contract CTEs, pivots", seed=20261008, count=6, model_count=24, dialects=("duckdb", "postgres", "snowflake", "bigquery"), expected_minimum_native=500, expected_minimum_expression_shapes=40, + expected_minimum_pivot_proofs=80, + expected_minimum_proven_pivots=12, + expected_maximum_enrichment_deferrals=120, + expected_minimum_native_enrichments=90, ) ], ids=lambda case: case.description, @@ -75,11 +89,20 @@ def test_given_generated_projects_when_analysing_natively_then_matches_python( assert mismatches(inputs=parity.names, expected=parity.python, actual=parity.native) == [] assert kinds["analysis_session:session"] == kinds["analysis_session:expression_shapes"] == 0 + assert kinds["analysis_session:dynamic_pivot"] == 0 assert ( parity.analysed_models - kinds["analysis_session:legacy_analysis"] >= test_case.expected_minimum_native ) assert parity.expression_shapes >= test_case.expected_minimum_expression_shapes + assert parity.pivot_proofs >= test_case.expected_minimum_pivot_proofs + assert parity.standalone_proofs >= test_case.expected_minimum_pivot_proofs + assert parity.proven_pivots >= test_case.expected_minimum_proven_pivots + assert ( + kinds["analysis_session:input_enrichment"] + <= test_case.expected_maximum_enrichment_deferrals + ) + assert parity.native_enrichments >= test_case.expected_minimum_native_enrichments @pytest.mark.parametrize( @@ -169,5 +192,46 @@ def test_given_session_failure_after_deferrals_when_analysing_then_records_only_ assert deferral_kinds(record_dir) == Counter(test_case.expected_kinds) +@pytest.mark.parametrize( + "test_case", + [ + StandalonePivotProofTestCase( + description="a selection that leaves the pivots out of model analysis", + analysed_models=frozenset({"orders_list"}), + expected_native_proofs=2, + expected_proven_by_model={ + "orders_list": None, + "status_amounts": True, + "status_passthrough": True, + }, + ) + ], + ids=lambda case: case.description, +) +def test_given_unanalysed_pivot_model_when_assembling_then_native_proves_it_without_the_wheel( + test_case: StandalonePivotProofTestCase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + inputs: CompileProjectInputs = compile_inputs( + project_dir=tmp_path / "project", files=pivot_project_files() + ) + monkeypatch.setenv(COMPILER_ENGINE_ENV_VAR, CompilerEngine.NATIVE_PREVIEW.value) + proofs: list[DynamicColumnContractProof | None] = native_pivot_proofs(monkeypatch=monkeypatch) + + project: CompiledProject = assemble_compiled_project( + inputs=inputs, + inference_profile=ExpressionInferenceProfile(sql_analysis_dialect="duckdb"), + analysis_model_names=test_case.analysed_models, + ) + + assert ( + all(proof is not None for proof in proofs), + len(proofs), + { + model.name: getattr(model.dynamic_column_contract, "output_proven", None) + for model in project.models + }, + ) == (True, test_case.expected_native_proofs, test_case.expected_proven_by_model) + + if __name__ == "__main__": pytest.main([__file__, "-vv"])