diff --git a/crates/sqlbuild-analysis/src/lineage/_helpers/mod.rs b/crates/sqlbuild-analysis/src/lineage/_helpers/mod.rs new file mode 100644 index 0000000000..b9f011aba7 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/_helpers/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod parsed_lineage; +pub(crate) mod projections; +pub(crate) mod python_payload; +pub(crate) mod references; +pub(crate) mod stars; diff --git a/crates/sqlbuild-analysis/src/lineage/_helpers/parsed_lineage.rs b/crates/sqlbuild-analysis/src/lineage/_helpers/parsed_lineage.rs new file mode 100644 index 0000000000..ad66de574c --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/_helpers/parsed_lineage.rs @@ -0,0 +1,231 @@ +//! Python's `_build_polyglot_fast_model_column_lineage`: the parse fallback without compact facts. + +use std::collections::{HashMap, HashSet}; + +use polyglot_sql::{ComplexityGuardOptions, Dialect, DialectType, Expression, ParseOptions}; + +use crate::lineage::_helpers::projections::{ + UnreadableExpression, classify_transform, is_star, projection_output_name, + projection_upstream_columns, set_expression_selects, table_alias_map, unaliased, +}; +use crate::lineage::_helpers::references::{PhysicalResource, normalized_sql, physical_resources}; +use crate::lineage::_helpers::stars::star_lineage; +use crate::lineage::constants::{KIND_SELECT, KIND_UNION, MAX_FUNCTION_CALL_DEPTH}; +use crate::lineage::models::{ + FastLineageOutcome, LineageColumn, LineageConfidence, LineageDeferral, LineageResourceType, + LineageSource, LineageTransformKind, +}; + +/// The parse options SQLBuild's Polyglot proxy gives `parse_one`, decoded as the wheel does. +pub(crate) fn proxy_parse_options() -> Result { + let guard: ComplexityGuardOptions = serde_json::from_value(serde_json::json!({ + "maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH, + }))?; + Ok(ParseOptions { + complexity_guard: Some(guard), + }) +} + +/// Python's `parse_one(dialect=dialect or "generic")` dialect, if this parser build carries it. +pub(crate) fn lineage_dialect(name: Option<&str>) -> Option { + let name = name.filter(|name| !name.is_empty()).unwrap_or("generic"); + let Ok(dialect) = name.parse::() else { + return None; + }; + matches!( + dialect, + DialectType::Generic + | DialectType::PostgreSQL + | DialectType::BigQuery + | DialectType::Snowflake + | DialectType::DuckDB + | DialectType::TSQL + | DialectType::Databricks + ) + .then_some(dialect) +} + +pub(crate) struct ParsedModel<'a> { + pub(crate) query_sql: &'a str, + pub(crate) inferred_names: &'a [String], + pub(crate) schema: &'a HashMap>, + pub(crate) dialect: Option, + pub(crate) options: &'a Result, +} + +pub(crate) fn parsed_model_lineage(model: &ParsedModel<'_>) -> FastLineageOutcome { + let Some(dialect) = model.dialect else { + return FastLineageOutcome::Deferred(LineageDeferral::UnsupportedDialect); + }; + let Ok(options) = model.options else { + return FastLineageOutcome::Deferred(LineageDeferral::NativeFailure); + }; + let physical = physical_resources(model.query_sql); + let sql = normalized_sql(model.query_sql); + let mut statements = match Dialect::get(dialect).parse_with_options(&sql, options) { + Ok(statements) => statements, + Err(error) => return FastLineageOutcome::Unparsed(error.to_string()), + }; + if statements.len() != 1 { + return FastLineageOutcome::Unparsed(format!( + "Expected 1 statement, found {}", + statements.len() + )); + } + let parsed = statements.remove(0); + let built = match parsed.variant_name() { + KIND_UNION => union_lineage(&parsed, model, &physical), + KIND_SELECT => select_lineage(&parsed, model, &physical), + _ => Ok(None), + }; + match built { + Ok(Some((columns, has_star))) => FastLineageOutcome::Built { columns, has_star }, + Ok(None) => FastLineageOutcome::Omitted, + Err(UnreadableExpression) => FastLineageOutcome::Deferred(LineageDeferral::NativeFailure), + } +} + +type BuiltLineage = Option<(Vec, bool)>; + +fn select_lineage( + parsed: &Expression, + model: &ParsedModel<'_>, + physical: &[PhysicalResource], +) -> Result { + let alias_map = table_alias_map(parsed, physical); + let unqualified = alias_map.single_resource(); + let mut lineages: Vec = Vec::new(); + let mut has_star = false; + for (index, projection) in parsed.get_expressions().iter().enumerate() { + if is_star(projection) { + has_star = true; + continue; + } + let inner = unaliased(projection)?; + if is_star(inner) { + has_star = true; + continue; + } + let Some(output_column) = projection_output_name(projection, index, model.inferred_names) + else { + continue; + }; + let (upstream_columns, confidence) = + projection_upstream_columns(projection, &alias_map, unqualified)?; + let transform_kind = classify_transform(inner, &upstream_columns); + let confidence = + if !upstream_columns.is_empty() || transform_kind == LineageTransformKind::Constant { + confidence + } else { + LineageConfidence::Unknown + }; + lineages.push(LineageColumn { + output_column, + transform_kind, + confidence, + upstream_columns, + }); + } + Ok(Some(with_star_columns(lineages, has_star, model, physical))) +} + +fn union_lineage( + parsed: &Expression, + model: &ParsedModel<'_>, + physical: &[PhysicalResource], +) -> Result { + let selects: Vec = set_expression_selects(parsed)?; + if selects.is_empty() { + return Ok(None); + } + let alias_maps: Vec<_> = selects + .iter() + .map(|select| table_alias_map(select, physical)) + .collect(); + let max_projection_count = selects + .iter() + .map(|select| select.get_expressions().len()) + .max() + .unwrap_or(0); + let mut lineages: Vec = Vec::new(); + let mut has_star = false; + for index in 0..max_projection_count { + let mut output_column = model.inferred_names.get(index).cloned(); + let mut upstream_columns: Vec = Vec::new(); + let mut seen: HashSet<(LineageResourceType, String, String)> = HashSet::new(); + let mut confidence = LineageConfidence::High; + let mut transform_kind = LineageTransformKind::Direct; + for (select, alias_map) in selects.iter().zip(&alias_maps) { + let Some(projection) = select.get_expressions().get(index) else { + continue; + }; + if is_star(projection) { + has_star = true; + continue; + } + let inner = unaliased(projection)?; + if is_star(inner) { + has_star = true; + continue; + } + if output_column.is_none() { + output_column = projection_output_name(projection, index, model.inferred_names); + } + let (branch_upstream, branch_confidence) = + projection_upstream_columns(projection, alias_map, alias_map.single_resource())?; + if branch_confidence == LineageConfidence::Unknown { + confidence = LineageConfidence::Unknown; + } else if branch_confidence == LineageConfidence::Medium + && confidence == LineageConfidence::High + { + confidence = LineageConfidence::Medium; + } + let branch_transform = classify_transform(inner, &branch_upstream); + if branch_transform != LineageTransformKind::Direct { + transform_kind = branch_transform; + } + for source in branch_upstream { + let key = ( + source.resource_type, + source.resource_name.clone(), + source.column_name.clone(), + ); + if seen.insert(key) { + upstream_columns.push(source); + } + } + } + let Some(output_column) = output_column else { + continue; + }; + if upstream_columns.is_empty() && transform_kind != LineageTransformKind::Constant { + confidence = LineageConfidence::Unknown; + } + lineages.push(LineageColumn { + output_column, + transform_kind, + confidence, + upstream_columns, + }); + } + Ok(Some(with_star_columns(lineages, has_star, model, physical))) +} + +fn with_star_columns( + mut lineages: Vec, + has_star: bool, + model: &ParsedModel<'_>, + physical: &[PhysicalResource], +) -> (Vec, bool) { + if has_star { + let star_columns = star_lineage( + model.schema, + physical, + lineages + .iter() + .map(|lineage| lineage.output_column.as_str()), + ); + lineages.extend(star_columns); + } + (lineages, has_star) +} diff --git a/crates/sqlbuild-analysis/src/lineage/_helpers/projections.rs b/crates/sqlbuild-analysis/src/lineage/_helpers/projections.rs new file mode 100644 index 0000000000..d13bb561dc --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/_helpers/projections.rs @@ -0,0 +1,315 @@ +//! Python's fast lineage reads of wheel expression accessors, over the same crate expressions. + +use std::collections::{HashMap, HashSet}; + +use polyglot_sql::{Expression, ExpressionWalk, ast_json}; +use serde_json::Value; + +use crate::lineage::_helpers::python_payload::{PythonValue, python_payload}; +use crate::lineage::_helpers::references::PhysicalResource; +use crate::lineage::constants::{ + AGGREGATE_KINDS, CAST_KINDS, KIND_ALIAS, KIND_COLUMN, KIND_SELECT, KIND_TABLE, PAYLOAD_COLUMN, + PAYLOAD_NAME, PAYLOAD_TABLE, SET_OPERATION_KINDS, SET_OPERATION_SIDES, STAR_COLUMN_NAME, +}; +use crate::lineage::models::{ + LineageConfidence, LineageResourceType, LineageSource, LineageTransformKind, +}; + +/// Payload nesting past which Python's recursive `to_dict` walk may exceed its recursion limit. +const MAX_PAYLOAD_DEPTH: usize = 400; + +/// A read Python would make differently, or fail on, so the model is deferred to Python. +#[derive(Debug)] +pub(crate) struct UnreadableExpression; + +/// Table names and aliases mapped to resources, in Python's dictionary insertion order. +pub(crate) struct AliasMap<'a> { + entries: Vec<(String, &'a PhysicalResource)>, +} + +impl<'a> AliasMap<'a> { + fn insert(&mut self, key: &str, resource: &'a PhysicalResource) { + match self.entries.iter_mut().find(|(name, _)| name == key) { + Some(entry) => entry.1 = resource, + None => self.entries.push((key.to_owned(), resource)), + } + } + + fn get(&self, key: &str) -> Option<&'a PhysicalResource> { + self.entries + .iter() + .find(|(name, _)| name == key) + .map(|(_, resource)| *resource) + } + + /// Python's `_single_alias_resource`: the one resource every entry names, by resource name. + pub(crate) fn single_resource(&self) -> Option<&'a PhysicalResource> { + let mut values = self.entries.iter().map(|(_, resource)| *resource); + let first = values.next()?; + values + .all(|candidate| candidate.resource_name == first.resource_name) + .then_some(first) + } +} + +pub(crate) fn is_star(expression: &Expression) -> bool { + match expression { + Expression::Star(_) => true, + Expression::Column(column) => column.name.name == STAR_COLUMN_NAME, + _ => false, + } +} + +/// `projection.this` for an alias, the projection itself otherwise. +pub(crate) fn unaliased(projection: &Expression) -> Result<&Expression, UnreadableExpression> { + if projection.variant_name() == KIND_ALIAS { + return projection.get_this().ok_or(UnreadableExpression); + } + Ok(projection) +} + +/// Python's `_polyglot_projection_output_name`. +pub(crate) fn projection_output_name( + projection: &Expression, + index: usize, + inferred_names: &[String], +) -> Option { + let output_name = projection.get_output_name(); + if !output_name.is_empty() && output_name != STAR_COLUMN_NAME { + return Some(output_name.to_owned()); + } + let alias = projection.get_alias(); + let alias_or_name = if alias.is_empty() { + projection.get_name() + } else { + alias + }; + if !alias_or_name.is_empty() && alias_or_name != STAR_COLUMN_NAME { + return Some(alias_or_name.to_owned()); + } + inferred_names.get(index).cloned() +} + +/// Python's `_polyglot_table_alias_map` over `find_all("table")`. +pub(crate) fn table_alias_map<'a>( + expression: &Expression, + physical_resources: &'a [PhysicalResource], +) -> AliasMap<'a> { + let mut by_name: HashMap<&str, &'a PhysicalResource> = HashMap::new(); + for resource in physical_resources { + let _ = by_name.insert(resource.physical_name.as_str(), resource); + } + for resource in physical_resources { + let _ = by_name.insert(resource.resource_name.as_str(), resource); + } + let mut alias_map = AliasMap { + entries: Vec::new(), + }; + for table in expression + .dfs() + .skip(1) + .filter(|node| node.variant_name() == KIND_TABLE) + { + let table_name = table.get_name(); + let Some(resource) = by_name.get(table_name) else { + continue; + }; + alias_map.insert(table_name, resource); + let alias = table.get_alias(); + let alias_or_name = if alias.is_empty() { table_name } else { alias }; + if !alias_or_name.is_empty() { + alias_map.insert(alias_or_name, resource); + } + } + alias_map +} + +/// Python's `_polyglot_projection_upstream_columns`. +pub(crate) fn projection_upstream_columns( + projection: &Expression, + alias_map: &AliasMap<'_>, + unqualified_resource: Option<&PhysicalResource>, +) -> Result<(Vec, LineageConfidence), UnreadableExpression> { + let mut columns: Vec = Vec::new(); + let mut seen: HashSet<(LineageResourceType, String, String)> = HashSet::new(); + let mut confidence = LineageConfidence::High; + for (column_name, table_name) in column_refs(projection)? { + if column_name.is_empty() { + continue; + } + let resource = if !table_name.is_empty() { + alias_map.get(&table_name) + } else if let Some(resource) = unqualified_resource { + confidence = LineageConfidence::Medium; + Some(resource) + } else { + confidence = LineageConfidence::Unknown; + None + }; + let Some(resource) = resource else { + continue; + }; + let key = ( + resource.resource_type, + resource.resource_name.clone(), + column_name, + ); + if seen.contains(&key) { + continue; + } + columns.push(LineageSource { + resource_type: key.0, + resource_name: key.1.clone(), + column_name: key.2.clone(), + }); + let _ = seen.insert(key); + } + Ok((columns, confidence)) +} + +/// Python's `_polyglot_classify_transform`. +pub(crate) fn classify_transform( + expression: &Expression, + upstream_columns: &[LineageSource], +) -> LineageTransformKind { + let kind = expression.variant_name(); + if is_star(expression) { + return LineageTransformKind::Star; + } + if CAST_KINDS.contains(&kind) { + return LineageTransformKind::Cast; + } + if expression + .dfs() + .any(|node| AGGREGATE_KINDS.contains(&node.variant_name())) + { + return LineageTransformKind::Aggregation; + } + if upstream_columns.is_empty() { + return LineageTransformKind::Constant; + } + if kind == KIND_COLUMN && upstream_columns.len() == 1 { + return LineageTransformKind::Direct; + } + LineageTransformKind::Expression +} + +/// Python's `_polyglot_set_expression_selects`, serialized once and walked iteratively, left first. +pub(crate) fn set_expression_selects( + expression: &Expression, +) -> Result, UnreadableExpression> { + let kind = expression.variant_name(); + if kind == KIND_SELECT { + return Ok(vec![expression.clone()]); + } + let mut selects: Vec = Vec::new(); + if !SET_OPERATION_KINDS.contains(&kind) { + return Ok(selects); + } + let payload: Value = serde_json::to_value(expression).map_err(|_| UnreadableExpression)?; + let mut pending: Vec<&Value> = vec![&payload]; + while let Some(node) = pending.pop() { + if set_operation_kind(node).is_some() { + let Some(Value::Object(arguments)) = payload_arguments(node) else { + continue; + }; + for side in SET_OPERATION_SIDES.iter().rev() { + if let Some(child) = arguments.get(*side) { + pending.push(child); + } + } + continue; + } + let Ok(child) = ast_json::expression_from_value(node.clone()) else { + continue; + }; + if child.variant_name() == KIND_SELECT { + selects.push(child); + } else if SET_OPERATION_KINDS.contains(&child.variant_name()) { + return Err(UnreadableExpression); + } + } + Ok(selects) +} + +/// The set-operation kind a serialized expression is tagged with, if any. +fn set_operation_kind(value: &Value) -> Option<&'static str> { + let Value::Object(map) = value else { + return None; + }; + let (key, _) = map.iter().next().filter(|_| map.len() == 1)?; + SET_OPERATION_KINDS + .iter() + .copied() + .find(|kind| *kind == key.as_str()) +} + +/// The wheel's `args`: the payload under an expression's single variant key. +fn payload_arguments(payload: &Value) -> Option<&Value> { + match payload { + Value::Object(map) => map.values().next(), + _ => None, + } +} + +/// Python's `_polyglot_column_refs_in_expression`: `(column, table)` pairs in payload order. +fn column_refs(expression: &Expression) -> Result, UnreadableExpression> { + let payload: PythonValue = python_payload(expression).map_err(|_| UnreadableExpression)?; + if expression.variant_name() == KIND_COLUMN { + let table = payload + .get(PAYLOAD_COLUMN) + .and_then(|column| column.get(PAYLOAD_TABLE)); + let table_name = match table.and_then(|table| table.get(PAYLOAD_NAME)) { + Some(PythonValue::Str(name)) => name.clone(), + _ => String::new(), + }; + return Ok(vec![(expression.get_name().to_owned(), table_name)]); + } + payload_column_refs(&payload, 0) +} + +/// Python's recursive `visit`: a `column` dict ends its branch; tuples are never entered. +fn payload_column_refs( + node: &PythonValue, + depth: usize, +) -> Result, UnreadableExpression> { + if depth > MAX_PAYLOAD_DEPTH { + return Err(UnreadableExpression); + } + let children: Vec<&PythonValue> = match node { + PythonValue::Dict(entries) => { + if let Some(column @ PythonValue::Dict(_)) = node.get(PAYLOAD_COLUMN) { + let table_name = match column.get(PAYLOAD_TABLE) { + Some(table @ PythonValue::Dict(_)) => { + name_payload_value(table.get(PAYLOAD_NAME)) + } + _ => String::new(), + }; + return Ok(vec![( + name_payload_value(column.get(PAYLOAD_NAME)), + table_name, + )]); + } + entries.iter().map(|(_, value)| value).collect() + } + PythonValue::List(values) => values.iter().collect(), + _ => Vec::new(), + }; + let mut refs: Vec<(String, String)> = Vec::new(); + for child in children { + refs.extend(payload_column_refs(child, depth + 1)?); + } + Ok(refs) +} + +/// Python's `_polyglot_name_payload_value`. +fn name_payload_value(payload: Option<&PythonValue>) -> String { + match payload { + Some(PythonValue::Str(name)) => name.clone(), + Some(map @ PythonValue::Dict(_)) => match map.get(PAYLOAD_NAME) { + Some(PythonValue::Str(name)) => name.clone(), + _ => String::new(), + }, + _ => String::new(), + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/_helpers/python_payload.rs b/crates/sqlbuild-analysis/src/lineage/_helpers/python_payload.rs new file mode 100644 index 0000000000..d30247b6c5 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/_helpers/python_payload.rs @@ -0,0 +1,340 @@ +//! The shape the wheel's `to_dict` gives an expression: pythonize's dicts, lists and tuples. + +use serde::ser::{ + Serialize, SerializeMap, SerializeSeq, SerializeStruct, SerializeStructVariant, SerializeTuple, + SerializeTupleStruct, SerializeTupleVariant, Serializer, +}; + +use crate::lineage::errors::PayloadError; + +/// A Python value as pythonize builds it, keeping only what the lineage walk reads. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum PythonValue { + /// Entries in insertion order; `None` keys are not strings. + Dict(Vec<(Option, PythonValue)>), + List(Vec), + /// Serde tuples, tuple structs and tuple variant bodies, which pythonize makes Python tuples. + Tuple(Vec), + Str(String), + Scalar, +} + +impl PythonValue { + /// Python's `dict.get` for a string key. + pub(crate) fn get(&self, key: &str) -> Option<&PythonValue> { + let PythonValue::Dict(entries) = self else { + return None; + }; + entries + .iter() + .find(|(name, _)| name.as_deref() == Some(key)) + .map(|(_, value)| value) + } +} + +/// The value pythonize would build for `value`. +pub(crate) fn python_payload( + value: &T, +) -> Result { + value.serialize(PayloadSerializer) +} + +fn variant(name: &str, value: PythonValue) -> PythonValue { + PythonValue::Dict(vec![(Some(name.to_owned()), value)]) +} + +struct PayloadSerializer; + +struct Collection { + tuple: bool, + items: Vec, +} + +impl Collection { + fn finish(self) -> PythonValue { + if self.tuple { + PythonValue::Tuple(self.items) + } else { + PythonValue::List(self.items) + } + } +} + +struct TupleVariant { + name: &'static str, + items: Vec, +} + +struct Dict { + entries: Vec<(Option, PythonValue)>, + key: Option>, +} + +impl Dict { + fn with_capacity(length: usize) -> Self { + Self { + entries: Vec::with_capacity(length), + key: None, + } + } + + /// Python's dictionary assignment: an existing string key keeps its position. + fn assign(&mut self, key: Option, value: PythonValue) { + let existing: Option = key.as_deref().and_then(|name| { + self.entries + .iter() + .position(|(other, _)| other.as_deref() == Some(name)) + }); + match existing { + Some(index) => self.entries[index].1 = value, + None => self.entries.push((key, value)), + } + } +} + +struct StructVariant { + name: &'static str, + dict: Dict, +} + +impl Serializer for PayloadSerializer { + type Ok = PythonValue; + type Error = PayloadError; + type SerializeSeq = Collection; + type SerializeTuple = Collection; + type SerializeTupleStruct = Collection; + type SerializeTupleVariant = TupleVariant; + type SerializeMap = Dict; + type SerializeStruct = Dict; + type SerializeStructVariant = StructVariant; + + fn serialize_bool(self, _: bool) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_i8(self, _: i8) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_i16(self, _: i16) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_i32(self, _: i32) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_i64(self, _: i64) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_u8(self, _: u8) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_u16(self, _: u16) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_u32(self, _: u32) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_u64(self, _: u64) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_f32(self, _: f32) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_f64(self, _: f64) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_char(self, value: char) -> Result { + Ok(PythonValue::Str(value.to_string())) + } + fn serialize_str(self, value: &str) -> Result { + Ok(PythonValue::Str(value.to_owned())) + } + fn serialize_bytes(self, _: &[u8]) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_none(self) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_some(self, value: &T) -> Result { + value.serialize(self) + } + fn serialize_unit(self) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_unit_struct(self, _: &'static str) -> Result { + Ok(PythonValue::Scalar) + } + fn serialize_unit_variant( + self, + _: &'static str, + _: u32, + name: &'static str, + ) -> Result { + Ok(PythonValue::Str(name.to_owned())) + } + fn serialize_newtype_struct( + self, + _: &'static str, + value: &T, + ) -> Result { + value.serialize(self) + } + fn serialize_newtype_variant( + self, + _: &'static str, + _: u32, + name: &'static str, + value: &T, + ) -> Result { + Ok(variant(name, value.serialize(PayloadSerializer)?)) + } + fn serialize_seq(self, length: Option) -> Result { + Ok(Collection { + tuple: false, + items: Vec::with_capacity(length.unwrap_or(0)), + }) + } + fn serialize_tuple(self, length: usize) -> Result { + Ok(Collection { + tuple: true, + items: Vec::with_capacity(length), + }) + } + fn serialize_tuple_struct( + self, + _: &'static str, + length: usize, + ) -> Result { + self.serialize_tuple(length) + } + fn serialize_tuple_variant( + self, + _: &'static str, + _: u32, + name: &'static str, + length: usize, + ) -> Result { + Ok(TupleVariant { + name, + items: Vec::with_capacity(length), + }) + } + fn serialize_map(self, length: Option) -> Result { + Ok(Dict::with_capacity(length.unwrap_or(0))) + } + fn serialize_struct(self, _: &'static str, length: usize) -> Result { + Ok(Dict::with_capacity(length)) + } + fn serialize_struct_variant( + self, + _: &'static str, + _: u32, + name: &'static str, + length: usize, + ) -> Result { + Ok(StructVariant { + name, + dict: Dict::with_capacity(length), + }) + } +} + +impl SerializeSeq for Collection { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_element(&mut self, value: &T) -> Result<(), PayloadError> { + self.items.push(value.serialize(PayloadSerializer)?); + Ok(()) + } + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeTuple for Collection { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_element(&mut self, value: &T) -> Result<(), PayloadError> { + SerializeSeq::serialize_element(self, value) + } + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeTupleStruct for Collection { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_field(&mut self, value: &T) -> Result<(), PayloadError> { + SerializeSeq::serialize_element(self, value) + } + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeTupleVariant for TupleVariant { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_field(&mut self, value: &T) -> Result<(), PayloadError> { + self.items.push(value.serialize(PayloadSerializer)?); + Ok(()) + } + fn end(self) -> Result { + Ok(variant(self.name, PythonValue::Tuple(self.items))) + } +} + +impl SerializeMap for Dict { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_key(&mut self, key: &T) -> Result<(), PayloadError> { + let key = match key.serialize(PayloadSerializer)? { + PythonValue::Str(name) => Some(name), + _ => None, + }; + self.key = Some(key); + Ok(()) + } + fn serialize_value(&mut self, value: &T) -> Result<(), PayloadError> { + let Some(key) = self.key.take() else { + return Err(PayloadError::new("map value without a key")); + }; + self.assign(key, value.serialize(PayloadSerializer)?); + Ok(()) + } + fn end(self) -> Result { + Ok(PythonValue::Dict(self.entries)) + } +} + +impl SerializeStruct for Dict { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_field( + &mut self, + key: &'static str, + value: &T, + ) -> Result<(), PayloadError> { + self.assign(Some(key.to_owned()), value.serialize(PayloadSerializer)?); + Ok(()) + } + fn end(self) -> Result { + Ok(PythonValue::Dict(self.entries)) + } +} + +impl SerializeStructVariant for StructVariant { + type Ok = PythonValue; + type Error = PayloadError; + fn serialize_field( + &mut self, + key: &'static str, + value: &T, + ) -> Result<(), PayloadError> { + self.dict + .assign(Some(key.to_owned()), value.serialize(PayloadSerializer)?); + Ok(()) + } + fn end(self) -> Result { + Ok(variant(self.name, PythonValue::Dict(self.dict.entries))) + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/_helpers/references.rs b/crates/sqlbuild-analysis/src/lineage/_helpers/references.rs new file mode 100644 index 0000000000..7ce7e43047 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/_helpers/references.rs @@ -0,0 +1,149 @@ +//! Python's `_normalize_sqlbuild_refs`, matching its backreferencing patterns exactly. + +use sqlbuild_core::text::main::is_python_space::is_python_space; + +use crate::lineage::constants::{ + PHYSICAL_NAME_REPLACEMENT, PHYSICAL_RESOURCE_PREFIX, REFERENCE_PREFIX, UDF_CALL, +}; +use crate::lineage::models::LineageResourceType; + +/// A physical SQL identifier mapped to its SQLBuild resource. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PhysicalResource { + pub(crate) resource_type: LineageResourceType, + pub(crate) resource_name: String, + pub(crate) physical_name: String, +} + +const REFERENCE_KINDS: [(&str, LineageResourceType); 3] = [ + ("ref(", LineageResourceType::Model), + ("source(", LineageResourceType::Source), + ("seed(", LineageResourceType::Seed), +]; + +/// The resource references in `sql`, in match order, repeats included. +pub(crate) fn physical_resources(sql: &str) -> Vec { + let mut resources: Vec = Vec::new(); + scan_calls(sql, reference_call, |_, (resource_type, name)| { + resources.push(PhysicalResource { + resource_type, + resource_name: name.to_owned(), + physical_name: physical_resource_name(resource_type, name), + }); + }); + resources +} + +/// `sql` with every reference and UDF call replaced by its physical name. +pub(crate) fn normalized_sql(sql: &str) -> String { + let references = replace_calls(sql, reference_call, |(resource_type, name)| { + physical_resource_name(resource_type, name) + }); + replace_calls(&references, udf_call, |name| { + format!("{}(", physical_name_text(name)) + }) +} + +/// Python's `_physical_resource_name`. +pub(crate) fn physical_resource_name(resource_type: LineageResourceType, name: &str) -> String { + format!( + "{PHYSICAL_RESOURCE_PREFIX}{}__{}", + resource_type.as_str(), + physical_name_text(name) + ) +} + +fn physical_name_text(name: &str) -> String { + let mut text = String::with_capacity(name.len()); + for character in name.chars() { + if character.is_ascii_alphanumeric() || character == '_' { + text.push(character); + } else { + text.push_str(PHYSICAL_NAME_REPLACEMENT); + } + } + text +} + +fn replace_calls<'a, T>( + sql: &'a str, + matcher: fn(&'a str) -> Option<(usize, T)>, + mut replacement: impl FnMut(T) -> String, +) -> String { + let mut text = String::with_capacity(sql.len()); + let mut copied = 0; + scan_calls(sql, matcher, |span, value| { + text.push_str(&sql[copied..span.0]); + text.push_str(&replacement(value)); + copied = span.1; + }); + text.push_str(&sql[copied..]); + text +} + +/// Python's `finditer`: non-overlapping matches, left to right. Every match starts with `__`. +fn scan_calls<'a, T>( + sql: &'a str, + matcher: fn(&'a str) -> Option<(usize, T)>, + mut on_match: impl FnMut((usize, usize), T), +) { + let mut position = 0; + while let Some(offset) = sql[position..].find(REFERENCE_PREFIX) { + let start = position + offset; + match matcher(&sql[start..]) { + Some((length, value)) => { + on_match((start, start + length), value); + position = start + length; + } + None => position = start + 1, + } + } +} + +/// `__(ref|source|seed)\(\s*(['"])(?P[^'"]+)\2\s*\)` at the start of `text`. +fn reference_call(text: &str) -> Option<(usize, (LineageResourceType, &str))> { + let rest = &text[REFERENCE_PREFIX.len()..]; + let (keyword, resource_type) = REFERENCE_KINDS + .iter() + .find(|(keyword, _)| rest.starts_with(keyword))?; + let opened = REFERENCE_PREFIX.len() + keyword.len(); + let (end, name) = quoted_argument(text, opened)?; + Some((end, (*resource_type, name))) +} + +/// `__udf\(\s*(['"])(?P[^'"]+)\1\s*\)\s*\(` at the start of `text`. +fn udf_call(text: &str) -> Option<(usize, &str)> { + if !text.starts_with(UDF_CALL) { + return None; + } + let (closed, name) = quoted_argument(text, UDF_CALL.len())?; + let opened = skip_spaces(text, closed); + text[opened..] + .starts_with('(') + .then_some((opened + 1, name)) +} + +/// `\s*(['"])(?P[^'"]+)\1\s*\)` from `start`; returns the end after `)` and the name. +fn quoted_argument(text: &str, start: usize) -> Option<(usize, &str)> { + let quote_start = skip_spaces(text, start); + let quote = text[quote_start..].chars().next()?; + if quote != '\'' && quote != '"' { + return None; + } + let name_start = quote_start + 1; + let name_length = text[name_start..].find(['\'', '"'])?; + if name_length == 0 || !text[name_start + name_length..].starts_with(quote) { + return None; + } + let close = skip_spaces(text, name_start + name_length + 1); + text[close..] + .starts_with(')') + .then_some((close + 1, &text[name_start..name_start + name_length])) +} + +fn skip_spaces(text: &str, start: usize) -> usize { + text[start..] + .char_indices() + .find(|(_, character)| !is_python_space(*character)) + .map_or(text.len(), |(offset, _)| start + offset) +} diff --git a/crates/sqlbuild-analysis/src/lineage/_helpers/stars.rs b/crates/sqlbuild-analysis/src/lineage/_helpers/stars.rs new file mode 100644 index 0000000000..8d8778a51b --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/_helpers/stars.rs @@ -0,0 +1,57 @@ +//! Python's `_build_schema_mapping` and `_build_star_lineage`. + +use std::collections::{HashMap, HashSet}; + +use crate::lineage::_helpers::references::{PhysicalResource, physical_resource_name}; +use crate::lineage::models::{ + LineageColumn, LineageConfidence, LineageSchemaResource, LineageSource, LineageTransformKind, +}; + +/// Known columns by physical name; a later resource with columns replaces an earlier one. +pub(crate) fn schema_mapping(resources: &[LineageSchemaResource]) -> HashMap> { + let mut schema = HashMap::with_capacity(resources.len()); + for resource in resources { + let mut seen = HashSet::with_capacity(resource.columns.len()); + let columns: Vec = resource + .columns + .iter() + .filter(|column| seen.insert(column.as_str())) + .cloned() + .collect(); + if !columns.is_empty() { + let _ = schema.insert( + physical_resource_name(resource.resource_type, &resource.name), + columns, + ); + } + } + schema +} + +/// Star columns for every referenced resource's known columns not already output. +pub(crate) fn star_lineage<'a>( + schema: &HashMap>, + physical_resources: &[PhysicalResource], + existing_columns: impl IntoIterator, +) -> Vec { + let mut seen: HashSet = existing_columns.into_iter().map(str::to_owned).collect(); + let mut lineages: Vec = Vec::new(); + for resource in physical_resources { + for column_name in schema.get(&resource.physical_name).into_iter().flatten() { + if !seen.insert(column_name.clone()) { + continue; + } + lineages.push(LineageColumn { + output_column: column_name.clone(), + transform_kind: LineageTransformKind::Star, + confidence: LineageConfidence::Medium, + upstream_columns: vec![LineageSource { + resource_type: resource.resource_type, + resource_name: resource.resource_name.clone(), + column_name: column_name.clone(), + }], + }); + } + } + lineages +} diff --git a/crates/sqlbuild-analysis/src/lineage/constants.rs b/crates/sqlbuild-analysis/src/lineage/constants.rs new file mode 100644 index 0000000000..10e8c04e7b --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/constants.rs @@ -0,0 +1,29 @@ +//! Polyglot expression kinds and SQLBuild reference spellings read by fast lineage. + +pub(crate) const STAR_COLUMN_NAME: &str = "*"; +pub(crate) const KIND_ALIAS: &str = "alias"; +pub(crate) const KIND_COLUMN: &str = "column"; +pub(crate) const KIND_SELECT: &str = "select"; +pub(crate) const KIND_TABLE: &str = "table"; +pub(crate) const KIND_UNION: &str = "union"; +pub(crate) const SET_OPERATION_KINDS: &[&str] = &["union", "intersect", "except"]; +pub(crate) const CAST_KINDS: &[&str] = &["cast", "try_cast"]; +pub(crate) const AGGREGATE_KINDS: &[&str] = &[ + "avg", + "count", + "max", + "min", + "sum", + "array_agg", + "string_agg", +]; +pub(crate) const PAYLOAD_COLUMN: &str = "column"; +pub(crate) const PAYLOAD_NAME: &str = "name"; +pub(crate) const PAYLOAD_TABLE: &str = "table"; +pub(crate) const SET_OPERATION_SIDES: [&str; 2] = ["left", "right"]; +/// The function-call depth SQLBuild's Polyglot proxy passes on every parse. +pub(crate) const MAX_FUNCTION_CALL_DEPTH: usize = 512; +pub(crate) const REFERENCE_PREFIX: &str = "__"; +pub(crate) const UDF_CALL: &str = "__udf("; +pub(crate) const PHYSICAL_RESOURCE_PREFIX: &str = "__sqlbuild_"; +pub(crate) const PHYSICAL_NAME_REPLACEMENT: &str = "__"; diff --git a/crates/sqlbuild-analysis/src/lineage/errors.rs b/crates/sqlbuild-analysis/src/lineage/errors.rs new file mode 100644 index 0000000000..677fe5e72f --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/errors.rs @@ -0,0 +1,29 @@ +//! Errors raised while reading expressions for fast lineage. + +use std::fmt::{self, Display}; + +use serde::ser; + +/// An expression pythonize could not have serialized. +#[derive(Debug)] +pub(crate) struct PayloadError(String); + +impl PayloadError { + pub(crate) fn new(message: &str) -> Self { + Self(message.to_owned()) + } +} + +impl Display for PayloadError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.0) + } +} + +impl std::error::Error for PayloadError {} + +impl ser::Error for PayloadError { + fn custom(message: T) -> Self { + Self(message.to_string()) + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/main/build_fast_lineage.rs b/crates/sqlbuild-analysis/src/lineage/main/build_fast_lineage.rs new file mode 100644 index 0000000000..f1266f2eff --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/main/build_fast_lineage.rs @@ -0,0 +1,62 @@ +//! Fast lineage for unresolved root stars and models without compact facts. + +use std::sync::Arc; + +use polyglot_sql::ParseOptions; +use rayon::ThreadPool; +use rayon::prelude::{IntoParallelRefIterator, ParallelIterator}; +use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic; + +use crate::lineage::_helpers::parsed_lineage::{ + ParsedModel, lineage_dialect, parsed_model_lineage, proxy_parse_options, +}; +use crate::lineage::_helpers::references::physical_resources; +use crate::lineage::_helpers::stars::{schema_mapping, star_lineage}; +use crate::lineage::models::{ + FastLineageModel, FastLineageOutcome, FastLineageRequest, LineageDeferral, +}; +use crate::semantic_validation::models::ProjectCatalog; + +/// One outcome per model, in request order, run on the compile's deep-stack analysis pool. +pub fn build_fast_lineage( + request: &FastLineageRequest, + catalog: &ProjectCatalog, +) -> Result, String> { + let pool: Arc = catalog.analysis_pool()?; + Ok(pool.install(|| model_outcomes(request))) +} + +/// A parser panic defers only its own model. +fn model_outcomes(request: &FastLineageRequest) -> Vec { + let schema = schema_mapping(&request.schema); + let dialect = lineage_dialect(request.dialect.as_deref()); + let options: Result = proxy_parse_options(); + request + .models + .par_iter() + .map(|model| match model { + FastLineageModel::StarExpansion { + query_sql, + existing_columns, + } => FastLineageOutcome::StarColumns(star_lineage( + &schema, + &physical_resources(query_sql), + existing_columns.iter().map(String::as_str), + )), + FastLineageModel::Parse { + query_sql, + inferred_columns, + } => { + let model = ParsedModel { + query_sql, + inferred_names: inferred_columns, + schema: &schema, + dialect, + options: &options, + }; + catch_compiler_panic(|| Ok(parsed_model_lineage(&model))) + .unwrap_or(FastLineageOutcome::Deferred(LineageDeferral::NativeFailure)) + } + }) + .collect() +} diff --git a/crates/sqlbuild-analysis/src/lineage/main/mod.rs b/crates/sqlbuild-analysis/src/lineage/main/mod.rs new file mode 100644 index 0000000000..f1d5cd565b --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/main/mod.rs @@ -0,0 +1 @@ +pub mod build_fast_lineage; diff --git a/crates/sqlbuild-analysis/src/lineage/mod.rs b/crates/sqlbuild-analysis/src/lineage/mod.rs index 9819e31334..6eb4e57e06 100644 --- a/crates/sqlbuild-analysis/src/lineage/mod.rs +++ b/crates/sqlbuild-analysis/src/lineage/mod.rs @@ -1 +1,10 @@ //! Fast column lineage facts for compiled models. + +pub(crate) mod _helpers; +pub(crate) mod constants; +pub(crate) mod errors; +pub mod main; +pub mod models; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-analysis/src/lineage/models.rs b/crates/sqlbuild-analysis/src/lineage/models.rs new file mode 100644 index 0000000000..7076c5dcaa --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/models.rs @@ -0,0 +1,171 @@ +//! Plain-data request and outcomes for fast column lineage. + +/// The SQLBuild resource kinds a `__ref`, `__source` or `__seed` call names. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum LineageResourceType { + Model, + Source, + Seed, +} + +impl LineageResourceType { + /// Python's `CompiledResourceType` value. + pub fn as_str(self) -> &'static str { + match self { + Self::Model => "model", + Self::Source => "source", + Self::Seed => "seed", + } + } + + /// Parse Python's `CompiledResourceType` value for the kinds lineage reads. + pub fn from_value(value: &str) -> Option { + match value { + "model" => Some(Self::Model), + "source" => Some(Self::Source), + "seed" => Some(Self::Seed), + _ => None, + } + } +} + +/// Python's `ColumnTransformKind` values that fast lineage produces. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LineageTransformKind { + Direct, + Cast, + Expression, + Aggregation, + Star, + Constant, +} + +impl LineageTransformKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Direct => "direct", + Self::Cast => "cast", + Self::Expression => "expression", + Self::Aggregation => "aggregation", + Self::Star => "star", + Self::Constant => "constant", + } + } +} + +/// Python's `ColumnLineageConfidence` values that fast lineage produces. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LineageConfidence { + High, + Medium, + Unknown, +} + +impl LineageConfidence { + pub fn as_str(self) -> &'static str { + match self { + Self::High => "high", + Self::Medium => "medium", + Self::Unknown => "unknown", + } + } +} + +/// One resource's known columns, in the order Python's schema mapping lists them. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LineageSchemaResource { + pub resource_type: LineageResourceType, + pub name: String, + /// Inferred then declared column names; repeats keep their first position. + pub columns: Vec, +} + +/// The work one model needs: star expansion over compact facts, or a parse of its SQL. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FastLineageModel { + /// A model with compact facts and an unresolved root star. + StarExpansion { + query_sql: String, + existing_columns: Vec, + }, + /// A model without compact facts. + Parse { + query_sql: String, + inferred_columns: Vec, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FastLineageRequest { + /// The adapter's analysis dialect; `None` parses as generic SQL. + pub dialect: Option, + pub schema: Vec, + pub models: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LineageSource { + pub resource_type: LineageResourceType, + pub resource_name: String, + pub column_name: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LineageColumn { + pub output_column: String, + pub transform_kind: LineageTransformKind, + pub confidence: LineageConfidence, + pub upstream_columns: Vec, +} + +/// Why a model is handed back to Python, which builds its lineage exactly as before. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LineageDeferral { + /// The dialect is not compiled into the native parser. + UnsupportedDialect, + /// The parsed SQL could not be read as Python reads it. + NativeFailure, +} + +impl LineageDeferral { + pub fn as_str(self) -> &'static str { + match self { + Self::UnsupportedDialect => "unsupported_dialect", + Self::NativeFailure => "native_failure", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FastLineageOutcome { + /// Star columns to append to the model's compact facts. + StarColumns(Vec), + /// Lineage built from the parsed SQL. + Built { + columns: Vec, + has_star: bool, + }, + /// Python records no lineage for this model. + Omitted, + /// The SQL did not parse; Python logs this message and records no lineage. + Unparsed(String), + Deferred(LineageDeferral), +} + +impl FastLineageOutcome { + /// `(status, columns, has_star, detail)`: detail is the parse error or the deferral kind. + pub fn into_parts(self) -> (&'static str, Vec, bool, Option) { + match self { + Self::StarColumns(columns) => ("star", columns, true, None), + Self::Built { columns, has_star } => ("built", columns, has_star, None), + Self::Omitted => ("omitted", Vec::new(), false, None), + Self::Unparsed(message) => ("unparsed", Vec::new(), false, Some(message)), + Self::Deferred(kind) => ( + "deferred", + Vec::new(), + false, + Some(kind.as_str().to_owned()), + ), + } + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/tests.rs b/crates/sqlbuild-analysis/src/lineage/tests.rs new file mode 100644 index 0000000000..191a850e53 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/tests.rs @@ -0,0 +1,8 @@ +#[path = "tests/test_build_fast_lineage.rs"] +mod build_fast_lineage; +#[path = "tests/helpers.rs"] +mod helpers; +#[path = "tests/test_references.rs"] +mod references; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-analysis/src/lineage/tests/helpers.rs b/crates/sqlbuild-analysis/src/lineage/tests/helpers.rs new file mode 100644 index 0000000000..c375293a7a --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/tests/helpers.rs @@ -0,0 +1,125 @@ +use polyglot_sql::{DialectType, SchemaValidationOptions}; + +use crate::lineage::main::build_fast_lineage::build_fast_lineage; +use crate::lineage::models::{ + FastLineageModel, FastLineageOutcome, FastLineageRequest, LineageColumn, LineageResourceType, + LineageSchemaResource, +}; +use crate::semantic_validation::models::ProjectCatalog; + +pub(crate) fn strings(values: &[&str]) -> Vec { + values.iter().map(|value| (*value).to_owned()).collect() +} + +/// Build one model against a fixed schema: `(status, column lines, has_star, detail)`. +pub(crate) fn outcome_lines( + dialect: Option<&str>, + model: FastLineageModel, +) -> Vec<(&'static str, Vec, bool, Option)> { + let request = FastLineageRequest { + dialect: dialect.map(str::to_owned), + schema: schema(), + models: vec![model], + }; + let outcomes: Vec = + build_fast_lineage(&request, &catalog()).expect("the analysis pool builds"); + outcomes.into_iter().map(described_outcome).collect() +} + +/// A `branches`-deep `UNION ALL` chain built from a caller thread with a `stack_bytes` stack. +pub(crate) fn deep_union_outcome_lines( + branches: usize, + stack_bytes: usize, +) -> Vec<(&'static str, Vec, bool, Option)> { + let sql: String = + vec!["SELECT order_id, amount FROM __ref('orders')"; branches].join(" UNION ALL "); + let worker = std::thread::Builder::new() + .stack_size(stack_bytes) + .spawn(move || { + outcome_lines( + Some("duckdb"), + FastLineageModel::Parse { + query_sql: sql, + inferred_columns: Vec::new(), + }, + ) + }) + .expect("the test thread starts"); + worker + .join() + .expect("the stage does not overflow the caller's stack") +} + +fn catalog() -> ProjectCatalog { + ProjectCatalog::with_options( + DialectType::DuckDB, + SchemaValidationOptions::default(), + false, + ) +} + +fn described_outcome( + outcome: FastLineageOutcome, +) -> (&'static str, Vec, bool, Option) { + let (status, columns, has_star, detail) = outcome.into_parts(); + (status, described(&columns), has_star, detail) +} + +fn described(columns: &[LineageColumn]) -> Vec { + columns.iter().map(described_column).collect() +} + +fn described_column(column: &LineageColumn) -> String { + let upstream: Vec = column + .upstream_columns + .iter() + .map(|source| { + format!( + "{}:{}.{}", + source.resource_type.as_str(), + source.resource_name, + source.column_name + ) + }) + .collect(); + format!( + "{} {} {} [{}]", + column.output_column, + column.transform_kind.as_str(), + column.confidence.as_str(), + upstream.join(" ") + ) +} + +fn schema() -> Vec { + vec![ + resource( + LineageResourceType::Model, + "orders", + &["order_id", "customer_id", "order_id", "amount"], + ), + resource( + LineageResourceType::Model, + "customers", + &["customer_id", "name"], + ), + resource(LineageResourceType::Model, "empty", &[]), + resource( + LineageResourceType::Seed, + "regions", + &["region_id", "label"], + ), + ] +} + +fn resource( + resource_type: LineageResourceType, + name: &str, + columns: &[&str], +) -> LineageSchemaResource { + LineageSchemaResource { + resource_type, + name: name.to_owned(), + columns: strings(columns), + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/tests/test_build_fast_lineage.rs b/crates/sqlbuild-analysis/src/lineage/tests/test_build_fast_lineage.rs new file mode 100644 index 0000000000..c166eb7c97 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/tests/test_build_fast_lineage.rs @@ -0,0 +1,330 @@ +use crate::lineage::models::FastLineageModel; +use crate::lineage::tests::helpers::{deep_union_outcome_lines, outcome_lines, strings}; +use crate::lineage::tests::test_types::{ + DeepUnionTestCase, ParsedLineageTestCase, StarExpansionTestCase, +}; + +#[test] +fn given_models_without_compact_facts_when_building_then_matches_python_lineage() { + let test_cases = [ + ParsedLineageTestCase { + description: "joins, aggregation, cast, constant and expression columns", + sql: "SELECT o.order_id, c.name AS customer_name, SUM(o.amount) AS total, \ + CAST(o.amount AS BIGINT) AS amount_big, 1 AS one, o.amount + c.customer_id AS mix \ + FROM __ref('orders') o JOIN __ref('customers') c ON o.customer_id = c.customer_id \ + GROUP BY 1, 2", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "order_id direct high [model:orders.order_id]", + "customer_name direct high [model:customers.name]", + "total aggregation high [model:orders.amount]", + "amount_big cast high [model:orders.amount]", + "one constant high []", + "mix expression high [model:orders.amount model:customers.customer_id]", + ], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "unqualified columns of one resource and an unknown qualifier", + sql: "SELECT order_id, amount * 2 AS doubled, missing.thing FROM __ref('orders')", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "order_id direct medium [model:orders.order_id]", + "doubled expression medium [model:orders.amount]", + "thing constant high []", + ], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "a root star expands known columns after explicit outputs", + sql: "SELECT *, amount AS amt FROM __ref('orders')", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "amt direct medium [model:orders.amount]", + "order_id star medium [model:orders.order_id]", + "customer_id star medium [model:orders.customer_id]", + "amount star medium [model:orders.amount]", + ], + expected_has_star: true, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "a qualified star expands every referenced resource", + sql: "SELECT o.* FROM __ref('orders') o JOIN __seed('regions') r ON TRUE", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "order_id star medium [model:orders.order_id]", + "customer_id star medium [model:orders.customer_id]", + "amount star medium [model:orders.amount]", + "region_id star medium [seed:regions.region_id]", + "label star medium [seed:regions.label]", + ], + expected_has_star: true, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "a three-branch union merges branch sources", + sql: "SELECT order_id FROM __ref('orders') UNION ALL SELECT customer_id \ + FROM __ref('customers') UNION ALL SELECT region_id + 1 FROM __seed('regions')", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "order_id expression medium [model:orders.order_id model:customers.customer_id \ + seed:regions.region_id]", + ], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "a union with a star branch takes inferred names", + sql: "SELECT order_id, amount FROM __ref('orders') UNION SELECT * FROM __ref('customers')", + dialect: None, + inferred_columns: &["first", "second"], + expected_status: "built", + expected_lines: &[ + "first direct medium [model:orders.order_id]", + "second direct medium [model:orders.amount]", + "order_id star medium [model:orders.order_id]", + "customer_id star medium [model:orders.customer_id]", + "amount star medium [model:orders.amount]", + "name star medium [model:customers.name]", + ], + expected_has_star: true, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "a CTE collapses onto its referenced model", + sql: "WITH base AS (SELECT order_id, amount FROM __ref('orders')) \ + SELECT order_id, amount FROM base", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "order_id direct medium [model:orders.order_id]", + "amount direct medium [model:orders.amount]", + ], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "unnamed projections take inferred names", + sql: "SELECT 1 + 1, COUNT(*) FROM __ref('orders')", + dialect: None, + inferred_columns: &["a", "b"], + expected_status: "built", + expected_lines: &["a constant high []", "b aggregation unknown []"], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "quoted and case-varied identifiers keep their spelling", + sql: "SELECT \"Order_ID\", O.Amount FROM __ref('orders') AS O", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "Order_ID direct medium [model:orders.Order_ID]", + "Amount direct high [model:orders.Amount]", + ], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "CASE branches are tuples Python's payload walk never enters", + sql: "SELECT CASE WHEN amount >= 9 THEN 'HIGH' ELSE 'LOW' END AS tier, \ + CASE order_id WHEN 1 THEN amount END AS picked, \ + COALESCE(amount, order_id) AS filled, amount BETWEEN 1 AND order_id AS ranged, \ + order_id IN (1, customer_id) AS listed, \ + ROW_NUMBER() OVER (PARTITION BY customer_id ORDER BY amount) AS rn, \ + EXTRACT(YEAR FROM amount) AS yr, order_id || '-' || amount AS joined \ + FROM __ref('orders')", + dialect: None, + inferred_columns: &[], + expected_status: "built", + expected_lines: &[ + "tier constant high []", + "picked expression medium [model:orders.order_id]", + "filled expression medium [model:orders.amount model:orders.order_id]", + "ranged expression medium [model:orders.amount model:orders.order_id]", + "listed expression medium [model:orders.order_id model:orders.customer_id]", + "rn expression medium [model:orders.customer_id model:orders.amount]", + "yr expression medium [model:orders.amount]", + "joined expression medium [model:orders.order_id model:orders.amount]", + ], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "a top-level EXCEPT has no fast lineage", + dialect: None, + sql: "SELECT order_id FROM __ref('orders') EXCEPT SELECT customer_id FROM __ref('customers')", + inferred_columns: &[], + expected_status: "omitted", + expected_lines: &[], + expected_has_star: false, + expected_detail: None, + }, + ParsedLineageTestCase { + description: "an unterminated query reports the parse error Python logs", + dialect: Some("duckdb"), + sql: "SELECT order_id FROM (", + inferred_columns: &[], + expected_status: "unparsed", + expected_lines: &[], + expected_has_star: false, + expected_detail: Some("Parse error at line 1, column 23: Unexpected end of input"), + }, + ParsedLineageTestCase { + description: "two statements report the statement count", + dialect: None, + sql: "SELECT 1; SELECT 2", + inferred_columns: &[], + expected_status: "unparsed", + expected_lines: &[], + expected_has_star: false, + expected_detail: Some("Expected 1 statement, found 2"), + }, + ParsedLineageTestCase { + description: "a parser panic on truncated T-SQL defers only that model", + sql: "SELECT IF(region > 1, re", + dialect: Some("tsql"), + inferred_columns: &[], + expected_status: "deferred", + expected_lines: &[], + expected_has_star: false, + expected_detail: Some("native_failure"), + }, + ParsedLineageTestCase { + description: "a dialect outside this parser build is deferred", + dialect: Some("mysql"), + sql: "SELECT 1", + inferred_columns: &[], + expected_status: "deferred", + expected_lines: &[], + expected_has_star: false, + expected_detail: Some("unsupported_dialect"), + }, + ParsedLineageTestCase { + description: "an unknown dialect is deferred", + dialect: Some("not-a-dialect"), + sql: "SELECT 1", + inferred_columns: &[], + expected_status: "deferred", + expected_lines: &[], + expected_has_star: false, + expected_detail: Some("unsupported_dialect"), + }, + ]; + for test_case in test_cases { + let outcomes: Vec<(&str, Vec, bool, Option)> = outcome_lines( + test_case.dialect, + FastLineageModel::Parse { + query_sql: test_case.sql.to_owned(), + inferred_columns: strings(test_case.inferred_columns), + }, + ); + assert_eq!( + outcomes, + vec![( + test_case.expected_status, + strings(test_case.expected_lines), + test_case.expected_has_star, + test_case.expected_detail.map(str::to_owned), + )], + "{}", + test_case.description + ); + } +} + +#[test] +fn given_compact_facts_with_unresolved_star_when_building_then_appends_unseen_columns() { + let test_cases = [ + StarExpansionTestCase { + description: "existing outputs and repeated references are skipped", + sql: "SELECT *, 1 AS amount FROM __ref('orders') JOIN __ref('orders') USING (order_id) \ + JOIN __seed('regions') USING (region_id)", + existing_columns: &["amount"], + expected_lines: &[ + "order_id star medium [model:orders.order_id]", + "customer_id star medium [model:orders.customer_id]", + "region_id star medium [seed:regions.region_id]", + "label star medium [seed:regions.label]", + ], + }, + StarExpansionTestCase { + description: "resources without known columns add nothing", + sql: "SELECT * FROM __ref('empty') JOIN __source('raw') ON TRUE", + existing_columns: &[], + expected_lines: &[], + }, + ]; + for test_case in test_cases { + let outcomes: Vec<(&str, Vec, bool, Option)> = outcome_lines( + Some("not-a-dialect"), + FastLineageModel::StarExpansion { + query_sql: test_case.sql.to_owned(), + existing_columns: strings(test_case.existing_columns), + }, + ); + assert_eq!( + outcomes, + vec![("star", strings(test_case.expected_lines), true, None)], + "{}", + test_case.description + ); + } +} + +#[test] +fn given_deep_union_chain_when_building_from_small_stack_then_runs_on_the_analysis_pool() { + let test_cases = [ + DeepUnionTestCase { + description: "the deepest chain the parser guard accepts", + branches: 512, + stack_bytes: 256 * 1024, + expected_status: "built", + expected_lines: &[ + "order_id direct medium [model:orders.order_id]", + "amount direct medium [model:orders.amount]", + ], + expected_detail: None, + }, + DeepUnionTestCase { + description: "one branch past the parser guard reports its parse error", + branches: 513, + stack_bytes: 256 * 1024, + expected_status: "unparsed", + expected_lines: &[], + expected_detail: Some( + "Parse error at line 0, column 0: E_GUARD_AST_DEPTH_EXCEEDED: value 513 exceeds configured limit 512", + ), + }, + ]; + for test_case in test_cases { + let outcomes: Vec<(&str, Vec, bool, Option)> = + deep_union_outcome_lines(test_case.branches, test_case.stack_bytes); + assert_eq!( + outcomes, + vec![( + test_case.expected_status, + strings(test_case.expected_lines), + false, + test_case.expected_detail.map(str::to_owned), + )], + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/tests/test_references.rs b/crates/sqlbuild-analysis/src/lineage/tests/test_references.rs new file mode 100644 index 0000000000..a673e3c7d0 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/tests/test_references.rs @@ -0,0 +1,59 @@ +use crate::lineage::_helpers::references::{normalized_sql, physical_resources}; +use crate::lineage::tests::test_types::ReferenceTestCase; + +#[test] +fn given_reference_calls_when_normalizing_then_matches_python_patterns() { + let test_cases = [ + ReferenceTestCase { + description: "each resource kind, spaced quotes and an underscore run", + sql: "SELECT * FROM __ref('orders') JOIN __source( \"raw.events\" ) ON 1=1 \ + JOIN ___seed('regions')", + expected_resources: &[ + ("model", "orders", "__sqlbuild_model__orders"), + ("source", "raw.events", "__sqlbuild_source__raw__events"), + ("seed", "regions", "__sqlbuild_seed__regions"), + ], + expected_normalized: "SELECT * FROM __sqlbuild_model__orders JOIN \ + __sqlbuild_source__raw__events ON 1=1 JOIN ___sqlbuild_seed__regions", + }, + ReferenceTestCase { + description: "a UDF call, mismatched and empty quotes, and a non-ASCII name", + sql: "SELECT __udf('my.fn') (x), __ref('a\") , __ref(''), __ref( '\u{e9} b' )\t)", + expected_resources: &[("model", "\u{e9} b", "__sqlbuild_model______b")], + expected_normalized: "SELECT my__fn(x), __ref('a\") , __ref(''), __sqlbuild_model______b\t)", + }, + ReferenceTestCase { + description: "Python whitespace beyond ASCII and an unknown call", + sql: "SELECT __ref('orders'\u{1c}) FROM x __unknown('a') __ref(\"q\")", + expected_resources: &[ + ("model", "orders", "__sqlbuild_model__orders"), + ("model", "q", "__sqlbuild_model__q"), + ], + expected_normalized: "SELECT __sqlbuild_model__orders FROM x __unknown('a') __sqlbuild_model__q", + }, + ]; + for test_case in test_cases { + let resources: Vec<(&str, String, String)> = physical_resources(test_case.sql) + .into_iter() + .map(|resource| { + ( + resource.resource_type.as_str(), + resource.resource_name, + resource.physical_name, + ) + }) + .collect(); + let expected: Vec<(&str, String, String)> = test_case + .expected_resources + .iter() + .map(|(kind, name, physical)| (*kind, (*name).to_owned(), (*physical).to_owned())) + .collect(); + assert_eq!(resources, expected, "{}", test_case.description); + assert_eq!( + normalized_sql(test_case.sql), + test_case.expected_normalized, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-analysis/src/lineage/tests/test_types.rs b/crates/sqlbuild-analysis/src/lineage/tests/test_types.rs new file mode 100644 index 0000000000..60235bf023 --- /dev/null +++ b/crates/sqlbuild-analysis/src/lineage/tests/test_types.rs @@ -0,0 +1,34 @@ +pub(crate) struct ReferenceTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) expected_resources: &'static [(&'static str, &'static str, &'static str)], + pub(crate) expected_normalized: &'static str, +} + +pub(crate) struct ParsedLineageTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: Option<&'static str>, + pub(crate) sql: &'static str, + pub(crate) inferred_columns: &'static [&'static str], + pub(crate) expected_status: &'static str, + /// One `output transform confidence [type:resource.column ...]` line per column. + pub(crate) expected_lines: &'static [&'static str], + pub(crate) expected_has_star: bool, + pub(crate) expected_detail: Option<&'static str>, +} + +pub(crate) struct StarExpansionTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) existing_columns: &'static [&'static str], + pub(crate) expected_lines: &'static [&'static str], +} + +pub(crate) struct DeepUnionTestCase { + pub(crate) description: &'static str, + pub(crate) branches: usize, + pub(crate) stack_bytes: usize, + pub(crate) expected_status: &'static str, + pub(crate) expected_lines: &'static [&'static str], + pub(crate) expected_detail: Option<&'static str>, +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/lineage/fast_lineage.rs b/crates/sqlbuild-python/src/bindings/_helpers/lineage/fast_lineage.rs new file mode 100644 index 0000000000..7d20c1eba4 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/lineage/fast_lineage.rs @@ -0,0 +1,118 @@ +//! Fast column lineage for the preview compiler engine. + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyRef, PyResult, Python}; +use pyo3::{pyfunction, wrap_pyfunction}; +use sqlbuild_analysis::lineage::main::build_fast_lineage::build_fast_lineage; +use sqlbuild_analysis::lineage::models::{ + FastLineageModel, FastLineageOutcome, FastLineageRequest, LineageColumn, LineageResourceType, + LineageSchemaResource, LineageSource, +}; +use sqlbuild_analysis::semantic_validation::models as validation; + +use crate::bindings::_helpers::boundary::panics::compiler_error; +use crate::bindings::models::ProjectCatalog; +use crate::bindings::types::CompilerDetach; + +type SourceRow = (&'static str, String, String); +type ColumnRow = (String, &'static str, &'static str, Vec); +/// `(dialect, [(resource_type, name, columns)], [(star_expansion, query_sql, column_names)])`. +type LineageRequestInput = ( + Option, + Vec<(String, String, Vec)>, + Vec<(bool, String, Vec)>, +); +/// `(status, columns, has_star, detail)`; see `FastLineageOutcome::into_parts`. +type OutcomeRow = (&'static str, Vec, bool, Option); + +/// Build fast lineage for the requested models on `catalog`'s analysis pool. +#[pyfunction] +fn build_fast_column_lineage( + py: Python<'_>, + catalog: PyRef<'_, ProjectCatalog>, + request: LineageRequestInput, +) -> PyResult> { + let (dialect, schema, models) = request; + let schema: Vec = schema + .into_iter() + .map(schema_resource) + .collect::>()?; + let models: Vec = models.into_iter().map(lineage_model).collect(); + let request = FastLineageRequest { + dialect, + schema, + models, + }; + let catalog: &validation::ProjectCatalog = &catalog.inner; + let outcomes: Vec = py + .compiler_detach(|| build_fast_lineage(&request, catalog)) + .map_err(compiler_error)?; + Ok(outcomes.into_iter().map(outcome_row).collect()) +} + +fn schema_resource( + (resource_type, name, columns): (String, String, Vec), +) -> PyResult { + let Some(kind) = LineageResourceType::from_value(&resource_type) else { + return Err(PyValueError::new_err(format!( + "unknown resource: {resource_type}" + ))); + }; + Ok(LineageSchemaResource { + resource_type: kind, + name, + columns, + }) +} + +fn lineage_model( + (star_expansion, query_sql, names): (bool, String, Vec), +) -> FastLineageModel { + if star_expansion { + FastLineageModel::StarExpansion { + query_sql, + existing_columns: names, + } + } else { + FastLineageModel::Parse { + query_sql, + inferred_columns: names, + } + } +} + +fn outcome_row(outcome: FastLineageOutcome) -> OutcomeRow { + let (status, columns, has_star, detail) = outcome.into_parts(); + (status, column_rows(columns), has_star, detail) +} + +fn column_rows(columns: Vec) -> Vec { + columns + .into_iter() + .map(|column| { + ( + column.output_column, + column.transform_kind.as_str(), + column.confidence.as_str(), + column + .upstream_columns + .into_iter() + .map(source_row) + .collect(), + ) + }) + .collect() +} + +fn source_row(source: LineageSource) -> SourceRow { + ( + source.resource_type.as_str(), + source.resource_name, + source.column_name, + ) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(build_fast_column_lineage, module)?)?; + Ok(()) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/lineage/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/lineage/mod.rs index 7444bd0374..768d0f4744 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/lineage/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/lineage/mod.rs @@ -1 +1,2 @@ +pub(crate) mod fast_lineage; pub(crate) mod registration; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/lineage/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/lineage/registration.rs index e747ba06c2..e22e2d12ee 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/lineage/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/lineage/registration.rs @@ -2,6 +2,8 @@ use pyo3::prelude::{Bound, PyModule, PyResult}; -pub(crate) fn register(_module: &Bound<'_, PyModule>) -> PyResult<()> { - Ok(()) +use crate::bindings::_helpers::lineage::fast_lineage; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + fast_lineage::register(module) } diff --git a/scripts/compiler_differential/_helpers/analysis_corpus/lineage.py b/scripts/compiler_differential/_helpers/analysis_corpus/lineage.py index fe6729e247..1508bda6de 100644 --- a/scripts/compiler_differential/_helpers/analysis_corpus/lineage.py +++ b/scripts/compiler_differential/_helpers/analysis_corpus/lineage.py @@ -1,9 +1,61 @@ """Failure cases owned by the native FAST lineage facts lane.""" +from scripts.compiler_differential._helpers.corpus.case_builder import failure_case +from scripts.compiler_differential.constants import FAILURE_MART_PATH from scripts.compiler_differential.models import FailureCase +_ORDERS: str = '__ref("stg_orders")' +_RAW: str = '__source("raw_orders")' +_OPTED_OUT_HEADER: str = ( + 'MODEL (\n description "Order totals per customer",\n sql_analysis false,\n);\n\n' +) +_BROKEN_PATH: str = "models/marts/broken_orders.sql" +_BROKEN_MODEL: str = ( + 'MODEL (\n description "Orders with an unknown qualifier",\n);\n\n' + f"SELECT x.order_id\nFROM {_ORDERS} o\n" +) + + +def _opted_out_mart(body: str) -> dict[str, str]: + """An opted-out mart whose lineage comes from the parse fallback, beside a failing model.""" + + return {FAILURE_MART_PATH: _OPTED_OUT_HEADER + body, _BROKEN_PATH: _BROKEN_MODEL} + def lineage_failure_cases() -> tuple[FailureCase, ...]: """Return this lane's failure cases; the lane appends here without editing shared lists.""" - return () + return ( + failure_case( + name="lineage-opted-out-star", + expected_code="B004", + files=_opted_out_mart(f"SELECT o.order_id AS first_order, o.*\nFROM {_ORDERS} o\n"), + ), + failure_case( + name="lineage-opted-out-union-of-stars", + expected_code="B004", + files=_opted_out_mart( + f"SELECT *, 1 AS copy_number\nFROM {_ORDERS}\n" + f'UNION ALL\nSELECT "R".*, 2 FROM {_RAW} AS "R" JOIN {_ORDERS} s ON TRUE\n' + ), + ), + failure_case( + name="lineage-opted-out-except-and-ctes", + expected_code="B004", + files=_opted_out_mart( + f"WITH kept AS (SELECT customer_id FROM {_ORDERS}),\n" + "again AS (SELECT * FROM kept)\n" + "SELECT customer_id, COUNT(*) AS total_amount FROM again GROUP BY 1\n" + f"EXCEPT\nSELECT customer_id, amount FROM {_ORDERS}\n" + ), + ), + failure_case( + name="lineage-opted-out-quoted-columns", + expected_code="B004", + files=_opted_out_mart( + 'SELECT "Customer_ID", SUM(O.Amount) AS "Total_Amount",\n' + " CAST(MAX(o.status) AS VARCHAR) AS last_status, 'fixed' AS label\n" + f"FROM {_ORDERS} AS O\nGROUP BY 1\n" + ), + ), + ) diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index cc0cc43c2f..68448e065b 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -270,6 +270,11 @@ def digest_files(paths: list[str]) -> list[str | None]: ... # Native analysis: contracts. # Native analysis: column lineage facts. +def build_fast_column_lineage( + catalog: object, + request: tuple[str | None, list[tuple[str, str, list[str]]], list[tuple[bool, str, list[str]]]], + /, +) -> list[tuple[str, list[tuple[str, str, str, list[tuple[str, str, str]]]], bool, str | None]]: ... # Native analysis: SQL test planning glue. diff --git a/src/sqlbuild/compiler/lineage/_helpers/native_deferrals.py b/src/sqlbuild/compiler/lineage/_helpers/native_deferrals.py new file mode 100644 index 0000000000..8d7e3cd387 --- /dev/null +++ b/src/sqlbuild/compiler/lineage/_helpers/native_deferrals.py @@ -0,0 +1,26 @@ +"""Debug-only record of fast lineage models the native engine hands back to Python.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +from sqlbuild.compiler.lineage.constants import ( + ANALYSIS_DEFERRAL_RECORD_PREFIX, + NATIVE_LINEAGE_DEFERRAL_SITE, +) +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR + + +def record_lineage_deferral(*, kind: str) -> None: + """Append one `{"kind", "site"}` line when the analysis record directory is set.""" + + directory: str | None = os.environ.get(ANALYSIS_RECORD_DIR_ENV_VAR) + if not directory: + return + target: Path = Path(directory) + target.mkdir(parents=True, exist_ok=True) + record: str = json.dumps({"kind": kind, "site": NATIVE_LINEAGE_DEFERRAL_SITE}) + with open(target / f"{ANALYSIS_DEFERRAL_RECORD_PREFIX}{os.getpid()}.jsonl", "a") as handle: + _ = handle.write(record + "\n") diff --git a/src/sqlbuild/compiler/lineage/_helpers/native_fast_columns.py b/src/sqlbuild/compiler/lineage/_helpers/native_fast_columns.py new file mode 100644 index 0000000000..61895dfc95 --- /dev/null +++ b/src/sqlbuild/compiler/lineage/_helpers/native_fast_columns.py @@ -0,0 +1,203 @@ +"""Fast column lineage built by the native engine, assembled exactly as Python assembles it.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence + +import sqlbuild._native as _native +from sqlbuild.compiler.compile.models import ( + CompactLineageFacts, + CompiledLineageColumnFact, + CompiledLineageSourceFact, + CompiledModel, + CompiledProject, +) +from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.lineage._helpers.columns import _build_schema_mapping +from sqlbuild.compiler.lineage._helpers.fast_columns import ( + _build_polyglot_fast_model_column_lineage, +) +from sqlbuild.compiler.lineage._helpers.native_deferrals import record_lineage_deferral +from sqlbuild.compiler.lineage.constants import ( + NATIVE_LINEAGE_BUILT, + NATIVE_LINEAGE_DEFERRED, + NATIVE_LINEAGE_NO_CATALOG, + NATIVE_LINEAGE_STAR, + NATIVE_LINEAGE_UNPARSED, +) +from sqlbuild.compiler.lineage.models import ( + ColumnLineage, + ColumnLineageSource, + ModelColumnLineage, + ProjectColumnLineage, +) +from sqlbuild.compiler.lineage.types import ( + ColumnLineageConfidence, + ColumnTransformKind, + InferredNullability, +) +from sqlbuild.diagnostics.main.log_debug_event import log_debug_event + +_DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.lineage") + +type _NativeColumn = tuple[str, str, str, list[tuple[str, str, str]]] +type _NativeOutcome = tuple[str, list[_NativeColumn], bool, str | None] + + +def build_native_fast_project_column_lineage( + *, + project: CompiledProject, + dialect: str | None, + model_names: frozenset[str] | None, +) -> ProjectColumnLineage | None: + """Build `build_fast_project_column_lineage`'s graph with native star expansion and parsing.""" + + if not project.settings.sql_analysis: + return None + models: tuple[CompiledModel, ...] = tuple( + model for model in project.models if model_names is None or model.name in model_names + ) + requests: list[tuple[bool, str, list[str]]] = [] + for model in models: + request: tuple[bool, str, list[str]] | None = _native_request(model) + if request is not None: + requests.append(request) + catalog: object | None = ( + project.binding_catalog.native if project.binding_catalog is not None else None + ) + if catalog is None and requests: + for _ in requests: + record_lineage_deferral(kind=NATIVE_LINEAGE_NO_CATALOG) + return None + outcomes: list[_NativeOutcome] = ( + _native.build_fast_column_lineage(catalog, (dialect, _schema_resources(project), requests)) + if requests + else [] + ) + outcome_index: int = 0 + python_schema: dict[str, dict[str, str]] | None = None + model_results: dict[str, ModelColumnLineage] = {} + compact_models: dict[str, tuple[Sequence[CompiledLineageColumnFact], bool]] = {} + model_order: list[str] = [] + for model in models: + if model.fast_lineage_columns is not None: + columns: Sequence[CompiledLineageColumnFact] = model.fast_lineage_columns + if _needs_star_expansion(model): + status, star_columns, _, _ = outcomes[outcome_index] + outcome_index += 1 + if status != NATIVE_LINEAGE_STAR: + return None + columns = (*columns, *(_column_fact(column) for column in star_columns)) + compact_models[model.name] = (columns, model.fast_lineage_has_star) + model_order.append(model.name) + continue + status, built_columns, has_star, detail = outcomes[outcome_index] + outcome_index += 1 + result: ModelColumnLineage | None = None + if status == NATIVE_LINEAGE_BUILT: + result = ModelColumnLineage( + model_name=model.name, + columns=tuple(_column_lineage(column) for column in built_columns), + has_star=has_star, + ) + elif status == NATIVE_LINEAGE_UNPARSED: + log_debug_event( + logger=_DEBUG_LOGGER, + message="fast column lineage parse failed; falling back", + sqlbuild_model=model.name, + sqlbuild_error=detail, + ) + elif status == NATIVE_LINEAGE_DEFERRED: + record_lineage_deferral(kind=str(detail)) + if python_schema is None: + python_schema = _build_schema_mapping(project) + result = _build_polyglot_fast_model_column_lineage( + model=model, schema=python_schema, dialect=dialect + ) + if result is None: + continue + model_results[model.name] = result + model_order.append(model.name) + return ProjectColumnLineage.from_fast_facts( + models=model_results, + compact_models=compact_models, + model_order=tuple(model_order), + ) + + +def _needs_star_expansion(model: CompiledModel) -> bool: + return model.fast_lineage_has_star and not model.fast_lineage_star_resolved + + +def _native_request(model: CompiledModel) -> tuple[bool, str, list[str]] | None: + columns: Sequence[CompiledLineageColumnFact] | None = model.fast_lineage_columns + if columns is None: + return (False, model.query_sql, [column.name for column in model.inferred_columns or ()]) + if not _needs_star_expansion(model): + return None + if isinstance(columns, CompactLineageFacts): + return (True, model.query_sql, [columns.string_pool[row[0]] for row in columns.rows]) + return (True, model.query_sql, [column.output_column for column in columns]) + + +def _schema_resources(project: CompiledProject) -> list[tuple[str, str, list[str]]]: + resources: list[tuple[str, str, list[str]]] = [] + for model in project.models: + columns: list[str] = [column.name for column in model.inferred_columns or ()] + if model.schema_entry is not None: + columns.extend(column.name for column in model.schema_entry.columns) + resources.append((CompiledResourceType.MODEL.value, model.name, columns)) + for source in project.sources: + resources.append( + ( + CompiledResourceType.SOURCE.value, + source.name, + [column.name for column in source.source_entry.columns], + ) + ) + for seed in project.seeds: + resources.append( + ( + CompiledResourceType.SEED.value, + seed.name, + [column.name for column in seed.schema_entry.columns], + ) + ) + return resources + + +def _column_fact(column: _NativeColumn) -> CompiledLineageColumnFact: + output_column, transform_kind, confidence, sources = column + return CompiledLineageColumnFact( + output_column=output_column, + upstream_columns=tuple( + CompiledLineageSourceFact( + resource_type=resource_type, + resource_name=resource_name, + column_name=column_name, + ) + for resource_type, resource_name, column_name in sources + ), + transform_kind=ColumnTransformKind(transform_kind), + confidence=ColumnLineageConfidence(confidence), + ) + + +def _column_lineage(column: _NativeColumn) -> ColumnLineage: + output_column, transform_kind, confidence, sources = column + return ColumnLineage( + output_column=output_column, + transform_kind=ColumnTransformKind(transform_kind), + expression_sql=None, + upstream_columns=tuple( + ColumnLineageSource( + resource_type=resource_type, + resource_name=resource_name, + column_name=column_name, + ) + for resource_type, resource_name, column_name in sources + ), + nullability=InferredNullability.UNKNOWN, + confidence=ColumnLineageConfidence(confidence), + ) diff --git a/src/sqlbuild/compiler/lineage/constants.py b/src/sqlbuild/compiler/lineage/constants.py index 7720007300..f6f49d295b 100644 --- a/src/sqlbuild/compiler/lineage/constants.py +++ b/src/sqlbuild/compiler/lineage/constants.py @@ -3,3 +3,10 @@ STAR_COLUMN_NAME: str = "*" POLYGLOT_RESOLVED_SOURCE_CONFIDENCE: str = "resolved" INDEXED_EDGE_RECORD_LENGTH: int = 6 +NATIVE_LINEAGE_STAR: str = "star" +NATIVE_LINEAGE_BUILT: str = "built" +NATIVE_LINEAGE_UNPARSED: str = "unparsed" +NATIVE_LINEAGE_DEFERRED: str = "deferred" +NATIVE_LINEAGE_NO_CATALOG: str = "no_analysis_catalog" +NATIVE_LINEAGE_DEFERRAL_SITE: str = "fast_columns.py" +ANALYSIS_DEFERRAL_RECORD_PREFIX: str = "analysis-deferrals-" diff --git a/src/sqlbuild/compiler/lineage/main/_build_native_column_lineage.py b/src/sqlbuild/compiler/lineage/main/_build_native_column_lineage.py index d19bf5b536..d5acf51a0d 100644 --- a/src/sqlbuild/compiler/lineage/main/_build_native_column_lineage.py +++ b/src/sqlbuild/compiler/lineage/main/_build_native_column_lineage.py @@ -3,6 +3,9 @@ from __future__ import annotations from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.compiler.lineage._helpers.native_fast_columns import ( + build_native_fast_project_column_lineage, +) from sqlbuild.compiler.lineage.models import ProjectColumnLineage @@ -11,4 +14,6 @@ def build_native_column_lineage( ) -> ProjectColumnLineage | None: """Return the fast lineage graph, or None where Python must build it.""" - return None + return build_native_fast_project_column_lineage( + project=project, dialect=dialect, model_names=model_names + ) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/__init__.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/_test_types.py new file mode 100644 index 0000000000..3b48472302 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/_test_types.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class NativeFastLineageCliTestCase: + """Every engine compiling and tracing fast lineage of one project through the CLI.""" + + description: str + files: dict[str, str] + lineage_targets: tuple[str, ...] + expected_lineage: dict[str, dict[str, object]] + expected_minimum_python_fallback_parses: int + expected_minimum_traced_edges: int diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/helpers.py new file mode 100644 index 0000000000..3e3231cb28 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/helpers.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import json +import re +import subprocess +from collections import Counter +from itertools import product +from pathlib import Path +from typing import NamedTuple, cast + +from sqlbuild.cli.compile_reuse.constants import REUSE_DISABLE_ENV_VAR +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR +from tests.e2e.src.sqlbuild.cli.commands.main.compile.helpers import run_installed_sqb + +_TIMING: re.Pattern[str] = re.compile(r" \(\d+\.\d+s\)") +_DIRECTIONS: tuple[str, ...] = ("upstream", "downstream") +FALLBACK_SITE: str = ( + "compiler/lineage/_helpers/fast_columns.py:_build_polyglot_fast_model_column_lineage" +) + + +class EngineLineageRun(NamedTuple): + """One engine's compile report, lineage traces, wheel fallback parses and deferral files.""" + + compile_report: str + compile_returncode: int + traces: list[tuple[int, str, str]] + fallback_parses: int + deferral_files: list[Path] + + +def engine_lineage_run( + *, project_dir: Path, files: dict[str, str], engine: str, targets: tuple[str, ...] +) -> EngineLineageRun: + """Write the project, then compile and trace every target both ways under one engine.""" + + for relative_path, contents in files.items(): + path: Path = project_dir / "project" / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + _ = path.write_text(contents, encoding="utf-8") + record_dir: Path = project_dir / "records" + environment: dict[str, str] = { + REUSE_DISABLE_ENV_VAR: "1", + ANALYSIS_RECORD_DIR_ENV_VAR: str(record_dir), + } + engine_args: tuple[str, ...] = ("--compiler-engine", engine) + compiled: subprocess.CompletedProcess[str] = run_installed_sqb( + project_dir=project_dir / "project", + args=(*engine_args, "compile", "--json", "--no-cache"), + env=environment, + ) + report: dict[str, object] = json.loads(compiled.stdout.replace(str(project_dir), "")) + _ = report.pop("compiler_engine", None) + _ = report.pop("compile_timings", None) + traces: list[tuple[int, str, str]] = [] + for target, direction in product(targets, _DIRECTIONS): + result: subprocess.CompletedProcess[str] = run_installed_sqb( + project_dir=project_dir / "project", + args=( + *engine_args, + "lineage", + target, + "--mode", + "fast", + "--direction", + direction, + "--json", + ), + env=environment, + ) + stdout: str = result.stdout.replace(str(project_dir), "") + traces.append((result.returncode, stdout, _TIMING.sub("", result.stderr))) + return EngineLineageRun( + compile_report=json.dumps(report, indent=2), + compile_returncode=compiled.returncode, + traces=traces, + fallback_parses=_site_calls(record_dir)[FALLBACK_SITE], + deferral_files=sorted(record_dir.glob("analysis-deferrals-*.jsonl")), + ) + + +def model_lineage_summaries(compile_report: str) -> dict[str, object]: + """The `lineage` summary of every model in a JSON compile report, by model name.""" + + payload: dict[str, object] = json.loads(compile_report) + models: list[dict[str, object]] = cast( + list[dict[str, object]], cast(dict[str, object], payload["resources"])["models"] + ) + return {str(model["name"]): model["lineage"] for model in models} + + +def _site_calls(record_dir: Path) -> Counter[str]: + calls: Counter[str] = Counter() + for path in record_dir.glob("polyglot-sites-*.json"): + for site, _, count in json.loads(path.read_text("utf-8"))["calls"]: + calls[site] += count + return calls diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/test_native_fast_lineage.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/test_native_fast_lineage.py new file mode 100644 index 0000000000..6ea7469c75 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/fast_lineage/test_native_fast_lineage.py @@ -0,0 +1,146 @@ +"""The preview engine builds fast column lineage natively with Python's exact CLI output.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.compile.fast_lineage._test_types import ( + NativeFastLineageCliTestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.main.compile.fast_lineage.helpers import ( + EngineLineageRun, + engine_lineage_run, + model_lineage_summaries, +) + +_ENGINES: tuple[str, ...] = ("python", "native", "native-preview") +_PROJECT_FILES: dict[str, str] = { + "sqlbuild_project.toml": ( + 'name = "orders_lineage"\nadapter = "duckdb"\n\n[connection]\ndatabase = "orders.duckdb"\n' + ), + "sources/raw.yml": ( + "sources:\n - name: raw_orders\n description: Orders feed.\n expression: >-\n" + " (SELECT 1 AS order_id, 10 AS customer_id, CAST(5 AS DOUBLE) AS amount)\n" + " columns:\n - name: order_id\n type: INTEGER\n" + " - name: customer_id\n type: INTEGER\n - name: amount\n type: DOUBLE\n" + ), + "models/staging/stg_orders.sql": ( + 'MODEL (description "Staged orders");\n\nSELECT * FROM __source("raw_orders")\n' + ), + "models/marts/order_union.sql": ( + 'MODEL (description "Orders twice", sql_analysis false);\n\n' + 'SELECT o.*, 1 AS copy_number FROM __ref("stg_orders") o\n' + 'UNION ALL\nSELECT *, 2 FROM __ref("stg_orders")\n' + ), + "models/marts/customer_totals.sql": ( + 'MODEL (description "Totals per customer", sql_analysis false);\n\n' + "WITH totals AS (\n" + ' SELECT customer_id, SUM(amount) AS "Total_Amount" FROM __ref("order_union")\n' + " GROUP BY customer_id\n)\n" + 'SELECT T.customer_id, "Total_Amount", CAST("Total_Amount" AS BIGINT) AS rounded\n' + "FROM totals AS T\n" + ), + "models/marts/order_copies.sql": ( + 'MODEL (description "Order copies");\n\n' + 'SELECT u.*, c.rounded FROM __ref("order_union") u\n' + 'JOIN __ref("customer_totals") c ON u.customer_id = c.customer_id\n' + ), +} + +_DEEP_UNION_BRANCHES: int = 480 +_DEEP_UNION_FILES: dict[str, str] = { + "sqlbuild_project.toml": _PROJECT_FILES["sqlbuild_project.toml"], + "sources/raw.yml": _PROJECT_FILES["sources/raw.yml"], + "models/marts/orders_union.sql": ( + 'MODEL (description "Orders repeated", sql_analysis false);\n\n' + + " UNION ALL ".join( + ['SELECT order_id, amount FROM __source("raw_orders")'] * _DEEP_UNION_BRANCHES + ) + + "\n" + ), + "models/marts/orders_second.sql": ( + 'MODEL (description "Order ids", sql_analysis false);\n\n' + 'SELECT order_id FROM __source("raw_orders")\n' + ), +} + + +@pytest.mark.parametrize( + "test_case", + [ + NativeFastLineageCliTestCase( + description="stars, opted-out unions, nested CTEs and quoted identifiers", + files=_PROJECT_FILES, + lineage_targets=( + "customer_totals.Total_Amount", + "order_union.amount", + "stg_orders.amount", + ), + expected_lineage={ + "order_union": { + "available": True, + "column_count": 0, + "edge_count": 3, + "has_star": True, + } + }, + expected_minimum_python_fallback_parses=2, + expected_minimum_traced_edges=4, + ), + NativeFastLineageCliTestCase( + description="a 480-branch UNION ALL chain beside another opted-out model", + files=_DEEP_UNION_FILES, + lineage_targets=("orders_union.amount", "orders_second.order_id"), + expected_lineage={ + "orders_union": { + "available": True, + "column_count": 0, + "edge_count": 2, + "has_star": False, + }, + "orders_second": { + "available": True, + "column_count": 0, + "edge_count": 1, + "has_star": False, + }, + }, + expected_minimum_python_fallback_parses=2, + expected_minimum_traced_edges=2, + ), + ], + ids=lambda case: case.description, +) +def test_given_project_when_compiling_with_each_engine_then_fast_lineage_output_matches( + test_case: NativeFastLineageCliTestCase, tmp_path: Path +) -> None: + runs: dict[str, EngineLineageRun] = { + engine: engine_lineage_run( + project_dir=tmp_path / engine, + files=test_case.files, + engine=engine, + targets=test_case.lineage_targets, + ) + for engine in _ENGINES + } + python: EngineLineageRun = runs["python"] + summaries: dict[str, object] = model_lineage_summaries(python.compile_report) + + assert python.compile_returncode == 0, python.compile_report + assert [code for code, _, _ in python.traces] == [0] * len(python.traces) + assert {name: summaries[name] for name in test_case.expected_lineage} == ( + test_case.expected_lineage + ) + assert sum('"source"' in stdout for _, stdout, _ in python.traces) >= ( + test_case.expected_minimum_traced_edges + ) + assert python.fallback_parses >= test_case.expected_minimum_python_fallback_parses + assert runs["native"] == python + assert runs["native-preview"]._replace(fallback_parses=python.fallback_parses) == python + assert runs["native-preview"].fallback_parses == 0 + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_native_analysis_seams.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_native_analysis_seams.py index 8cd7acd8a6..3ec873e7e6 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_native_analysis_seams.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_native_analysis_seams.py @@ -3,9 +3,11 @@ from __future__ import annotations from pathlib import Path +from unittest.mock import ANY import pytest +from sqlbuild.compiler.lineage.models import ProjectColumnLineage from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import ( NativeAnalysisSeamTestCase, ) @@ -32,7 +34,7 @@ "assemble_native_sql_tests": [None], "complete_native_semantic_diagnostics": [None], "evaluate_native_model_contracts": [None], - "build_native_column_lineage": [None], + "build_native_column_lineage": [ANY], "plan_native_sql_test_artifacts": [None], }, ) @@ -69,6 +71,7 @@ def test_given_project_when_compiling_with_each_engine_then_seams_defer_and_outp {}, test_case.expected_preview_returns, ) + assert isinstance(preview_seams["build_native_column_lineage"][0], ProjectColumnLineage) assert (native_run.returncode, preview_run.returncode) == (0, 0) assert report_without_engine(native_run) == report_without_engine(python_run) assert report_without_engine(preview_run) == report_without_engine(python_run) diff --git a/tests/integration/src/sqlbuild/compiler/lineage/__init__.py b/tests/integration/src/sqlbuild/compiler/lineage/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integration/src/sqlbuild/compiler/lineage/_test_types.py b/tests/integration/src/sqlbuild/compiler/lineage/_test_types.py new file mode 100644 index 0000000000..e56e9b4f5e --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/lineage/_test_types.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class GeneratedLineageParityTestCase: + """Seeded generated projects whose fast lineage Python and the native engine must agree on.""" + + description: str + seed: int + count: int + dialects: tuple[str | None, ...] + expected_minimum_native: int + expected_minimum_parsed: int + expected_minimum_star_expansions: int + + +@dataclass(frozen=True) +class DeferredLineageTestCase: + """Parsed models the native engine hands back to Python, which records each deferral.""" + + description: str + dialect: str + keeps_catalog: bool + expected_native_statuses: dict[str, int] + expected_kind: str + + +@dataclass(frozen=True) +class UnparsedLineageTestCase: + """A model without compact facts whose SQL does not parse, so it has no fast lineage.""" + + description: str + query_sql: str + expected_status: str + expected_messages: tuple[str, ...] diff --git a/tests/integration/src/sqlbuild/compiler/lineage/helpers.py b/tests/integration/src/sqlbuild/compiler/lineage/helpers.py new file mode 100644 index 0000000000..4a6af823e5 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/lineage/helpers.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import json +import random +from collections import Counter +from collections.abc import Callable, Sequence +from dataclasses import replace +from pathlib import Path +from typing import Any, cast + +import pytest + +import sqlbuild._native as native_module +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs +from sqlbuild.compiler.lineage._helpers.fast_columns import build_fast_project_column_lineage +from sqlbuild.compiler.lineage.main._build_native_column_lineage import ( + build_native_column_lineage, +) +from sqlbuild.compiler.lineage.models import ProjectColumnLineage +from sqlbuild.compiler.pipeline.main.graph import build_project_graph + +_PROJECT_TOML: str = ( + 'name = "orders_lineage"\nadapter = "duckdb"\n\n[connection]\ndatabase = "orders.duckdb"\n' +) +_SOURCES: str = """sources: + - name: raw_orders + description: Orders feed. + expression: >- + (SELECT 1 AS order_id, 10 AS customer_id, CAST(5 AS DOUBLE) AS amount, 'placed' AS status) + columns: + - name: order_id + type: INTEGER + - name: customer_id + type: INTEGER + - name: amount + type: DOUBLE + - name: status + type: VARCHAR + - name: raw_customers + description: Customers feed. + expression: >- + (SELECT 10 AS customer_id, 'Ada' AS Customer_Name, 'east' AS region) + columns: + - name: customer_id + type: INTEGER + - name: Customer_Name + type: VARCHAR + - name: region + type: VARCHAR +""" +_SEEDS: str = """seeds: + - name: regions + description: Region labels. + columns: + - name: region + type: VARCHAR + - name: region_label + type: VARCHAR +""" +_INITIAL_RELATIONS: dict[str, tuple[str, ...]] = { + '__source("raw_orders")': ("order_id", "customer_id", "amount", "status"), + '__source("raw_customers")': ("customer_id", "Customer_Name", "region"), + '__seed("regions")': ("region", "region_label"), +} +_OPT_OUT_SHARE: float = 0.25 +_OPT_OUTS: tuple[str, str] = ("", ", sql_analysis false") + +type _Template = Callable[ + [random.Random, list[tuple[str, tuple[str, ...]]]], tuple[str, tuple[str, ...]] +] + + +def _star( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + return f"SELECT * FROM {relation}", columns + + +def _qualified_star( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + (left, left_columns), (right, right_columns) = inputs[0], inputs[-1] + extra: str = rng.choice(right_columns) + sql: str = ( + f"SELECT a.*, b.{extra} AS joined_{extra.lower()}, 1 AS flag\n" + f"FROM {left} a\nJOIN {right} b ON TRUE" + ) + return sql, (*left_columns, f"joined_{extra.lower()}", "flag") + + +def _nested_ctes( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + picked: list[str] = rng.sample(list(columns), k=min(2, len(columns))) + sql: str = ( + f"WITH base AS (\n SELECT {', '.join(picked)} FROM {relation}\n),\n" + "layered AS (\n SELECT * FROM base\n)\n" + f"SELECT {picked[0]} AS key_value, COUNT(*) AS row_count\nFROM layered\nGROUP BY 1" + ) + return sql, ("key_value", "row_count") + + +def _union( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + (left, left_columns), (right, right_columns) = inputs[0], inputs[-1] + operator: str = rng.choice(("UNION ALL", "UNION", "EXCEPT", "INTERSECT")) + sql: str = ( + f"SELECT {left_columns[0]} AS shared_key, 'left' AS side FROM {left}\n" + f"{operator}\nSELECT {right_columns[0]}, 'right' FROM {right}" + ) + return sql, ("shared_key", "side") + + +def _union_stars( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + sql: str = f"SELECT * FROM {relation}\nUNION ALL\nSELECT t.* FROM {relation} t" + return sql, columns + + +def _quoted( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + column: str = rng.choice(columns) + sql: str = ( + f'SELECT "{column}" AS "Mixed_Case", UPPER(CAST(Src.{column.upper()} AS VARCHAR)) AS Up\n' + f"FROM {relation} AS Src" + ) + return sql, ("Mixed_Case", "Up") + + +def _expressions( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + column: str = rng.choice(columns) + sql: str = ( + f"SELECT CAST({column} AS VARCHAR) AS as_text, {column} IS NULL AS is_missing,\n" + f" 'fixed' AS literal_value, COUNT({column}) OVER () AS total_rows,\n" + f" (SELECT 1) AS scalar_one\nFROM {relation}" + ) + return sql, ("as_text", "is_missing", "literal_value", "total_rows", "scalar_one") + + +def _branching( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + first, last = columns[0], columns[-1] + sql: str = ( + f"SELECT CASE WHEN {first} IS NULL THEN 'none' ELSE 'some' END AS presence,\n" + f" CASE {first} WHEN {last} THEN {first} END AS matched,\n" + f" COALESCE({first}, {last}) AS filled, {first} IN ({last}, {first}) AS listed,\n" + f" ROW_NUMBER() OVER (PARTITION BY {last} ORDER BY {first}) AS position\n" + f"FROM {relation}" + ) + return sql, ("presence", "matched", "filled", "listed", "position") + + +def _subquery( + rng: random.Random, inputs: list[tuple[str, tuple[str, ...]]] +) -> tuple[str, tuple[str, ...]]: + relation, columns = inputs[0] + column: str = rng.choice(columns) + sql: str = f"SELECT s.{column}, s.* FROM (SELECT * FROM {relation}) s" + return sql, columns + + +_TEMPLATES: tuple[_Template, ...] = ( + _star, + _qualified_star, + _nested_ctes, + _union, + _union_stars, + _quoted, + _expressions, + _branching, + _subquery, +) + + +def generated_lineage_files(*, rng: random.Random, model_count: int) -> dict[str, str]: + """A project of sources, a seed and models drawing on every lineage template.""" + + files: dict[str, str] = { + "sqlbuild_project.toml": _PROJECT_TOML, + "sources/raw.yml": _SOURCES, + "seeds/regions.yml": _SEEDS, + "seeds/regions.csv": "region,region_label\neast,East\n", + } + relations: list[tuple[str, tuple[str, ...]]] = list(_INITIAL_RELATIONS.items()) + for index in range(model_count): + template: _Template = _TEMPLATES[index % len(_TEMPLATES)] + inputs: list[tuple[str, tuple[str, ...]]] = rng.sample(relations, k=2) + sql, columns = template(rng, inputs) + name: str = f"orders_model_{index}" + opt_out: str = _OPT_OUTS[rng.random() < _OPT_OUT_SHARE] + files[f"models/layer_{index % 3}/{name}.sql"] = ( + f'MODEL (description "Generated orders model {index}"{opt_out});\n\n{sql}\n' + ) + relations.append((f'__ref("{name}")', columns)) + return files + + +def compiled_project(*, project_dir: Path, files: dict[str, str]) -> CompiledProject: + """Write and compile a project with the Python compiler, keeping any diagnostics.""" + + for relative_path, contents in files.items(): + path: Path = project_dir / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + _ = path.write_text(contents, encoding="utf-8") + return build_project_graph( + discovered_inputs=discover_project_inputs(project_dir=project_dir), + adapter=DuckDbAdapter(), + ).project + + +def without_compact_facts(project: CompiledProject) -> CompiledProject: + """The project with every model sent through the parse fallback.""" + + return replace( + project, + models=tuple(replace(model, fast_lineage_columns=None) for model in project.models), + ) + + +def lineage_views( + *, + project: CompiledProject, + dialect: str | None, + model_names: frozenset[str] | None, +) -> tuple[list[object], list[object], list[object]]: + """Each public lineage view's name, Python's value and the native stage's value.""" + + python: ProjectColumnLineage = cast( + ProjectColumnLineage, + build_fast_project_column_lineage( + project=project, dialect=dialect, model_names=model_names + ), + ) + native: ProjectColumnLineage = cast( + ProjectColumnLineage, + build_native_column_lineage(project=project, dialect=dialect, model_names=model_names) + or build_fast_project_column_lineage( + project=project, dialect=dialect, model_names=model_names + ), + ) + names, python_values = _public_views(project=project, lineage=python) + return names, python_values, _public_views(project=project, lineage=native)[1] + + +def _public_views( + *, project: CompiledProject, lineage: ProjectColumnLineage +) -> tuple[list[object], list[object]]: + names: list[object] = [] + values: list[object] = [] + for model in project.models: + name: str = model.name + views: tuple[tuple[str, object], ...] = ( + ("has_model", lineage.has_model(name)), + ("has_star", lineage.model_has_star(name)), + ("edge_count", lineage.edge_count_targeting(name)), + ("targeting", lineage.edges_targeting(name)), + ("sourced", lineage.edges_sourced_from(name)), + ) + for view, value in views: + names.append(f"{name}.{view}") + values.append(value) + names.extend(("models", "edges")) + values.extend((list(lineage.models.items()), lineage.edges)) + return names, values + + +def lineage_log_records(caplog: pytest.LogCaptureFixture) -> list[tuple[str, str, object]]: + """Captured records as logger, message and the logged parse error.""" + + return [ + (record.name, record.getMessage(), record.__dict__.get("sqlbuild_error")) + for record in caplog.records + ] + + +def deferral_records(record_dir: Path) -> list[object]: + """Every deferral line the native lineage stage wrote below `record_dir`.""" + + records: list[object] = [] + for path in sorted(record_dir.glob("analysis-deferrals-*.jsonl")): + for line in path.read_text("utf-8").splitlines(): + records.append(json.loads(line)) + return records + + +def record_native_outcomes(*, monkeypatch: pytest.MonkeyPatch) -> Counter[str]: + """Count the native engine's per-model outcome statuses for the rest of the test.""" + + statuses: Counter[str] = Counter() + build: Callable[..., Sequence[tuple[str, object, bool, str | None]]] = ( + native_module.build_fast_column_lineage + ) + + def counted(*arguments: Any) -> Sequence[tuple[str, object, bool, str | None]]: + outcomes: Sequence[tuple[str, object, bool, str | None]] = build(*arguments) + statuses.update(outcome[0] for outcome in outcomes) + return outcomes + + monkeypatch.setattr(native_module, "build_fast_column_lineage", counted) + return statuses diff --git a/tests/integration/src/sqlbuild/compiler/lineage/test_native_fast_lineage.py b/tests/integration/src/sqlbuild/compiler/lineage/test_native_fast_lineage.py new file mode 100644 index 0000000000..9894d30670 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/lineage/test_native_fast_lineage.py @@ -0,0 +1,191 @@ +"""Native fast column lineage equals Python's on generated projects, model by model.""" + +from __future__ import annotations + +import logging +import random +from collections import Counter +from dataclasses import replace +from pathlib import Path + +import pytest + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.compiler.lineage._helpers.fast_columns import build_fast_project_column_lineage +from sqlbuild.compiler.lineage.main._build_native_column_lineage import ( + build_native_column_lineage, +) +from sqlbuild.compiler.lineage.models import ProjectColumnLineage +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR +from tests.integration.src.sqlbuild.compiler.helpers import mismatches +from tests.integration.src.sqlbuild.compiler.lineage._test_types import ( + DeferredLineageTestCase, + GeneratedLineageParityTestCase, + UnparsedLineageTestCase, +) +from tests.integration.src.sqlbuild.compiler.lineage.helpers import ( + compiled_project, + deferral_records, + generated_lineage_files, + lineage_log_records, + lineage_views, + record_native_outcomes, + without_compact_facts, +) + +_MODEL_COUNT: int = 16 + + +@pytest.mark.parametrize( + "test_case", + [ + GeneratedLineageParityTestCase( + description="stars, nested CTEs, set operations, quoting and opted-out models", + seed=20261008, + count=4, + dialects=(None, "duckdb", "postgres", "snowflake"), + expected_minimum_native=400, + expected_minimum_parsed=250, + expected_minimum_star_expansions=20, + ) + ], + ids=lambda case: case.description, +) +def test_given_generated_projects_when_building_fast_lineage_natively_then_matches_python( + test_case: GeneratedLineageParityTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rng: random.Random = random.Random(test_case.seed) + statuses: Counter[str] = record_native_outcomes(monkeypatch=monkeypatch) + differences: list[tuple[object, object, object]] = [] + for index in range(test_case.count): + project: CompiledProject = compiled_project( + project_dir=tmp_path / f"project_{index}", + files=generated_lineage_files(rng=rng, model_count=_MODEL_COUNT), + ) + selected: frozenset[str] = frozenset( + rng.sample([model.name for model in project.models], k=_MODEL_COUNT // 2) + ) + for variant in (project, without_compact_facts(project)): + for dialect in test_case.dialects: + for model_names in (None, selected): + names, python_views, native_views = lineage_views( + project=variant, dialect=dialect, model_names=model_names + ) + differences.extend( + mismatches(inputs=names, expected=python_views, actual=native_views) + ) + + assert differences == [] + assert statuses["deferred"] == 0 + assert sum(statuses.values()) >= test_case.expected_minimum_native + assert statuses["built"] + statuses["omitted"] >= test_case.expected_minimum_parsed + assert statuses["star"] >= test_case.expected_minimum_star_expansions + + +@pytest.mark.parametrize( + "test_case", + [ + UnparsedLineageTestCase( + description="an unterminated subquery", + query_sql='SELECT order_id FROM (SELECT * FROM __ref("orders_model_0")', + expected_status="unparsed", + expected_messages=("fast column lineage parse failed; falling back",), + ), + UnparsedLineageTestCase( + description="two statements", + query_sql="SELECT 1; SELECT 2", + expected_status="unparsed", + expected_messages=("fast column lineage parse failed; falling back",), + ), + ], + ids=lambda case: case.description, +) +def test_given_unparsable_model_when_building_fast_lineage_then_logs_and_omits_like_python( + test_case: UnparsedLineageTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + statuses: Counter[str] = record_native_outcomes(monkeypatch=monkeypatch) + project: CompiledProject = compiled_project( + project_dir=tmp_path / "project", + files=generated_lineage_files(rng=random.Random(11), model_count=_MODEL_COUNT), + ) + broken: CompiledProject = replace( + project, + models=( + replace(project.models[0], query_sql=test_case.query_sql, fast_lineage_columns=None), + *project.models[1:], + ), + ) + + with caplog.at_level(logging.DEBUG, logger="sqlbuild.lineage"): + python: ProjectColumnLineage | None = build_fast_project_column_lineage(project=broken) + python_records: list[tuple[str, str, object]] = lineage_log_records(caplog) + caplog.clear() + native: ProjectColumnLineage | None = build_native_column_lineage( + project=broken, dialect=None, model_names=None + ) + native_records: list[tuple[str, str, object]] = lineage_log_records(caplog) + + assert native == python + assert native is not None + assert not native.has_model(project.models[0].name) + assert statuses[test_case.expected_status] == 1 + assert native_records == python_records + assert tuple(message for _, message, _ in native_records) == test_case.expected_messages + + +@pytest.mark.parametrize( + "test_case", + [ + DeferredLineageTestCase( + description="a dialect outside the native parser build", + dialect="mysql", + keeps_catalog=True, + expected_native_statuses={"deferred": _MODEL_COUNT}, + expected_kind="unsupported_dialect", + ), + DeferredLineageTestCase( + description="a project without an analysis catalog to run on", + dialect="duckdb", + keeps_catalog=False, + expected_native_statuses={}, + expected_kind="no_analysis_catalog", + ), + ], + ids=lambda case: case.description, +) +def test_given_deferred_models_when_building_fast_lineage_then_python_builds_and_records( + test_case: DeferredLineageTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + record_dir: Path = tmp_path / "records" + monkeypatch.setenv(ANALYSIS_RECORD_DIR_ENV_VAR, str(record_dir)) + statuses: Counter[str] = record_native_outcomes(monkeypatch=monkeypatch) + compiled: CompiledProject = without_compact_facts( + compiled_project( + project_dir=tmp_path / "project", + files=generated_lineage_files(rng=random.Random(7), model_count=_MODEL_COUNT), + ) + ) + project: CompiledProject = replace( + compiled, binding_catalog=(None, compiled.binding_catalog)[test_case.keeps_catalog] + ) + + names, python_views, native_views = lineage_views( + project=project, dialect=test_case.dialect, model_names=None + ) + + assert mismatches(inputs=names, expected=python_views, actual=native_views) == [] + assert statuses == Counter(test_case.expected_native_statuses) + assert deferral_records(record_dir) == ( + [{"kind": test_case.expected_kind, "site": "fast_columns.py"}] * _MODEL_COUNT + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"])