diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/cte_facts.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/cte_facts.rs new file mode 100644 index 0000000000..dbb9d925c9 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/cte_facts.rs @@ -0,0 +1,1297 @@ +//! Python's CTE fact recovery and filtered non-null outputs, over the wheel's expression view. + +use std::cell::Cell; +use std::collections::HashSet; +use std::rc::Rc; + +use polyglot_sql::Expression; +use polyglot_sql::traversal::ExpressionWalk; +use serde_json::{Map, Value}; + +use crate::assembly::analysis_session::_helpers::dict_walk::{ + PyValue, alias_or_name, arg, children, column_table_name, direct_tables, expression_args, + find_all, is_star, kind, name, name_payload, output_name, parse_one, projected_expression, + py_truthy, python_int, raw_name_payload, set_operation, to_dict, top_level_ctes, + unwrap_annotations, where_clause, +}; +use crate::assembly::analysis_session::constants::{ + BINARY_OPERAND_COUNT, BOOLEAN_RESULT_AST_KINDS, BOOLEAN_TYPE, CASE_AST_KIND, CAST_AST_KIND, + CAST_AST_KINDS, COALESCE_AST_KIND, COLUMN_AST_KIND, CONCAT_AST_KIND, CONDITIONAL_RESULT_RULE, + COUNT_AST_KIND, CUSTOM_TYPE_NAME, DECIMAL_TYPE, FIRST_ARG_RULE, FUNCTION_AST_KIND, + IF_FUNC_AST_KIND, IS_NULL_AST_KIND, JOIN_FULL, JOIN_LEFT, JOIN_RIGHT, LITERAL_AST_KIND, + NON_NULL_NULLABILITY, NULL_AST_KIND, NULL_SET_OPERATION_TYPE, NULLABLE_NULLABILITY, + NULLIF_FUNCTION_NAME, POLYGLOT_TYPE_NAMES, SELECT_AST_KIND, SET_OPERATION_AST_KINDS, + STRING_LITERAL_TYPE, SUBSTRING_AST_KIND, TABLE_AST_KIND, TEXT_TYPE, TIMESTAMP_TYPE_NAME, + TIMESTAMP_TZ_TYPE, TRY_CAST_AST_KIND, TYPE_PASSTHROUGH_AST_KINDS, UNKNOWN_NULLABILITY, + VARCHAR_DATA_TYPES, WILDCARD, +}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::type_system::main::normalize_type::normalize_type; +use crate::type_system::models::{NormalizedType, TypeFamily}; + +/// A relation's column types; the `Rc` is the identity Python's `id()` dedup compares. +type Shape = Rc; +/// Relation facts in Python dict order. +type Relations = Vec<(String, Shape)>; +/// A relation's column nullabilities. +type NullShape = Rc>; +type NullRelations = Vec<(String, NullShape)>; +/// Python's `NonNullFilterContext`: the selected relations and the proven non-null columns. +type FilterContext = (Vec<(String, NullShape)>, HashSet<(String, String)>); +/// The result of one step, or the reason Python's result is not reproduced. +type Fact = Result; + +/// The inference profile facts the recovery reads. +pub(crate) struct RecoveryProfile<'a> { + pub(crate) dialect: &'a str, + pub(crate) function_return_types: &'a Pairs, + /// Adapter nullability rules by function name; None where a rule is not one Python ships. + pub(crate) rules: Option<&'a Pairs>, +} + +/// One enrichment's recovery request. +pub(crate) struct RecoveryInput<'a> { + pub(crate) cleaned_sql: &'a str, + pub(crate) input_schemas: &'a Shapes, + /// Whether Python runs `_polyglot_cte_passthrough_facts` past its early return. + pub(crate) recover: bool, + /// Whether a filter reads NULL, so Python looks for filtered non-null outputs. + pub(crate) null_filter: bool, + pub(crate) profile: RecoveryProfile<'a>, +} + +/// What Python's compact re-analysis takes from the parsed query. +#[derive(Debug, Default)] +pub(crate) struct Recovery { + pub(crate) types: Vec<(String, String)>, + pub(crate) nullability: Vec<(String, &'static str)>, + pub(crate) direct_outputs: HashSet, + pub(crate) non_null_outputs: HashSet, +} + +/// Python's CTE pass-through facts and filtered non-null outputs, or why they defer. +pub(crate) fn recovery(input: &RecoveryInput<'_>) -> Fact { + NON_ASCII_FOLDED.set(false); + let recovered: Fact = recovered(input); + if NON_ASCII_FOLDED.replace(false) { + return Err("non-ASCII text Python casefolds or upper-cases differently".to_owned()); + } + recovered +} + +thread_local! { + /// Whether this recovery folded non-ASCII text, where Python's `casefold`/`upper` may differ. + static NON_ASCII_FOLDED: Cell = const { Cell::new(false) }; +} + +fn recovered(input: &RecoveryInput<'_>) -> Fact { + let mut recovery: Recovery = Recovery::default(); + let mut parsed: Option> = None; + if input.recover { + let root: Option = parse_one(input.cleaned_sql, input.profile.dialect)?; + if let Some(root) = &root { + recovery = passthrough_facts(root, input, recovery)?; + } + parsed = Some(root); + } + if input.null_filter { + let root: Option = match parsed { + Some(root) => root, + None => parse_one(input.cleaned_sql, input.profile.dialect)?, + }; + if let Some(root) = &root { + recovery.non_null_outputs = filtered_non_null(root, &input_nullability(input))?; + } + } + Ok(recovery) +} + +/// Python's enrichment types: the input shapes. +fn input_types(input: &RecoveryInput<'_>) -> Relations { + input + .input_schemas + .iter() + .map(|(table, shape)| (table.clone(), Rc::new(shape.clone()))) + .collect() +} + +/// Python's enrichment nullability: every input column unknown. +fn input_nullability(input: &RecoveryInput<'_>) -> NullRelations { + input + .input_schemas + .iter() + .map(|(table, shape)| (table.clone(), Rc::new(unknown_columns(shape)))) + .collect() +} + +fn unknown_columns(shape: &Pairs) -> Vec<(String, &'static str)> { + shape + .iter() + .map(|(column, _)| (column.clone(), UNKNOWN_NULLABILITY)) + .collect() +} + +/// Python's `_polyglot_cte_passthrough_facts` after a successful parse. +fn passthrough_facts( + root: &Expression, + input: &RecoveryInput<'_>, + mut recovery: Recovery, +) -> Fact { + let ctes: Vec<(&str, bool, &Expression)> = top_level_ctes(root); + if ctes.is_empty() { + return Ok(recovery); + } + if input.profile.rules.is_none() { + return Err("an adapter nullability rule Python does not ship".to_owned()); + } + let context = Context { + profile: &input.profile, + }; + recovery.types = context.passthrough_types(root, &input_types(input))?; + recovery.nullability = context.passthrough_nullability(root, &input_nullability(input))?; + let names: HashSet = ctes.iter().map(|(name, _, _)| casefold(name)).collect(); + recovery.direct_outputs = direct_cte_outputs(root, &names); + Ok(recovery) +} + +/// Python's `_polyglot_direct_cte_output_names`. +fn direct_cte_outputs(root: &Expression, cte_names: &HashSet) -> HashSet { + let reads_cte: bool = direct_tables(root) + .iter() + .any(|table| cte_names.contains(&casefold(table.get_name()))); + if !reads_cte { + return HashSet::new(); + } + let mut outputs: HashSet = HashSet::new(); + for projection in root.get_expressions() { + let projection: Option<&Expression> = unwrap_annotations(Some(projection)); + if kind(projected_expression(projection)) != COLUMN_AST_KIND { + continue; + } + let output: &str = output_name(projection); + if !output.is_empty() && output != WILDCARD { + outputs.insert(output.to_owned()); + } + } + outputs +} + +/// The profile every resolver reads. +struct Context<'a, 'b> { + profile: &'a RecoveryProfile<'b>, +} + +impl Context<'_, '_> { + /// Python's `_polyglot_cte_passthrough_types_from_parsed`. + fn passthrough_types( + &self, + root: &Expression, + facts: &Relations, + ) -> Fact> { + let ctes: Vec<(&str, bool, &Expression)> = top_level_ctes(root); + if ctes.is_empty() { + return Ok(Vec::new()); + } + let mut relations: Relations = referenced_facts(root, facts); + for (cte_name, has_column_aliases, body) in ctes { + if has_column_aliases { + continue; + } + let inferred: Vec<(String, String)> = + self.select_output_types(Some(body), &relations)?; + if !inferred.is_empty() { + relations = dict_set(relations, cte_name, Rc::new(inferred)); + } + } + if SET_OPERATION_AST_KINDS.contains(&kind(Some(root))) { + return self.set_operation_types(root, &relations); + } + self.direct_select_types(Some(root), &relations) + } + + /// Python's `_polyglot_select_output_types`. + fn select_output_types( + &self, + select: Option<&Expression>, + relations: &Relations, + ) -> Fact> { + let Some(select) = unwrap_annotations(select) else { + return Err("Python reads CTEs of a missing expression".to_owned()); + }; + if !top_level_ctes(select).is_empty() { + return self.passthrough_types(select, relations); + } + if SET_OPERATION_AST_KINDS.contains(&kind(Some(select))) { + return self.set_operation_types(select, relations); + } + self.direct_select_types(Some(select), relations) + } + + /// Python's `_polyglot_set_operation_output_types`. + fn set_operation_types( + &self, + operation: &Expression, + relations: &Relations, + ) -> Fact> { + let Some((left, right, by_name)) = set_operation(operation) else { + return Ok(Vec::new()); + }; + let mut output: Vec<(String, String)> = Vec::new(); + if by_name { + let left_types: Vec<(String, String)> = + self.select_output_types(Some(left), relations)?; + let right_types: Vec<(String, String)> = + self.select_output_types(Some(right), relations)?; + for (name, left_type) in &left_types { + let right_type: Option<&str> = ci_get(&right_types, name).map(String::as_str); + if let Some(common) = self.common_type(Some(left_type), right_type)? { + output = dict_set(output, name, common); + } + } + return Ok(output); + } + let left_slots: Vec = self.type_slots(Some(left), relations)?; + if left_slots.is_empty() { + return Ok(Vec::new()); + } + let right_slots: Vec = self.type_slots(Some(right), relations)?; + if right_slots.is_empty() || left_slots.len() != right_slots.len() { + return Ok(Vec::new()); + } + for ((name, left_type), (_, right_type)) in left_slots.iter().zip(&right_slots) { + if let Some(common) = self.common_type(left_type.as_deref(), right_type.as_deref())? { + output = dict_set(output, name, common); + } + } + Ok(output) + } + + /// Python's `_polyglot_select_output_type_slots`. + fn type_slots(&self, select: Option<&Expression>, relations: &Relations) -> Fact> { + let select: Option<&Expression> = unwrap_annotations(select); + if let Some((left, right, by_name)) = select.and_then(set_operation) { + let left_slots: Vec = self.type_slots(Some(left), relations)?; + if left_slots.is_empty() { + return Ok(Vec::new()); + } + let right_slots: Vec = self.type_slots(Some(right), relations)?; + if right_slots.is_empty() { + return Ok(Vec::new()); + } + return self.combined_slots(&left_slots, &right_slots, by_name); + } + let Some(select) = select.filter(|select| kind(Some(select)) == SELECT_AST_KIND) else { + return Ok(Vec::new()); + }; + let inferred: Vec<(String, String)> = self.direct_select_types(Some(select), relations)?; + let mut slots: Vec = Vec::new(); + for projection in select.get_expressions() { + let projection: Option<&Expression> = unwrap_annotations(Some(projection)); + let output: &str = output_name(projection); + if is_star(projection) || output.is_empty() || output == WILDCARD { + return Ok(Vec::new()); + } + let mut inferred_type: Option = ci_get(&inferred, output).cloned(); + if inferred_type.is_none() && kind(projected_expression(projection)) == NULL_AST_KIND { + inferred_type = Some(NULL_SET_OPERATION_TYPE.to_owned()); + } + slots.push((output.to_owned(), inferred_type)); + } + Ok(slots) + } + + fn combined_slots(&self, left: &[Slot], right: &[Slot], by_name: bool) -> Fact> { + if by_name { + let mut right_types: Vec = Vec::new(); + for (name, slot_type) in right { + right_types = dict_set(right_types, name, slot_type.clone()); + } + let mut slots: Vec = Vec::with_capacity(left.len()); + for (name, left_type) in left { + let right_type: Option<&str> = ci_get_optional(&right_types, name); + slots.push(( + name.clone(), + self.common_type(left_type.as_deref(), right_type)?, + )); + } + return Ok(slots); + } + if left.len() != right.len() { + return Ok(Vec::new()); + } + let mut slots: Vec = Vec::with_capacity(left.len()); + for ((name, left_type), (_, right_type)) in left.iter().zip(right) { + slots.push(( + name.clone(), + self.common_type(left_type.as_deref(), right_type.as_deref())?, + )); + } + Ok(slots) + } + + /// Python's `_polyglot_common_set_operation_type`. + fn common_type(&self, left: Option<&str>, right: Option<&str>) -> Fact> { + let null: Option<&str> = Some(NULL_SET_OPERATION_TYPE); + if left == null && right == null { + return Ok(None); + } + if left == null { + return Ok(right.map(str::to_owned)); + } + if right == null { + return Ok(left.map(str::to_owned)); + } + let (Some(left), Some(right)) = (left, right) else { + return Ok(None); + }; + Ok(self.types_equal(left, right)?.then(|| left.to_owned())) + } + + /// Python's `_polyglot_types_equal`. + fn types_equal(&self, left: &str, right: &str) -> Fact { + Ok(left == right || self.normalized(left)? == self.normalized(right)?) + } + + /// Python's `normalize_type`, or a deferral where the native normalization defers. + fn normalized(&self, type_sql: &str) -> Fact { + normalize_type(type_sql, self.profile.dialect) + .map(|normalization| normalization.normalized) + .ok_or_else(|| format!("type normalization of {type_sql} defers")) + } + + /// Python's `_polyglot_direct_select_output_types`. + fn direct_select_types( + &self, + select: Option<&Expression>, + relations: &Relations, + ) -> Fact> { + let Some(select) = select.filter(|select| kind(Some(select)) == SELECT_AST_KIND) else { + return Ok(Vec::new()); + }; + let mut by_alias: Relations = Vec::new(); + for table in direct_tables(select) { + let table_name: &str = table.get_name(); + let Some(column_types) = ci_get(relations, table_name) else { + continue; + }; + let column_types: Shape = Rc::clone(column_types); + by_alias = dict_set(by_alias, table_name, Rc::clone(&column_types)); + let alias: &str = alias_or_name(table); + if !alias.is_empty() { + by_alias = dict_set(by_alias, alias, column_types); + } + } + let mut inferred: Vec<(String, String)> = Vec::new(); + for projection in select.get_expressions() { + let projection: Option<&Expression> = unwrap_annotations(Some(projection)); + if let Some(star) = projection.filter(|projection| is_star(Some(projection))) { + for (column, column_type) in star_types(star, &by_alias)? { + inferred = dict_set(inferred, &column, column_type); + } + continue; + } + let output: &str = output_name(projection); + if output.is_empty() || output == WILDCARD { + continue; + } + let expression: Option<&Expression> = projected_expression(projection); + if let Some(inferred_type) = self.expression_type_in(expression, &by_alias)? { + inferred = dict_set(inferred, output, inferred_type); + } + } + Ok(inferred) + } + + /// Python's `_polyglot_direct_expression_type`. + fn expression_type_in( + &self, + expression: Option<&Expression>, + by_alias: &Relations, + ) -> Fact> { + let expression: Option<&Expression> = unwrap_annotations(expression); + let node_kind: &str = kind(expression); + if let Some(inferred) = self.expression_type(expression)? { + return Ok(Some(inferred)); + } + let Some(node) = expression else { + return Ok(None); + }; + if node_kind == COLUMN_AST_KIND { + return direct_column_type(node, by_alias); + } + if BOOLEAN_RESULT_AST_KINDS.contains(&node_kind) { + return Ok(Some(BOOLEAN_TYPE.to_owned())); + } + if node_kind == CONCAT_AST_KIND { + return self.concat_type(node, by_alias); + } + if node_kind == SUBSTRING_AST_KIND { + return self.value_type(arg(Some(node), "this")?, by_alias); + } + if TYPE_PASSTHROUGH_AST_KINDS.contains(&node_kind) { + if let Some(inner) = node.get_this() { + return self.expression_type_in(Some(inner), by_alias); + } + return self.value_type(arg(Some(node), "this")?, by_alias); + } + let results: Vec> = result_expressions(node)?; + if results.is_empty() { + return Ok(None); + } + self.common_result_type(&results, by_alias) + } + + /// The type of a converted payload value; a non-expression has none. + fn value_type(&self, value: Option, by_alias: &Relations) -> Fact> { + match value { + Some(PyValue::Expr(inner)) => self.expression_type_in(Some(&inner), by_alias), + _ => Ok(None), + } + } + + fn concat_type(&self, node: &Expression, by_alias: &Relations) -> Fact> { + let operands: Vec = [arg(Some(node), "left")?, arg(Some(node), "right")?] + .into_iter() + .flatten() + .collect(); + if operands.len() != BINARY_OPERAND_COUNT { + return Ok(None); + } + let mut operand_types: Vec = Vec::with_capacity(BINARY_OPERAND_COUNT); + for operand in operands { + let PyValue::Expr(operand) = operand else { + return Ok(None); + }; + if string_literal(&operand)? { + operand_types.push(TEXT_TYPE.to_owned()); + continue; + } + let Some(operand_type) = self.expression_type_in(Some(&operand), by_alias)? else { + return Ok(None); + }; + operand_types.push(operand_type); + } + for operand_type in &operand_types { + if self.normalized(operand_type)?.family != TypeFamily::String { + return Ok(None); + } + } + Ok(Some(TEXT_TYPE.to_owned())) + } + + /// Python's `_polyglot_common_result_type`. + fn common_result_type( + &self, + expressions: &[Option], + by_alias: &Relations, + ) -> Fact> { + let mut resolved: Vec = Vec::new(); + for expression in expressions { + let Some(expression) = expression else { + return Ok(None); + }; + let expression: Option<&Expression> = unwrap_annotations(Some(expression)); + if kind(expression) == NULL_AST_KIND { + continue; + } + let Some(resolved_type) = self.expression_type_in(expression, by_alias)? else { + return Ok(None); + }; + resolved.push(resolved_type); + } + let Some((first, others)) = resolved.split_first() else { + return Ok(None); + }; + for other in others { + if !self.types_equal(first, other)? { + return Ok(None); + } + } + Ok(Some(first.clone())) + } + + /// Python's `_polyglot_expression_type`. + fn expression_type(&self, expression: Option<&Expression>) -> Fact> { + let node_kind: &str = kind(expression); + let node_name: &str = name(expression); + let function_name: &str = if node_name.is_empty() { + node_kind + } else { + node_name + }; + if node_kind != COLUMN_AST_KIND + && let Some(declared) = self.function_return_type(function_name)? + { + return Ok(Some(declared)); + } + if BOOLEAN_RESULT_AST_KINDS.contains(&node_kind) { + return Ok(Some(BOOLEAN_TYPE.to_owned())); + } + let Some(node) = expression.filter(|_| CAST_AST_KINDS.contains(&node_kind)) else { + return Ok(None); + }; + let value: Value = to_dict(node)?; + let target: Option<&Map> = value + .get(node_kind) + .and_then(Value::as_object) + .and_then(|payload| payload.get("to")) + .and_then(Value::as_object); + Ok(target.and_then(cast_target_type)) + } + + /// Python's `ExpressionInferenceProfile.function_return_type`. + fn function_return_type(&self, function_name: &str) -> Fact> { + let upper: String = upper(function_name); + Ok(self + .profile + .function_return_types + .iter() + .find(|(declared, _)| *declared == upper) + .map(|(_, return_type)| return_type.clone())) + } + + /// Python's `_polyglot_cte_passthrough_nullability_from_parsed`. + fn passthrough_nullability( + &self, + root: &Expression, + facts: &NullRelations, + ) -> Fact> { + let ctes: Vec<(&str, bool, &Expression)> = top_level_ctes(root); + if ctes.is_empty() { + return Ok(Vec::new()); + } + let mut relations: NullRelations = referenced_facts(root, facts); + for (cte_name, has_column_aliases, body) in ctes { + if has_column_aliases { + continue; + } + let inferred: Vec<(String, &'static str)> = + self.direct_select_nullability(body, &relations)?; + if !inferred.is_empty() { + relations = dict_set(relations, cte_name, Rc::new(inferred)); + } + } + self.direct_select_nullability(root, &relations) + } + + /// Python's `_polyglot_direct_select_output_nullability`. + fn direct_select_nullability( + &self, + select: &Expression, + relations: &NullRelations, + ) -> Fact> { + if kind(Some(select)) != SELECT_AST_KIND { + return Ok(Vec::new()); + } + let mut scoped: NullRelations = Vec::new(); + for table in direct_tables(select) { + let table_name: &str = table.get_name(); + if let Some(facts) = ci_get(relations, table_name) { + scoped = dict_set(scoped, table_name, Rc::new(facts.as_ref().clone())); + } + } + let known: bool = scoped.iter().any(|(_, facts)| has_known_nullability(facts)); + let (aliases, scoped): (Vec<(String, &'static str)>, NullRelations) = if known { + alias_nullability(select, scoped)? + } else { + (Vec::new(), scoped) + }; + let filter: Option = filter_context(select, &scoped)?; + let mut inferred: Vec<(String, &'static str)> = Vec::new(); + for projection in select.get_expressions() { + let projection: Option<&Expression> = unwrap_annotations(Some(projection)); + if is_star(projection) { + continue; + } + let output: &str = output_name(projection); + if output.is_empty() || output == WILDCARD { + continue; + } + let expression: Option<&Expression> = projected_expression(projection); + let nullability: &'static str = if non_null_after_filter(expression, filter.as_ref())? { + NON_NULL_NULLABILITY + } else if known { + self.nullability(expression, &aliases, &scoped)? + } else { + self.shallow_nullability(expression)? + }; + if nullability != UNKNOWN_NULLABILITY { + inferred = dict_set(inferred, output, nullability); + } + } + Ok(inferred) + } + + /// Python's `_infer_polyglot_nullability`. + fn nullability( + &self, + expression: Option<&Expression>, + aliases: &[(String, &'static str)], + scoped: &NullRelations, + ) -> Fact<&'static str> { + let node_kind: &str = kind(expression); + match node_kind { + NULL_AST_KIND => return Ok(NULLABLE_NULLABILITY), + LITERAL_AST_KIND | COUNT_AST_KIND | IS_NULL_AST_KIND => { + return Ok(NON_NULL_NULLABILITY); + } + TRY_CAST_AST_KIND => return Ok(UNKNOWN_NULLABILITY), + _ => {} + } + if let Some(column) = expression.filter(|_| node_kind == COLUMN_AST_KIND) { + return column_nullability(column, aliases, scoped); + } + if node_kind == CAST_AST_KIND { + return match expression.and_then(Expression::get_this) { + Some(inner) => self.nullability(Some(inner), aliases, scoped), + None => Ok(UNKNOWN_NULLABILITY), + }; + } + let rule: Option<&str> = self.rule(node_kind)?; + if node_kind != COALESCE_AST_KIND && rule.is_none() { + return Ok(UNKNOWN_NULLABILITY); + } + let mut arguments: Vec<&'static str> = Vec::new(); + for child in expression_args(expression) { + arguments.push(self.nullability(Some(child), aliases, scoped)?); + } + Ok(combined(node_kind, rule, &arguments)) + } + + /// Python's `_infer_polyglot_shallow_nullability`. + fn shallow_nullability(&self, expression: Option<&Expression>) -> Fact<&'static str> { + let node_kind: &str = kind(expression); + match node_kind { + NULL_AST_KIND => return Ok(NULLABLE_NULLABILITY), + LITERAL_AST_KIND | COUNT_AST_KIND | IS_NULL_AST_KIND => { + return Ok(NON_NULL_NULLABILITY); + } + COLUMN_AST_KIND | TRY_CAST_AST_KIND => return Ok(UNKNOWN_NULLABILITY), + _ => {} + } + if node_kind == CAST_AST_KIND { + return match expression.and_then(Expression::get_this) { + Some(inner) => self.shallow_nullability(Some(inner)), + None => Ok(UNKNOWN_NULLABILITY), + }; + } + let rule: Option<&str> = self.rule(node_kind)?; + if node_kind != COALESCE_AST_KIND && rule.is_none() { + return Ok(UNKNOWN_NULLABILITY); + } + let mut arguments: Vec<&'static str> = Vec::new(); + for child in expression_args(expression) { + arguments.push(self.shallow_nullability(Some(child))?); + } + Ok(combined(node_kind, rule, &arguments)) + } + + /// `inference_profile.function_nullability_rule(kind)`. + fn rule(&self, node_kind: &str) -> Fact> { + let rules: &Pairs = self + .profile + .rules + .ok_or("an adapter nullability rule Python does not ship")?; + let upper: String = upper(node_kind); + Ok(rules + .iter() + .find(|(declared, _)| *declared == upper) + .map(|(_, rule)| rule.as_str())) + } +} + +/// One set operation output slot: its name and type, None where unproven. +type Slot = (String, Option); + +/// COALESCE's or an adapter rule's nullability over the argument nullabilities. +fn combined(node_kind: &str, rule: Option<&str>, arguments: &[&'static str]) -> &'static str { + if node_kind == COALESCE_AST_KIND { + if arguments.contains(&NON_NULL_NULLABILITY) { + return NON_NULL_NULLABILITY; + } + if !arguments.is_empty() && arguments.iter().all(|value| *value == NULLABLE_NULLABILITY) { + return NULLABLE_NULLABILITY; + } + return UNKNOWN_NULLABILITY; + } + match rule { + Some(FIRST_ARG_RULE) => arguments.first().copied().unwrap_or(UNKNOWN_NULLABILITY), + Some(CONDITIONAL_RESULT_RULE) => conditional_result(arguments), + _ => UNKNOWN_NULLABILITY, + } +} + +/// Python's `conditional_result_nullability`. +fn conditional_result(arguments: &[&'static str]) -> &'static str { + let [_, when_true, when_false, ..] = arguments else { + return UNKNOWN_NULLABILITY; + }; + if *when_true == NON_NULL_NULLABILITY && *when_false == NON_NULL_NULLABILITY { + return NON_NULL_NULLABILITY; + } + if *when_true == NULLABLE_NULLABILITY || *when_false == NULLABLE_NULLABILITY { + return NULLABLE_NULLABILITY; + } + UNKNOWN_NULLABILITY +} + +/// Python's `_infer_polyglot_column_nullability`. +fn column_nullability( + column: &Expression, + aliases: &[(String, &'static str)], + scoped: &NullRelations, +) -> Fact<&'static str> { + let column_name: &str = column.get_name(); + if column_name.is_empty() { + return Ok(UNKNOWN_NULLABILITY); + } + let table_name: String = column_table_name(column)?; + if !table_name.is_empty() { + if dict_get(aliases, &table_name) == Some(&NULLABLE_NULLABILITY) { + return Ok(NULLABLE_NULLABILITY); + } + return Ok(dict_get(scoped, &table_name) + .and_then(|facts| dict_get(facts, column_name)) + .copied() + .unwrap_or(UNKNOWN_NULLABILITY)); + } + let matches: Vec<&'static str> = scoped + .iter() + .filter_map(|(_, facts)| dict_get(facts, column_name).copied()) + .collect(); + Ok(match matches.as_slice() { + [only] => only, + _ => UNKNOWN_NULLABILITY, + }) +} + +/// Python's `_polyglot_alias_nullability_from_select`, copying facts to aliases in `scoped`. +fn alias_nullability( + select: &Expression, + mut scoped: NullRelations, +) -> Fact<(Vec<(String, &'static str)>, NullRelations)> { + let mut aliases: Vec<(String, &'static str)> = Vec::new(); + let mut current: Vec = Vec::new(); + let value: Value = to_dict(select)?; + let Some(payload) = value.get("select").and_then(Value::as_object) else { + return Ok((aliases, scoped)); + }; + let from: Option<&Vec> = payload + .get("from") + .and_then(Value::as_object) + .and_then(|from| from.get("expressions")) + .and_then(Value::as_array); + if let Some([relation]) = from.map(Vec::as_slice) + && let Some(table) = relation + .get(TABLE_AST_KIND) + .filter(|table| table.is_object()) + { + let (alias, table_name) = table_alias_and_name(table); + current.push(alias.clone()); + aliases = dict_set(aliases, &alias, UNKNOWN_NULLABILITY); + scoped = copy_facts_to_alias(&alias, &table_name, scoped); + } + let Some(joins) = payload.get("joins").and_then(Value::as_array) else { + return Ok((aliases, scoped)); + }; + for join in joins.iter().filter_map(Value::as_object) { + let Some(table) = join + .get("this") + .filter(|this| this.is_object()) + .and_then(|this| this.get(TABLE_AST_KIND)) + .filter(|table| table.is_object()) + else { + continue; + }; + let (joined, joined_name) = table_alias_and_name(table); + let side: String = match join.get("kind") { + None | Some(Value::Null) => String::new(), + Some(Value::String(side)) => upper(side), + Some(_) => return Err("a join side Python stringifies".to_owned()), + }; + let current_side: Option<&'static str> = match side.as_str() { + JOIN_RIGHT | JOIN_FULL => Some(NULLABLE_NULLABILITY), + _ => None, + }; + if let Some(nullability) = current_side { + for alias in ¤t { + aliases = dict_set(aliases, alias, nullability); + } + } + let joined_nullability: &'static str = match side.as_str() { + JOIN_LEFT | JOIN_FULL => NULLABLE_NULLABILITY, + _ => UNKNOWN_NULLABILITY, + }; + aliases = dict_set(aliases, &joined, joined_nullability); + current.push(joined.clone()); + scoped = copy_facts_to_alias(&joined, &joined_name, scoped); + } + Ok((aliases, scoped)) +} + +/// Python's `_polyglot_table_payload_alias_and_name`. +fn table_alias_and_name(table: &Value) -> (String, String) { + let table_name: String = raw_name_payload(table.get("name")); + let alias: String = raw_name_payload(table.get("alias")); + ( + if alias.is_empty() { + table_name.clone() + } else { + alias + }, + table_name, + ) +} + +/// Python's `_copy_table_facts_to_alias`. +fn copy_facts_to_alias(alias: &str, table_name: &str, mut scoped: NullRelations) -> NullRelations { + if alias == table_name || dict_get(&scoped, alias).is_some() { + return scoped; + } + if let Some(facts) = dict_get(&scoped, table_name).cloned() { + scoped.push((alias.to_owned(), facts)); + } + scoped +} + +fn has_known_nullability(facts: &[(String, &'static str)]) -> bool { + facts.iter().any(|(_, value)| *value != UNKNOWN_NULLABILITY) +} + +/// Python's `_polyglot_non_null_filter_context`. +fn filter_context(select: &Expression, nullability: &NullRelations) -> Fact> { + let Some(clause) = where_clause(select)? else { + return Ok(None); + }; + let Some(clause) = clause.as_object() else { + return Ok(None); + }; + let references: Vec<(String, String)> = non_null_conjuncts(clause.get("this")); + if references.is_empty() { + return Ok(None); + } + let mut relations: Vec<(String, NullShape)> = Vec::new(); + for table in direct_tables(select) { + let table_name: &str = table.get_name(); + let alias: &str = alias_or_name(table); + if let Some(columns) = ci_get(nullability, table_name) { + let key: &str = if alias.is_empty() { table_name } else { alias }; + relations.push((key.to_owned(), Rc::clone(columns))); + } + } + let columns: HashSet<(String, String)> = references + .iter() + .filter_map(|reference| resolved_column(reference, &relations)) + .collect(); + Ok((!columns.is_empty()).then_some((relations, columns))) +} + +/// Python's `_polyglot_non_null_conjuncts`. +fn non_null_conjuncts(payload: Option<&Value>) -> Vec<(String, String)> { + let Some(payload) = payload.and_then(Value::as_object) else { + return Vec::new(); + }; + if let Some(paren) = payload.get("paren").and_then(Value::as_object) { + return non_null_conjuncts(paren.get("this")); + } + if let Some(and) = payload.get("and").and_then(Value::as_object) { + let mut references: Vec<(String, String)> = non_null_conjuncts(and.get("left")); + references.extend(non_null_conjuncts(and.get("right"))); + return references; + } + let Some(is_null) = payload.get(IS_NULL_AST_KIND).and_then(Value::as_object) else { + return Vec::new(); + }; + if is_null.is_empty() || is_null.get("not") != Some(&Value::Bool(true)) { + return Vec::new(); + } + let Some(column) = is_null + .get("this") + .and_then(Value::as_object) + .and_then(|this| this.get(COLUMN_AST_KIND)) + .and_then(Value::as_object) + else { + return Vec::new(); + }; + let column_name: String = raw_name_payload(column.get("name")); + let table_name: String = raw_name_payload(column.get("table")); + if column_name.is_empty() { + return Vec::new(); + } + vec![(table_name, column_name)] +} + +/// Python's `_polyglot_resolved_relation_column`. +fn resolved_column( + reference: &(String, String), + relations: &[(String, NullShape)], +) -> Option<(String, String)> { + let (table_name, column_name) = reference; + let matching: Vec<(String, String)> = relations + .iter() + .filter(|(alias, _)| table_name.is_empty() || casefold(alias) == casefold(table_name)) + .filter_map(|(alias, columns)| folded_column(alias, columns, column_name)) + .collect(); + match <[(String, String); 1]>::try_from(matching) { + Ok([only]) => Some(only), + Err(_) => None, + } +} + +fn folded_column(alias: &str, columns: &NullShape, column_name: &str) -> Option<(String, String)> { + let key: &str = ci_key(columns, column_name)?; + Some((casefold(alias), casefold(key))) +} + +/// Python's `_polyglot_expression_is_non_null_after_filter`. +fn non_null_after_filter( + expression: Option<&Expression>, + context: Option<&FilterContext>, +) -> Fact { + let Some((relations, columns)) = context else { + return Ok(false); + }; + let Some(reference) = direct_column_reference(expression)? else { + return Ok(false); + }; + Ok(resolved_column(&reference, relations).is_some_and(|resolved| columns.contains(&resolved))) +} + +/// Python's `_polyglot_direct_column_reference`. +fn direct_column_reference(expression: Option<&Expression>) -> Fact> { + let mut expression: Option<&Expression> = unwrap_annotations(expression); + if kind(expression) == CAST_AST_KIND { + expression = unwrap_annotations(expression.and_then(Expression::get_this)); + } + if kind(expression) != COLUMN_AST_KIND { + return Ok(None); + } + let column_name: &str = name(expression); + if column_name.is_empty() { + return Ok(None); + } + let table_name: String = name_payload(arg(expression, "table")?.as_ref()); + Ok(Some((table_name, column_name.to_owned()))) +} + +/// Python's `_polyglot_filtered_non_null_outputs` once its filter test passed. +fn filtered_non_null(root: &Expression, nullability: &NullRelations) -> Fact> { + let select: Option<&Expression> = if kind(Some(root)) == SELECT_AST_KIND { + Some(root) + } else { + root.dfs() + .skip(1) + .find(|node| node.variant_name() == SELECT_AST_KIND) + }; + let Some(select) = select else { + return Ok(HashSet::new()); + }; + let context: Option = filter_context(select, nullability)?; + let mut outputs: HashSet = HashSet::new(); + for projection in select.get_expressions() { + let projection: Option<&Expression> = unwrap_annotations(Some(projection)); + let output: &str = output_name(projection); + if output.is_empty() || output == WILDCARD { + continue; + } + if non_null_after_filter(projected_expression(projection), context.as_ref())? { + outputs.insert(output.to_owned()); + } + } + Ok(outputs) +} + +/// Python's `_polyglot_referenced_relation_facts`. +fn referenced_facts(root: &Expression, facts: &[(String, V)]) -> Vec<(String, V)> { + let mut referenced: Vec<(String, V)> = Vec::new(); + for table in find_all(root, TABLE_AST_KIND) { + let table_name: &str = table.get_name(); + if let Some(column_facts) = dict_get(facts, table_name) { + referenced = dict_set(referenced, table_name, column_facts.clone()); + } + } + referenced +} + +/// Python's `_polyglot_direct_column_type`. +fn direct_column_type(column: &Expression, by_alias: &Relations) -> Fact> { + let column_name: &str = column.get_name(); + let table_name: String = column_table_name(column)?; + if !table_name.is_empty() { + return Ok(ci_get(by_alias, &table_name) + .and_then(|column_types| ci_get(column_types, column_name)) + .cloned()); + } + let matching: Vec<&String> = distinct_shapes(by_alias) + .into_iter() + .filter_map(|column_types| ci_get(column_types, column_name)) + .collect(); + Ok(match matching.as_slice() { + [only] => Some((*only).clone()), + _ => None, + }) +} + +/// The relations of `by_alias` in order, each `Rc` once, as Python's `id()` dedup keeps them. +fn distinct_shapes(by_alias: &Relations) -> Vec<&Shape> { + let mut distinct: Vec<&Shape> = Vec::new(); + for (_, shape) in by_alias { + if !distinct.iter().any(|seen| Rc::ptr_eq(seen, shape)) { + distinct.push(shape); + } + } + distinct +} + +/// Python's `_polyglot_star_output_types`. +fn star_types(star: &Expression, by_alias: &Relations) -> Fact> { + if py_truthy(arg(Some(star), "rename")?.as_ref()) { + return Ok(Vec::new()); + } + let table_name: String = name_payload(arg(Some(star), "table")?.as_ref()); + let sources: Vec<&Shape> = if table_name.is_empty() { + distinct_shapes(by_alias) + } else { + ci_get(by_alias, &table_name).into_iter().collect() + }; + let excluded: HashSet = listed_names(arg(Some(star), "except")?, false)?; + let replaced: HashSet = listed_names(arg(Some(star), "replace")?, true)?; + let mut inferred: Vec<(String, String)> = Vec::new(); + let mut ambiguous: Vec = Vec::new(); + for column_types in sources { + for (column_name, column_type) in column_types.iter() { + let folded: String = casefold(column_name); + if excluded.contains(&folded) || replaced.contains(&folded) { + continue; + } + if let Some((existing, _)) = inferred.iter().find(|(name, _)| casefold(name) == folded) + { + ambiguous.push(existing.clone()); + continue; + } + inferred.push((column_name.clone(), column_type.clone())); + } + } + inferred.retain(|(name, _)| !ambiguous.contains(name)); + Ok(inferred) +} + +/// The folded names of a star's `except` list, or of its `replace` aliases. +fn listed_names(value: Option, aliases: bool) -> Fact> { + let items: Vec = match value { + None => Vec::new(), + Some(PyValue::Exprs(_)) => return Ok(HashSet::new()), + Some(PyValue::Raw(Value::Array(items))) => items, + Some(PyValue::Raw(raw)) if !py_truthy(Some(&PyValue::Raw(raw.clone()))) => Vec::new(), + Some(_) => return Err("a star modifier Python iterates unusually".to_owned()), + }; + Ok(items + .iter() + .filter(|item| !aliases || item.is_object()) + .map(|item| match aliases { + true => raw_name_payload(item.get("alias")), + false => raw_name_payload(Some(item)), + }) + .filter(|name| !name.is_empty()) + .map(|name| casefold(&name)) + .collect()) +} + +/// Whether `node` is a literal whose `literal_type` arg is Python's `"string"`. +fn string_literal(node: &Expression) -> Fact { + if node.variant_name() != LITERAL_AST_KIND { + return Ok(false); + } + Ok(matches!( + arg(Some(node), "literal_type")?, + Some(PyValue::Raw(Value::String(literal_type))) if literal_type == STRING_LITERAL_TYPE + )) +} + +/// Python's `_polyglot_result_expressions`; None stands for a value that is not an expression. +fn result_expressions(node: &Expression) -> Fact>> { + let node_kind: &str = node.variant_name(); + if node_kind == COALESCE_AST_KIND { + return Ok(match arg(Some(node), "expressions")? { + Some(PyValue::Exprs(values)) => values.into_iter().map(Some).collect(), + Some(PyValue::Raw(Value::Array(values))) => values.iter().map(|_| None).collect(), + _ => Vec::new(), + }); + } + if node_kind == IF_FUNC_AST_KIND { + let mut values: Vec> = Vec::new(); + for key in ["true_value", "false_value"] { + if let Some(value) = arg(Some(node), key)? { + values.push(expression_value(value)); + } + } + return Ok(values); + } + if node_kind == CASE_AST_KIND { + return case_results(node); + } + if node_kind == FUNCTION_AST_KIND && upper(node.get_name()) == NULLIF_FUNCTION_NAME { + return Ok(match arg(Some(node), "args")? { + Some(PyValue::Exprs(mut values)) if !values.is_empty() => { + vec![Some(values.swap_remove(0))] + } + Some(PyValue::Raw(Value::Array(values))) if !values.is_empty() => vec![None], + _ => Vec::new(), + }); + } + + Ok(Vec::new()) +} + +fn expression_value(value: PyValue) -> Option { + match value { + PyValue::Expr(expression) => Some(expression), + _ => None, + } +} + +/// CASE's result branches by child position, as Python indexes `children()`. +fn case_results(node: &Expression) -> Fact>> { + let node_children: Vec<&Expression> = children(node); + let mut values: Vec> = Vec::new(); + let when_count: Option = match arg(Some(node), "whens")? { + Some(PyValue::Exprs(whens)) => Some(whens.len()), + Some(PyValue::Raw(Value::Array(whens))) => Some(whens.len()), + _ => None, + }; + if let Some(when_count) = when_count { + let first: usize = usize::from(arg(Some(node), "operand")?.is_some()); + for index in 0..when_count { + let child: &Expression = node_children + .get(first + index * 2 + 1) + .ok_or("Python indexes past CASE's children")?; + values.push(Some(child.clone())); + } + } + if arg(Some(node), "else_")?.is_some() + && let Some(last) = node_children.last() + { + values.push(Some((*last).clone())); + } + Ok(values) +} + +/// Python's cast-target rendering in `_polyglot_expression_type`. +fn cast_target_type(target: &Map) -> Option { + let raw_type: &str = target + .get("data_type") + .and_then(Value::as_str) + .filter(|raw| !raw.is_empty())?; + if upper(raw_type) == CUSTOM_TYPE_NAME + && let Some(custom) = target + .get("name") + .and_then(Value::as_str) + .filter(|custom| !custom.is_empty()) + { + return Some(upper(custom)); + } + let lower: String = casefold(raw_type); + if lower == TIMESTAMP_TYPE_NAME && target.get("timezone") == Some(&Value::Bool(true)) { + return Some(TIMESTAMP_TZ_TYPE.to_owned()); + } + let type_name: String = POLYGLOT_TYPE_NAMES + .iter() + .find(|(raw, _)| *raw == raw_type) + .map_or_else( + || upper(&raw_type.replace('_', " ")), + |(_, rendered)| (*rendered).to_owned(), + ); + if VARCHAR_DATA_TYPES.contains(&lower.as_str()) + && let Some(length) = python_int(target.get("length")) + { + return Some(format!("VARCHAR({length})")); + } + if type_name == DECIMAL_TYPE + && let Some(precision) = python_int(target.get("precision")) + { + return Some(match python_int(target.get("scale")) { + Some(scale) => format!("DECIMAL({precision}, {scale})"), + None => format!("DECIMAL({precision})"), + }); + } + Some(type_name) +} + +/// Python's `str.casefold`, exact for ASCII; non-ASCII text marks the recovery deferred. +fn casefold(text: &str) -> String { + flag_non_ascii(text); + text.to_ascii_lowercase() +} + +/// Python's `str.upper`, exact for ASCII; non-ASCII text marks the recovery deferred. +fn upper(text: &str) -> String { + flag_non_ascii(text); + text.to_ascii_uppercase() +} + +fn flag_non_ascii(text: &str) { + if !text.is_ascii() { + NON_ASCII_FOLDED.set(true); + } +} + +fn dict_get<'a, V>(dict: &'a [(String, V)], key: &str) -> Option<&'a V> { + dict.iter() + .find(|(candidate, _)| candidate == key) + .map(|(_, value)| value) +} + +/// Python's `dict[key] = value`: replace in place, otherwise append. +fn dict_set(mut dict: Vec<(String, V)>, key: &str, value: V) -> Vec<(String, V)> { + match dict.iter_mut().find(|(candidate, _)| candidate == key) { + Some((_, existing)) => *existing = value, + None => dict.push((key.to_owned(), value)), + } + dict +} + +/// Python's `_case_insensitive_mapping_get` over values that are never None. +fn ci_get<'a, V>(dict: &'a [(String, V)], key: &str) -> Option<&'a V> { + if let Some(direct) = dict_get(dict, key) { + return Some(direct); + } + unique_folded(dict, key) +} + +/// Python's `_case_insensitive_mapping_get` over values that may be None. +fn ci_get_optional<'a>(dict: &'a [(String, Option)], key: &str) -> Option<&'a str> { + if let Some(Some(direct)) = dict_get(dict, key) { + return Some(direct); + } + unique_folded(dict, key).and_then(Option::as_deref) +} + +fn unique_folded<'a, V>(dict: &'a [(String, V)], key: &str) -> Option<&'a V> { + let folded: String = casefold(key); + let matches: Vec<&V> = dict + .iter() + .filter(|(candidate, _)| casefold(candidate) == folded) + .map(|(_, value)| value) + .collect(); + match matches.as_slice() { + [only] => Some(only), + _ => None, + } +} + +/// Python's `_case_insensitive_mapping_key`. +fn ci_key<'a, V>(dict: &'a [(String, V)], key: &str) -> Option<&'a str> { + if let Some((candidate, _)) = dict.iter().find(|(candidate, _)| candidate == key) { + return Some(candidate); + } + let folded: String = casefold(key); + let matches: Vec<&str> = dict + .iter() + .filter(|(candidate, _)| casefold(candidate) == folded) + .map(|(candidate, _)| candidate.as_str()) + .collect(); + match matches.as_slice() { + [only] => Some(only), + _ => None, + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs index d11f30c0ae..96b1a60d02 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs @@ -1,10 +1,14 @@ -//! Python's dict walk over a parsed query's `to_dict()` payload, on the same serde values. +//! Python's view of a parsed query: the wheel's expression accessors and its dict walk. -use serde_json::{Map, Value}; +use polyglot_sql::expressions::Cte; +use polyglot_sql::traversal::ExpressionWalk; +use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions, ast_json}; +use serde_json::{Map, Value, json}; use crate::assembly::analysis_session::constants::{ - COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES, FUNCTION_AST_KIND, RENDERED_TYPE_NAMES, - STAR_AST_KIND, TABLE_AST_KIND, + ALIAS_AST_KIND, ANNOTATED_AST_KIND, COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES, + FUNCTION_AST_KIND, MAX_FUNCTION_CALL_DEPTH, RENDERED_TYPE_NAMES, STAR_AST_KIND, TABLE_AST_KIND, + WILDCARD, }; /// Python's `bool(value)`. @@ -142,7 +146,7 @@ pub(crate) fn render_type(node: Option<&Value>) -> Option { } /// Python's `str(value)` for a value `isinstance(value, int)` accepts. -fn python_int(value: Option<&Value>) -> Option { +pub(crate) 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()), @@ -154,3 +158,236 @@ fn python_int(value: Option<&Value>) -> Option { pub(crate) fn casefold(text: &str) -> String { text.to_ascii_lowercase() } + +/// A value the wheel hands Python for one payload key: an expression, a list of them, or data. +#[derive(Debug, Clone)] +pub(crate) enum PyValue { + Expr(Expression), + Exprs(Vec), + Raw(Value), +} + +/// The wheel's guarded `parse_one`; Ok(None) is the `PolyglotError` Python catches. +pub(crate) fn parse_one(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 mut expressions: Vec = match parsed { + Ok(expressions) => expressions, + Err(_) => return Ok(None), + }; + if expressions.len() != 1 { + return Ok(None); + } + Ok(expressions.pop()) +} + +/// `str(getattr(node, "kind", ""))`. +pub(crate) fn kind(node: Option<&Expression>) -> &'static str { + node.map_or("", Expression::variant_name) +} + +/// `node.name`, "" for None. +pub(crate) fn name(node: Option<&Expression>) -> &str { + node.map_or("", Expression::get_name) +} + +/// `node.output_name`, "" for None. +pub(crate) fn output_name(node: Option<&Expression>) -> &str { + node.map_or("", Expression::get_output_name) +} + +/// `node.alias_or_name`. +pub(crate) fn alias_or_name(node: &Expression) -> &str { + let alias: &str = node.get_alias(); + if alias.is_empty() { + node.get_name() + } else { + alias + } +} + +/// `node.is_star`. +pub(crate) fn is_star(node: Option<&Expression>) -> bool { + match node { + Some(Expression::Star(_)) => true, + Some(Expression::Column(column)) => column.name.name == WILDCARD, + _ => false, + } +} + +/// Python's `_unwrap_polyglot_annotations`. +pub(crate) fn unwrap_annotations(mut node: Option<&Expression>) -> Option<&Expression> { + while kind(node) == ANNOTATED_AST_KIND { + node = node.and_then(Expression::get_this); + } + node +} + +/// A projection's expression: an alias's `this`, otherwise the projection. +pub(crate) fn projected_expression(projection: Option<&Expression>) -> Option<&Expression> { + if kind(projection) == ALIAS_AST_KIND { + return projection.and_then(Expression::get_this); + } + projection +} + +/// `node.this` then `node.expressions`: Python's `_polyglot_expression_args`. +pub(crate) fn expression_args(node: Option<&Expression>) -> Vec<&Expression> { + let Some(node) = node else { + return Vec::new(); + }; + node.get_this() + .into_iter() + .chain(node.get_expressions()) + .collect() +} + +/// `node.children()`. +pub(crate) fn children(node: &Expression) -> Vec<&Expression> { + ExpressionWalk::children(node) +} + +/// Python's `_polyglot_direct_select_tables`: the table children. +pub(crate) fn direct_tables(node: &Expression) -> Vec<&Expression> { + children(node) + .into_iter() + .filter(|child| child.variant_name() == TABLE_AST_KIND) + .collect() +} + +/// `node.find_all(kind)`: descendants in depth-first order, the node itself excluded. +pub(crate) fn find_all<'a>(node: &'a Expression, wanted: &str) -> Vec<&'a Expression> { + node.dfs() + .skip(1) + .filter(|descendant| descendant.variant_name() == wanted) + .collect() +} + +/// `node.to_dict()`. +pub(crate) fn to_dict(node: &Expression) -> Result { + serde_json::to_value(node).map_err(|error| error.to_string()) +} + +/// The wheel's `expression_payload`: the tagged value's fields. +pub(crate) fn expression_payload(node: &Expression) -> Result, String> { + Ok(match to_dict(node)? { + Value::Object(map) => match map.into_iter().next() { + Some((_, Value::Object(payload))) => payload, + _ => Map::new(), + }, + _ => Map::new(), + }) +} + +/// `node.args.get(key)` (and `node.arg(key)` outside its special keys); None for Python's None. +pub(crate) fn arg(node: Option<&Expression>, key: &str) -> Result, String> { + let Some(node) = node else { + return Ok(None); + }; + Ok(expression_payload(node)?.remove(key).and_then(py_value)) +} + +/// The wheel's `value_to_python_object`; None where Python sees None. +fn py_value(value: Value) -> Option { + if let Ok(expression) = ast_json::expression_from_value(value.clone()) { + return Some(PyValue::Expr(expression)); + } + if let Ok(expressions) = ast_json::expressions_from_value(value.clone()) { + return Some(PyValue::Exprs(expressions)); + } + (!value.is_null()).then_some(PyValue::Raw(value)) +} + +/// Python's `bool()` of a converted value: expressions are always truthy. +pub(crate) fn py_truthy(value: Option<&PyValue>) -> bool { + match value { + None => false, + Some(PyValue::Expr(_)) => true, + Some(PyValue::Exprs(items)) => !items.is_empty(), + Some(PyValue::Raw(raw)) => truthy(Some(raw)), + } +} + +/// Python's `_polyglot_name_payload_value` of a converted value. +pub(crate) fn name_payload(value: Option<&PyValue>) -> String { + match value { + Some(PyValue::Raw(raw)) => raw_name_payload(Some(raw)), + _ => String::new(), + } +} + +/// Python's `_polyglot_name_payload_value` of plain data. +pub(crate) fn raw_name_payload(value: Option<&Value>) -> String { + match value { + Some(Value::String(text)) => text.clone(), + Some(Value::Object(object)) => object + .get("name") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + _ => String::new(), + } +} + +/// Python's `_polyglot_column_table_name`: `to_dict()["column"]["table"]["name"]`. +pub(crate) fn column_table_name(node: &Expression) -> Result { + let value: Value = to_dict(node)?; + let table: Option<&Value> = value + .get("column") + .filter(|column| column.is_object()) + .and_then(|column| column.get("table")) + .filter(|table| table.is_object()); + Ok(raw_name_payload(table)) +} + +/// The top-level CTEs `with_ctes()` returns: name, whether it has column aliases, and its body. +pub(crate) fn top_level_ctes(node: &Expression) -> Vec<(&str, bool, &Expression)> { + let with = match node { + Expression::Select(select) => select.with.as_ref(), + Expression::Union(union) => union.with.as_ref(), + Expression::Intersect(intersect) => intersect.with.as_ref(), + Expression::Except(except) => except.with.as_ref(), + _ => None, + }; + with.map(|with| with.ctes.iter().map(cte_entry).collect()) + .unwrap_or_default() +} + +fn cte_entry(cte: &Cte) -> (&str, bool, &Expression) { + (cte.alias.name.as_str(), !cte.columns.is_empty(), &cte.this) +} + +/// A set operation's `arg("left")`, `arg("right")` and `arg("by_name")`. +pub(crate) fn set_operation(node: &Expression) -> Option<(&Expression, &Expression, bool)> { + match node { + Expression::Union(operation) => { + Some((&operation.left, &operation.right, operation.by_name)) + } + Expression::Intersect(operation) => { + Some((&operation.left, &operation.right, operation.by_name)) + } + Expression::Except(operation) => { + Some((&operation.left, &operation.right, operation.by_name)) + } + _ => None, + } +} + +/// A select's `arg("where_clause")` as plain data. +pub(crate) fn where_clause(node: &Expression) -> Result, String> { + let Expression::Select(select) = node else { + return Ok(None); + }; + select + .where_clause + .as_ref() + .map(|clause| serde_json::to_value(clause).map_err(|error| error.to_string())) + .transpose() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs index 78f0b3bdba..40386cc58b 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs @@ -2,17 +2,16 @@ use std::collections::{HashMap, HashSet}; -use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions}; -use serde_json::{Map, Value, json}; +use serde_json::{Map, Value}; 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, + parse_one, payload, relation_name, render_type, to_dict, 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, + MOTHERDUCK_DIALECT, PIVOT_AST_KIND, QUERY_AST_KINDS, SELECT_AST_KIND, SIMPLIFIED_PIVOT_DIALECTS, SUPPORTED_PIVOT_DIALECTS, TYPE_PRESERVING_AGGREGATES, UNKNOWN_NULLABILITY, UNKNOWN_TYPE, }; @@ -80,25 +79,9 @@ pub(crate) fn pivot_outcome( /// 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()) + parse_one(sql, dialect)? + .map(|expression| to_dict(&expression)) + .transpose() } /// Casefolding equals Python's only for ASCII text. diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs index 2e176ff81e..4a2da05dde 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs @@ -6,11 +6,14 @@ 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::cte_facts::{ + Recovery, RecoveryInput, RecoveryProfile, recovery, +}; 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, + MAX_FUNCTION_CALL_DEPTH, NON_NULL_NULLABILITY, 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, }; @@ -25,6 +28,8 @@ pub(crate) struct EnrichmentInput<'a> { pub(crate) input_schemas: &'a Shapes, pub(crate) dialect: &'a str, pub(crate) function_return_types: &'a Pairs, + /// Adapter nullability rules by function name; None where one is not a rule Python ships. + pub(crate) nullability_rules: Option<&'a Pairs>, } /// Python's successful re-analysis: columns, plain lineage facts and the star flag. @@ -75,7 +80,7 @@ fn analysed(input: &EnrichmentInput<'_>) -> Result, String> { Err(_) => return Ok(None), }; let analysis: Value = serde_json::to_value(analysis).map_err(|error| error.to_string())?; - Ok(projected(input, &analysis)) + Ok(projected(input, &analysis, &cleaned)) } /// Python's `_compact_analysis_schema` over the input shapes, nullability unknown. @@ -109,15 +114,26 @@ fn analysis_schema(input: &EnrichmentInput<'_>) -> Option { } /// Python's projection of `analyze_query` facts into columns and lineage. -fn projected(input: &EnrichmentInput<'_>, analysis: &Value) -> Option { +fn projected(input: &EnrichmentInput<'_>, analysis: &Value, cleaned: &str) -> 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 recovery: Recovery = match recovery(&RecoveryInput { + cleaned_sql: cleaned, + input_schemas: input.input_schemas, + recover: recovers_cte_facts(input, analysis), + null_filter: has_null_filter(analysis), + profile: RecoveryProfile { + dialect: input.dialect, + function_return_types: input.function_return_types, + rules: input.nullability_rules, + }, + }) { + Ok(recovery) => recovery, + Err(_deferred) => return None, + }; let resources: HashMap<&str, (&str, &str)> = input .model .lineage_references @@ -136,10 +152,24 @@ fn projected(input: &EnrichmentInput<'_>, analysis: &Value) -> Option = if direct { + recovered(&recovery.types, name).cloned() + } else { + projection_type(projection, input.function_return_types)? + }; + let nullability: &str = if recovery.non_null_outputs.contains(name) { + NON_NULL_NULLABILITY + } else { + recovered(&recovery.nullability, name) + .filter(|_| direct) + .copied() + .unwrap_or_else(|| projection_nullability(projection, infer_nullability)) + }; columns.push(ColumnFact { name: name.to_owned(), - data_type: projection_type(projection, input.function_return_types)?, - nullability: projection_nullability(projection, infer_nullability).to_owned(), + data_type, + nullability: nullability.to_owned(), }); let (sources, confidence) = upstream(projection, &resources); let transform: u8 = transform_code(projection, !sources.is_empty()); @@ -176,7 +206,15 @@ fn projected(input: &EnrichmentInput<'_>, analysis: &Value) -> Option(facts: &'a [(String, V)], name: &str) -> Option<&'a V> { + facts + .iter() + .find(|(output, _)| output == name) + .map(|(_, value)| value) +} + +/// Whether Python's CTE pass-through recovery runs past its early return. fn recovers_cte_facts(input: &EnrichmentInput<'_>, analysis: &Value) -> bool { input.model.recover_cte_facts && truthy(analysis.get("cteFacts")) 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 1ff9c1dda6..3534a64196 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,6 @@ pub(crate) mod catalog_state; pub(crate) mod compact_batch; +pub(crate) mod cte_facts; pub(crate) mod dict_walk; pub(crate) mod dynamic_pivot; pub(crate) mod enrichment; 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 cf945f7c71..43418c7881 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs @@ -352,6 +352,7 @@ impl AnalysisSession { input_schemas, dialect: &self.request.dialect, function_return_types: &self.request.function_return_types, + nullability_rules: self.request.nullability_rules.as_ref(), }) .collect(); let pool = self.catalog.native.analysis_pool()?; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs index d0d6fa5520..ccdf67f459 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs @@ -95,3 +95,79 @@ pub(crate) const RENDERED_TYPE_NAMES: [(&str, &str); 12] = [ ]; /// The stack a standalone dynamic pivot proof parses on, as the analysis workers have. pub(crate) const PIVOT_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; +/// Stack for one standalone CTE fact recovery, whose AST walks recurse with query depth. +pub(crate) const CTE_FACT_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; +/// The wheel's node kinds Python's CTE fact recovery reads. +pub(crate) const ANNOTATED_AST_KIND: &str = "annotated"; +pub(crate) const LITERAL_AST_KIND: &str = "literal"; +pub(crate) const NULL_AST_KIND: &str = "null"; +pub(crate) const CONCAT_AST_KIND: &str = "concat"; +pub(crate) const SUBSTRING_AST_KIND: &str = "substring"; +pub(crate) const COALESCE_AST_KIND: &str = "coalesce"; +pub(crate) const IF_FUNC_AST_KIND: &str = "if_func"; +pub(crate) const CASE_AST_KIND: &str = "case"; +pub(crate) const TRY_CAST_AST_KIND: &str = "try_cast"; +pub(crate) const COUNT_AST_KIND: &str = "count"; +pub(crate) const IS_NULL_AST_KIND: &str = "is_null"; +pub(crate) const SET_OPERATION_AST_KINDS: [&str; 3] = ["union", "intersect", "except"]; +pub(crate) const CAST_AST_KINDS: [&str; 2] = ["cast", "try_cast"]; +pub(crate) const TYPE_PASSTHROUGH_AST_KINDS: [&str; 4] = + ["max", "min", "window_function", "within_group"]; +/// Python's `POLYGLOT_BOOLEAN_RESULT_KINDS`. +pub(crate) const BOOLEAN_RESULT_AST_KINDS: [&str; 18] = [ + "boolean", + "and", + "between", + "eq", + "exists", + "gt", + "gte", + "ilike", + "in", + "is", + "is_null", + "like", + "lt", + "lte", + "neq", + "not", + "or", + "regexp_like", +]; +pub(crate) const NULLIF_FUNCTION_NAME: &str = "NULLIF"; +pub(crate) const STRING_LITERAL_TYPE: &str = "string"; +pub(crate) const BINARY_OPERAND_COUNT: usize = 2; +/// The type Python's set operation slots give a bare NULL. +pub(crate) const NULL_SET_OPERATION_TYPE: &str = "__SQLBUILD_NULL_SET_OPERATION_TYPE__"; +pub(crate) const BOOLEAN_TYPE: &str = "BOOLEAN"; +pub(crate) const TEXT_TYPE: &str = "TEXT"; +pub(crate) const CUSTOM_TYPE_NAME: &str = "CUSTOM"; +pub(crate) const TIMESTAMP_TYPE_NAME: &str = "timestamp"; +pub(crate) const TIMESTAMP_TZ_TYPE: &str = "TIMESTAMPTZ"; +pub(crate) const DECIMAL_TYPE: &str = "DECIMAL"; +pub(crate) const VARCHAR_DATA_TYPES: [&str; 2] = ["var_char", "varchar"]; +/// Python's `_polyglot_type_name` names, matched case-sensitively as Python's dict is. +pub(crate) const POLYGLOT_TYPE_NAMES: [(&str, &str); 13] = [ + ("big_int", "BIGINT"), + ("bool", "BOOLEAN"), + ("boolean", "BOOLEAN"), + ("date", "DATE"), + ("decimal", "DECIMAL"), + ("double", "DOUBLE"), + ("float", "FLOAT"), + ("int", "INT"), + ("integer", "INT"), + ("text", "TEXT"), + ("timestamp", "TIMESTAMP"), + ("var_char", "TEXT"), + ("varchar", "TEXT"), +]; +/// Python's join sides that make relations nullable. +pub(crate) const JOIN_LEFT: &str = "LEFT"; +pub(crate) const JOIN_RIGHT: &str = "RIGHT"; +pub(crate) const JOIN_FULL: &str = "FULL"; +/// Ids of the adapter nullability rules Python registers, by the function it calls. +pub(crate) const FIRST_ARG_RULE: &str = "first_arg"; +pub(crate) const CONDITIONAL_RESULT_RULE: &str = "conditional_result"; +pub(crate) const NON_NULL_NULLABILITY: &str = "non_null"; +pub(crate) const NULLABLE_NULLABILITY: &str = "nullable"; 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 6c5394355c..2e74f7a45c 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs @@ -2,5 +2,6 @@ pub mod expression_shapes; pub mod finish_analysis_session; pub mod prove_dynamic_contract; pub mod provide_deferred_analyses; +pub mod recover_cte_facts; pub mod run_analysis_session; pub mod start_analysis_session; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/recover_cte_facts.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/recover_cte_facts.rs new file mode 100644 index 0000000000..84aace7ca4 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/recover_cte_facts.rs @@ -0,0 +1,47 @@ +//! Python's CTE pass-through facts and filtered non-null outputs for one query. + +use crate::assembly::analysis_session::_helpers::cte_facts::{ + Recovery, RecoveryInput, RecoveryProfile, recovery, +}; +use crate::assembly::analysis_session::constants::CTE_FACT_WORKER_STACK_BYTES; +use crate::assembly::analysis_session::models::{CteFactRequest, CteFacts}; + +/// The facts enrichment recovers for the query, or why they are left to Python. +pub fn recover_cte_facts(request: &CteFactRequest) -> Result { + std::thread::scope(|scope| { + std::thread::Builder::new() + .stack_size(CTE_FACT_WORKER_STACK_BYTES) + .spawn_scoped(scope, || recovered_facts(request)) + .map_err(|error| error.to_string())? + .join() + .map_err(|_| "native CTE fact recovery panicked".to_owned())? + }) +} + +fn recovered_facts(request: &CteFactRequest) -> Result { + let recovered: Recovery = recovery(&RecoveryInput { + cleaned_sql: &request.cleaned_sql, + input_schemas: &request.input_schemas, + recover: request.recover, + null_filter: request.null_filter, + profile: RecoveryProfile { + dialect: &request.dialect, + function_return_types: &request.function_return_types, + rules: request.nullability_rules.as_ref(), + }, + })?; + let mut direct_outputs: Vec = recovered.direct_outputs.into_iter().collect(); + direct_outputs.sort_unstable(); + let mut non_null_outputs: Vec = recovered.non_null_outputs.into_iter().collect(); + non_null_outputs.sort_unstable(); + Ok(CteFacts { + types: recovered.types, + nullability: recovered + .nullability + .into_iter() + .map(|(name, nullability)| (name, nullability.to_owned())) + .collect(), + direct_outputs, + non_null_outputs, + }) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs index d10d56c0ad..30657d426e 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs @@ -89,6 +89,8 @@ pub struct SessionRequest { /// Whether published shapes keep authored quoting, Python's `inferred_binding_shape` test. pub case_sensitive_shapes: bool, pub function_return_types: Pairs, + /// Adapter nullability rules as `(function name, rule id)`; None where one is not Python's. + pub nullability_rules: Option, pub rich_type_inference: bool, pub column_types: Shapes, pub column_nullability: Shapes, @@ -257,3 +259,30 @@ pub struct AnalysisSession { pub(crate) publications: Shapes, pub(crate) failures: Vec, } + +/// One query's CTE fact recovery, as Python's compact enrichment runs it. +#[derive(Debug, Clone)] +pub struct CteFactRequest { + pub cleaned_sql: String, + pub dialect: String, + /// Input relation types; every input column's nullability is unknown, as in enrichment. + pub input_schemas: Shapes, + pub function_return_types: Pairs, + /// Adapter nullability rules as `(function name, rule id)`; None where one is not Python's. + pub nullability_rules: Option, + /// Whether Python runs `_polyglot_cte_passthrough_facts` past its early return. + pub recover: bool, + /// Whether a filter reads NULL, so Python looks for filtered non-null outputs. + pub null_filter: bool, +} + +/// Python's recovered CTE pass-through facts and filtered non-null outputs for one query. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CteFacts { + pub types: Pairs, + pub nullability: Pairs, + /// Sorted direct CTE output names. + pub direct_outputs: Vec, + /// Sorted outputs a filter proves non-null. + pub non_null_outputs: Vec, +} 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 1f274fdfc1..33eabbad61 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs @@ -1,3 +1,6 @@ +use crate::assembly::analysis_session::_helpers::cte_facts::{ + Recovery, RecoveryInput, RecoveryProfile, recovery, +}; use crate::assembly::analysis_session::_helpers::mappings::catalog_relations; use crate::assembly::analysis_session::main::finish_analysis_session::finish_analysis_session; use crate::assembly::analysis_session::main::provide_deferred_analyses::provide_deferred_analyses; @@ -8,7 +11,7 @@ use crate::assembly::analysis_session::models::{ ModelOutcome, ModelReference, ModelRequest, PivotOutcome, PivotRequest, SessionRequest, SessionStep, }; -use crate::assembly::analysis_session::tests::test_types::ModelSpec; +use crate::assembly::analysis_session::tests::test_types::{ModelSpec, RecoveredFacts}; use crate::assembly::analysis_session::types::{Pairs, Shapes}; use crate::semantic_validation::models::{CatalogInput, ProjectCatalog}; @@ -95,6 +98,7 @@ pub(crate) fn orders_request(models: Vec) -> SessionRequest { dialect: "duckdb".to_owned(), case_sensitive_shapes: false, function_return_types: Vec::new(), + nullability_rules: Some(Vec::new()), rich_type_inference: true, column_types: raw_orders.clone(), column_nullability: shapes(&[( @@ -297,3 +301,52 @@ pub(crate) fn proven( bare_dynamic_pivot: bare, }) } + +/// Python's recovery over `orders` for a contract-enforced model whose filter reads NULL. +pub(crate) fn recovered_facts(sql: &str) -> Option<(Pairs, Pairs, Vec, Vec)> { + let input_schemas: Shapes = shapes(&[( + "orders", + &[ + ("order_id", "INTEGER"), + ("amount", "DOUBLE"), + ("status", "VARCHAR"), + ], + )]); + let rules: Pairs = pairs(&[("UPPER", "first_arg")]); + let recovered: Recovery = recovery(&RecoveryInput { + cleaned_sql: sql, + input_schemas: &input_schemas, + recover: true, + null_filter: true, + profile: RecoveryProfile { + dialect: "duckdb", + function_return_types: &Vec::new(), + rules: Some(&rules), + }, + }) + .ok()?; + let mut direct: Vec = recovered.direct_outputs.into_iter().collect(); + direct.sort(); + let mut non_null: Vec = recovered.non_null_outputs.into_iter().collect(); + non_null.sort(); + let nullability: Pairs = recovered + .nullability + .into_iter() + .map(|(name, value)| (name, value.to_owned())) + .collect(); + Some((recovered.types, nullability, direct, non_null)) +} + +/// The expected recovery as owned values. +pub(crate) fn expected_facts( + facts: Option, +) -> Option<(Pairs, Pairs, Vec, Vec)> { + facts.map(|(types, nullability, direct, non_null)| { + ( + pairs(types), + pairs(nullability), + direct.iter().map(|name| (*name).to_owned()).collect(), + non_null.iter().map(|name| (*name).to_owned()).collect(), + ) + }) +} 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 c56625f012..ff31a21f99 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 @@ -2,10 +2,11 @@ use crate::assembly::analysis_session::main::prove_dynamic_contract::prove_dynam 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, failed, model_requests, orders_request, pivot_request, proven, session_lines, + catalog, expected_facts, failed, model_requests, orders_request, pivot_request, proven, + recovered_facts, session_lines, }; use crate::assembly::analysis_session::tests::test_types::{ - PivotTestCase, SessionTestCase, UnscheduledTestCase, + CteRecoveryTestCase, PivotTestCase, SessionTestCase, UnscheduledTestCase, }; #[test] @@ -291,3 +292,68 @@ fn given_dynamic_pivots_when_proving_then_matches_python_or_defers() { ); } } + +#[test] +fn given_cte_reads_when_recovering_facts_then_matches_python_or_defers() { + let test_cases = [ + CteRecoveryTestCase { + description: "direct CTE reads keep input and cast types", + sql: "WITH base AS (SELECT order_id, CAST(amount AS DECIMAL(10, 2)) AS amount \ + FROM orders) SELECT order_id, amount FROM base", + expected_facts: Some(( + &[("order_id", "INTEGER"), ("amount", "DECIMAL(10, 2)")], + &[], + &["amount", "order_id"], + &[], + )), + }, + CteRecoveryTestCase { + description: "literals and filtered columns are non-null through an aliased CTE", + sql: "WITH base AS (SELECT 'x' AS tag, UPPER(status) AS up, order_id FROM orders \ + WHERE order_id IS NOT NULL) SELECT b.tag, b.up, b.order_id FROM base AS b", + expected_facts: Some(( + &[("order_id", "INTEGER")], + &[("tag", "non_null"), ("order_id", "non_null")], + &["order_id", "tag", "up"], + &[], + )), + }, + CteRecoveryTestCase { + description: "a set operation branch with an unnamed NULL proves nothing", + sql: "WITH u AS (SELECT order_id AS k, status AS s FROM orders UNION ALL \ + SELECT NULL, 'y') SELECT k, s FROM u", + expected_facts: Some((&[], &[], &["k", "s"], &[])), + }, + CteRecoveryTestCase { + description: "a star with EXCLUDE and a joined relation", + sql: "WITH a AS (SELECT * EXCLUDE (amount) FROM orders) SELECT o.status, a.order_id \ + FROM a LEFT JOIN orders AS o ON a.order_id = o.order_id", + expected_facts: Some(( + &[("status", "VARCHAR"), ("order_id", "INTEGER")], + &[], + &["order_id", "status"], + &[], + )), + }, + CteRecoveryTestCase { + description: "no CTE leaves only the filtered non-null outputs", + sql: "SELECT order_id, status FROM orders WHERE order_id IS NOT NULL", + expected_facts: Some((&[], &[], &[], &["order_id"])), + }, + CteRecoveryTestCase { + description: "non-ASCII names Python casefolds differently", + sql: "WITH b\u{e4}se AS (SELECT order_id FROM orders) SELECT order_id FROM b\u{e4}se", + expected_facts: None, + }, + ]; + for test_case in test_cases { + let facts = recovered_facts(test_case.sql); + + assert_eq!( + facts, + expected_facts(test_case.expected_facts), + "{}", + 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 c09b783048..8e4d0b84e1 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 @@ -51,3 +51,18 @@ pub(crate) struct PivotTestCase { pub(crate) family: Option<(&'static str, &'static str, &'static str)>, pub(crate) expected_outcome: PivotOutcome, } + +/// Types, nullabilities, direct CTE outputs and filtered non-null outputs, in Python's order. +pub(crate) type RecoveredFacts = ( + &'static [(&'static str, &'static str)], + &'static [(&'static str, &'static str)], + &'static [&'static str], + &'static [&'static str], +); + +pub(crate) struct CteRecoveryTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + /// None where the recovery defers to Python. + pub(crate) expected_facts: Option, +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs index 54e3fdf710..feba290be6 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs @@ -1,2 +1,3 @@ +pub(crate) mod oracle; pub(crate) mod registration; pub(crate) mod session; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/oracle.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/oracle.rs new file mode 100644 index 0000000000..e631bc54ea --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/oracle.rs @@ -0,0 +1,62 @@ +//! Test-only hook comparing native CTE fact recovery with Python's, function for function. + +use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyResult, Python}; +use pyo3::{pyfunction, wrap_pyfunction}; +use sqlbuild_analysis::assembly::analysis_session::main::recover_cte_facts::recover_cte_facts; +use sqlbuild_analysis::assembly::analysis_session::models::{CteFactRequest, CteFacts}; +use sqlbuild_analysis::assembly::analysis_session::types::{Pairs, Shapes}; + +use crate::bindings::_helpers::boundary::panics::value_error; +use crate::bindings::types::CompilerDetach; + +/// `(cleaned SQL, dialect, input schemas, return types, nullability rules, recover, null filter)`. +type RecoveryRow = (String, String, Shapes, Pairs, Option, bool, bool); +/// `(deferral, types, nullability, sorted direct outputs, sorted non-null outputs)`. +type FactsRow = (Option, Pairs, Pairs, Vec, Vec); + +/// The facts native enrichment recovers for one query, or the reason it defers to Python. +#[pyfunction] +fn _oracle_cte_fact_recovery(py: Python<'_>, request: RecoveryRow) -> PyResult { + let ( + cleaned_sql, + dialect, + input_schemas, + function_return_types, + nullability_rules, + recover, + null_filter, + ) = request; + let request = CteFactRequest { + cleaned_sql, + dialect, + input_schemas, + function_return_types, + nullability_rules, + recover, + null_filter, + }; + let facts: Result = py + .compiler_detach(|| Ok(recover_cte_facts(&request))) + .map_err(value_error)?; + Ok(match facts { + Ok(facts) => ( + None, + facts.types, + facts.nullability, + facts.direct_outputs, + facts.non_null_outputs, + ), + Err(deferral) => ( + Some(deferral), + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + ), + }) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(_oracle_cte_fact_recovery, module)?)?; + Ok(()) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs index bfe82cb5d1..426d0f43eb 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs @@ -2,8 +2,9 @@ use pyo3::prelude::{Bound, PyModule, PyResult}; -use crate::bindings::_helpers::analysis_session::session; +use crate::bindings::_helpers::analysis_session::{oracle, session}; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - session::register(module) + session::register(module)?; + oracle::register(module) } 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 87af8f9b33..7d04738e92 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs @@ -47,6 +47,7 @@ type RequestRow = ( String, bool, Pairs, + Option, bool, Shapes, Shapes, @@ -226,6 +227,7 @@ fn session_request(request: RequestRow) -> SessionRequest { dialect, case_sensitive_shapes, function_return_types, + nullability_rules, rich_type_inference, column_types, column_nullability, @@ -238,6 +240,7 @@ fn session_request(request: RequestRow) -> SessionRequest { dialect, case_sensitive_shapes, function_return_types, + nullability_rules, rich_type_inference, column_types, column_nullability, diff --git a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py index 4a7c34a23e..c769c944bd 100644 --- a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py +++ b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py @@ -26,6 +26,13 @@ def _pivot_model(*, pivot_column: str, body: str) -> str: ) +def _contract_header(columns: str) -> str: + return ( + 'MODEL (\n description "Order totals per customer",\n contract enforced,\n' + f" columns (\n{columns} ),\n);\n\n" + ) + + def analysis_session_failure_cases() -> tuple[FailureCase, ...]: """Return this lane's failure cases; the lane appends here without editing shared lists.""" @@ -70,6 +77,29 @@ def analysis_session_failure_cases() -> tuple[FailureCase, ...]: f"SELECT customer_id, amount\nFROM (SELECT * FROM {_ORDERS}) staged\n", }, ), + failure_case( + name="analysis-session-cte-fact-type-mismatch", + expected_code="K002", + files={ + FAILURE_MART_PATH: _contract_header( + " customer_id (type INTEGER),\n total_amount (type VARCHAR),\n" + ) + + "WITH totals AS (\n SELECT customer_id, CAST(amount AS DOUBLE) AS total_amount\n" + f" FROM {_ORDERS}\n)\n" + "SELECT customer_id + 0 AS customer_id, total_amount FROM totals\n", + }, + ), + failure_case( + name="analysis-session-cte-fact-nullability", + expected_code="K004", + files={ + FAILURE_MART_PATH: _contract_header( + " customer_id (type INTEGER),\n note (type VARCHAR, nullable false),\n" + ) + + "WITH base AS (\n SELECT customer_id, CAST(NULL AS VARCHAR) AS note\n" + f" FROM {_ORDERS}\n)\nSELECT customer_id + 0 AS customer_id, note FROM base\n", + }, + ), failure_case( name="analysis-session-pivot-values-on-clause", expected_code="K011", diff --git a/scripts/compiler_differential/classes/analysis_features.py b/scripts/compiler_differential/classes/analysis_features.py index e45a65edaa..9240dc67dc 100644 --- a/scripts/compiler_differential/classes/analysis_features.py +++ b/scripts/compiler_differential/classes/analysis_features.py @@ -331,6 +331,28 @@ def _python_sql(self) -> None: def _analysis_modes(self) -> None: self._features.add("analysis_modes") base: ModelPlan = self._base() + self._model( + name=f"cte_facts_{base.name}", + description=f"Contracted CTE reads of {base.name}", + header=( + " contract enforced,\n columns (\n" + + "".join( + f' {column} (description "Recovered {column}"),\n' + for column in ("id", "amount", "kind", "status") + ) + + " ),\n" + ), + body=( + "WITH base AS (\n" + " SELECT id, CAST(amount AS DECIMAL(12, 2)) AS amount, 'event' AS kind,\n" + " COALESCE(status, 'unknown') AS status\n" + f' FROM __ref("{base.name}")\n' + " WHERE id IS NOT NULL\n" + ")\n" + "SELECT id + 0 AS id, amount, UPPER(kind) AS kind, UPPER(status) AS status\n" + "FROM base\n" + ), + ) self.commands = tuple( _selected_command(label=label, arguments=arguments, model=base.name) for label, arguments in GENERATOR_ANALYSIS_MODE_COMMANDS diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index b98b504587..6042e56cb7 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -452,6 +452,18 @@ def prove_dynamic_column_contract( 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]: ... +def _oracle_cte_fact_recovery( + request: tuple[ + str, + str, + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, str]], + list[tuple[str, str]] | None, + bool, + bool, + ], + /, +) -> tuple[str | None, list[tuple[str, str]], list[tuple[str, str]], list[str], list[str]]: ... # Native analysis: semantic completion. class SemanticTypeRecovery: diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py index 093c361e6e..99cd11b33f 100644 --- a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py +++ b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py @@ -5,6 +5,15 @@ from collections.abc import Mapping from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.adapter.contract.types import FunctionNullabilityRule +from sqlbuild.adapter.type_system.main.conditional_result_nullability import ( + conditional_result_nullability, +) +from sqlbuild.adapter.type_system.main.first_arg_nullability import first_arg_nullability +from sqlbuild.compiler.analysis_session.constants import ( + NULLABILITY_RULE_CONDITIONAL_RESULT, + NULLABILITY_RULE_FIRST_ARG, +) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest from sqlbuild.compiler.analysis_session.types import ( ColumnRow, @@ -30,6 +39,11 @@ from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily +_NULLABILITY_RULE_IDS: dict[FunctionNullabilityRule, str] = { + first_arg_nullability: NULLABILITY_RULE_FIRST_ARG, + conditional_result_nullability: NULLABILITY_RULE_CONDITIONAL_RESULT, +} + def shape_rows(shapes: Mapping[str, Mapping[str, object]]) -> ShapeRows | None: """`{relation: {column: value}}` as ordered rows, or None when a value is not text.""" @@ -66,6 +80,7 @@ def session_request( profile.sql_analysis_dialect or "generic", case_sensitive_shapes(profile=profile, dialect=profile.sql_analysis_dialect), list(profile.function_return_types.items()), + nullability_rule_rows(profile), request.rich_type_inference, *shapes, family_rows(request.dynamic_families_by_table), @@ -73,6 +88,18 @@ def session_request( ) +def nullability_rule_rows(profile: ExpressionInferenceProfile) -> list[tuple[str, str]] | None: + """Adapter nullability rules as `(name, rule id)`, or None for a rule SQLBuild lacks.""" + + rows: list[tuple[str, str]] = [] + for name, rule in profile.function_nullability_rules.items(): + rule_id: str | None = _NULLABILITY_RULE_IDS.get(rule) + if rule_id is None: + return None + rows.append((name, rule_id)) + return rows + + def contract_proof(row: ProofRow | None) -> DynamicColumnContractProof | None: """A native dynamic pivot proof as Python's, None when absent or left to Python.""" diff --git a/src/sqlbuild/compiler/analysis_session/constants.py b/src/sqlbuild/compiler/analysis_session/constants.py index 9d819d795d..4a4e269c40 100644 --- a/src/sqlbuild/compiler/analysis_session/constants.py +++ b/src/sqlbuild/compiler/analysis_session/constants.py @@ -16,3 +16,5 @@ 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" +NULLABILITY_RULE_FIRST_ARG: str = "first_arg" +NULLABILITY_RULE_CONDITIONAL_RESULT: str = "conditional_result" diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_helpers.py new file mode 100644 index 0000000000..7622a8d2bf --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_helpers.py @@ -0,0 +1,440 @@ +"""Seeded CTE queries and both CTE fact recoveries, compared function for function.""" + +from __future__ import annotations + +import random +from collections import Counter +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import sqlbuild._native as native_module +import sqlbuild.compiler.compile._helpers.analysis.compact as compact_analysis +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.adapter.type_system.main.conditional_result_nullability import ( + conditional_result_nullability, +) +from sqlbuild.adapter.type_system.main.first_arg_nullability import first_arg_nullability +from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter +from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter +from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.adapters.sqlserver.classes.sqlserver_adapter import SqlServerAdapter +from sqlbuild.compiler.analysis_session._helpers.session_rows import nullability_rule_rows +from sqlbuild.compiler.compile._helpers.analysis.cte_facts import ( + _polyglot_cte_passthrough_facts, + _polyglot_filtered_non_null_outputs, +) +from sqlbuild.compiler.compile.models import CteFactResolvers +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.compiler.sql_analysis.main.import_polyglot import import_polyglot + +type Generate = Callable[[random.Random, list[str], int], str] +type CteFactView = ( + tuple[list[tuple[str, str]], list[tuple[str, str]], list[str], list[str]] | tuple[str, str] +) + + +@dataclass(frozen=True) +class CteFactQuery: + """One generated query and the enrichment facts both recoveries read.""" + + sql: str + profile: ExpressionInferenceProfile + recover: bool + null_filter: bool + + +@dataclass +class CteFactParity: + """Python's and native's recovered facts for every query native answered.""" + + queries: list[object] + python: list[object] + native: list[object] + counts: Counter[str] + + +_SCHEMAS: dict[str, dict[str, str]] = { + "orders": { + "order_id": "INTEGER", + "customer_id": "BIGINT", + "qty": "INT", + "price": "DECIMAL(10, 2)", + "big_price": "NUMERIC(12,2)", + "ratio": "DOUBLE", + "status": "VARCHAR", + "code": "VARCHAR(20)", + "note": "TEXT", + "flag": "BOOLEAN", + "placed_on": "DATE", + "placed_at": "TIMESTAMP", + "MixedCase": "INTEGER", + "mixedcase": "BIGINT", + "Price": "FLOAT", + }, + "customers": { + "customer_id": "BIGINT", + "order_id": "VARCHAR", + "region": "STRING", + "price": "DOUBLE", + "Region": "CHAR(2)", + "empty_type": "", + }, + "Events": {"event_id": "INT", "kind": "VARCHAR(10)"}, + "events": {"event_id": "BIGINT", "payload": "JSON"}, +} +_INPUT_NULLABILITY: dict[str, dict[str, InferredNullability]] = { + table: dict.fromkeys(columns, InferredNullability.UNKNOWN) + for table, columns in _SCHEMAS.items() +} +_INPUT_SHAPES: list[tuple[str, list[tuple[str, str]]]] = [ + (table, list(columns.items())) for table, columns in _SCHEMAS.items() +] +_NULL_FILTER_USE: dict[str, str] = {"context": "filter", "expressionSql": "x IS NOT NULL"} +_RESOLVERS: CteFactResolvers = CteFactResolvers( + expression_type=compact_analysis._polyglot_expression_type, + nullability=compact_analysis._infer_polyglot_nullability, + shallow_nullability=compact_analysis._infer_polyglot_shallow_nullability, + alias_nullability=compact_analysis._polyglot_alias_nullability_from_select, +) +_ADAPTERS: tuple[BaseAdapter, ...] = ( + DuckDbAdapter(), + SnowflakeAdapter(), + BigQueryAdapter(), + DatabricksAdapter(), + PostgresAdapter(), + SqlServerAdapter(), +) +# Rules apply by node kind. The wheel parses IF, IFF and IIF as `if_func`, whose arguments +# Python never collects, so adapters' IF/IFF rules stay unknown here; the profile keyed by +# node kinds with collected arguments reaches the conditional_result and first_arg rules. +PROFILES: tuple[ExpressionInferenceProfile, ...] = ( + *(adapter.expression_inference_profile() for adapter in _ADAPTERS), + ExpressionInferenceProfile( + sql_analysis_dialect="duckdb", + function_nullability_rules=DuckDbAdapter() + .expression_inference_profile() + .function_nullability_rules, + function_return_types={"LENGTH": "BIGINT", "UPPER": "VARCHAR(5)", "ABS": "DECIMAL(10,2)"}, + ), + ExpressionInferenceProfile( + sql_analysis_dialect="snowflake", + function_nullability_rules={ + "FUNCTION": conditional_result_nullability, + "IN": conditional_result_nullability, + "ARRAY_FUNC": first_arg_nullability, + "CONCAT": first_arg_nullability, + "UPPER": first_arg_nullability, + "LOWER": first_arg_nullability, + }, + function_return_types=SnowflakeAdapter() + .expression_inference_profile() + .function_return_types, + ), + ExpressionInferenceProfile(sql_analysis_dialect="generic"), + ExpressionInferenceProfile(sql_analysis_dialect=None), +) +_TYPES: tuple[str, ...] = ( + "INT", "INTEGER", "BIGINT", "SMALLINT", "TINYINT", "DECIMAL(10, 2)", "DECIMAL(12,2)", + "DECIMAL", "DECIMAL(9)", "NUMERIC(10, 2)", "NUMERIC", "NUMBER", "NUMBER(10, 2)", "VARCHAR", + "VARCHAR(10)", "TEXT", "STRING", "CHAR(3)", "DOUBLE", "FLOAT", "REAL", "FLOAT64", "INT64", + "BOOLEAN", "BOOL", "DATE", "TIMESTAMP", "TIMESTAMPTZ", "TIMESTAMP WITH TIME ZONE", + "TIMESTAMP_NTZ", "DATETIME", "JSON", "VARIANT", "BINARY", "HUGEINT", "BIGNUMERIC", + "NVARCHAR(5)", "DOUBLE PRECISION", "INTERVAL", "UUID", "ARRAY", "INT[]", + "STRUCT", "MAP(VARCHAR, INT)", +) # fmt: skip +_LITERALS: tuple[str, ...] = ("1", "'x'", "NULL", "1.5", "TRUE", "DATE '2020-01-01'") +_MAX_DEPTH: int = 3 +_ALIASES: tuple[str, ...] = ("b", "c", "o2", "X") +_JOINS: tuple[tuple[str, str], ...] = ( + ("JOIN", " ON TRUE"), + ("LEFT JOIN", " ON TRUE"), + ("RIGHT JOIN", " ON TRUE"), + ("FULL OUTER JOIN", " ON TRUE"), + ("LEFT OUTER JOIN", " ON TRUE"), + ("INNER JOIN", " ON TRUE"), + ("CROSS JOIN", ""), +) +_OUTPUT_NAMES: tuple[str, ...] = ( + "c{index}", "C{index}", '"Mx{index}"', "status", "qty", "dup", "Dup", '"Größe"', + '"STRASSE"', '"straße"', "ıd", +) # fmt: skip +_STARS: tuple[str, ...] = ( + "*", + "{alias}.*", + "* EXCLUDE (status)", + "* EXCLUDE (order_id, qty)", + "* REPLACE (1 AS qty)", +) +_SET_OPERATIONS: tuple[str, ...] = ( + "UNION ALL", + "UNION", + "UNION ALL BY NAME", + "EXCEPT", + "INTERSECT", +) +_FACT_KINDS: tuple[str, ...] = ("types", "nullability", "direct", "non_null") +_CTE_NAMES: tuple[str, ...] = ("cte{index}", "Cte{index}", '"Cte{index}"', "base", "orders") + + +def _quoted(rng: random.Random, name: str) -> str: + return rng.choices((f'"{name}"', name.upper(), name.lower(), name), weights=(10, 10, 5, 75))[0] + + +def _column(rng: random.Random, columns: list[str]) -> str: + return _quoted(rng, rng.choice(columns)) + + +def _nested(rng: random.Random, columns: list[str], depth: int) -> str: + return _expression(rng=rng, columns=columns, depth=depth + 1) + + +_LEAVES: tuple[Generate, ...] = ( + lambda rng, columns, depth: _column(rng, columns), + lambda rng, columns, depth: rng.choice(_LITERALS), +) +_COMPOUNDS: tuple[Generate, ...] = ( + lambda rng, columns, depth: _column(rng, columns), + lambda rng, columns, depth: _column(rng, columns), + lambda r, c, d: f"CAST({_nested(r, c, d)} AS {r.choice(_TYPES)})", + lambda r, c, d: f"TRY_CAST({_nested(r, c, d)} AS {r.choice(_TYPES)})", + lambda r, c, d: f"{_nested(r, c, d)}::{r.choice(('INT', 'DECIMAL(5, 1)', 'VARCHAR', 'TEXT'))}", + lambda rng, columns, depth: rng.choice(_LITERALS), + lambda r, c, d: ( + f"CASE WHEN {_nested(r, c, d)} THEN {_nested(r, c, d)} ELSE {_nested(r, c, d)} END" + ), + lambda r, c, d: ( + f"CASE {_nested(r, c, d)} WHEN 1 THEN {_nested(r, c, d)} WHEN 2 THEN {_nested(r, c, d)} END" + ), + lambda r, c, d: f"CASE WHEN {_nested(r, c, d)} THEN {_nested(r, c, d)} END", + lambda r, c, d: f"COALESCE({_nested(r, c, d)}, {_nested(r, c, d)})", + lambda r, c, d: f"COALESCE({_nested(r, c, d)}, {_nested(r, c, d)}, {_nested(r, c, d)})", + lambda r, c, d: f"NULLIF({_nested(r, c, d)}, {_nested(r, c, d)})", + lambda r, c, d: f"IF({_nested(r, c, d)}, {_nested(r, c, d)}, {_nested(r, c, d)})", + lambda r, c, d: f"IFF({_nested(r, c, d)}, {_nested(r, c, d)}, {_nested(r, c, d)})", + lambda r, c, d: f"UPPER({_nested(r, c, d)})", + lambda r, c, d: f"LOWER({_nested(r, c, d)})", + lambda r, c, d: f"{_nested(r, c, d)} || {_nested(r, c, d)}", + lambda r, c, d: f"SUBSTRING({_nested(r, c, d)}, 1, 2)", + lambda r, c, d: f"{r.choice(('MAX', 'MIN', 'SUM', 'AVG', 'COUNT'))}({_nested(r, c, d)})", + lambda r, c, d: f"MAX({_nested(r, c, d)}) OVER (PARTITION BY {_column(r, c)})", + lambda r, c, d: f"ROW_NUMBER() OVER (ORDER BY {_column(r, c)})", + lambda r, c, d: f"{_nested(r, c, d)} {r.choice(('>', '=', '<>', 'LIKE'))} {_nested(r, c, d)}", + lambda r, c, d: f"{_nested(r, c, d)} IS NULL", + lambda r, c, d: f"{_nested(r, c, d)} IS NOT NULL", + lambda r, c, d: f"{_nested(r, c, d)} {r.choice(('+', '*', '/'))} {_nested(r, c, d)}", + lambda r, c, d: ( + f"{r.choice(('TO_VARCHAR', 'ABS', 'LENGTH', 'REPLACE', 'CONCAT', 'TO_DATE'))}" + f"({_nested(r, c, d)}, {_nested(r, c, d)})" + ), + lambda r, c, d: f"LENGTH({_nested(r, c, d)})", + lambda r, c, d: ( + f"{r.choice(('GREATEST', 'CONCAT_WS', 'DECODE'))}({_nested(r, c, d)}, " + f"{_nested(r, c, d)}, {_nested(r, c, d)})" + ), + lambda r, c, d: f"{_nested(r, c, d)} IN ({_nested(r, c, d)}, {_nested(r, c, d)})", + lambda r, c, d: f"[{_nested(r, c, d)}, {_nested(r, c, d)}]", +) + + +def _expression(*, rng: random.Random, columns: list[str], depth: int) -> str: + builders: tuple[Generate, ...] = (_LEAVES, _COMPOUNDS)[depth < _MAX_DEPTH] + return rng.choice(builders)(rng, columns, depth) + + +def _star_projection( + rng: random.Random, aliases: list[str], columns: list[str], index: int +) -> tuple[str, list[str]]: + return rng.choice(_STARS).format(alias=aliases[0]), list(columns) + + +def _column_projection( + rng: random.Random, aliases: list[str], columns: list[str], index: int +) -> tuple[str, list[str]]: + column: str = rng.choice(columns) + qualifier: str = rng.choices(("", f"{rng.choice(aliases)}."), weights=(6, 4))[0] + return f"{qualifier}{_quoted(rng, column)}", [column] + + +def _expression_projection( + rng: random.Random, aliases: list[str], columns: list[str], index: int +) -> tuple[str, list[str]]: + name: str = rng.choice(_OUTPUT_NAMES).format(index=index) + return f"{_expression(rng=rng, columns=columns, depth=0)} AS {name}", [name.strip('"')] + + +_PROJECTIONS: tuple[Callable[..., tuple[str, list[str]]], ...] = ( + _star_projection, + _column_projection, + _expression_projection, +) + + +def _predicate(rng: random.Random, aliases: list[str], columns: list[str]) -> str: + qualifier: str = rng.choice(("", f"{rng.choice(aliases)}.")) + column: str = f"{qualifier}{_column(rng, columns)}" + return rng.choice( + ( + f"{column} IS NOT NULL", + f"({column} IS NOT NULL)", + f"{column} > 1", + f"NOT {column} IS NULL", + ) + ) + + +def _select( + *, rng: random.Random, relations: dict[str, list[str]], nests: bool +) -> tuple[str, list[str]]: + first: str = rng.choice(list(relations)) + alias_sql, alias = rng.choice(((("", first)), (" a", "a"), (" AS A", "A"), (" t1", "t1"))) + columns: list[str] = list(relations[first]) + aliases: list[str] = [alias] + joins: list[str] = [] + for _ in range(rng.choice((0, 0, 1, 1, 2))): + other: str = rng.choice(list(relations)) + join_alias: str = rng.choice(_ALIASES) + join, condition = rng.choice(_JOINS) + joins.append(f"\n{join} {_quoted(rng, other)} AS {join_alias}{condition}") + columns.extend(relations[other]) + aliases.append(join_alias) + projections: list[str] = [] + outputs: list[str] = [] + for index in range(rng.randint(1, 5)): + projection, names = rng.choices(_PROJECTIONS, weights=(12, 33, 55))[0]( + rng, aliases, columns, index + ) + projections.append(projection) + outputs.extend(names) + predicates: list[str] = [_predicate(rng, aliases, columns) for _ in range(rng.randint(1, 3))] + where: str = rng.choices( + ("", "\nWHERE " + rng.choice((" AND ", " OR ", " AND ")).join(predicates)), + weights=(3, 7), + )[0] + group: str = rng.choices(("", "\nGROUP BY 1"), weights=(9, 1))[0] + body: str = ( + f"SELECT {', '.join(projections)}\nFROM {_quoted(rng, first)}{alias_sql}" + f"{''.join(joins)}{where}{group}" + ) + operation: Callable[[], str] = rng.choices( + (lambda: "", lambda: _set_operation(rng=rng, relations=relations)), weights=(85, 15) + )[0] + selected: str = f"{body}{operation()}" + nested: bool = nests and rng.random() < 0.15 + return ( + lambda: (selected, outputs), + lambda: _nested_select(rng=rng, relations=relations), + )[nested]() + + +def _set_operation(*, rng: random.Random, relations: dict[str, list[str]]) -> str: + other, _ = _select(rng=rng, relations=relations, nests=False) + return f"\n{rng.choice(_SET_OPERATIONS)}\n{other}" + + +def _nested_select(*, rng: random.Random, relations: dict[str, list[str]]) -> tuple[str, list[str]]: + inner, outputs = _select(rng=rng, relations=relations, nests=False) + return f"WITH inner_cte AS ({inner})\nSELECT * FROM inner_cte", outputs + + +def generated_cte_query(*, rng: random.Random) -> str: + """A query of one to three CTEs over typed inputs and an outer select reading them.""" + + relations: dict[str, list[str]] = {table: list(columns) for table, columns in _SCHEMAS.items()} + ctes: list[str] = [] + for index in range(rng.randint(1, 3)): + body, outputs = _select(rng=rng, relations=relations, nests=True) + name: str = rng.choice(_CTE_NAMES).format(index=index) + column_aliases: str = rng.choices(("", " (x, y)"), weights=(93, 7))[0] + ctes.append(f"{name}{column_aliases} AS (\n{body}\n)") + relations[name.strip('"')] = outputs or ["order_id"] + cte_relations: dict[str, list[str]] = dict( + filter(lambda item: item[0] not in _SCHEMAS or rng.random() < 0.2, relations.items()) + ) + outer, _ = _select(rng=rng, relations=cte_relations or relations, nests=False) + return "WITH " + ",\n".join(ctes) + "\n" + outer + + +def generated_cte_fact_query(*, rng: random.Random) -> CteFactQuery: + """A generated query under one adapter's profile, with enrichment's recovery gates.""" + + return CteFactQuery( + sql=generated_cte_query(rng=rng), + profile=rng.choice(PROFILES), + recover=rng.random() < 0.9, + null_filter=rng.random() < 0.8, + ) + + +def python_cte_facts(*, query: CteFactQuery) -> CteFactView: + """Python's recovered facts, as enrichment computes them, or what it raises.""" + + analysis: dict[str, Any] = { + "cteFacts": [1][: int(query.recover)], + "columnUses": [_NULL_FILTER_USE][: int(query.null_filter)], + } + dialect: str | None = query.profile.sql_analysis_dialect + try: + types, nullability, direct, parsed = _polyglot_cte_passthrough_facts( + polyglot_module=import_polyglot(), + cleaned_sql=query.sql, + dialect=dialect, + column_types_by_table=_SCHEMAS, + column_nullability_by_table=_INPUT_NULLABILITY, + inference_profile=query.profile, + analysis=analysis, + resolvers=_RESOLVERS, + ) + non_null: frozenset[str] = _polyglot_filtered_non_null_outputs( + polyglot_module=import_polyglot(), + cleaned_sql=query.sql, + dialect=dialect, + analysis=analysis, + column_nullability_by_table=_INPUT_NULLABILITY, + parsed=parsed, + ) + except Exception as error: # noqa: BLE001 - native must not answer where Python raises + return (type(error).__name__, str(error)) + return ( + list(types.items()), + [(name, value.value) for name, value in nullability.items()], + sorted(direct), + sorted(non_null), + ) + + +def native_cte_facts(*, query: CteFactQuery) -> tuple[str | None, CteFactView]: + """Native recovery's deferral reason, or None and its facts.""" + + deferral, types, nullability, direct, non_null = native_module._oracle_cte_fact_recovery( + ( + query.sql, + query.profile.sql_analysis_dialect or "generic", + _INPUT_SHAPES, + list(query.profile.function_return_types.items()), + nullability_rule_rows(query.profile), + query.recover, + query.null_filter, + ) + ) + return deferral, (types, nullability, direct, non_null) + + +def compare_cte_facts(*, query: CteFactQuery, parity: CteFactParity) -> None: + """Record both recoveries of `query` where native answers, and what each recovered.""" + + deferral, native = native_cte_facts(query=query) + answered: list[CteFactView] = [native][: int(deferral is None)] + python: list[CteFactView] = [python_cte_facts(query=query) for _ in answered] + parity.counts.update( + [f"deferred:{str(deferral).partition(' ')[0]}"][: int(deferral is not None)] + ) + parity.counts["compared"] += len(answered) + for index, kind in enumerate(_FACT_KINDS): + parity.counts[f"recovered_{kind}"] += sum(bool(view[index]) for view in answered) + parity.queries.extend((query.sql, query.profile.sql_analysis_dialect) for _ in answered) + parity.python.extend(python) + parity.native.extend(answered) 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 f3ea3a0aac..e9958740b8 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py @@ -15,6 +15,7 @@ class GeneratedAnalysisParityTestCase: expected_minimum_native: int expected_minimum_expression_shapes: int expected_minimum_pivot_proofs: int + expected_minimum_python_cte_recoveries: int expected_minimum_proven_pivots: int expected_maximum_enrichment_deferrals: int expected_minimum_native_enrichments: int @@ -49,3 +50,14 @@ class StandalonePivotProofTestCase: analysed_models: frozenset[str] expected_native_proofs: int expected_proven_by_model: dict[str, bool | None] + + +@dataclass(frozen=True) +class CteFactRecoveryParityTestCase: + """Seeded CTE queries whose native and Python CTE fact recoveries must agree.""" + + description: str + seed: int + count: int + expected_minimum_compared: int + expected_minimum_recovered: dict[str, int] diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py index 47095774c0..6c6a9d89be 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py @@ -15,6 +15,7 @@ import sqlbuild._native as native_module import sqlbuild.compiler.analysis_session.classes.native_model_analysis as native_model_analysis +import sqlbuild.compiler.compile._helpers.analysis.compact as compact_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 @@ -171,6 +172,8 @@ def _pivot_models(rng: random.Random) -> dict[str, str]: return files +_PYTHON_CTE_RECOVERY: Callable[..., Any] = compact_analysis._polyglot_cte_passthrough_facts + type _Relation = tuple[str, tuple[str, ...]] type _Template = Callable[[random.Random, list[_Relation]], tuple[str, str, tuple[str, ...]]] @@ -258,6 +261,35 @@ def _contract(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tu return header, sql, ("contract_key", "contract_label") +_CTE_CONTRACT_BODIES: tuple[str, ...] = ( + "WITH base AS (\n SELECT {a}, {b} FROM {left}\n)\nSELECT {a}, {b} FROM base", + "WITH base AS (\n SELECT CAST({a} AS VARCHAR(12)) AS label, COALESCE({b}, {a}) AS filled,\n" + " 'fixed' AS tag, CAST(1.5 AS DECIMAL(10, 2)) AS ratio, {a} IS NULL AS missing\n" + " FROM {left}\n WHERE {a} IS NOT NULL\n)\n" + "SELECT label, filled, tag, ratio, missing FROM base", + "WITH base AS (\n SELECT CASE WHEN {a} IS NULL THEN 'none' ELSE 'some' END AS state,\n" + " UPPER(CAST({b} AS VARCHAR)) AS upper_b, CAST({a} AS TEXT) || '-x' AS joined\n" + " FROM {left}\n)\nSELECT state, upper_b, joined FROM base", + "WITH unioned AS (\n SELECT {a} AS shared FROM {left}\n UNION ALL\n" + " SELECT {c} FROM {right}\n)\nSELECT shared FROM unioned", + "WITH base AS (\n SELECT * EXCLUDE ({a}) FROM {left}\n)\nSELECT {b} FROM base", + "WITH lhs AS (\n SELECT {a}, {b} FROM {left}\n),\nrhs AS (\n SELECT {c} FROM {right}\n)\n" + "SELECT l.{a}, r.{c}\nFROM lhs l\nLEFT JOIN rhs r ON l.{a} = r.{c}", + "WITH outer_cte AS (\n WITH inner_cte AS (SELECT {a}, {b} FROM {left})\n" + " SELECT {a}, {b} FROM inner_cte\n)\nSELECT {a}, {b} FROM outer_cte", + "WITH base AS (\n SELECT {a}, NULLIF({b}, {b}) AS cleared, MAX({b}) AS latest FROM {left}\n" + " GROUP BY {a}\n)\nSELECT {a}, cleared, latest FROM base WHERE {a} IS NOT NULL", +) + + +def _cte_contract(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tuple[str, ...]]: + (left, left_columns), (right, right_columns) = inputs[0], inputs[-1] + first, second = rng.sample(list(left_columns), k=2) + body: str = rng.choice(_CTE_CONTRACT_BODIES) + sql: str = body.format(a=first, b=second, c=right_columns[0], left=left, right=right) + return " contract enforced,\n", sql, (first, second) + + def _unknown_column( rng: random.Random, inputs: list[_Relation] ) -> tuple[str, str, tuple[str, ...]]: @@ -282,6 +314,7 @@ def _subquery(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tu _aggregate, _quoted, _contract, + _cte_contract, _unknown_column, _subquery, ) @@ -373,6 +406,7 @@ class AnalysisParity: analysed_models: int = 0 expression_shapes: int = 0 pivot_proofs: int = 0 + python_cte_recoveries: int = 0 standalone_proofs: int = 0 proven_pivots: int = 0 native_enrichments: int = 0 @@ -404,7 +438,13 @@ def models_by_both( dynamic_families_by_table=dynamic_families_by_table, ) ) - python: dict[str, ModelSqlAnalysis] = python_analysis.func(**keywords) + with monkeypatch.context() as patch: + patch.setattr( + compact_analysis, + "_polyglot_cte_passthrough_facts", + partial(_counted_cte_recovery, parity=parity), + ) + python: dict[str, ModelSqlAnalysis] = python_analysis.func(**keywords) parity.analysed_models += len(python) _append(parity, "model analyses", _analysis_views(python), _analysis_views(native)) _append( @@ -515,6 +555,12 @@ def deferral_kinds(record_dir: Path) -> Counter[str]: return Counter(f"{record['site']}:{record['kind']}" for record in records) +def _counted_cte_recovery(*, parity: AnalysisParity, **arguments: Any) -> Any: + recovered: Any = _PYTHON_CTE_RECOVERY(**arguments) + parity.python_cte_recoveries += bool(recovered[2]) + return recovered + + def _append(parity: AnalysisParity, name: str, python: object, native: object) -> None: parity.names.append(name) parity.python.append(python) diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_cte_fact_recovery.py b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_cte_fact_recovery.py new file mode 100644 index 0000000000..119e8e0c71 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_cte_fact_recovery.py @@ -0,0 +1,56 @@ +"""Native CTE fact recovery equals Python's recovery on seeded queries, query by query.""" + +from __future__ import annotations + +import random +from collections import Counter + +import pytest + +from tests.integration.src.sqlbuild.compiler.analysis_session._test_helpers import ( + CteFactParity, + compare_cte_facts, + generated_cte_fact_query, +) +from tests.integration.src.sqlbuild.compiler.analysis_session._test_types import ( + CteFactRecoveryParityTestCase, +) +from tests.integration.src.sqlbuild.compiler.helpers import mismatches + + +@pytest.mark.parametrize( + "test_case", + [ + CteFactRecoveryParityTestCase( + description="casts, CASE, COALESCE, IF/IFF, adapter rules, set operations, stars, joins, filters", + seed=20261009, + count=5000, + expected_minimum_compared=4319, + expected_minimum_recovered={ + "types": 1243, + "nullability": 417, + "direct": 859, + "non_null": 38, + }, + ) + ], + ids=lambda case: case.description, +) +def test_given_seeded_cte_queries_when_recovering_facts_natively_then_matches_python( + test_case: CteFactRecoveryParityTestCase, +) -> None: + rng: random.Random = random.Random(test_case.seed) + parity: CteFactParity = CteFactParity(queries=[], python=[], native=[], counts=Counter()) + for _ in range(test_case.count): + compare_cte_facts(query=generated_cte_fact_query(rng=rng), parity=parity) + + assert mismatches(inputs=parity.queries, expected=parity.python, actual=parity.native) == [] + assert parity.counts["compared"] >= test_case.expected_minimum_compared + assert { + kind: parity.counts[f"recovered_{kind}"] >= minimum + for kind, minimum in test_case.expected_minimum_recovered.items() + } == dict.fromkeys(test_case.expected_minimum_recovered, True) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) 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 4e9527ea64..f0d116fe9d 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 @@ -54,7 +54,7 @@ "test_case", [ GeneratedAnalysisParityTestCase( - description="stars, CTEs, set operations, untyped inputs, contracts, contract CTEs, pivots", + description="stars, CTEs, CTE facts, set operations, untyped inputs, contracts, pivots", seed=20261008, count=6, model_count=24, @@ -62,8 +62,9 @@ expected_minimum_native=500, expected_minimum_expression_shapes=40, expected_minimum_pivot_proofs=80, + expected_minimum_python_cte_recoveries=40, expected_minimum_proven_pivots=12, - expected_maximum_enrichment_deferrals=120, + expected_maximum_enrichment_deferrals=0, expected_minimum_native_enrichments=90, ) ], @@ -90,8 +91,11 @@ 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.python_cte_recoveries >= test_case.expected_minimum_python_cte_recoveries assert ( - parity.analysed_models - kinds["analysis_session:legacy_analysis"] + parity.analysed_models + - kinds["analysis_session:legacy_analysis"] + - kinds["analysis_session:input_enrichment"] >= test_case.expected_minimum_native ) assert parity.expression_shapes >= test_case.expected_minimum_expression_shapes