diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/catalog_state.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/catalog_state.rs new file mode 100644 index 0000000000..a72ed4906b --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/catalog_state.rs @@ -0,0 +1,106 @@ +//! A session copy of the compile's binding catalog that records what Python's would gain. + +use std::collections::{HashMap, HashSet}; + +use crate::assembly::analysis_session::_helpers::mappings::{ + ShapeTable, catalog_columns, catalog_relations, same_mapping, +}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::models::{Columns, ProjectCatalog}; +use crate::semantic_validation::types::{BindingRequest, Relations}; + +/// Python's `BindingCatalog` state that model analysis reads and changes. +#[derive(Debug)] +pub(crate) struct SessionCatalog { + pub(crate) native: ProjectCatalog, + schemas: HashMap, + additions: Shapes, + analysis_names: Vec, + recorded_analysis: HashSet, +} + +impl SessionCatalog { + /// Copy `native`, whose Python wrapper holds `schemas`. + pub(crate) fn new(native: &ProjectCatalog, schemas: &Shapes) -> Self { + Self { + native: native.with_relations(Relations::new()), + schemas: schemas.iter().cloned().collect(), + additions: Vec::new(), + analysis_names: Vec::new(), + recorded_analysis: HashSet::new(), + } + } + + /// Python's `BindingCatalog.prepare`: add unknown closed relations, then override the rest. + pub(crate) fn prepare(&mut self, requests: &[(&str, &Shapes)]) -> Vec { + let mut additions: Shapes = Vec::new(); + for (_, schema) in requests { + for (name, columns) in schema.iter() { + if !columns.is_empty() && !self.schemas.contains_key(name) { + self.schemas.insert(name.clone(), columns.clone()); + additions.push((name.clone(), columns.clone())); + } + } + } + if !additions.is_empty() { + self.native.update_relations(catalog_relations(&additions)); + self.additions.extend(additions); + } + requests + .iter() + .map(|(sql, schema)| self.binding_request(sql, schema)) + .collect() + } + + fn binding_request(&self, sql: &str, schema: &Shapes) -> BindingRequest { + let mut overrides: Relations = Relations::new(); + let mut references: Vec<(String, bool)> = Vec::with_capacity(schema.len()); + for (name, columns) in schema { + references.push((name.clone(), !columns.is_empty())); + let known: bool = self + .schemas + .get(name) + .is_some_and(|known| same_mapping(known, columns)); + if !columns.is_empty() && !known { + overrides.insert(name.clone(), catalog_columns(columns)); + } + } + references.sort(); + (sql.to_owned(), references, overrides) + } + + /// Python's `BindingCatalog.prepare_analysis` for the relations `names` a batch reads. + pub(crate) fn prepare_analysis( + &mut self, + names: &[String], + types: &ShapeTable, + nullability: &ShapeTable, + ) { + let mut updates: HashMap = HashMap::new(); + for name in names { + let Some(columns) = nullability.get(name) else { + continue; + }; + let column_types: Pairs = types.get(name).cloned().unwrap_or_default(); + updates.insert( + name.clone(), + (catalog_columns(&column_types), catalog_columns(columns)), + ); + if self.recorded_analysis.insert(name.clone()) { + self.analysis_names.push(name.clone()); + } + } + if !updates.is_empty() { + self.native.update_analysis(updates); + } + } + + pub(crate) fn register_override(&mut self, overrides: Relations) -> usize { + self.native.register_override(overrides) + } + + /// The `schemas` additions and analysis-shape names Python's catalog records. + pub(crate) fn into_changes(self) -> (Shapes, Vec) { + (self.additions, self.analysis_names) + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/compact_batch.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/compact_batch.rs new file mode 100644 index 0000000000..5a3b0bd450 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/compact_batch.rs @@ -0,0 +1,603 @@ +//! One compact analysis batch built, run and projected exactly as Python's unshared batch. + +use std::collections::{HashMap, HashSet}; + +use serde_json::{Map, Value, json}; + +use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; +use crate::assembly::analysis_session::_helpers::mappings::{ShapeTable, native_dialect}; +use crate::assembly::analysis_session::constants::{ + BINDING_SEVERITIES, COMPACT_WORKERS, CONFIDENCE_CODES, DEFAULT_BINDING_MESSAGE, FACT_LENGTH, + INVALID_NATIVE_RESPONSE, LEGACY_FALLBACK, LEGACY_RESPONSE_LENGTH, NULLABILITY_BY_CODE, + RESPONSE_LENGTH, SOURCE_LENGTH, TRANSFORM_CODES, +}; +use crate::assembly::analysis_session::models::{ColumnFact, LineageRow}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::main::diagnostics::binding_diagnostics; +use crate::semantic_validation::types::{DiagnosticRow, NormalizationRequest}; + +const INVALID_BATCH: &str = "native compact query analysis returned an invalid batch response"; +const INVALID_TEMPLATE: &str = "native compact query analysis returned an invalid template"; +const INVALID_COLUMN: &str = "native compact analysis returned invalid column facts"; + +/// One query of a batch with the facts Python's preparation reads. +pub(crate) struct BatchMember<'a> { + pub(crate) query_sql: &'a str, + pub(crate) placeholders: &'a Pairs, + /// Every reference's analysis name, in reference order. + pub(crate) analysis_names: Vec<&'a str>, + pub(crate) lineage_references: &'a [(String, String, String)], + pub(crate) recover_cte_facts: bool, + pub(crate) binding_schema: Option<&'a Shapes>, +} + +/// One batch with the inference profile settings and relation facts its analysis reads. +pub(crate) struct Batch<'a> { + pub(crate) dialect: &'a str, + pub(crate) function_return_types: &'a Pairs, + pub(crate) rich_type_inference: bool, + pub(crate) members: Vec>, + pub(crate) types: &'a ShapeTable, + pub(crate) nullability: &'a ShapeTable, +} + +/// What the native engine returned for one member. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum MemberAnalysis { + Projected { + columns: Vec, + lineage: Vec, + has_star: bool, + star_resolved: bool, + }, + /// Analysis failed; Python records an unsuccessful analysis. + Failed, + /// The engine asked for Python's legacy analysis, keeping lineage rows when it had them. + Legacy { lineage: Option> }, +} + +/// One member's projected native analysis. +#[derive(Debug, Clone)] +pub(crate) struct MemberResult { + pub(crate) cleaned_sql: String, + pub(crate) analysis: MemberAnalysis, + pub(crate) binding_diagnostics: Option>, + /// The failure Python logs at debug level. + pub(crate) failure: Option, +} + +/// The batch's response tables Python's projection reads. +struct BatchResponse<'a> { + strings: Vec<&'a str>, + facts: &'a [Value], + templates: &'a [Value], +} + +/// The prepared request and each member's query index. +struct BatchPayload { + request: Value, + query_indexes: Vec, +} + +impl SessionCatalog { + /// Normalize, prepare, analyse and project a batch; an error is the one Python raises. + pub(crate) fn analyze_batch(&mut self, batch: &Batch<'_>) -> Result, String> { + if batch.members.is_empty() { + return Ok(Vec::new()); + } + let mut required: Vec = Vec::new(); + let mut seen: HashSet<&str> = HashSet::new(); + for name in batch + .members + .iter() + .flat_map(|member| member.analysis_names.iter()) + { + if seen.insert(name) { + required.push((*name).to_owned()); + } + } + self.prepare_analysis(&required, batch.types, batch.nullability); + let cleaned: Vec = self.normalize_members(batch)?; + let payload: BatchPayload = self.batch_payload(batch, &cleaned)?; + let job = self.native.prepare_compact(&payload.request.to_string())?; + let analysis = job + .take_analysis() + .ok_or("compact analysis job already ran")?; + let response: Value = + serde_json::from_str(&job.run(analysis)?).map_err(|error| error.to_string())?; + let results: Vec = attach_validations( + &response, + &payload.query_indexes, + batch.dialect, + project_response(&response, cleaned)?, + )?; + self.with_separate_validation(&batch.members, results) + } + + fn normalize_members(&self, batch: &Batch<'_>) -> Result, String> { + let requests: Vec = batch + .members + .iter() + .map(|member| { + ( + member.query_sql.to_owned(), + HashMap::new(), + member.placeholders.iter().cloned().collect(), + ) + }) + .collect(); + self.native + .normalize_analysis_sqls(batch.dialect, requests)? + .into_iter() + .collect() + } + + /// Python's `_prepare_compact_analysis_batch` request and each member's query index. + fn batch_payload( + &mut self, + batch: &Batch<'_>, + cleaned: &[String], + ) -> Result { + let function_return_types: Map = batch + .function_return_types + .iter() + .map(|(name, value)| (name.clone(), Value::String(value.clone()))) + .collect(); + let mut queries: Vec = Vec::new(); + let mut query_keys: HashMap = HashMap::new(); + let mut templates: Vec = Vec::new(); + let mut template_keys: HashMap = HashMap::new(); + let mut projections: Vec = Vec::with_capacity(batch.members.len()); + let mut query_indexes: Vec = Vec::with_capacity(batch.members.len()); + for (member, sql) in batch.members.iter().zip(cleaned) { + let query: Value = self.member_query(batch.dialect, member, sql)?; + let query_index: usize = *query_keys.entry(query.to_string()).or_insert_with(|| { + queries.push(query); + queries.len() - 1 + }); + query_indexes.push(query_index); + let template: Value = json!({ + "queryIndex": query_index, + "references": lineage_resources(member), + "functionReturnTypes": function_return_types, + "declaredColumnOrder": declared_column_order(member, batch.nullability), + "recoverCteFacts": member.recover_cte_facts, + "richTypeInference": batch.rich_type_inference, + }); + let template_index: usize = + *template_keys + .entry(template.to_string()) + .or_insert_with(|| { + templates.push(template); + templates.len() - 1 + }); + let resource_names: Map = member + .lineage_references + .iter() + .map(|(name, _, resource)| (name.clone(), Value::String(resource.clone()))) + .collect(); + projections + .push(json!({"templateIndex": template_index, "resourceNames": resource_names})); + } + Ok(BatchPayload { + request: json!({ + "queries": queries, + "templates": templates, + "projections": projections, + "workers": COMPACT_WORKERS, + }), + query_indexes, + }) + } + + /// One member's query: analysis references and, when bound, its catalog binding. + fn member_query( + &mut self, + dialect: &str, + member: &BatchMember<'_>, + sql: &str, + ) -> Result { + let mut lineage_names: Vec<&str> = member + .lineage_references + .iter() + .map(|(name, _, _)| name.as_str()) + .collect(); + lineage_names.sort_unstable(); + let analysis_references: Vec<[&str; 2]> = + lineage_names.iter().map(|name| [*name, *name]).collect(); + let mut query: Value = json!({ + "sql": sql, + "dialect": dialect, + "analysis_references": analysis_references, + }); + match member.binding_schema { + Some(schema) => { + let (_, references, overrides) = self + .prepare(&[(sql, schema)]) + .pop() + .ok_or("binding preparation returned no request")?; + query["binding_references"] = json!(references); + if !overrides.is_empty() { + query["binding_override"] = json!(self.register_override(overrides)); + } + } + None if !member.lineage_references.is_empty() => { + return Err("unbound members with relations are not supported".to_owned()); + } + None => {} + } + Ok(query) + } + + /// Python's separate schema validation for bound members the batch did not validate. + fn with_separate_validation( + &mut self, + members: &[BatchMember<'_>], + mut results: Vec, + ) -> Result, String> { + let mut indexes: Vec = Vec::new(); + let mut requests: Vec<(&str, &Shapes)> = Vec::new(); + for (index, (member, result)) in members.iter().zip(&results).enumerate() { + if let (Some(schema), None) = (member.binding_schema, &result.binding_diagnostics) { + indexes.push(index); + requests.push((result.cleaned_sql.as_str(), schema)); + } + } + if indexes.is_empty() { + return Ok(results); + } + let prepared: Vec<_> = self.prepare(&requests); + let rows: Vec> = self.native.binding_results(prepared)?; + if rows.len() != indexes.len() { + return Err( + "native SQL schema validation returned an invalid batch response".to_owned(), + ); + } + for (index, diagnostics) in indexes.into_iter().zip(rows) { + results[index].binding_diagnostics = Some(diagnostics); + } + Ok(results) + } +} + +fn lineage_resources(member: &BatchMember<'_>) -> Map { + member + .lineage_references + .iter() + .map(|(name, resource_type, _)| { + ( + name.clone(), + json!({"resourceType": resource_type, "resourceName": name}), + ) + }) + .collect() +} + +/// A single-reference member's declared column order, as Python's template records it. +fn declared_column_order<'a>( + member: &BatchMember<'_>, + nullability: &'a ShapeTable, +) -> Vec<&'a str> { + let [single] = member.analysis_names.as_slice() else { + return Vec::new(); + }; + let Some(columns) = nullability.get(single) else { + return Vec::new(); + }; + columns.iter().map(|(name, _)| name.as_str()).collect() +} + +fn to_index(value: u64) -> Result { + usize::try_from(value).map_err(|error| error.to_string()) +} + +/// Python's `_project_compact_analysis_batch`. +fn project_response(response: &Value, cleaned: Vec) -> Result, String> { + let strings: Vec<&str> = response + .get("strings") + .and_then(Value::as_array) + .ok_or(INVALID_BATCH)? + .iter() + .map(|value| value.as_str().ok_or(INVALID_BATCH)) + .collect::>()?; + let batch = BatchResponse { + strings, + facts: response + .get("facts") + .and_then(Value::as_array) + .ok_or(INVALID_BATCH)?, + templates: response + .get("templates") + .and_then(Value::as_array) + .ok_or(INVALID_BATCH)?, + }; + let analyses = response + .get("analyses") + .and_then(Value::as_array) + .filter(|analyses| analyses.len() == cleaned.len()) + .ok_or(INVALID_BATCH)?; + let mut results: Vec = Vec::with_capacity(cleaned.len()); + for (cleaned_sql, analysis) in cleaned.into_iter().zip(analyses) { + let (analysis, failure) = project_member(&batch, analysis)?; + results.push(MemberResult { + cleaned_sql, + analysis, + binding_diagnostics: None, + failure, + }); + } + Ok(results) +} + +fn project_member( + batch: &BatchResponse<'_>, + analysis: &Value, +) -> Result<(MemberAnalysis, Option), String> { + let entry: Option<&Vec> = analysis + .as_array() + .filter(|entry| entry.len() == RESPONSE_LENGTH); + let (Some(template_index), Some(mappings)) = ( + entry.and_then(|entry| entry[0].as_i64()), + entry.and_then(|entry| entry[1].as_array()), + ) else { + return Ok(( + MemberAnalysis::Failed, + Some(INVALID_NATIVE_RESPONSE.to_owned()), + )); + }; + let template: &Value = u64::try_from(template_index) + .map_err(|error| error.to_string()) + .and_then(to_index) + .map(|index| batch.templates.get(index))? + .ok_or("native compact query analysis returned an invalid template index")?; + if let Some(error) = template.as_str() { + let analysis = if error == LEGACY_FALLBACK { + MemberAnalysis::Legacy { lineage: None } + } else { + MemberAnalysis::Failed + }; + return Ok((analysis, Some(error.to_owned()))); + } + let parts: &Vec = template.as_array().ok_or(INVALID_TEMPLATE)?; + let flagged: bool = parts.len() == LEGACY_RESPONSE_LENGTH; + let legacy: bool = flagged && parts[2].as_str() == Some(LEGACY_FALLBACK); + let (Some(rows), Some(has_star)) = ( + parts.first().and_then(Value::as_array), + parts.get(1).and_then(Value::as_bool), + ) else { + return Err(INVALID_TEMPLATE.to_owned()); + }; + if !(parts.len() == RESPONSE_LENGTH || legacy || (flagged && parts[2].is_boolean())) { + return Err(INVALID_TEMPLATE.to_owned()); + } + let resource_indexes: HashMap = resource_name_indexes(batch, mappings)?; + let (columns, lineage) = project_rows(batch, rows, &resource_indexes)?; + if legacy { + return Ok(( + MemberAnalysis::Legacy { + lineage: Some(lineage), + }, + None, + )); + } + Ok(( + MemberAnalysis::Projected { + columns, + lineage, + has_star, + star_resolved: flagged && parts[2].as_bool() == Some(true), + }, + None, + )) +} + +fn resource_name_indexes( + batch: &BatchResponse<'_>, + mappings: &[Value], +) -> Result, String> { + let mut indexes: HashMap = HashMap::new(); + for mapping in mappings { + let pair: &Vec = mapping + .as_array() + .filter(|pair| pair.len() == RESPONSE_LENGTH) + .ok_or("native compact query analysis returned an invalid resource mapping")?; + let (Some(canonical), Some(resource)) = (pair[0].as_u64(), pair[1].as_u64()) else { + return Err( + "native compact query analysis returned an invalid resource mapping".to_owned(), + ); + }; + ensure_pooled(batch, canonical)?; + ensure_pooled(batch, resource)?; + indexes.insert(canonical, resource); + } + Ok(indexes) +} + +fn pooled<'a>(batch: &BatchResponse<'a>, index: u64) -> Result<&'a str, String> { + batch + .strings + .get(to_index(index)?) + .copied() + .ok_or_else(|| "native compact analysis returned an out-of-range index".to_owned()) +} + +fn ensure_pooled(batch: &BatchResponse<'_>, index: u64) -> Result<(), String> { + pooled(batch, index).map(drop) +} + +/// Python's `_compact_projected_analysis_result` columns and resolved lineage rows. +fn project_rows( + batch: &BatchResponse<'_>, + rows: &[Value], + resource_indexes: &HashMap, +) -> Result<(Vec, Vec), String> { + let mut columns: Vec = Vec::with_capacity(rows.len()); + let mut lineage: Vec = Vec::with_capacity(rows.len()); + for row in rows { + let fact: &Vec = batch + .facts + .get(to_index(row.as_u64().ok_or(INVALID_COLUMN)?)?) + .and_then(Value::as_array) + .filter(|fact| fact.len() == FACT_LENGTH) + .ok_or("native compact analysis returned an invalid column")?; + let (column, lineage_row) = project_fact(batch, fact, resource_indexes)?; + columns.push(column); + lineage.push(lineage_row); + } + Ok((columns, lineage)) +} + +fn project_fact( + batch: &BatchResponse<'_>, + fact: &[Value], + resource_indexes: &HashMap, +) -> Result<(ColumnFact, LineageRow), String> { + let (Some(name), Some(nullability), Some(transform), Some(confidence), Some(sources)) = ( + fact[0].as_u64(), + fact[2].as_u64(), + fact[3].as_u64(), + fact[4].as_u64(), + fact[5].as_array(), + ) else { + return Err(INVALID_COLUMN.to_owned()); + }; + let nullability: &str = NULLABILITY_BY_CODE + .get(to_index(nullability)?) + .copied() + .filter(|_| transform < TRANSFORM_CODES && confidence < CONFIDENCE_CODES) + .ok_or("native compact analysis returned an invalid fact code")?; + let data_type: Option = match &fact[1] { + Value::Null => None, + value => Some( + pooled( + batch, + value + .as_u64() + .ok_or("native compact analysis returned an invalid type index")?, + )? + .to_owned(), + ), + }; + let output_column: String = pooled(batch, name)?.to_owned(); + let mut resolved: Vec<(String, String, String)> = Vec::with_capacity(sources.len()); + for source in sources { + resolved.push(project_source(batch, source, resource_indexes)?); + } + Ok(( + ColumnFact { + name: output_column.clone(), + data_type, + nullability: nullability.to_owned(), + }, + LineageRow { + output_column, + transform_code: u8::try_from(transform).map_err(|error| error.to_string())?, + confidence_code: u8::try_from(confidence).map_err(|error| error.to_string())?, + sources: resolved, + }, + )) +} + +fn project_source( + batch: &BatchResponse<'_>, + source: &Value, + resource_indexes: &HashMap, +) -> Result<(String, String, String), String> { + const INVALID_SOURCE: &str = "native compact analysis returned invalid upstream lineage"; + let parts: &Vec = source + .as_array() + .filter(|parts| parts.len() == SOURCE_LENGTH) + .ok_or(INVALID_SOURCE)?; + let (Some(resource_type), Some(resource), Some(column)) = + (parts[0].as_u64(), parts[1].as_u64(), parts[2].as_u64()) + else { + return Err(INVALID_SOURCE.to_owned()); + }; + ensure_pooled(batch, resource)?; + Ok(( + pooled(batch, resource_type)?.to_owned(), + pooled( + batch, + resource_indexes.get(&resource).copied().unwrap_or(resource), + )? + .to_owned(), + pooled(batch, column)?.to_owned(), + )) +} + +/// Python's `_attach_compiled_bindings`: decode the batch's fused binding validations. +fn attach_validations( + response: &Value, + query_indexes: &[usize], + dialect: &str, + mut results: Vec, +) -> Result, String> { + let Some(validations) = response.get("validations").filter(|value| !value.is_null()) else { + return Ok(results); + }; + let query_count: usize = query_indexes.iter().max().map_or(0, |index| index + 1); + let validations: &Vec = validations + .as_array() + .filter(|validations| validations.len() == query_count) + .ok_or("native compilation returned an invalid binding batch")?; + for (result, query_index) in results.iter_mut().zip(query_indexes) { + let validation: &Value = &validations[*query_index]; + if !validation.is_null() { + result.binding_diagnostics = + Some(decode_validation(&result.cleaned_sql, dialect, validation)?); + } + } + Ok(results) +} + +/// Python's `_binding_result` over one native validation. +fn decode_validation( + sql: &str, + dialect: &str, + validation: &Value, +) -> Result, String> { + let errors: &Vec = validation + .get("errors") + .and_then(Value::as_array) + .ok_or("native SQL schema validation returned an invalid response")?; + let mut rows: Vec = Vec::new(); + for error in errors { + if let Some(row) = diagnostic_row(error)? { + rows.push(row); + } + } + binding_diagnostics(sql, native_dialect(dialect), rows) +} + +fn diagnostic_row(error: &Value) -> Result, String> { + let (Some(severity), Some(code)) = ( + error.get("severity").and_then(Value::as_str), + error.get("code").and_then(Value::as_str), + ) else { + return Ok(None); + }; + if !BINDING_SEVERITIES.contains(&severity) { + return Ok(None); + } + let message: String = match error.get("message") { + Some(Value::String(message)) if !message.is_empty() => message.clone(), + _ => DEFAULT_BINDING_MESSAGE.to_owned(), + }; + Ok(Some(( + code.to_owned(), + message, + position(error, "line")?, + position(error, "column")?, + position(error, "start")?, + position(error, "end")?, + severity.to_owned(), + ))) +} + +fn position(error: &Value, key: &str) -> Result, String> { + error + .get(key) + .and_then(Value::as_u64) + .map(to_index) + .transpose() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mappings.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mappings.rs new file mode 100644 index 0000000000..792c6a0766 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mappings.rs @@ -0,0 +1,126 @@ +//! Python dict semantics over ordered pairs. + +use std::collections::HashMap; + +use crate::assembly::analysis_session::constants::DIALECT_ALIASES; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::models::Columns; +use crate::semantic_validation::types::Relations; + +/// Python's `dict(pairs)`: the first position of each key with its last value. +pub(crate) fn dict_from_pairs(pairs: impl IntoIterator) -> Pairs { + let mut positions: HashMap = HashMap::new(); + let mut result: Pairs = Vec::new(); + for (key, value) in pairs { + if let Some(position) = positions.get(&key) { + result[*position].1 = value; + continue; + } + positions.insert(key.clone(), result.len()); + result.push((key, value)); + } + result +} + +/// Python's `dict ==`: the same items in any order. +pub(crate) fn same_mapping(left: &Pairs, right: &Pairs) -> bool { + let right: HashMap<&str, &str> = right + .iter() + .map(|(key, value)| (key.as_str(), value.as_str())) + .collect(); + left.len() == right.len() + && left + .iter() + .all(|(key, value)| right.get(key.as_str()) == Some(&value.as_str())) +} + +/// Python's `dict.keys() ==`: the same keys in any order. +pub(crate) fn same_keys(left: &Pairs, right: &Pairs) -> bool { + let right: HashMap<&str, &str> = right + .iter() + .map(|(key, value)| (key.as_str(), value.as_str())) + .collect(); + left.len() == right.len() && left.iter().all(|(key, _)| right.contains_key(key.as_str())) +} + +/// Python's `dict ==` over `{relation: {column: type}}`. +pub(crate) fn same_relations(left: &Shapes, right: &Shapes) -> bool { + let right: HashMap<&str, &Pairs> = right + .iter() + .map(|(name, shape)| (name.as_str(), shape)) + .collect(); + if left.len() != right.len() { + return false; + } + for (name, shape) in left { + match right.get(name.as_str()) { + Some(other) if same_mapping(shape, other) => {} + _ => return false, + } + } + true +} + +/// Python's `dict[name] = value` over ordered pairs. +pub(crate) fn with_pair(mut pairs: Pairs, key: &str, value: &str) -> Pairs { + match pairs.iter().position(|(existing, _)| existing == key) { + Some(position) => value.clone_into(&mut pairs[position].1), + None => pairs.push((key.to_owned(), value.to_owned())), + } + pairs +} + +/// Relation shapes by name, keeping Python's `setdefault` semantics. +#[derive(Debug, Clone, Default)] +pub(crate) struct ShapeTable { + shapes: HashMap, +} + +impl ShapeTable { + pub(crate) fn from_shapes(shapes: &Shapes) -> Self { + Self { + shapes: shapes.iter().cloned().collect(), + } + } + + pub(crate) fn get(&self, name: &str) -> Option<&Pairs> { + self.shapes.get(name) + } + + pub(crate) fn contains(&self, name: &str) -> bool { + self.shapes.contains_key(name) + } + + /// Python's `setdefault`. + pub(crate) fn set_default(&mut self, name: &str, shape: Pairs) { + self.shapes.entry(name.to_owned()).or_insert(shape); + } +} + +/// Native catalog columns for a `{name: type}` shape. +pub(crate) fn catalog_columns(shape: &Pairs) -> Columns { + Columns( + shape + .iter() + .map(|(name, value)| (name.clone(), Some(value.clone()))) + .collect(), + ) +} + +/// Native catalog relations for ordered shapes. +pub(crate) fn catalog_relations<'a>( + shapes: impl IntoIterator, +) -> Relations { + shapes + .into_iter() + .map(|(name, shape)| (name.clone(), catalog_columns(shape))) + .collect() +} + +/// Python's `NATIVE_DIALECT_ALIASES.get(dialect, dialect)`. +pub(crate) fn native_dialect(dialect: &str) -> &str { + DIALECT_ALIASES + .iter() + .find(|(alias, _)| *alias == dialect) + .map_or(dialect, |(_, native)| native) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs new file mode 100644 index 0000000000..f992aa5b95 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs @@ -0,0 +1,6 @@ +pub(crate) mod catalog_state; +pub(crate) mod compact_batch; +pub(crate) mod mappings; +pub(crate) mod publication; +pub(crate) mod schedule; +pub(crate) mod session; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/publication.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/publication.rs new file mode 100644 index 0000000000..5032309904 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/publication.rs @@ -0,0 +1,80 @@ +//! Python's `published_model_shape` and `inferred_binding_shape`. + +use std::collections::HashMap; + +use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; +use crate::assembly::analysis_session::_helpers::mappings::{ + catalog_columns, catalog_relations, with_pair, +}; +use crate::assembly::analysis_session::constants::{QUOTED_IDENTIFIER_DELIMITER, UNKNOWN_TYPE}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::main::normalize::normalize_analysis_sql; +use crate::semantic_validation::models::NormalizationInput; + +/// How a session publishes relation shapes. +pub(crate) struct ShapeOptions<'a> { + pub(crate) dialect: &'a str, + pub(crate) case_sensitive: bool, +} + +/// One relation's inferred columns and the SQL and inputs they came from. +pub(crate) struct ShapeSource<'a> { + pub(crate) sql: &'a str, + pub(crate) columns: Pairs, + pub(crate) inputs: &'a Shapes, + pub(crate) snapshot_columns: Option<&'a (String, String)>, +} + +impl SessionCatalog { + /// A model relation's shape, including snapshot validity columns. + pub(crate) fn published_model_shape( + &self, + options: &ShapeOptions<'_>, + source: ShapeSource<'_>, + ) -> Result { + let snapshot_columns: Option<&(String, String)> = source.snapshot_columns; + let mut shape: Pairs = self.inferred_binding_shape(options, source)?; + if let Some((valid_from, valid_to)) = snapshot_columns { + shape = with_pair(shape, valid_from, UNKNOWN_TYPE); + shape = with_pair(shape, valid_to, UNKNOWN_TYPE); + } + Ok(shape) + } + + /// Column names keep authored quoting only where the dialect binds quoted names exactly. + pub(crate) fn inferred_binding_shape( + &self, + options: &ShapeOptions<'_>, + source: ShapeSource<'_>, + ) -> Result { + if !options.case_sensitive { + return Ok(source.columns); + } + let normalized: String = normalize_analysis_sql(NormalizationInput { + sql: source.sql.to_owned(), + dialect: options.dialect.to_owned(), + stubs: HashMap::new(), + placeholders: HashMap::new(), + })?; + if !normalized.contains(QUOTED_IDENTIFIER_DELIMITER) && !has_exact_input(source.inputs) { + return Ok(source.columns); + } + Ok(self + .native + .inferred_schema( + &normalized, + catalog_columns(&source.columns), + catalog_relations(source.inputs), + )? + .into_iter() + .map(|(name, data_type)| (name, data_type.unwrap_or_default())) + .collect()) + } +} + +fn has_exact_input(inputs: &Shapes) -> bool { + inputs + .iter() + .flat_map(|(_, shape)| shape.iter()) + .any(|(name, _)| name.contains(QUOTED_IDENTIFIER_DELIMITER)) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/schedule.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/schedule.rs new file mode 100644 index 0000000000..2d39e878d0 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/schedule.rs @@ -0,0 +1,61 @@ +//! Python's topological analysis waves over the models a compile analyses. + +use std::collections::{BTreeSet, HashMap}; + +use crate::assembly::analysis_session::models::ModelRequest; + +/// Each model's analysed `ref` producers, or None when two models share a name. +pub(crate) fn producers(models: &[ModelRequest]) -> Option>> { + let mut indexes: HashMap<&str, usize> = HashMap::with_capacity(models.len()); + for (index, model) in models.iter().enumerate() { + if indexes.insert(model.name.as_str(), index).is_some() { + return None; + } + } + let mut producers: Vec> = Vec::with_capacity(models.len()); + for model in models { + let mut model_producers: BTreeSet = BTreeSet::new(); + for reference in model + .references + .iter() + .filter(|reference| reference.model_ref) + { + if let Some(index) = indexes.get(reference.analysis_name.as_str()) { + model_producers.insert(*index); + } + } + producers.push(model_producers.into_iter().collect()); + } + Some(producers) +} + +/// `TopologicalSorter` waves in request order, or None when the graph has a cycle. +pub(crate) fn waves(producers: &[Vec]) -> Option>> { + let mut pending: Vec = producers.iter().map(Vec::len).collect(); + let mut consumers: Vec> = vec![Vec::new(); producers.len()]; + for (model, model_producers) in producers.iter().enumerate() { + for producer in model_producers { + consumers[*producer].push(model); + } + } + let mut ready: Vec = (0..producers.len()) + .filter(|model| pending[*model] == 0) + .collect(); + let mut waves: Vec> = Vec::new(); + let mut scheduled: usize = 0; + while !ready.is_empty() { + scheduled += ready.len(); + let mut next: Vec = Vec::new(); + for model in &ready { + for consumer in &consumers[*model] { + pending[*consumer] -= 1; + if pending[*consumer] == 0 { + next.push(*consumer); + } + } + } + next.sort_unstable(); + waves.push(std::mem::replace(&mut ready, next)); + } + (scheduled == producers.len()).then_some(waves) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs new file mode 100644 index 0000000000..de099cbc93 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs @@ -0,0 +1,510 @@ +//! The session's phases: Python's uncached model analysis, wave by wave. + +use std::collections::HashMap; + +use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; +use crate::assembly::analysis_session::_helpers::compact_batch::{ + Batch, BatchMember, MemberAnalysis, MemberResult, +}; +use crate::assembly::analysis_session::_helpers::mappings::{ + ShapeTable, dict_from_pairs, same_keys, same_relations, +}; +use crate::assembly::analysis_session::_helpers::publication::{ShapeOptions, ShapeSource}; +use crate::assembly::analysis_session::_helpers::schedule::{producers, waves}; +use crate::assembly::analysis_session::constants::{ + DEFERRAL_ANALYSIS, DEFERRAL_ENRICHMENT, UNKNOWN_NULLABILITY, UNKNOWN_TYPE, +}; +use crate::assembly::analysis_session::models::{ + AnalysisSession, Awaiting, ColumnFact, Deferral, DeferredAnalysis, LineageFacts, ModelAnalysis, + ModelOutcome, ModelRequest, Phase, SessionOutcome, SessionRequest, SessionStep, +}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::models::ProjectCatalog; +use crate::semantic_validation::types::DiagnosticRow; + +/// The enrichment deferrals of one wave and whether each model's star was pending. +type Enrichments = (Vec, HashMap); + +impl AnalysisSession { + /// A session over `request`, or None where Python's cycle handling must analyse. + pub(crate) fn start(request: SessionRequest, catalog: &ProjectCatalog) -> Option { + let waves: Vec> = waves(&producers(&request.models)?)?; + let complete_shapes: ShapeTable = ShapeTable::from_shapes(&request.complete_schemas); + let dependency_ordered: bool = request + .models + .iter() + .flat_map(|model| model.required_names.iter()) + .any(|name| complete_shapes.get(name).is_none_or(Vec::is_empty)); + Some(Self { + catalog: SessionCatalog::new(catalog, &request.catalog_schemas), + available_types: ShapeTable::from_shapes(&request.column_types), + available_nullability: ShapeTable::from_shapes(&request.column_nullability), + complete_shapes, + outcomes: vec![None; request.models.len()], + waves: if dependency_ordered || waves.is_empty() { + waves + } else { + vec![(0..request.models.len()).collect()] + }, + dependency_ordered, + next_wave: 0, + phase: Phase::Analyze, + publications: Vec::new(), + failures: Vec::new(), + request, + }) + } + + /// Advance until Python must answer deferrals; a step without deferrals means done. + pub(crate) fn run(&mut self) -> Result { + loop { + match std::mem::replace(&mut self.phase, Phase::Done) { + Phase::Await(awaiting) => { + self.phase = Phase::Await(awaiting); + return Err("the session is waiting for deferred analyses".to_owned()); + } + Phase::Done => return Ok(self.step(Vec::new())), + Phase::Analyze => { + let Some(models) = self.waves.get(self.next_wave).cloned() else { + return Ok(self.step(Vec::new())); + }; + let wave: usize = self.next_wave; + self.next_wave += 1; + let deferrals: Vec = self.analyze_wave(&models)?; + if deferrals.is_empty() { + self.phase = Phase::Complete(wave); + continue; + } + let deferred: Vec = deferrals.iter().map(Deferral::model).collect(); + self.phase = Phase::Await(Awaiting::Analyses { + wave, + models: deferred, + }); + return Ok(self.step(deferrals)); + } + Phase::Complete(wave) => { + let (deferrals, star_pending) = self.complete_wave(wave)?; + if deferrals.is_empty() { + self.phase = Phase::Finish(wave); + continue; + } + self.phase = Phase::Await(Awaiting::Enrichments { wave, star_pending }); + return Ok(self.step(deferrals)); + } + Phase::Finish(wave) => { + self.finish_wave(wave)?; + self.phase = Phase::Analyze; + } + } + } + } + + /// Answer the deferrals of the last step. + pub(crate) fn provide(&mut self, results: Vec) -> Result<(), String> { + let Phase::Await(awaiting) = std::mem::replace(&mut self.phase, Phase::Done) else { + return Err("the session is not waiting for deferred analyses".to_owned()); + }; + match awaiting { + Awaiting::Analyses { wave, models } => { + let mut answered: Vec = results.iter().map(|result| result.model).collect(); + answered.sort_unstable(); + let mut expected: Vec = models; + expected.sort_unstable(); + if answered != expected { + return Err("deferred analyses do not match the deferrals".to_owned()); + } + for result in results { + let model: usize = result.model; + self.outcome_mut(model)?.analysis = python_analysis(result); + } + self.phase = Phase::Complete(wave); + } + Awaiting::Enrichments { wave, star_pending } => { + if results.len() != star_pending.len() { + return Err("deferred enrichments do not match the deferrals".to_owned()); + } + for result in results { + let pending: bool = *star_pending + .get(&result.model) + .ok_or("an enrichment answers no deferral")?; + let outcome: &mut ModelOutcome = self.outcome_mut(result.model)?; + outcome.analysis = enriched_analysis(&outcome.analysis, result, pending); + } + self.phase = Phase::Finish(wave); + } + } + Ok(()) + } + + /// Every model's outcome and the binding catalog changes Python records. + pub(crate) fn finish(self) -> Result { + if !matches!(self.phase, Phase::Done) { + return Err("the session has not finished".to_owned()); + } + let models: Vec = self + .outcomes + .into_iter() + .collect::>() + .ok_or("the session left a model unanalysed")?; + let (schema_additions, analysis_names) = self.catalog.into_changes(); + Ok(SessionOutcome { + models, + schema_additions, + analysis_names, + }) + } + + fn step(&mut self, deferrals: Vec) -> SessionStep { + SessionStep { + publications: std::mem::take(&mut self.publications), + deferrals, + failures: std::mem::take(&mut self.failures), + } + } + + fn outcome_mut(&mut self, model: usize) -> Result<&mut ModelOutcome, String> { + self.outcomes + .get_mut(model) + .and_then(Option::as_mut) + .ok_or_else(|| "a deferral names an unanalysed model".to_owned()) + } + + /// Python's `_analyze_model_sql_requests` for one wave. + fn analyze_wave(&mut self, models: &[usize]) -> Result, String> { + let schemas: Vec = models + .iter() + .map(|model| self.binding_schema(&self.request.models[*model])) + .collect(); + let batch = Batch { + dialect: &self.request.dialect, + function_return_types: &self.request.function_return_types, + rich_type_inference: self.request.rich_type_inference, + members: models + .iter() + .zip(&schemas) + .map(|(model, schema)| batch_member(&self.request.models[*model], schema)) + .collect(), + types: &self.available_types, + nullability: &self.available_nullability, + }; + let results: Vec = self.catalog.analyze_batch(&batch)?; + let mut deferrals: Vec = Vec::new(); + for ((model, schema), result) in models.iter().zip(schemas).zip(results) { + if let Some(failure) = result.failure { + self.failures.push(failure); + } + let diagnostics: Vec = result.binding_diagnostics.unwrap_or_default(); + let analysis: ModelAnalysis = match result.analysis { + MemberAnalysis::Projected { + columns, + lineage, + has_star, + star_resolved, + } => ModelAnalysis { + analysis_succeeded: true, + columns: Some(columns), + lineage: LineageFacts::Native(lineage), + has_star, + star_resolved, + binding_diagnostics: diagnostics, + binding_validated: true, + }, + MemberAnalysis::Failed => { + failed_analysis(LineageFacts::Native(Vec::new()), diagnostics) + } + MemberAnalysis::Legacy { lineage } => { + deferrals.push(Deferral::Analysis { + model: *model, + cleaned_sql: result.cleaned_sql.clone(), + binding_schema: schema.clone(), + binding_diagnostics: diagnostics.clone(), + lineage, + }); + failed_analysis(LineageFacts::PythonAnalysis, diagnostics) + } + }; + self.outcomes[*model] = Some(ModelOutcome { + analysis, + cleaned_sql: result.cleaned_sql, + validated_schema: schema, + fused_binding_validated: true, + }); + } + Ok(deferrals) + } + + /// The closed shape of each relation the model binds, open relations empty. + fn binding_schema(&self, model: &ModelRequest) -> Shapes { + model + .required_names + .iter() + .map(|name| { + ( + name.clone(), + self.complete_shapes.get(name).cloned().unwrap_or_default(), + ) + }) + .collect() + } + + /// Python's `_complete_inferred_bindings` star and type pass for one wave. + fn complete_wave(&mut self, wave: usize) -> Result { + let mut deferrals: Vec = Vec::new(); + let mut star_pending_by_model: HashMap = HashMap::new(); + for model in self.waves[wave].clone() { + let request: &ModelRequest = &self.request.models[model]; + let required: Vec = request.required_names.clone(); + let has_set_operation: bool = request.has_set_operation; + let inputs_known: bool = !required.is_empty() + && required + .iter() + .all(|name| self.complete_shapes.contains(name)); + let stars_expanded_natively: bool = + required.iter().all(|name| self.supplied_relation(name)); + let input_schemas: Shapes = self.input_schemas(&required); + let analysis: &mut ModelAnalysis = &mut self.outcome_mut(model)?.analysis; + let mut star_pending: bool = + analysis.has_star && !analysis.star_resolved && inputs_known; + if star_pending && has_columns(analysis.columns.as_ref()) && stars_expanded_natively { + analysis.star_resolved = true; + star_pending = false; + } + let untyped: bool = analysis + .columns + .iter() + .flatten() + .any(|column| column.data_type.is_none()); + if !required.is_empty() && (untyped || star_pending) && !has_set_operation { + star_pending_by_model.insert(model, star_pending); + deferrals.push(Deferral::Enrichment { + model, + input_schemas, + }); + } + } + Ok((deferrals, star_pending_by_model)) + } + + fn input_schemas(&self, required: &[String]) -> Shapes { + let mut schemas: Shapes = Vec::with_capacity(required.len()); + for name in required { + if let Some(shape) = self.complete_shapes.get(name) { + schemas.push((name.clone(), shape.clone())); + } + } + schemas + } + + /// Whether the dataflow supplied `name`'s closed shape to the wave's analysis. + fn supplied_relation(&self, name: &str) -> bool { + let Some(shape) = self.complete_shapes.get(name) else { + return false; + }; + self.dependency_ordered + && same_keys(self.available_types.get(name).unwrap_or(&Vec::new()), shape) + } + + /// Python's deferred binding validation, then the dataflow's shape publication. + fn finish_wave(&mut self, wave: usize) -> Result<(), String> { + let models: Vec = self.waves[wave].clone(); + let mut validations: Vec<(usize, String, Shapes)> = Vec::new(); + for model in &models { + let schema: Shapes = self.binding_schema(&self.request.models[*model]); + let outcome: &ModelOutcome = self.outcomes[*model] + .as_ref() + .ok_or("a wave model was not analysed")?; + if !outcome.fused_binding_validated + || !same_relations(&outcome.validated_schema, &schema) + { + validations.push((*model, outcome.cleaned_sql.clone(), schema)); + } + } + if !validations.is_empty() { + let requests: Vec<(&str, &Shapes)> = validations + .iter() + .map(|(_, sql, schema)| (sql.as_str(), schema)) + .collect(); + let prepared: Vec<_> = self.catalog.prepare(&requests); + let rows: Vec> = self.catalog.native.binding_results(prepared)?; + for ((model, _, _), diagnostics) in validations.iter().zip(rows) { + let analysis: &mut ModelAnalysis = &mut self.outcome_mut(*model)?.analysis; + analysis.binding_diagnostics = diagnostics; + analysis.binding_validated = true; + } + } + if self.dependency_ordered { + for model in models { + self.publish(model)?; + } + } + Ok(()) + } + + /// The dataflow's `_publish` for one analysed model. + fn publish(&mut self, model: usize) -> Result<(), String> { + let request: &ModelRequest = &self.request.models[model]; + let outcome: &ModelOutcome = self.outcomes[model] + .as_ref() + .ok_or("a published model was not analysed")?; + let analysis: &ModelAnalysis = &outcome.analysis; + let required_known: bool = request + .required_names + .iter() + .all(|name| self.complete_shapes.contains(name)); + let star_known: bool = !analysis.has_star || (analysis.star_resolved && required_known); + let Some(columns) = analysis + .columns + .as_ref() + .filter(|columns| !columns.is_empty()) + else { + return Ok(()); + }; + if !star_known { + return Ok(()); + } + let inputs: Shapes = self.binding_schema(request); + let shape: Pairs = self.catalog.published_model_shape( + &ShapeOptions { + dialect: &self.request.dialect, + case_sensitive: self.request.case_sensitive_shapes, + }, + ShapeSource { + sql: if outcome.cleaned_sql.is_empty() { + &request.query_sql + } else { + &outcome.cleaned_sql + }, + columns: shape_columns(columns), + inputs: &inputs, + snapshot_columns: request.snapshot_columns.as_ref(), + }, + )?; + let name: String = request.name.clone(); + let nullability: Pairs = shape + .iter() + .map(|(column, _)| (column.clone(), UNKNOWN_NULLABILITY.to_owned())) + .collect(); + self.complete_shapes.set_default(&name, shape.clone()); + self.available_types.set_default(&name, shape.clone()); + self.available_nullability.set_default(&name, nullability); + self.publications.push((name, shape)); + Ok(()) + } +} + +fn batch_member<'a>(model: &'a ModelRequest, schema: &'a Shapes) -> BatchMember<'a> { + BatchMember { + query_sql: &model.query_sql, + placeholders: &model.placeholders, + analysis_names: model + .references + .iter() + .map(|reference| reference.analysis_name.as_str()) + .collect(), + lineage_references: &model.lineage_references, + recover_cte_facts: model.recover_cte_facts, + binding_schema: Some(schema), + } +} + +/// Python's `{column.name: column.type or "UNKNOWN"}`. +pub(crate) fn shape_columns(columns: &[ColumnFact]) -> Pairs { + dict_from_pairs(columns.iter().map(shape_column)) +} + +fn shape_column(column: &ColumnFact) -> (String, String) { + let data_type: &str = match column.data_type.as_deref() { + Some(data_type) if !data_type.is_empty() => data_type, + _ => UNKNOWN_TYPE, + }; + (column.name.clone(), data_type.to_owned()) +} + +fn has_columns(columns: Option<&Vec>) -> bool { + columns.is_some_and(|columns| !columns.is_empty()) +} + +fn failed_analysis(lineage: LineageFacts, diagnostics: Vec) -> ModelAnalysis { + ModelAnalysis { + analysis_succeeded: false, + columns: None, + lineage, + has_star: false, + star_resolved: false, + binding_diagnostics: diagnostics, + binding_validated: true, + } +} + +fn python_analysis(result: DeferredAnalysis) -> ModelAnalysis { + ModelAnalysis { + analysis_succeeded: result.analysis_succeeded, + columns: result.columns, + lineage: LineageFacts::PythonAnalysis, + has_star: result.has_star, + star_resolved: result.star_resolved, + binding_diagnostics: result.binding_diagnostics, + binding_validated: result.binding_validated, + } +} + +impl Deferral { + /// The deferred model's request index. + pub fn model(&self) -> usize { + match self { + Self::Analysis { model, .. } | Self::Enrichment { model, .. } => *model, + } + } + + /// Python's name for the deferral's kind. + pub fn kind(&self) -> &'static str { + match self { + Self::Analysis { .. } => DEFERRAL_ANALYSIS, + Self::Enrichment { .. } => DEFERRAL_ENRICHMENT, + } + } +} + +/// Python's merge of a re-analysis with known inputs into the model's analysis. +fn enriched_analysis( + analysis: &ModelAnalysis, + enriched: DeferredAnalysis, + star_pending: bool, +) -> ModelAnalysis { + let mut recovered: HashMap = HashMap::new(); + for column in enriched + .columns + .iter() + .flatten() + .chain(analysis.columns.iter().flatten()) + { + if let Some(data_type) = &column.data_type { + recovered.insert(column.name.clone(), data_type.clone()); + } + } + let star_expanded: bool = star_pending && enriched.analysis_succeeded; + let replaced: bool = + (!has_columns(analysis.columns.as_ref()) || star_expanded) && enriched.analysis_succeeded; + let mut merged: ModelAnalysis = if replaced { + ModelAnalysis { + analysis_succeeded: enriched.analysis_succeeded, + columns: enriched.columns, + lineage: LineageFacts::PythonEnrichment, + has_star: enriched.has_star, + star_resolved: enriched.star_resolved, + binding_diagnostics: analysis.binding_diagnostics.clone(), + binding_validated: analysis.binding_validated, + } + } else { + analysis.clone() + }; + let mut columns: Vec = merged.columns.take().unwrap_or_default(); + for column in &mut columns { + if column.data_type.as_deref().is_none_or(str::is_empty) { + column.data_type = recovered.get(&column.name).cloned(); + } + } + merged.columns = Some(columns); + merged.star_resolved = merged.star_resolved || star_expanded; + merged +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs new file mode 100644 index 0000000000..735e39dbc8 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs @@ -0,0 +1,40 @@ +//! Wire values shared with Python's model analysis. + +/// The native engine's request for Python's legacy type recovery. +pub(crate) const LEGACY_FALLBACK: &str = "native project type recovery requires legacy fallback"; +/// Python logs this for a malformed per-model response. +pub(crate) const INVALID_NATIVE_RESPONSE: &str = "invalid native response"; +/// Python's `UNKNOWN` column type. +pub(crate) const UNKNOWN_TYPE: &str = "UNKNOWN"; +/// Python's `InferredNullability.UNKNOWN`. +pub(crate) const UNKNOWN_NULLABILITY: &str = "unknown"; +/// Python's `InferredNullability` values by native code. +pub(crate) const NULLABILITY_BY_CODE: [&str; 3] = ["unknown", "non_null", "nullable"]; +/// The transform and confidence code counts Python's enum tuples accept. +pub(crate) const TRANSFORM_CODES: u64 = 6; +pub(crate) const CONFIDENCE_CODES: u64 = 3; +/// Python's `COMPACT_ANALYSIS_RESPONSE_LENGTH`: a member's `[template, mappings]`. +pub(crate) const RESPONSE_LENGTH: usize = 2; +/// Python's `COMPACT_ANALYSIS_LEGACY_RESPONSE_LENGTH`: `[rows, has_star, flag]`. +pub(crate) const LEGACY_RESPONSE_LENGTH: usize = 3; +/// Python's `COMPACT_ANALYSIS_FACT_LENGTH`. +pub(crate) const FACT_LENGTH: usize = 6; +/// Python's `COMPACT_ANALYSIS_SOURCE_LENGTH`. +pub(crate) const SOURCE_LENGTH: usize = 3; +/// Python's `BINDING_SEVERITIES`. +pub(crate) const BINDING_SEVERITIES: [&str; 2] = ["error", "warning"]; +/// Python's default binding message. +pub(crate) const DEFAULT_BINDING_MESSAGE: &str = "SQL binding failed"; +/// Analysis workers per compact batch, as Python requests. +pub(crate) const COMPACT_WORKERS: usize = 4; +/// Python's `NATIVE_DIALECT_ALIASES`. +pub(crate) const DIALECT_ALIASES: [(&str, &str); 3] = [ + ("postgres", "postgresql"), + ("motherduck", "duckdb"), + ("sqlserver", "tsql"), +]; +/// The quote that marks an exact identifier. +pub(crate) const QUOTED_IDENTIFIER_DELIMITER: char = '"'; +/// Python's names for a legacy-analysis and an enrichment deferral. +pub(crate) const DEFERRAL_ANALYSIS: &str = "analysis"; +pub(crate) const DEFERRAL_ENRICHMENT: &str = "enrichment"; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/expression_shapes.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/expression_shapes.rs new file mode 100644 index 0000000000..973018ff0f --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/expression_shapes.rs @@ -0,0 +1,68 @@ +//! Python's `get_expression_source_shapes` for expressions its catalog has not inferred. + +use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; +use crate::assembly::analysis_session::_helpers::compact_batch::{ + Batch, BatchMember, MemberAnalysis, MemberResult, +}; +use crate::assembly::analysis_session::_helpers::mappings::ShapeTable; +use crate::assembly::analysis_session::_helpers::publication::{ShapeOptions, ShapeSource}; +use crate::assembly::analysis_session::_helpers::session::shape_columns; +use crate::assembly::analysis_session::models::{ExpressionShape, ExpressionShapeRequest}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::models::ProjectCatalog; + +/// One shape per expression, in order; an error is the one Python raises. +pub fn expression_shapes( + catalog: &ProjectCatalog, + request: &ExpressionShapeRequest, +) -> Result, String> { + let mut session_catalog: SessionCatalog = SessionCatalog::new(catalog, &Vec::new()); + let no_placeholders: Pairs = Vec::new(); + let empty: ShapeTable = ShapeTable::default(); + let batch = Batch { + dialect: &request.dialect, + function_return_types: &request.function_return_types, + rich_type_inference: true, + members: request + .expressions + .iter() + .map(|expression| BatchMember { + query_sql: expression, + placeholders: &no_placeholders, + analysis_names: Vec::new(), + lineage_references: &[], + recover_cte_facts: true, + binding_schema: None, + }) + .collect(), + types: &empty, + nullability: &empty, + }; + let results: Vec = session_catalog.analyze_batch(&batch)?; + let options = ShapeOptions { + dialect: &request.dialect, + case_sensitive: request.case_sensitive_shapes, + }; + let no_inputs: Shapes = Vec::new(); + let mut shapes: Vec = Vec::with_capacity(results.len()); + for (expression, result) in request.expressions.iter().zip(results) { + shapes.push(match result.analysis { + MemberAnalysis::Legacy { .. } => ExpressionShape::Deferred, + MemberAnalysis::Projected { + columns, has_star, .. + } if !columns.is_empty() && !has_star => { + ExpressionShape::Inferred(session_catalog.inferred_binding_shape( + &options, + ShapeSource { + sql: expression, + columns: shape_columns(&columns), + inputs: &no_inputs, + snapshot_columns: None, + }, + )?) + } + MemberAnalysis::Projected { .. } | MemberAnalysis::Failed => ExpressionShape::Absent, + }); + } + Ok(shapes) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/finish_analysis_session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/finish_analysis_session.rs new file mode 100644 index 0000000000..f7a4fec8b3 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/finish_analysis_session.rs @@ -0,0 +1,8 @@ +//! Collect a finished native model analysis session's outcomes. + +use crate::assembly::analysis_session::models::{AnalysisSession, SessionOutcome}; + +/// Every model's outcome and the binding catalog changes Python records. +pub fn finish_analysis_session(session: AnalysisSession) -> Result { + session.finish() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs new file mode 100644 index 0000000000..5893bfd5e8 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs @@ -0,0 +1,5 @@ +pub mod expression_shapes; +pub mod finish_analysis_session; +pub mod provide_deferred_analyses; +pub mod run_analysis_session; +pub mod start_analysis_session; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/provide_deferred_analyses.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/provide_deferred_analyses.rs new file mode 100644 index 0000000000..ba205cd7c0 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/provide_deferred_analyses.rs @@ -0,0 +1,11 @@ +//! Resume a native model analysis session with Python's answers. + +use crate::assembly::analysis_session::models::{AnalysisSession, DeferredAnalysis}; + +/// Answer the deferrals of the session's last step. +pub fn provide_deferred_analyses( + session: &mut AnalysisSession, + results: Vec, +) -> Result<(), String> { + session.provide(results) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/run_analysis_session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/run_analysis_session.rs new file mode 100644 index 0000000000..0c3c2a5c68 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/run_analysis_session.rs @@ -0,0 +1,8 @@ +//! Advance a native model analysis session to its next deferrals. + +use crate::assembly::analysis_session::models::{AnalysisSession, SessionStep}; + +/// Advance until Python must answer deferrals; a step without deferrals means done. +pub fn run_analysis_session(session: &mut AnalysisSession) -> Result { + session.run() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/start_analysis_session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/start_analysis_session.rs new file mode 100644 index 0000000000..3d716c5fe3 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/start_analysis_session.rs @@ -0,0 +1,12 @@ +//! Start Python's uncached model analysis as a native session. + +use crate::assembly::analysis_session::models::{AnalysisSession, SessionRequest}; +use crate::semantic_validation::models::ProjectCatalog; + +/// A session over `request` on a copy of `catalog`, or None where Python must analyse. +pub fn start_analysis_session( + request: SessionRequest, + catalog: &ProjectCatalog, +) -> Option { + AnalysisSession::start(request, catalog) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/mod.rs new file mode 100644 index 0000000000..536cf2088d --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/mod.rs @@ -0,0 +1,10 @@ +//! Python's uncached model analysis, run natively with deferrals back to Python. + +pub(crate) mod _helpers; +pub(crate) mod constants; +pub mod main; +pub mod models; +pub mod types; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs new file mode 100644 index 0000000000..a52f3e38ce --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs @@ -0,0 +1,207 @@ +//! Plain-data request, deferrals and outcomes of the native model analysis session. + +use std::collections::HashMap; + +use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; +use crate::assembly::analysis_session::_helpers::mappings::ShapeTable; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::types::DiagnosticRow; + +/// One `ref`, `source`, `seed`, `table_fn` or `udf` call a model's SQL makes. +#[derive(Debug, Clone)] +pub struct ModelReference { + /// Python's `_analysis_reference_name`. + pub analysis_name: String, + /// Whether the call is a `ref`, so names a model the analysis orders by. + pub model_ref: bool, +} + +/// One analysed model, with the facts Python derives from its compile input. +#[derive(Debug, Clone)] +pub struct ModelRequest { + pub name: String, + /// Cursor-intrinsic analysis SQL. + pub query_sql: String, + pub placeholders: Pairs, + pub references: Vec, + /// Python's `lineage_reference_map`: `(analysis name, resource type, resource name)`. + pub lineage_references: Vec<(String, String, String)>, + /// Sorted `binding_relation_names`. + pub required_names: Vec, + pub recover_cte_facts: bool, + /// Whether the SQL names `UNION`, `INTERSECT` or `EXCEPT`, as Python's search decides it. + pub has_set_operation: bool, + /// The snapshot validity columns a published shape appends. + pub snapshot_columns: Option<(String, String)>, +} + +/// Everything one compile's model analysis reads. +#[derive(Debug, Clone, Default)] +pub struct SessionRequest { + /// The analysis dialect, `generic` when the adapter names none. + pub dialect: String, + /// Whether published shapes keep authored quoting, Python's `inferred_binding_shape` test. + pub case_sensitive_shapes: bool, + pub function_return_types: Pairs, + pub rich_type_inference: bool, + pub column_types: Shapes, + pub column_nullability: Shapes, + pub complete_schemas: Shapes, + /// The Python binding catalog's `schemas` before analysis. + pub catalog_schemas: Shapes, + pub models: Vec, +} + +/// One inferred output column; nullability is Python's `InferredNullability` value. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ColumnFact { + pub name: String, + pub data_type: Option, + pub nullability: String, +} + +/// One compact lineage row with resolved names; codes index Python's enum tuples. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LineageRow { + pub output_column: String, + pub transform_code: u8, + pub confidence_code: u8, + pub sources: Vec<(String, String, String)>, +} + +/// Where a model's lineage facts live. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LineageFacts { + /// Native compact rows. + Native(Vec), + /// The lineage of the analysis Python returned for this model's analysis deferral. + PythonAnalysis, + /// The lineage of the analysis Python returned for this model's enrichment deferral. + PythonEnrichment, +} + +/// Python's `PolyglotAnalysisResult` for one model. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModelAnalysis { + pub analysis_succeeded: bool, + pub columns: Option>, + pub lineage: LineageFacts, + pub has_star: bool, + pub star_resolved: bool, + pub binding_diagnostics: Vec, + pub binding_validated: bool, +} + +/// Python's `ModelSqlAnalysis` for one model, without the placeholders Python keeps. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModelOutcome { + pub analysis: ModelAnalysis, + pub cleaned_sql: String, + pub validated_schema: Shapes, + pub fused_binding_validated: bool, +} + +/// Work native analysis hands back to Python for one model. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Deferral { + /// The native engine asked for Python's legacy analysis of this model. + Analysis { + model: usize, + cleaned_sql: String, + binding_schema: Shapes, + binding_diagnostics: Vec, + /// Lineage Python projects from native rows, when the engine kept them. + lineage: Option>, + }, + /// Re-analysis with the model's known input shapes. + Enrichment { model: usize, input_schemas: Shapes }, +} + +/// Python's answer to one deferral, without the lineage Python keeps. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeferredAnalysis { + pub model: usize, + pub analysis_succeeded: bool, + pub columns: Option>, + pub has_star: bool, + pub star_resolved: bool, + pub binding_diagnostics: Vec, + pub binding_validated: bool, +} + +/// One phase of the session: the deferrals Python must answer, or none when it is done. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct SessionStep { + /// Shapes published since the previous step, in publication order. + pub publications: Shapes, + pub deferrals: Vec, + /// Native analysis failures Python logs at debug level. + pub failures: Vec, +} + +/// The session's results once every deferral is answered. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct SessionOutcome { + pub models: Vec, + /// Relations the Python binding catalog adds to `schemas`, in order. + pub schema_additions: Shapes, + /// Relations whose analysis shapes the Python binding catalog records. + pub analysis_names: Vec, +} + +/// One expression source's shape: inferred, absent, or left to Python. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExpressionShape { + Inferred(Pairs), + Absent, + Deferred, +} + +/// What expression-source shape inference reads from the inference profile. +#[derive(Debug, Clone, Default)] +pub struct ExpressionShapeRequest { + pub dialect: String, + pub case_sensitive_shapes: bool, + pub function_return_types: Pairs, + pub expressions: Vec, +} + +/// Deferrals the session waits on before it continues a wave. +#[derive(Debug)] +pub(crate) enum Awaiting { + Analyses { + wave: usize, + models: Vec, + }, + Enrichments { + wave: usize, + star_pending: HashMap, + }, +} + +/// Where the session is within its waves. +#[derive(Debug)] +pub(crate) enum Phase { + Analyze, + Complete(usize), + Finish(usize), + Await(Awaiting), + Done, +} + +/// One compile's model analysis, advanced phase by phase and resumed with Python's answers. +#[derive(Debug)] +pub struct AnalysisSession { + pub(crate) request: SessionRequest, + pub(crate) catalog: SessionCatalog, + pub(crate) dependency_ordered: bool, + pub(crate) waves: Vec>, + pub(crate) next_wave: usize, + pub(crate) available_types: ShapeTable, + pub(crate) available_nullability: ShapeTable, + pub(crate) complete_shapes: ShapeTable, + pub(crate) outcomes: Vec>, + pub(crate) phase: Phase, + pub(crate) publications: Shapes, + pub(crate) failures: Vec, +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests.rs new file mode 100644 index 0000000000..d3cf656166 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests.rs @@ -0,0 +1,10 @@ +#[path = "tests/test_analysis_session.rs"] +mod analysis_session; +#[path = "tests/test_expression_shapes.rs"] +mod expression_shapes; +#[path = "tests/helpers.rs"] +mod helpers; +#[path = "tests/test_mappings.rs"] +mod mappings; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs new file mode 100644 index 0000000000..fccee1d92b --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs @@ -0,0 +1,212 @@ +use crate::assembly::analysis_session::_helpers::mappings::catalog_relations; +use crate::assembly::analysis_session::main::finish_analysis_session::finish_analysis_session; +use crate::assembly::analysis_session::main::provide_deferred_analyses::provide_deferred_analyses; +use crate::assembly::analysis_session::main::run_analysis_session::run_analysis_session; +use crate::assembly::analysis_session::main::start_analysis_session::start_analysis_session; +use crate::assembly::analysis_session::models::{ + AnalysisSession, Deferral, DeferredAnalysis, ModelOutcome, ModelReference, ModelRequest, + SessionRequest, SessionStep, +}; +use crate::assembly::analysis_session::tests::test_types::ModelSpec; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::models::{CatalogInput, ProjectCatalog}; + +pub(crate) fn pairs(values: &[(&str, &str)]) -> Pairs { + values + .iter() + .map(|(name, value)| ((*name).to_owned(), (*value).to_owned())) + .collect() +} + +pub(crate) fn shapes(values: &[(&str, &[(&str, &str)])]) -> Shapes { + values + .iter() + .map(|(name, shape)| ((*name).to_owned(), pairs(shape))) + .collect() +} + +/// The compile catalog holding the closed `relations` Python registered before analysis. +pub(crate) fn catalog(dialect: &str, relations: &Shapes) -> ProjectCatalog { + ProjectCatalog::new(CatalogInput { + dialect: dialect.to_owned(), + quoted_ignore_case: false, + known_functions: Vec::new(), + known_types: Vec::new(), + relations: catalog_relations(relations), + }) + .expect("the catalog dialect is known") +} + +/// A model reading `sources` (source names) and `refs` (model names). +pub(crate) fn model(name: &str, sql: &str, sources: &[&str], refs: &[&str]) -> ModelRequest { + let mut references: Vec = Vec::new(); + let mut lineage_references: Vec<(String, String, String)> = Vec::new(); + for source in sources { + references.push(ModelReference { + analysis_name: (*source).to_owned(), + model_ref: false, + }); + lineage_references.push(( + (*source).to_owned(), + "source".to_owned(), + (*source).to_owned(), + )); + } + for model_ref in refs { + references.push(ModelReference { + analysis_name: (*model_ref).to_owned(), + model_ref: true, + }); + lineage_references.push(( + (*model_ref).to_owned(), + "model".to_owned(), + (*model_ref).to_owned(), + )); + } + let mut required_names: Vec = sources + .iter() + .chain(refs) + .map(|name| (*name).to_owned()) + .collect(); + required_names.sort(); + ModelRequest { + name: name.to_owned(), + query_sql: sql.to_owned(), + placeholders: Vec::new(), + references, + lineage_references, + required_names, + recover_cte_facts: false, + has_set_operation: sql.contains("UNION"), + snapshot_columns: None, + } +} + +/// A DuckDB request whose `raw_orders` source is closed and typed. +pub(crate) fn orders_request(models: Vec) -> SessionRequest { + let raw_orders: Shapes = shapes(&[( + "raw_orders", + &[("order_id", "INTEGER"), ("amount", "DOUBLE")], + )]); + SessionRequest { + dialect: "duckdb".to_owned(), + case_sensitive_shapes: false, + function_return_types: Vec::new(), + rich_type_inference: true, + column_types: raw_orders.clone(), + column_nullability: shapes(&[( + "raw_orders", + &[("order_id", "non_null"), ("amount", "unknown")], + )]), + complete_schemas: raw_orders.clone(), + catalog_schemas: raw_orders, + models, + } +} + +pub(crate) fn started(request: SessionRequest) -> AnalysisSession { + let catalog: ProjectCatalog = catalog(&request.dialect, &request.catalog_schemas); + start_analysis_session(request, &catalog).expect("the models form no cycle") +} + +/// Run to completion, answering every deferral with `answer`; return the steps and outcomes. +pub(crate) fn completed( + mut session: AnalysisSession, + answer: fn(&Deferral) -> DeferredAnalysis, +) -> (Vec, Vec) { + let mut done: bool = false; + let steps: Vec = std::iter::from_fn(|| { + (!done).then(|| { + let step: SessionStep = run_analysis_session(&mut session).expect("the session runs"); + let answers: Vec = step.deferrals.iter().map(answer).collect(); + done = answers.is_empty(); + (!done).then(|| { + provide_deferred_analyses(&mut session, answers).expect("the answers match") + }); + step + }) + }) + .collect(); + let outcome = finish_analysis_session(session).expect("the session finished"); + (steps, outcome.models) +} + +/// `name type nullability` per column, or `failed` when analysis did not succeed. +pub(crate) fn described(outcome: &ModelOutcome) -> Vec { + let analysis = &outcome.analysis; + let mut lines: Vec = analysis + .columns + .iter() + .flatten() + .map(|column| { + format!( + "{} {} {}", + column.name, + column.data_type.as_deref().unwrap_or("-"), + column.nullability + ) + }) + .collect(); + lines.push(format!( + "succeeded={} star={}/{} diagnostics={:?}", + analysis.analysis_succeeded, + analysis.has_star, + analysis.star_resolved, + analysis + .binding_diagnostics + .iter() + .map(|row| row.0.as_str()) + .collect::>() + )); + lines +} + +/// A successful Python answer with no columns and no star. +pub(crate) fn empty_answer(deferral: &Deferral) -> DeferredAnalysis { + DeferredAnalysis { + model: deferral.model(), + analysis_succeeded: true, + columns: None, + has_star: false, + star_resolved: false, + binding_diagnostics: Vec::new(), + binding_validated: true, + } +} + +/// `publish name col:TYPE ...` per publication, then `defer kind model` per deferral. +pub(crate) fn step_lines(step: &SessionStep) -> Vec { + let mut lines: Vec = Vec::new(); + for (name, shape) in &step.publications { + let columns: Vec = shape + .iter() + .map(|(column, data_type)| format!("{column}:{data_type}")) + .collect(); + lines.push(format!("publish {name} {}", columns.join(" "))); + } + lines.extend( + step.deferrals + .iter() + .map(|deferral| format!("defer {} {}", deferral.kind(), deferral.model())), + ); + lines +} + +pub(crate) fn model_requests(models: &[ModelSpec]) -> Vec { + models + .iter() + .map(|(name, sql, sources, refs)| model(name, sql, sources, refs)) + .collect() +} + +/// Run `models` over the orders request: each step's lines and each model's description. +pub(crate) fn session_lines(models: &[ModelSpec]) -> (Vec>, Vec>) { + let (steps, outcomes) = completed( + started(orders_request(model_requests(models))), + empty_answer, + ); + ( + steps.iter().map(step_lines).collect(), + outcomes.iter().map(described).collect(), + ) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs new file mode 100644 index 0000000000..4f53c3093f --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs @@ -0,0 +1,166 @@ +use crate::assembly::analysis_session::main::start_analysis_session::start_analysis_session; +use crate::assembly::analysis_session::tests::helpers::{ + catalog, model_requests, orders_request, session_lines, +}; +use crate::assembly::analysis_session::tests::test_types::{SessionTestCase, UnscheduledTestCase}; + +#[test] +fn given_models_when_running_the_session_then_defers_and_publishes_as_python_does() { + let test_cases = [ + SessionTestCase { + description: "typed producers publish closed shapes their consumers bind", + models: &[ + ( + "stg_orders", + "SELECT order_id, amount * 2 AS doubled FROM __source(\"raw_orders\")", + &["raw_orders"], + &[], + ), + ( + "orders_mart", + "SELECT doubled, order_id + 1 AS next_id FROM __ref(\"stg_orders\")", + &[], + &["stg_orders"], + ), + ( + "orders_star", + "SELECT * FROM __ref(\"stg_orders\")", + &[], + &["stg_orders"], + ), + ], + expected_steps: &[&[ + "publish stg_orders order_id:INTEGER doubled:DOUBLE", + "publish orders_mart doubled:DOUBLE next_id:INT", + "publish orders_star order_id:INT doubled:DOUBLE", + ]], + expected_outcomes: &[ + &[ + "order_id INTEGER non_null", + "doubled DOUBLE unknown", + "succeeded=true star=false/true diagnostics=[]", + ], + &[ + "doubled DOUBLE unknown", + "next_id INT unknown", + "succeeded=true star=false/true diagnostics=[]", + ], + &[ + "order_id INT unknown", + "doubled DOUBLE unknown", + "succeeded=true star=true/true diagnostics=[]", + ], + ], + }, + SessionTestCase { + description: "an unknown column is a binding diagnostic and an untyped output defers", + models: &[( + "orders_bad", + "SELECT missing_column FROM __source(\"raw_orders\")", + &["raw_orders"], + &[], + )], + expected_steps: &[&["defer enrichment 0"], &[]], + expected_outcomes: &[&[ + "missing_column - unknown", + "succeeded=true star=false/false diagnostics=[\"B002\"]", + ]], + }, + SessionTestCase { + description: "an open source's consumers defer to Python", + models: &[ + ( + "events", + "SELECT event_id FROM __source(\"raw_events\")", + &["raw_events"], + &[], + ), + ( + "events_mart", + "SELECT event_id, 1 AS one FROM __ref(\"events\")", + &[], + &["events"], + ), + ], + expected_steps: &[ + &["defer enrichment 0"], + &["publish events event_id:UNKNOWN", "defer analysis 1"], + &[], + ], + expected_outcomes: &[ + &[ + "event_id - unknown", + "succeeded=true star=false/false diagnostics=[]", + ], + &["succeeded=true star=false/false diagnostics=[]"], + ], + }, + SessionTestCase { + description: "a set operation never re-analyses with its inputs", + models: &[( + "orders_union", + "SELECT order_id FROM __source(\"raw_events\") UNION ALL SELECT 1", + &["raw_events"], + &[], + )], + expected_steps: &[&["publish orders_union order_id:UNKNOWN"]], + expected_outcomes: &[&[ + "order_id - unknown", + "succeeded=true star=false/false diagnostics=[]", + ]], + }, + ]; + for test_case in test_cases { + let (steps, outcomes) = session_lines(test_case.models); + + assert_eq!(steps, test_case.expected_steps, "{}", test_case.description); + assert_eq!( + outcomes, test_case.expected_outcomes, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_unschedulable_models_when_starting_then_python_analyses() { + let test_cases = [ + UnscheduledTestCase { + description: "models that reference each other", + models: &[ + ( + "orders", + "SELECT order_id FROM __ref(\"returns\")", + &[], + &["returns"], + ), + ( + "returns", + "SELECT order_id FROM __ref(\"orders\")", + &[], + &["orders"], + ), + ], + expected_started: false, + }, + UnscheduledTestCase { + description: "two models with one name", + models: &[ + ("orders", "SELECT 1 AS order_id", &[], &[]), + ("orders", "SELECT 2 AS order_id", &[], &[]), + ], + expected_started: false, + }, + ]; + for test_case in test_cases { + let request = orders_request(model_requests(test_case.models)); + let catalog = catalog(&request.dialect, &request.catalog_schemas); + + assert_eq!( + start_analysis_session(request, &catalog).is_some(), + test_case.expected_started, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_expression_shapes.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_expression_shapes.rs new file mode 100644 index 0000000000..6811ffd806 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_expression_shapes.rs @@ -0,0 +1,70 @@ +use crate::assembly::analysis_session::main::expression_shapes::expression_shapes; +use crate::assembly::analysis_session::models::ExpressionShapeRequest; +use crate::assembly::analysis_session::tests::helpers::catalog; +use crate::assembly::analysis_session::tests::test_types::ExpressionShapeTestCase; + +#[test] +fn given_source_expressions_when_inferring_shapes_then_matches_python_shapes() { + let test_cases = [ + ExpressionShapeTestCase { + description: "literal columns are typed", + dialect: "duckdb", + case_sensitive_shapes: false, + expression: "(SELECT 1 AS order_id, 'placed' AS status)", + expected_shape: r#"Inferred([("order_id", "INT"), ("status", "TEXT")])"#, + }, + ExpressionShapeTestCase { + description: "a star has no closed shape", + dialect: "duckdb", + case_sensitive_shapes: false, + expression: "SELECT * FROM orders", + expected_shape: "Absent", + }, + ExpressionShapeTestCase { + description: "unparseable SQL has no shape", + dialect: "duckdb", + case_sensitive_shapes: false, + expression: "SELECT FROM WHERE", + expected_shape: "Absent", + }, + ExpressionShapeTestCase { + description: "an unknown input column is UNKNOWN", + dialect: "duckdb", + case_sensitive_shapes: false, + expression: "SELECT order_id + 1 AS next_id FROM orders", + expected_shape: r#"Inferred([("next_id", "UNKNOWN")])"#, + }, + ExpressionShapeTestCase { + description: "case-insensitive dialects drop authored quoting", + dialect: "duckdb", + case_sensitive_shapes: false, + expression: "SELECT 1 AS \"Mixed_Case\", CAST(2 AS BIGINT) AS total", + expected_shape: r#"Inferred([("Mixed_Case", "INT"), ("total", "BIGINT")])"#, + }, + ExpressionShapeTestCase { + description: "case-sensitive dialects keep authored quoting", + dialect: "postgresql", + case_sensitive_shapes: true, + expression: "SELECT 1 AS \"Mixed_Case\", CAST(2 AS BIGINT) AS total", + expected_shape: r#"Inferred([("\"Mixed_Case\"", "INT"), ("total", "BIGINT")])"#, + }, + ]; + for test_case in test_cases { + let request = ExpressionShapeRequest { + dialect: test_case.dialect.to_owned(), + case_sensitive_shapes: test_case.case_sensitive_shapes, + function_return_types: Vec::new(), + expressions: vec![test_case.expression.to_owned()], + }; + + let shapes = expression_shapes(&catalog(test_case.dialect, &Vec::new()), &request) + .expect("the batch runs"); + + assert_eq!( + format!("{shapes:?}"), + format!("[{}]", test_case.expected_shape), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_mappings.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_mappings.rs new file mode 100644 index 0000000000..1f510017ce --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_mappings.rs @@ -0,0 +1,71 @@ +use crate::assembly::analysis_session::_helpers::mappings::{dict_from_pairs, same_mapping}; +use crate::assembly::analysis_session::_helpers::schedule::waves; +use crate::assembly::analysis_session::tests::helpers::pairs; +use crate::assembly::analysis_session::tests::test_types::{DictTestCase, WavesTestCase}; + +#[test] +fn given_pairs_when_building_a_dict_then_keeps_first_position_and_last_value() { + let test_cases = [ + DictTestCase { + description: "distinct keys keep their order", + pairs: &[("b", "INT"), ("a", "TEXT")], + expected_dict: &[("b", "INT"), ("a", "TEXT")], + }, + DictTestCase { + description: "a repeated key keeps its first position and takes the last value", + pairs: &[("up", "INT"), ("other", "TEXT"), ("up", "BIGINT")], + expected_dict: &[("up", "BIGINT"), ("other", "TEXT")], + }, + ]; + for test_case in test_cases { + let built = dict_from_pairs(pairs(test_case.pairs)); + + assert_eq!( + built, + pairs(test_case.expected_dict), + "{}", + test_case.description + ); + assert!( + same_mapping(&built, &built.iter().rev().cloned().collect()), + "{}", + test_case.description + ); + } +} + +#[test] +fn given_producers_when_scheduling_then_returns_topological_waves() { + let test_cases = [ + WavesTestCase { + description: "independent models share one wave", + producers: &[&[], &[]], + expected_waves: Some(&[&[0, 1]]), + }, + WavesTestCase { + description: "a consumer waits for every producer", + producers: &[&[2], &[], &[1], &[0, 1]], + expected_waves: Some(&[&[1], &[2], &[0], &[3]]), + }, + WavesTestCase { + description: "a cycle cannot be scheduled", + producers: &[&[1], &[0]], + expected_waves: None, + }, + ]; + for test_case in test_cases { + let producers: Vec> = test_case + .producers + .iter() + .map(|producers| producers.to_vec()) + .collect(); + let scheduled: Option>> = waves(&producers); + + assert_eq!( + format!("{scheduled:?}"), + format!("{:?}", test_case.expected_waves), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs new file mode 100644 index 0000000000..f94d532933 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs @@ -0,0 +1,42 @@ +/// A model as `(name, sql, source names, model refs)`. +pub(crate) type ModelSpec = ( + &'static str, + &'static str, + &'static [&'static str], + &'static [&'static str], +); + +pub(crate) struct SessionTestCase { + pub(crate) description: &'static str, + pub(crate) models: &'static [ModelSpec], + /// `step_lines` of each step, in order. + pub(crate) expected_steps: &'static [&'static [&'static str]], + /// `described` of each model's outcome, in request order. + pub(crate) expected_outcomes: &'static [&'static [&'static str]], +} + +pub(crate) struct UnscheduledTestCase { + pub(crate) description: &'static str, + pub(crate) models: &'static [ModelSpec], + pub(crate) expected_started: bool, +} + +pub(crate) struct ExpressionShapeTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) case_sensitive_shapes: bool, + pub(crate) expression: &'static str, + pub(crate) expected_shape: &'static str, +} + +pub(crate) struct DictTestCase { + pub(crate) description: &'static str, + pub(crate) pairs: &'static [(&'static str, &'static str)], + pub(crate) expected_dict: &'static [(&'static str, &'static str)], +} + +pub(crate) struct WavesTestCase { + pub(crate) description: &'static str, + pub(crate) producers: &'static [&'static [usize]], + pub(crate) expected_waves: Option<&'static [&'static [usize]]>, +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/types.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/types.rs new file mode 100644 index 0000000000..5c36906409 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/types.rs @@ -0,0 +1,6 @@ +//! Ordered mappings shared with Python's dicts. + +/// An ordered `name -> value` mapping in Python dict order. +pub type Pairs = Vec<(String, String)>; +/// Ordered relation shapes in Python dict order. +pub type Shapes = Vec<(String, Pairs)>; diff --git a/crates/sqlbuild-analysis/src/assembly/mod.rs b/crates/sqlbuild-analysis/src/assembly/mod.rs index 079a1bcfca..e3c29d7550 100644 --- a/crates/sqlbuild-analysis/src/assembly/mod.rs +++ b/crates/sqlbuild-analysis/src/assembly/mod.rs @@ -1,3 +1,4 @@ //! The native model analysis session: catalog, dataflow scheduling, projection and shape publication. +pub mod analysis_session; pub mod project; diff --git a/crates/sqlbuild-analysis/src/semantic_validation/_helpers/catalog.rs b/crates/sqlbuild-analysis/src/semantic_validation/_helpers/catalog.rs index 2c2e139df8..3ce72f609a 100644 --- a/crates/sqlbuild-analysis/src/semantic_validation/_helpers/catalog.rs +++ b/crates/sqlbuild-analysis/src/semantic_validation/_helpers/catalog.rs @@ -26,13 +26,13 @@ const ANALYSIS_WORKERS: usize = 4; const ANALYSIS_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; impl ProjectCatalog { - /// Map inferred column types onto the binding names the query projects. + /// Map inferred column types onto the binding names the query projects, in column order. pub fn inferred_schema( &self, sql: &str, columns: Columns, inputs: Relations, - ) -> Result>, String> { + ) -> Result)>, String> { use polyglot_sql::{Expression, ExpressionWalk}; let parsed = Dialect::get(self.dialect) .parse(sql) @@ -70,14 +70,21 @@ impl ProjectCatalog { } } } - let mut result: HashMap> = HashMap::new(); + let mut result: Vec<(String, Option)> = Vec::with_capacity(columns.0.len()); + let mut positions: HashMap = HashMap::new(); for (name, column_type) in columns.0 { let binding_name = identifiers .get(&name) .or_else(|| exact_inputs.get(&name)) .cloned() .unwrap_or(name); - result.insert(binding_name, column_type); + match positions.get(&binding_name) { + Some(position) => result[*position].1 = column_type, + None => { + positions.insert(binding_name.clone(), result.len()); + result.push((binding_name, column_type)); + } + } } Ok(result) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis/catalog.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis/catalog.rs index 70ef2a70d0..37314dacda 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis/catalog.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis/catalog.rs @@ -23,7 +23,7 @@ impl ProjectCatalog { sql: &str, #[pyo3(from_py_with = conversions::columns)] columns: validation::Columns, #[pyo3(from_py_with = conversions::relations)] inputs: Relations, - ) -> PyResult>> { + ) -> PyResult)>> { py.compiler_detach(|| self.inner.inferred_schema(sql, columns, inputs)) .map_err(compiler_error) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs index 7444bd0374..54e3fdf710 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/mod.rs @@ -1 +1,2 @@ pub(crate) mod registration; +pub(crate) mod session; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs index bbd9c13b6b..bfe82cb5d1 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/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::analysis_session::session; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + session::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs new file mode 100644 index 0000000000..2e3c48a906 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs @@ -0,0 +1,358 @@ +//! The native model analysis session and expression-source shapes for the preview engine. + +use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyRef, PyResult, Python}; +use pyo3::{pyclass, pyfunction, pymethods, wrap_pyfunction}; +use sqlbuild_analysis::assembly::analysis_session::main::expression_shapes::expression_shapes; +use sqlbuild_analysis::assembly::analysis_session::main::finish_analysis_session::finish_analysis_session; +use sqlbuild_analysis::assembly::analysis_session::main::provide_deferred_analyses::provide_deferred_analyses; +use sqlbuild_analysis::assembly::analysis_session::main::run_analysis_session::run_analysis_session; +use sqlbuild_analysis::assembly::analysis_session::main::start_analysis_session::start_analysis_session; +use sqlbuild_analysis::assembly::analysis_session::models::{ + AnalysisSession, ColumnFact, Deferral, DeferredAnalysis, ExpressionShape, + ExpressionShapeRequest, LineageFacts, LineageRow, ModelOutcome, ModelReference, ModelRequest, + SessionOutcome, SessionRequest, SessionStep, +}; +use sqlbuild_analysis::assembly::analysis_session::types::{Pairs, Shapes}; +use sqlbuild_analysis::semantic_validation::types::DiagnosticRow; + +use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic; + +use crate::bindings::_helpers::boundary::panics::compiler_error; +use crate::bindings::models::ProjectCatalog; +use crate::bindings::types::CompilerDetach; + +/// `(name, type, nullability)`. +type ColumnRow = (String, Option, String); +/// `(output column, transform code, confidence code, [(resource type, resource, column)])`. +type LineageItem = (String, u8, u8, Vec<(String, String, String)>); +/// A model's name, SQL, placeholders, references, lineage, required names and flags. +type ModelRow = ( + String, + String, + Pairs, + Vec<(String, bool)>, + Vec<(String, String, String)>, + Vec, + bool, + bool, + Option<(String, String)>, +); +/// The profile settings, relation shapes and models one session analyses. +type RequestRow = ( + String, + bool, + Pairs, + bool, + Shapes, + Shapes, + Shapes, + Shapes, + Vec, +); +/// `(kind, model, cleaned_sql, schemas, binding diagnostics, lineage)`. +type DeferralRow = ( + &'static str, + usize, + Option, + Shapes, + Vec, + Option>, +); +/// `(publications, deferrals, failures)`. +type StepRow = (Shapes, Vec, Vec); +/// `(model, succeeded, columns, has_star, star_resolved, binding diagnostics, validated)`. +type DeferredRow = ( + usize, + bool, + Option>, + bool, + bool, + Vec, + bool, +); +/// One model's analysis, lineage source and cleaned SQL. +type OutcomeRow = ( + bool, + Option>, + &'static str, + Vec, + bool, + bool, + Vec, + bool, + String, +); +/// `(models, catalog schema additions, analysis-shape names)`. +type FinishRow = (Vec, Shapes, Vec); + +/// One compile's native model analysis; any failure hands the whole analysis back to Python. +#[pyclass(module = "sqlbuild._native")] +pub(crate) struct NativeModelAnalysisSession { + inner: Option, + failure: Option, +} + +impl NativeModelAnalysisSession { + /// Keep `result`'s value, or end the session and keep its failure for Python's debug log. + fn kept(&mut self, result: Result) -> Option { + match result { + Ok(value) => Some(value), + Err(error) => { + self.inner = None; + self.failure = Some(error); + None + } + } + } +} + +#[pymethods] +impl NativeModelAnalysisSession { + /// Advance; None means Python must analyse every model, no deferrals means done. + fn run(&mut self, py: Python<'_>) -> Option { + let mut session: AnalysisSession = self.inner.take()?; + let result: Result<(AnalysisSession, SessionStep), String> = + py.compiler_detach(move || { + run_analysis_session(&mut session).map(|step| (session, step)) + }); + let (session, step) = self.kept(result)?; + self.inner = Some(session); + Some(step_row(step)) + } + + /// Answer the last step's deferrals; False means Python must analyse every model. + fn provide(&mut self, results: Vec) -> bool { + let Some(mut session) = self.inner.take() else { + return false; + }; + let answers: Vec = results.into_iter().map(deferred_analysis).collect(); + let provided: Result<(), String> = + catch_compiler_panic(|| provide_deferred_analyses(&mut session, answers)); + self.inner = Some(session); + self.kept(provided).is_some() + } + + /// Every model's outcome once the session is done, or None to analyse in Python. + fn finish(&mut self) -> Option { + let session: AnalysisSession = self.inner.take()?; + let outcome: Result = + catch_compiler_panic(|| finish_analysis_session(session).map(finish_row)); + self.kept(outcome) + } + + /// Why the session handed the analysis back to Python, when it did. + #[getter] + fn failure(&self) -> Option { + self.failure.clone() + } +} + +/// Start a session on `catalog`, or None where Python must analyse. +#[pyfunction] +fn start_model_analysis_session( + py: Python<'_>, + catalog: PyRef<'_, ProjectCatalog>, + request: RequestRow, +) -> PyResult> { + let request: SessionRequest = session_request(request); + let catalog = &catalog.inner; + let session: Option = py + .compiler_detach(|| Ok(start_analysis_session(request, catalog))) + .map_err(compiler_error)?; + Ok(session.map(|session| NativeModelAnalysisSession { + inner: Some(session), + failure: None, + })) +} + +/// Each expression's shape, with `False` where Python must infer it. +#[pyfunction] +fn infer_expression_source_shapes( + py: Python<'_>, + catalog: PyRef<'_, ProjectCatalog>, + request: (String, bool, Pairs, Vec), +) -> (Vec<(bool, Option)>, Option) { + let (dialect, case_sensitive_shapes, function_return_types, expressions) = request; + let request = ExpressionShapeRequest { + dialect, + case_sensitive_shapes, + function_return_types, + expressions, + }; + let catalog = &catalog.inner; + match py.compiler_detach(|| expression_shapes(catalog, &request)) { + Ok(shapes) => (shapes.into_iter().map(expression_shape_row).collect(), None), + Err(error) => (Vec::new(), Some(error)), + } +} + +fn expression_shape_row(shape: ExpressionShape) -> (bool, Option) { + match shape { + ExpressionShape::Inferred(shape) => (true, Some(shape)), + ExpressionShape::Absent => (true, None), + ExpressionShape::Deferred => (false, None), + } +} + +fn session_request(request: RequestRow) -> SessionRequest { + let ( + dialect, + case_sensitive_shapes, + function_return_types, + rich_type_inference, + column_types, + column_nullability, + complete_schemas, + catalog_schemas, + models, + ) = request; + SessionRequest { + dialect, + case_sensitive_shapes, + function_return_types, + rich_type_inference, + column_types, + column_nullability, + complete_schemas, + catalog_schemas, + models: models.into_iter().map(model_request).collect(), + } +} + +fn model_request(row: ModelRow) -> ModelRequest { + let ( + name, + query_sql, + placeholders, + references, + lineage_references, + required_names, + recover_cte_facts, + has_set_operation, + snapshot_columns, + ) = row; + ModelRequest { + name, + query_sql, + placeholders, + references: references + .into_iter() + .map(|(analysis_name, model_ref)| ModelReference { + analysis_name, + model_ref, + }) + .collect(), + lineage_references, + required_names, + recover_cte_facts, + has_set_operation, + snapshot_columns, + } +} + +fn step_row(step: SessionStep) -> StepRow { + ( + step.publications, + step.deferrals.into_iter().map(deferral_row).collect(), + step.failures, + ) +} + +fn deferral_row(deferral: Deferral) -> DeferralRow { + let kind: &'static str = deferral.kind(); + match deferral { + Deferral::Analysis { + model, + cleaned_sql, + binding_schema, + binding_diagnostics, + lineage, + } => ( + kind, + model, + Some(cleaned_sql), + binding_schema, + binding_diagnostics, + lineage.map(lineage_items), + ), + Deferral::Enrichment { + model, + input_schemas, + } => (kind, model, None, input_schemas, Vec::new(), None), + } +} + +fn deferred_analysis(row: DeferredRow) -> DeferredAnalysis { + let (model, analysis_succeeded, columns, has_star, star_resolved, diagnostics, validated) = row; + DeferredAnalysis { + model, + analysis_succeeded, + columns: columns.map(|columns| columns.into_iter().map(column_fact).collect()), + has_star, + star_resolved, + binding_diagnostics: diagnostics, + binding_validated: validated, + } +} + +fn column_fact((name, data_type, nullability): ColumnRow) -> ColumnFact { + ColumnFact { + name, + data_type, + nullability, + } +} + +fn column_row(column: ColumnFact) -> ColumnRow { + (column.name, column.data_type, column.nullability) +} + +fn lineage_items(rows: Vec) -> Vec { + rows.into_iter() + .map(|row| { + ( + row.output_column, + row.transform_code, + row.confidence_code, + row.sources, + ) + }) + .collect() +} + +fn finish_row(outcome: SessionOutcome) -> FinishRow { + ( + outcome.models.into_iter().map(outcome_row).collect(), + outcome.schema_additions, + outcome.analysis_names, + ) +} + +fn outcome_row(outcome: ModelOutcome) -> OutcomeRow { + let analysis = outcome.analysis; + let (source, lineage) = match analysis.lineage { + LineageFacts::Native(rows) => ("native", lineage_items(rows)), + LineageFacts::PythonAnalysis => ("analysis", Vec::new()), + LineageFacts::PythonEnrichment => ("enrichment", Vec::new()), + }; + ( + analysis.analysis_succeeded, + analysis + .columns + .map(|columns| columns.into_iter().map(column_row).collect()), + source, + lineage, + analysis.has_star, + analysis.star_resolved, + analysis.binding_diagnostics, + analysis.binding_validated, + outcome.cleaned_sql, + ) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_function(wrap_pyfunction!(start_model_analysis_session, module)?)?; + module.add_function(wrap_pyfunction!(infer_expression_source_shapes, module)?)?; + Ok(()) +} diff --git a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py index 9faaa527da..2c00e8a6d6 100644 --- a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py +++ b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py @@ -1,9 +1,50 @@ """Failure cases owned by the native model analysis session lane.""" +from scripts.compiler_differential._helpers.corpus.case_builder import failure_case +from scripts.compiler_differential.constants import FAILURE_MART_PATH, FAILURE_STAGING_PATH from scripts.compiler_differential.models import FailureCase +_ORDERS: str = '__ref("stg_orders")' +_MART_HEADER: str = 'MODEL (\n description "Order totals per customer",\n);\n\n' +_STAGING_HEADER: str = 'MODEL (\n description "Staged orders",\n);\n\n' +_LAYERED_PATH: str = "models/marts/order_layers.sql" +_LAYERED_MODEL: str = ( + 'MODEL (\n description "Orders layered over the staged star",\n);\n\n' + f"SELECT s.*, s.amount * 2 AS doubled_amount\nFROM {_ORDERS} s\n" +) + def analysis_session_failure_cases() -> tuple[FailureCase, ...]: """Return this lane's failure cases; the lane appends here without editing shared lists.""" - return () + return ( + failure_case( + name="analysis-session-star-published-unknown-column", + expected_code="B002", + files={ + FAILURE_STAGING_PATH: _STAGING_HEADER + 'SELECT *\nFROM __source("raw_orders")\n', + _LAYERED_PATH: _LAYERED_MODEL, + FAILURE_MART_PATH: _MART_HEADER + + "SELECT customer_id, SUM(missing_amount) AS total_amount\n" + 'FROM __ref("order_layers")\nGROUP BY customer_id\n', + }, + ), + failure_case( + name="analysis-session-set-operation-unknown-column", + expected_code="B002", + files={ + FAILURE_MART_PATH: _MART_HEADER + + f"SELECT customer_id, amount AS total_amount FROM {_ORDERS}\n" + f"UNION ALL\nSELECT customer_id, refund_amount FROM {_ORDERS}\n", + }, + ), + failure_case( + name="analysis-session-cte-chain-unknown-column", + expected_code="B002", + files={ + FAILURE_MART_PATH: _MART_HEADER + f"WITH base AS (SELECT * FROM {_ORDERS}),\n" + "totals AS (SELECT customer_id, SUM(amount) AS total_amount FROM base GROUP BY 1)\n" + "SELECT customer_id, total_amount, order_count FROM totals\n", + }, + ), + ) diff --git a/scripts/compiler_differential/_helpers/comparing/capture.py b/scripts/compiler_differential/_helpers/comparing/capture.py index 4f35198c9c..acd1d261c7 100644 --- a/scripts/compiler_differential/_helpers/comparing/capture.py +++ b/scripts/compiler_differential/_helpers/comparing/capture.py @@ -17,8 +17,10 @@ from scripts.compiler_differential.models import Divergence -def first_capture_difference(*, left: Path, right: Path) -> Divergence | None: - """Return the first logical JSON pointer, or line for invalid JSON, where captures differ.""" +def first_capture_difference( + *, left: Path, right: Path, ignored: frozenset[str] = frozenset() +) -> Divergence | None: + """Return the first JSON pointer outside `ignored`, or line for invalid JSON, that differs.""" try: left_capture: CaptureFile = CaptureFile(left) @@ -27,6 +29,7 @@ def first_capture_difference(*, left: Path, right: Path) -> Divergence | None: left=left_capture.root, right=right_capture.root, resolve=(left_capture.resolve, right_capture.resolve), + ignored=ignored, ) except (json.JSONDecodeError, UnicodeDecodeError): return _first_line_difference(left=left, right=right) diff --git a/scripts/compiler_differential/_helpers/comparing/compare.py b/scripts/compiler_differential/_helpers/comparing/compare.py index 4d6af02808..cb8d5a17c6 100644 --- a/scripts/compiler_differential/_helpers/comparing/compare.py +++ b/scripts/compiler_differential/_helpers/comparing/compare.py @@ -13,9 +13,14 @@ def first_json_difference( - *, left: object, right: object, pointer: str = "", resolve: NodeResolvers | None = None + *, + left: object, + right: object, + pointer: str = "", + resolve: NodeResolvers | None = None, + ignored: frozenset[str] = frozenset(), ) -> Divergence | None: - """Return the first JSON pointer whose value, type, or key order differs, via `resolve`.""" + """Return the first JSON pointer outside `ignored` whose value, type or key order differs.""" if resolve is not None: left_digest: str | None = shared_digest(left) @@ -28,12 +33,16 @@ def first_json_difference( right_object: dict[str, object] | None = as_json_object(right) if left_object is not None and right_object is not None: return _first_object_difference( - left=left_object, right=right_object, pointer=pointer, resolve=resolve + left=left_object, right=right_object, pointer=pointer, resolve=resolve, ignored=ignored ) if isinstance(left, list) and isinstance(right, list): for index, (left_item, right_item) in enumerate(zip(left, right, strict=False)): found: Divergence | None = first_json_difference( - left=left_item, right=right_item, pointer=f"{pointer}/{index}", resolve=resolve + left=left_item, + right=right_item, + pointer=f"{pointer}/{index}", + resolve=resolve, + ignored=ignored, ) if found is not None: return found @@ -99,9 +108,12 @@ def _first_object_difference( right: dict[str, object], pointer: str, resolve: NodeResolvers | None, + ignored: frozenset[str], ) -> Divergence | None: for key in (*left, *(key for key in right if key not in left)): location: str = f"{pointer}/{json_pointer_token(key)}" + if location in ignored: + continue if key not in left or key not in right: return _divergence( location=location, @@ -109,7 +121,7 @@ def _first_object_difference( right=_shown(value=right.get(key, MISSING_VALUE), resolve=resolve, side=1), ) found: Divergence | None = first_json_difference( - left=left[key], right=right[key], pointer=location, resolve=resolve + left=left[key], right=right[key], pointer=location, resolve=resolve, ignored=ignored ) if found is not None: return found diff --git a/scripts/compiler_differential/_helpers/comparing/comparison.py b/scripts/compiler_differential/_helpers/comparing/comparison.py index 53c5976c0b..20fd71c31a 100644 --- a/scripts/compiler_differential/_helpers/comparing/comparison.py +++ b/scripts/compiler_differential/_helpers/comparing/comparison.py @@ -19,6 +19,8 @@ strip_report_fields, ) from scripts.compiler_differential.constants import ( + ANALYSIS_CACHE_STATE_POINTERS, + COLD_CAPTURE_COMMAND, MISSING_VALUE, MODEL_RESOURCE_TYPE, PLAN_LABEL, @@ -266,7 +268,11 @@ def _capture_differences( differences.append( ( artifact, - first_capture_difference(left=left_files[name], right=right_files[name]), + first_capture_difference( + left=left_files[name], + right=right_files[name], + ignored=_cache_state_pointers(command), + ), ) ) return differences @@ -278,3 +284,9 @@ def _presence_divergence(*, left_present: bool, right_present: bool) -> Divergen left="present" if left_present else MISSING_VALUE, right="present" if right_present else MISSING_VALUE, ) + + +def _cache_state_pointers(command: str) -> frozenset[str]: + """Catalog state Python's warm analysis cache skips; compared only on the cold compile.""" + + return frozenset() if command == COLD_CAPTURE_COMMAND else ANALYSIS_CACHE_STATE_POINTERS diff --git a/scripts/compiler_differential/constants.py b/scripts/compiler_differential/constants.py index 03da9883f1..008e546c06 100644 --- a/scripts/compiler_differential/constants.py +++ b/scripts/compiler_differential/constants.py @@ -86,6 +86,10 @@ def short_model_names(*, model: Model, ctx: RuleContext) -> list[Finding]: COLD_COMPILE: DifferentialCommand = DifferentialCommand( label="compile", arguments=("compile", "--json", "--manifest", "--dag") ) +COLD_CAPTURE_COMMAND: str = f"0-{COLD_COMPILE.label}" +ANALYSIS_CACHE_STATE_POINTERS: frozenset[str] = frozenset( + {"/binding_catalog/schemas", "/binding_catalog/analysis_shapes"} +) WARM_COMPILE: DifferentialCommand = DifferentialCommand( label="compile-warm", arguments=("compile", "--json") ) diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 38f4bd7514..75d48ddb6c 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -341,6 +341,74 @@ def normalize_type( ) -> tuple[tuple[str, str, int | None, int | None, int | None], str | None] | None: ... # Native analysis: model analysis session. +type _AnalysisDiagnosticRow = tuple[str, str, int | None, int | None, int | None, int | None, str] + +class NativeModelAnalysisSession: + def run( + self, + ) -> ( + tuple[ + list[tuple[str, list[tuple[str, str]]]], + list[ + tuple[ + str, + int, + str | None, + list[tuple[str, list[tuple[str, str]]]], + list[_AnalysisDiagnosticRow], + list[tuple[str, int, int, list[tuple[str, str, str]]]] | None, + ] + ], + list[str], + ] + | None + ): ... + def provide( + self, + results: list[ + tuple[ + int, + bool, + list[tuple[str, str | None, str]] | None, + bool, + bool, + list[_AnalysisDiagnosticRow], + bool, + ] + ], + /, + ) -> bool: ... + def finish( + self, + ) -> ( + tuple[ + list[ + tuple[ + bool, + list[tuple[str, str | None, str]] | None, + str, + list[tuple[str, int, int, list[tuple[str, str, str]]]], + bool, + bool, + list[_AnalysisDiagnosticRow], + bool, + str, + ] + ], + list[tuple[str, list[tuple[str, str]]]], + list[str], + ] + | None + ): ... + @property + def failure(self) -> str | None: ... + +def start_model_analysis_session( + catalog: object, request: tuple[object, ...], / +) -> NativeModelAnalysisSession | None: ... +def infer_expression_source_shapes( + catalog: object, request: tuple[str, bool, list[tuple[str, str]], list[str]], / +) -> tuple[list[tuple[bool, list[tuple[str, str]] | None]], str | None]: ... # Native analysis: semantic completion. diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/__init__.py b/src/sqlbuild/compiler/analysis_session/_helpers/__init__.py new file mode 100644 index 0000000000..64fb62aa62 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/__init__.py @@ -0,0 +1 @@ +"""Native model analysis session helpers.""" diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/deferral_records.py b/src/sqlbuild/compiler/analysis_session/_helpers/deferral_records.py new file mode 100644 index 0000000000..c5ce9b8f4b --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/deferral_records.py @@ -0,0 +1,24 @@ +"""Debug-only record of the model analysis work native hands back to Python.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +from sqlbuild.compiler.analysis_session.constants import ANALYSIS_DEFERRAL_SITE +from sqlbuild.compiler.lineage.constants import ANALYSIS_DEFERRAL_RECORD_PREFIX +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR + + +def record_analysis_deferral(*, kind: str, count: int = 1) -> None: + """Append `count` `{"kind", "site"}` lines when the analysis record directory is set.""" + + directory: str | None = os.environ.get(ANALYSIS_RECORD_DIR_ENV_VAR) + if not directory or count <= 0: + return + target: Path = Path(directory) + target.mkdir(parents=True, exist_ok=True) + record: str = json.dumps({"kind": kind, "site": ANALYSIS_DEFERRAL_SITE}) + with open(target / f"{ANALYSIS_DEFERRAL_RECORD_PREFIX}{os.getpid()}.jsonl", "a") as handle: + _ = handle.write((record + "\n") * count) diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/expression_shapes.py b/src/sqlbuild/compiler/analysis_session/_helpers/expression_shapes.py new file mode 100644 index 0000000000..6a70d30bc3 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/expression_shapes.py @@ -0,0 +1,72 @@ +"""Infer expression-source shapes natively, filling the binding catalog's shape cache.""" + +from __future__ import annotations + +import logging +from typing import Any + +import sqlbuild._native as _native +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.analysis_session._helpers.deferral_records import record_analysis_deferral +from sqlbuild.compiler.analysis_session._helpers.session_rows import case_sensitive_shapes +from sqlbuild.compiler.analysis_session.constants import ( + DEFERRAL_EXPRESSION_SHAPES, + DEFERRAL_NO_CATALOG, + NATIVE_SHAPES_FAILURE_MESSAGE, +) +from sqlbuild.compiler.sql_analysis.constants import NATIVE_DIALECT_ALIASES +from sqlbuild.diagnostics.main.log_debug_event import log_debug_event + +_DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile") + + +def native_expression_source_shapes( + *, expressions: tuple[str, ...], profile: ExpressionInferenceProfile +) -> tuple[dict[str, str] | None, ...] | None: + """One shape per expression, or None (recorded) where Python must infer the shapes.""" + + catalog: Any = profile.binding_catalog + if catalog is None: + record_analysis_deferral(kind=DEFERRAL_NO_CATALOG) + return None + cache: dict[str, dict[str, str] | None] = catalog.expression_shapes + pending: tuple[str, ...] = tuple( + dict.fromkeys(expression for expression in expressions if expression not in cache) + ) + inferred: dict[str, dict[str, str] | None] | None = ( + _inferred_shapes(catalog=catalog, pending=pending, profile=profile) if pending else {} + ) + if inferred is None: + return None + cache.update(inferred) + return tuple( + None if cache[expression] is None else dict(cache[expression] or {}) + for expression in expressions + ) + + +def _inferred_shapes( + *, catalog: Any, pending: tuple[str, ...], profile: ExpressionInferenceProfile +) -> dict[str, dict[str, str] | None] | None: + dialect: str = profile.sql_analysis_dialect or "generic" + dialect = NATIVE_DIALECT_ALIASES.get(dialect, dialect) + shapes, failure = _native.infer_expression_source_shapes( + catalog.native, + ( + dialect, + case_sensitive_shapes(profile=profile, dialect=dialect), + list(profile.function_return_types.items()), + list(pending), + ), + ) + deferred: int = len(pending) if failure is not None else sum(not done for done, _ in shapes) + if deferred: + record_analysis_deferral(kind=DEFERRAL_EXPRESSION_SHAPES, count=deferred) + log_debug_event( + logger=_DEBUG_LOGGER, message=NATIVE_SHAPES_FAILURE_MESSAGE, sqlbuild_error=failure + ) + return None + return { + expression: dict(shape) if shape is not None else None + for expression, (_, shape) in zip(pending, shapes, strict=True) + } diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/model_analysis.py b/src/sqlbuild/compiler/analysis_session/_helpers/model_analysis.py new file mode 100644 index 0000000000..c0ec38e947 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/model_analysis.py @@ -0,0 +1,49 @@ +"""Start the native model analysis session, or hand the whole analysis to Python.""" + +from __future__ import annotations + +from typing import Any + +import sqlbuild._native as _native +from sqlbuild.compiler.analysis_session._helpers.deferral_records import record_analysis_deferral +from sqlbuild.compiler.analysis_session._helpers.session_rows import session_request +from sqlbuild.compiler.analysis_session.classes.native_model_analysis import NativeModelAnalysis +from sqlbuild.compiler.analysis_session.constants import ( + DEFERRAL_NO_CATALOG, + DEFERRAL_NO_COMPACT_ANALYSIS, + DEFERRAL_SESSION, +) +from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest +from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis +from sqlbuild.compiler.compile.models import ModelSqlAnalysis + + +def native_model_analyses( + *, request: NativeModelAnalysisRequest +) -> dict[str, ModelSqlAnalysis] | None: + """Each model's analysis by name, or None (recorded) where Python must analyse.""" + + catalog: Any = request.inference_profile.binding_catalog + if not request.allow_compact_analysis: + record_analysis_deferral(kind=DEFERRAL_NO_COMPACT_ANALYSIS) + return None + if catalog is None: + record_analysis_deferral(kind=DEFERRAL_NO_CATALOG) + return None + if not request.model_inputs: + return {} + python: PythonModelAnalysis = PythonModelAnalysis( + model_inputs=request.model_inputs, + inference_profile=request.inference_profile, + complete_binding_schemas=request.complete_binding_schemas, + ) + row: tuple[object, ...] | None = session_request( + request=request, python=python, schemas=catalog.schemas + ) + session: _native.NativeModelAnalysisSession | None = ( + _native.start_model_analysis_session(catalog.native, row) if row is not None else None + ) + if session is None: + record_analysis_deferral(kind=DEFERRAL_SESSION) + return None + return NativeModelAnalysis(request=request, python=python, catalog=catalog).analyses(session) diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py new file mode 100644 index 0000000000..88b2dab191 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py @@ -0,0 +1,145 @@ +"""Plain rows crossing into and out of the native model analysis session.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest +from sqlbuild.compiler.analysis_session.types import ( + ColumnRow, + DeferredRow, + DiagnosticRow, + LineageItem, + ShapeRows, +) +from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis +from sqlbuild.compiler.compile.models import ( + CompactLineageFacts, + InferredColumn, + PolyglotAnalysisResult, +) +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.compiler.sql_analysis.constants import CASE_SENSITIVE_BINDING_DIALECTS +from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic + + +def shape_rows(shapes: Mapping[str, Mapping[str, object]]) -> ShapeRows | None: + """`{relation: {column: value}}` as ordered rows, or None when a value is not text.""" + + rows: ShapeRows = [] + for name, shape in shapes.items(): + columns: list[tuple[str, str]] = [] + for column, value in shape.items(): + if not isinstance(column, str) or not isinstance(value, str): + return None + columns.append((column, str(value))) + rows.append((name, columns)) + return rows + + +def session_request( + *, + request: NativeModelAnalysisRequest, + python: PythonModelAnalysis, + schemas: Mapping[str, Mapping[str, str]], +) -> tuple[object, ...] | None: + """The session request, or None when a shape cannot cross as text.""" + + profile: ExpressionInferenceProfile = request.inference_profile + shapes: tuple[ShapeRows | None, ...] = ( + shape_rows(request.column_types_by_table), + shape_rows(request.column_nullability_by_table), + shape_rows(request.complete_binding_schemas), + shape_rows(schemas), + ) + if any(rows is None for rows in shapes): + return None + return ( + profile.sql_analysis_dialect or "generic", + case_sensitive_shapes(profile=profile, dialect=profile.sql_analysis_dialect), + list(profile.function_return_types.items()), + request.rich_type_inference, + *shapes, + python.model_rows(), + ) + + +def case_sensitive_shapes(*, profile: ExpressionInferenceProfile, dialect: str | None) -> bool: + """Python's `inferred_binding_shape` test for keeping authored identifier quoting.""" + + return not profile.quoted_identifiers_ignore_case and dialect in CASE_SENSITIVE_BINDING_DIALECTS + + +def column_rows(columns: tuple[InferredColumn, ...] | None) -> list[ColumnRow] | None: + if columns is None: + return None + return [(column.name, column.type, str(column.nullability)) for column in columns] + + +def diagnostic_rows(diagnostics: tuple[SqlBindingDiagnostic, ...]) -> list[DiagnosticRow]: + return [ + ( + diagnostic.code, + diagnostic.message, + diagnostic.line, + diagnostic.column, + diagnostic.start, + diagnostic.end, + diagnostic.severity, + ) + for diagnostic in diagnostics + ] + + +def deferred_row(*, model: int, analysis: PolyglotAnalysisResult) -> DeferredRow: + """Python's answer to one deferral, without the lineage Python keeps.""" + + return ( + model, + analysis.analysis_succeeded, + column_rows(analysis.columns), + analysis.has_star, + analysis.star_resolved, + diagnostic_rows(analysis.binding_diagnostics), + analysis.binding_validated, + ) + + +def inferred_columns(rows: list[ColumnRow] | None) -> tuple[InferredColumn, ...] | None: + if rows is None: + return None + return tuple( + InferredColumn(name=name, type=data_type, nullability=InferredNullability(nullability)) + for name, data_type, nullability in rows + ) + + +def binding_diagnostics(rows: list[DiagnosticRow]) -> tuple[SqlBindingDiagnostic, ...]: + return tuple(SqlBindingDiagnostic(*row) for row in rows) + + +def compact_lineage(rows: list[LineageItem]) -> CompactLineageFacts: + """Native lineage rows as Python's compact facts over a per-model string pool.""" + + pool: dict[str, int] = {} + compact_rows: list[tuple[int, int, int, tuple[tuple[int, int, int], ...]]] = [] + for output_column, transform_code, confidence_code, sources in rows: + source_indexes: list[tuple[int, int, int]] = [] + for resource_type, resource_name, column_name in sources: + source_indexes.append( + ( + pool.setdefault(resource_type, len(pool)), + pool.setdefault(resource_name, len(pool)), + pool.setdefault(column_name, len(pool)), + ) + ) + compact_rows.append( + ( + pool.setdefault(output_column, len(pool)), + transform_code, + confidence_code, + tuple(source_indexes), + ) + ) + return CompactLineageFacts(string_pool=tuple(pool), rows=tuple(compact_rows)) diff --git a/src/sqlbuild/compiler/analysis_session/classes/__init__.py b/src/sqlbuild/compiler/analysis_session/classes/__init__.py new file mode 100644 index 0000000000..74909fb82d --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/classes/__init__.py @@ -0,0 +1 @@ +"""Native model analysis session classes.""" diff --git a/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py b/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py new file mode 100644 index 0000000000..7674467755 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py @@ -0,0 +1,203 @@ +"""One compile's native model analysis, with Python answering the session's deferrals.""" + +from __future__ import annotations + +import logging +from collections import Counter +from collections.abc import Sequence +from typing import Any + +import sqlbuild._native as _native +from sqlbuild.compiler.analysis_session._helpers.deferral_records import record_analysis_deferral +from sqlbuild.compiler.analysis_session._helpers.session_rows import ( + binding_diagnostics, + compact_lineage, + deferred_row, + inferred_columns, +) +from sqlbuild.compiler.analysis_session.constants import ( + DEFERRAL_ANALYSIS, + DEFERRAL_ENRICHMENT, + DEFERRAL_LEGACY_ANALYSIS, + DEFERRAL_SESSION, + LINEAGE_NATIVE, + NATIVE_ANALYSIS_FAILURE_MESSAGE, + NATIVE_SESSION_FAILURE_MESSAGE, +) +from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest +from sqlbuild.compiler.analysis_session.types import ( + DeferralRow, + DeferredRow, + FinishRow, + OutcomeRow, + ShapeRows, + StepRow, +) +from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis +from sqlbuild.compiler.compile.models import ( + CompiledLineageColumnFact, + ModelSqlAnalysis, + ModelSqlAnalysisRequest, + PolyglotAnalysisResult, +) +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.diagnostics.main.log_debug_event import log_debug_event + +_DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile") + + +class NativeModelAnalysis: + """Drive one native session, answering deferrals with today's Python analysis.""" + + def __init__( + self, + *, + request: NativeModelAnalysisRequest, + python: PythonModelAnalysis, + catalog: Any, + ) -> None: + self._python: PythonModelAnalysis = python + self._catalog: Any = catalog + self._types: dict[str, dict[str, str]] = dict(request.column_types_by_table) + self._nullability: dict[str, dict[str, InferredNullability]] = dict( + request.column_nullability_by_table + ) + self._kept: dict[tuple[int, str], PolyglotAnalysisResult] = {} + self._deferrals: Counter[str] = Counter() + + def analyses( + self, session: _native.NativeModelAnalysisSession + ) -> dict[str, ModelSqlAnalysis] | None: + """Every model's analysis by name, or None when Python must analyse them all.""" + + finished: FinishRow | None = self._finished(session) + if finished is None: + record_analysis_deferral(kind=DEFERRAL_SESSION) + log_debug_event( + logger=_DEBUG_LOGGER, + message=NATIVE_SESSION_FAILURE_MESSAGE, + sqlbuild_error=session.failure, + ) + return None + outcomes, schema_additions, analysis_names = finished + self._record_catalog_changes( + schema_additions=schema_additions, analysis_names=analysis_names + ) + for kind, count in self._deferrals.items(): + record_analysis_deferral(kind=kind, count=count) + self._python.record_uncached() + return { + request.model_input.model_file.file_path.stem: self._model_analysis( + index=index, request=request, outcome=outcome + ) + for index, (request, outcome) in enumerate( + zip(self._python.requests, outcomes, strict=True) + ) + } + + def _finished(self, session: _native.NativeModelAnalysisSession) -> FinishRow | None: + step: StepRow | None = session.run() + while step is not None: + publications, deferrals, failures = step + for failure in failures: + log_debug_event( + logger=_DEBUG_LOGGER, + message=NATIVE_ANALYSIS_FAILURE_MESSAGE, + sqlbuild_error=failure, + ) + self._publish(publications) + if not deferrals: + return session.finish() + answers: list[DeferredRow] = [self._answer(deferral) for deferral in deferrals] + if not session.provide(answers): + return None + step = session.run() + return None + + def _publish(self, publications: ShapeRows) -> None: + for name, shape in publications: + published: dict[str, str] = dict(shape) + self._types.setdefault(name, published) + self._nullability.setdefault( + name, dict.fromkeys(published, InferredNullability.UNKNOWN) + ) + + def _answer(self, deferral: DeferralRow) -> DeferredRow: + kind, model, cleaned_sql, schemas, diagnostics, lineage = deferral + shapes: dict[str, dict[str, str]] = {name: dict(shape) for name, shape in schemas} + analysis: PolyglotAnalysisResult + if kind == DEFERRAL_ANALYSIS: + self._deferrals[DEFERRAL_LEGACY_ANALYSIS] += 1 + analysis = self._python.analyze_deferred( + model=model, + precomputed=PythonModelAnalysis.legacy_precomputed( + cleaned_sql=cleaned_sql or "", + binding_diagnostics=binding_diagnostics(diagnostics), + lineage=lineage, + ), + binding_schema=shapes, + column_types_by_table=self._types, + column_nullability_by_table=self._nullability, + ) + else: + self._deferrals[DEFERRAL_ENRICHMENT] += 1 + analysis = self._python.enrich(model=model, input_schemas=shapes) + self._kept[(model, kind)] = analysis + return deferred_row(model=model, analysis=analysis) + + def _record_catalog_changes( + self, *, schema_additions: ShapeRows, analysis_names: list[str] + ) -> None: + """Leave Python's binding catalog as its own analysis would have.""" + + additions: dict[str, dict[str, str]] = { + name: dict(shape) + for name, shape in schema_additions + if name not in self._catalog.schemas + } + self._catalog.schemas.update(additions) + if additions: + self._catalog.native.update_relations(additions) + self._catalog.prepare_analysis( + types={name: self._types[name] for name in analysis_names if name in self._types}, + nullability={ + name: self._nullability[name] + for name in analysis_names + if name in self._nullability + }, + ) + + def _model_analysis( + self, *, index: int, request: ModelSqlAnalysisRequest, outcome: OutcomeRow + ) -> ModelSqlAnalysis: + ( + succeeded, + columns, + lineage_source, + lineage_rows, + has_star, + star_resolved, + diagnostics, + validated, + cleaned_sql, + ) = outcome + lineage: Sequence[CompiledLineageColumnFact] = ( + (compact_lineage(lineage_rows) if succeeded else ()) + if lineage_source == LINEAGE_NATIVE + else self._kept[(index, lineage_source)].lineage_columns + ) + return ModelSqlAnalysis( + polyglot_analysis=PolyglotAnalysisResult( + analysis_succeeded=succeeded, + columns=inferred_columns(columns), + lineage_columns=lineage, + has_star=has_star, + star_resolved=star_resolved, + binding_diagnostics=binding_diagnostics(diagnostics), + binding_validated=validated, + ), + placeholders=request.placeholders, + fused_binding_validated=True, + cleaned_sql=cleaned_sql, + validated_schema=request.binding_schema, + ) diff --git a/src/sqlbuild/compiler/analysis_session/constants.py b/src/sqlbuild/compiler/analysis_session/constants.py new file mode 100644 index 0000000000..51a32d88cf --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/constants.py @@ -0,0 +1,15 @@ +"""Native model analysis session constants.""" + +ANALYSIS_DEFERRAL_SITE: str = "analysis_session" +DEFERRAL_NO_COMPACT_ANALYSIS: str = "no_compact_analysis" +DEFERRAL_NO_CATALOG: str = "no_analysis_catalog" +DEFERRAL_SESSION: str = "session" +DEFERRAL_LEGACY_ANALYSIS: str = "legacy_analysis" +DEFERRAL_ENRICHMENT: str = "input_enrichment" +DEFERRAL_EXPRESSION_SHAPES: str = "expression_shapes" +DEFERRAL_ANALYSIS: str = "analysis" +DEFERRAL_ENRICH: str = "enrichment" +LINEAGE_NATIVE: str = "native" +NATIVE_ANALYSIS_FAILURE_MESSAGE: str = "native compact query analysis failed" +NATIVE_SESSION_FAILURE_MESSAGE: str = "native model analysis failed; falling back" +NATIVE_SHAPES_FAILURE_MESSAGE: str = "native expression-source shape inference failed; falling back" diff --git a/src/sqlbuild/compiler/analysis_session/main/_analyze_native_model_sql.py b/src/sqlbuild/compiler/analysis_session/main/_analyze_native_model_sql.py index ea5cdadba4..f7807824f2 100644 --- a/src/sqlbuild/compiler/analysis_session/main/_analyze_native_model_sql.py +++ b/src/sqlbuild/compiler/analysis_session/main/_analyze_native_model_sql.py @@ -2,6 +2,7 @@ from __future__ import annotations +from sqlbuild.compiler.analysis_session._helpers.model_analysis import native_model_analyses from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest from sqlbuild.compiler.compile.models import ModelSqlAnalysis @@ -11,4 +12,4 @@ def analyze_native_model_sql( ) -> dict[str, ModelSqlAnalysis] | None: """Return each model's analysis (and any pivot proof) by name, or None to defer to Python.""" - return None + return native_model_analyses(request=request) diff --git a/src/sqlbuild/compiler/analysis_session/main/_infer_native_expression_source_shapes.py b/src/sqlbuild/compiler/analysis_session/main/_infer_native_expression_source_shapes.py index a39c7a6609..53cb602fc9 100644 --- a/src/sqlbuild/compiler/analysis_session/main/_infer_native_expression_source_shapes.py +++ b/src/sqlbuild/compiler/analysis_session/main/_infer_native_expression_source_shapes.py @@ -3,6 +3,9 @@ from __future__ import annotations from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.analysis_session._helpers.expression_shapes import ( + native_expression_source_shapes, +) def infer_native_expression_source_shapes( @@ -10,4 +13,4 @@ def infer_native_expression_source_shapes( ) -> tuple[dict[str, str] | None, ...] | None: """Return one shape per expression, or None where Python must infer the shapes.""" - return None + return native_expression_source_shapes(expressions=expressions, profile=profile) diff --git a/src/sqlbuild/compiler/analysis_session/types.py b/src/sqlbuild/compiler/analysis_session/types.py new file mode 100644 index 0000000000..746ce29cc0 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/types.py @@ -0,0 +1,25 @@ +"""Native model analysis session row types.""" + +from __future__ import annotations + +type ShapeRows = list[tuple[str, list[tuple[str, str]]]] +type ColumnRow = tuple[str, str | None, str] +type DiagnosticRow = tuple[str, str, int | None, int | None, int | None, int | None, str] +type LineageItem = tuple[str, int, int, list[tuple[str, str, str]]] +type DeferredRow = tuple[int, bool, list[ColumnRow] | None, bool, bool, list[DiagnosticRow], bool] +type DeferralRow = tuple[ + str, int, str | None, ShapeRows, list[DiagnosticRow], list[LineageItem] | None +] +type StepRow = tuple[ShapeRows, list[DeferralRow], list[str]] +type OutcomeRow = tuple[ + bool, + list[ColumnRow] | None, + str, + list[LineageItem], + bool, + bool, + list[DiagnosticRow], + bool, + str, +] +type FinishRow = tuple[list[OutcomeRow], ShapeRows, list[str]] diff --git a/src/sqlbuild/compiler/compile/classes/python_model_analysis.py b/src/sqlbuild/compiler/compile/classes/python_model_analysis.py new file mode 100644 index 0000000000..20a7c1b100 --- /dev/null +++ b/src/sqlbuild/compiler/compile/classes/python_model_analysis.py @@ -0,0 +1,213 @@ +"""Python's half of native model analysis: per-model request rows and deferred analyses.""" + +from __future__ import annotations + +import re + +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.compile._helpers.analysis.cache import record_analysis_cache_metrics +from sqlbuild.compiler.compile._helpers.analysis.columns import table_function_analysis_name +from sqlbuild.compiler.compile._helpers.analysis.compact import ( + analyze_columns_and_lineage_with_polyglot, +) +from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( + binding_relation_names, + binding_schema_for_model, +) +from sqlbuild.compiler.compile._helpers.render.cursor_intrinsics import ( + cursor_intrinsics_analysis_sql, +) +from sqlbuild.compiler.compile._helpers.sharing.binding import lineage_reference_map +from sqlbuild.compiler.compile.models import ( + CompileModelInput, + CompileSqlReference, + ModelSqlAnalysisRequest, + NativeCompactAnalysis, + PolyglotAnalysisResult, +) +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.compiler.planner.constants import ( + SNAPSHOT_DEFAULT_VALID_FROM_COLUMN, + SNAPSHOT_DEFAULT_VALID_TO_COLUMN, +) +from sqlbuild.compiler.planner.types import ContractPolicy, MaterializationType +from sqlbuild.compiler.references.types import SqlReferenceKind +from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic + +_SET_OPERATION_PATTERN: re.Pattern[str] = re.compile( + r"\b(?:UNION|INTERSECT|EXCEPT)\b", re.IGNORECASE +) + + +class PythonModelAnalysis: + """The analysed models' request facts, and today's Python analysis where native defers.""" + + def __init__( + self, + *, + model_inputs: tuple[CompileModelInput, ...], + inference_profile: ExpressionInferenceProfile, + complete_binding_schemas: dict[str, dict[str, str]], + ) -> None: + self._profile: ExpressionInferenceProfile = inference_profile + self.requests: tuple[ModelSqlAnalysisRequest, ...] = tuple( + ModelSqlAnalysisRequest( + model_input=model_input, + query_sql=cursor_intrinsics_analysis_sql( + sql=model_input.query_sql, + cursor_type=model_input.config.values.get("cursor_type"), + ), + placeholders=_placeholders(model_input), + cache_key=None, + binding_schema=binding_schema_for_model( + model_input=model_input, complete_binding_schemas=complete_binding_schemas + ), + ) + for model_input in model_inputs + ) + + def model_rows(self) -> list[tuple[object, ...]]: + """Each model's native request row, in analysis order.""" + + return [_model_row(request) for request in self.requests] + + def analyze_deferred( + self, + *, + model: int, + precomputed: NativeCompactAnalysis, + binding_schema: dict[str, dict[str, str]], + column_types_by_table: dict[str, dict[str, str]], + column_nullability_by_table: dict[str, dict[str, InferredNullability]], + ) -> PolyglotAnalysisResult: + """Python's legacy analysis of a model the native engine handed back.""" + + request: ModelSqlAnalysisRequest = self.requests[model] + return analyze_columns_and_lineage_with_polyglot( + query_sql=request.query_sql, + references=request.model_input.references, + placeholders=request.placeholders, + column_nullability_by_table=column_nullability_by_table, + column_types_by_table=column_types_by_table, + inference_profile=self._profile, + allow_compact_analysis=True, + binding_schema=binding_schema, + recover_cte_facts=_recovers_cte_facts(request.model_input), + precomputed=precomputed, + ) + + def enrich( + self, *, model: int, input_schemas: dict[str, dict[str, str]] + ) -> PolyglotAnalysisResult: + """Python's re-analysis of a model with its known input shapes.""" + + request: ModelSqlAnalysisRequest = self.requests[model] + return analyze_columns_and_lineage_with_polyglot( + query_sql=request.query_sql, + references=request.model_input.references, + placeholders=request.placeholders, + column_types_by_table=input_schemas, + column_nullability_by_table={ + table: dict.fromkeys(columns, InferredNullability.UNKNOWN) + for table, columns in input_schemas.items() + }, + inference_profile=self._profile, + allow_compact_analysis=True, + recover_cte_facts=_recovers_cte_facts(request.model_input), + ) + + @staticmethod + def legacy_precomputed( + *, + cleaned_sql: str, + binding_diagnostics: tuple[SqlBindingDiagnostic, ...], + lineage: list[tuple[str, int, int, list[tuple[str, str, str]]]] | None, + ) -> NativeCompactAnalysis: + """The native result Python's legacy analysis reads, with lineage rows when kept.""" + + if lineage is None: + return NativeCompactAnalysis( + cleaned_sql=cleaned_sql, + analysis=None, + projected=False, + binding_diagnostics=binding_diagnostics, + ) + pool: dict[str, int] = {} + facts: list[object] = [] + for output_column, transform_code, confidence_code, sources in lineage: + name_index: int = pool.setdefault(output_column, len(pool)) + source_indexes: list[list[int]] = [] + for source in sources: + source_indexes.append([pool.setdefault(value, len(pool)) for value in source]) + facts.append([name_index, None, 0, transform_code, confidence_code, source_indexes]) + return NativeCompactAnalysis( + cleaned_sql=cleaned_sql, + analysis=None, + projected=False, + binding_diagnostics=binding_diagnostics, + compact_rows=list(range(len(facts))), + compact_fact_rows=facts, + string_pool=tuple(pool), + compact_column_cache={}, + compact_template_index=0, + compact_fact_cache={}, + compact_decoded_fact_cache={}, + ) + + def record_uncached(self) -> None: + """Record that every analysed model bypassed the analysis cache.""" + + record_analysis_cache_metrics( + batch_hits=0, entry_hits=0, misses=0, bypasses=len(self.requests) + ) + + +def _model_row(request: ModelSqlAnalysisRequest) -> tuple[object, ...]: + model_input: CompileModelInput = request.model_input + references: tuple[CompileSqlReference, ...] = model_input.references + values: dict[str, object] = model_input.config.values + return ( + model_input.model_file.file_path.stem, + request.query_sql, + list((request.placeholders or {}).items()), + [ + (_analysis_name(reference), reference.ref_kind == SqlReferenceKind.REF) + for reference in references + ], + [ + (name, resource_type.value, resource_name) + for name, (resource_type, resource_name) in lineage_reference_map(references).items() + ], + sorted(binding_relation_names(references)), + _recovers_cte_facts(model_input), + _SET_OPERATION_PATTERN.search(request.query_sql) is not None, + ( + ( + str(values.get("valid_from_column") or SNAPSHOT_DEFAULT_VALID_FROM_COLUMN), + str(values.get("valid_to_column") or SNAPSHOT_DEFAULT_VALID_TO_COLUMN), + ) + if values.get("materialized") == MaterializationType.SNAPSHOT + else None + ), + ) + + +def _analysis_name(reference: CompileSqlReference) -> str: + if reference.ref_kind == SqlReferenceKind.TABLE_FUNCTION: + return table_function_analysis_name(reference.ref_name) + return reference.ref_name + + +def _placeholders(model_input: CompileModelInput) -> dict[str, str] | None: + raw_placeholders: object | None = model_input.config.values.get("placeholders") + return ( + {str(key): str(value) for key, value in raw_placeholders.items()} + if isinstance(raw_placeholders, dict) + else None + ) + + +def _recovers_cte_facts(model_input: CompileModelInput) -> bool: + return model_input.config.values.get("contract") == ContractPolicy.ENFORCED or ( + model_input.schema_entry is not None and bool(model_input.schema_entry.type_enforcement) + ) diff --git a/src/sqlbuild/compiler/sql_analysis/classes/binding_catalog.py b/src/sqlbuild/compiler/sql_analysis/classes/binding_catalog.py index 0e0afc4255..e5df48b406 100644 --- a/src/sqlbuild/compiler/sql_analysis/classes/binding_catalog.py +++ b/src/sqlbuild/compiler/sql_analysis/classes/binding_catalog.py @@ -80,7 +80,7 @@ def inferred_schema( break if SQL_QUOTED_IDENTIFIER_DELIMITER not in sql and not has_exact_input: return columns - return self.native.inferred_schema(sql=sql, columns=columns, inputs=inputs) + return dict(self.native.inferred_schema(sql=sql, columns=columns, inputs=inputs)) def prepare(self, requests: Sequence[SqlSchemaValidationRequest]) -> list[NativeBindingRequest]: additions: dict[str, Mapping[str, str]] = {} diff --git a/src/sqlbuild/compiler/sql_analysis/types.py b/src/sqlbuild/compiler/sql_analysis/types.py index 7e69e03f36..f8fa8d5ac4 100644 --- a/src/sqlbuild/compiler/sql_analysis/types.py +++ b/src/sqlbuild/compiler/sql_analysis/types.py @@ -17,7 +17,7 @@ def run(self) -> bytes: ... class NativeProjectCatalog(Protocol): def inferred_schema( self, *, sql: str, columns: Mapping[str, str], inputs: Mapping[str, Mapping[str, str]] - ) -> dict[str, str]: ... + ) -> list[tuple[str, str]]: ... def with_relations( self, relations: Mapping[str, Mapping[str, str]] ) -> NativeProjectCatalog: ... diff --git a/tests/e2e/scripts/compiler_differential/helpers.py b/tests/e2e/scripts/compiler_differential/helpers.py index dff1c5eb7f..407113d85c 100644 --- a/tests/e2e/scripts/compiler_differential/helpers.py +++ b/tests/e2e/scripts/compiler_differential/helpers.py @@ -73,6 +73,20 @@ def _perturbed(*, until, **stages): frontier.native_frontier = _perturbed """ +CATALOG_PERTURBATION: str = """ +import sqlbuild.compiler.analysis_session.classes.native_model_analysis as native_analysis + +_original = native_analysis.NativeModelAnalysis._record_catalog_changes + + +def _perturbed(self, **changes): + _original(self, **changes) + self._catalog.schemas["perturbed_by_test"] = {"order_id": "INTEGER"} + + +native_analysis.NativeModelAnalysis._record_catalog_changes = _perturbed +""" + DEFERRAL_PERTURBATION: str = """ import json import os diff --git a/tests/e2e/scripts/compiler_differential/test_compiler_differential.py b/tests/e2e/scripts/compiler_differential/test_compiler_differential.py index 9ab1471e8e..e0dde5cd4c 100644 --- a/tests/e2e/scripts/compiler_differential/test_compiler_differential.py +++ b/tests/e2e/scripts/compiler_differential/test_compiler_differential.py @@ -18,6 +18,7 @@ WheelSiteReportTestCase, ) from tests.e2e.scripts.compiler_differential.helpers import ( + CATALOG_PERTURBATION, COMPILED_PROJECT_CAPTURE, DEFERRAL_PERTURBATION, DISCOVERY_PERTURBATION, @@ -181,6 +182,39 @@ def test_given_perturbed_native_render_when_comparing_then_compile_inputs_captur assert not any(text in output for text in test_case.expected_absent), output +@pytest.mark.parametrize( + "test_case", + [ + HarnessRunTestCase( + description="wrong_native_catalog_shape", + extra_arguments=(), + expected_exit_code=1, + expected_lines=("DIFF project/waffle_shop",), + expected_patterns=( + r"- stage capture 0-compile/\d+-compiled_project\.json at /binding_catalog/schemas", + ), + expected_absent=("warm/", "3-plan/"), + ) + ], + ids=lambda case: case.description, +) +def test_given_wrong_native_catalog_shape_when_comparing_then_only_the_cold_capture_differs( + test_case: HarnessRunTestCase, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + exit_code: int = run_compiler_differential( + harness_arguments( + work_dir=tmp_path / "work", + extra=perturbation_arguments(tmp_path / "perturbation", source=CATALOG_PERTURBATION), + ) + ) + + output: str = capsys.readouterr().out + assert exit_code == test_case.expected_exit_code, output + assert all(line in output for line in test_case.expected_lines), output + assert all(re.search(pattern, output) for pattern in test_case.expected_patterns), output + assert not any(text in output for text in test_case.expected_absent), output + + @pytest.mark.parametrize( "test_case", [ @@ -402,13 +436,13 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness "test_case", [ WheelSiteReportTestCase( - description="python_and_preview_wheel_calls_without_deferrals", + description="python_and_preview_wheel_calls_with_input_enrichment_deferrals", perturbation="", expected_lines=( "Polyglot wheel calls (python):", "Polyglot wheel calls (native-preview):", "Analysis deferrals (python): none recorded", - "Analysis deferrals (native-preview): none recorded", + " input_enrichment analysis_session (project ", "Compiler differential passed: 1 projects identical", ), expected_sites=frozenset( @@ -417,7 +451,7 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness "_validate_sql_syntax_with_message parse_one", } ), - expected_deferrals=frozenset(), + expected_deferrals=frozenset({"input_enrichment analysis_session"}), ), WheelSiteReportTestCase( description="preview_deferrals_are_counted_per_kind_and_site", @@ -428,7 +462,9 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness " legacy_fallback orders.sql (project ", ), expected_sites=frozenset(), - expected_deferrals=frozenset({"legacy_fallback orders.sql"}), + expected_deferrals=frozenset( + {"legacy_fallback orders.sql", "input_enrichment analysis_session"} + ), ), ], ids=lambda case: case.description, 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 index 3e3231cb28..e8d1fe7758 100644 --- 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 @@ -9,6 +9,7 @@ from typing import NamedTuple, cast from sqlbuild.cli.compile_reuse.constants import REUSE_DISABLE_ENV_VAR +from sqlbuild.compiler.lineage.constants import NATIVE_LINEAGE_DEFERRAL_SITE 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 @@ -20,13 +21,13 @@ class EngineLineageRun(NamedTuple): - """One engine's compile report, lineage traces, wheel fallback parses and deferral files.""" + """One engine's compile report, lineage traces, wheel fallback parses and lineage deferrals.""" compile_report: str compile_returncode: int traces: list[tuple[int, str, str]] fallback_parses: int - deferral_files: list[Path] + lineage_deferrals: int def engine_lineage_run( @@ -75,7 +76,10 @@ def engine_lineage_run( compile_returncode=compiled.returncode, traces=traces, fallback_parses=_site_calls(record_dir)[FALLBACK_SITE], - deferral_files=sorted(record_dir.glob("analysis-deferrals-*.jsonl")), + lineage_deferrals=sum( + path.read_text("utf-8").count(f'"site": "{NATIVE_LINEAGE_DEFERRAL_SITE}"') + for path in record_dir.glob("analysis-deferrals-*.jsonl") + ), ) 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 3ec873e7e6..ea06219cb7 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 @@ -1,4 +1,4 @@ -"""The analysis stage seams defer to Python, so every engine compiles a project identically.""" +"""Every engine compiles a project identically, whichever analysis stages run natively.""" from __future__ import annotations @@ -29,8 +29,8 @@ description="models_contracts_lineage_and_sql_tests", expected_preview_returns={ "assemble_native_project": [None], - "infer_native_expression_source_shapes": [None], - "analyze_native_model_sql": [None], + "infer_native_expression_source_shapes": [ANY], + "analyze_native_model_sql": [ANY], "assemble_native_sql_tests": [None], "complete_native_semantic_diagnostics": [None], "evaluate_native_model_contracts": [None], @@ -41,7 +41,7 @@ ], ids=lambda case: case.description, ) -def test_given_project_when_compiling_with_each_engine_then_seams_defer_and_outputs_match( +def test_given_project_when_compiling_with_each_engine_then_native_seams_answer_and_outputs_match( test_case: NativeAnalysisSeamTestCase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -72,6 +72,8 @@ 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 isinstance(preview_seams["analyze_native_model_sql"][0], dict) + assert isinstance(preview_seams["infer_native_expression_source_shapes"][0], tuple) 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/analysis_session/__init__.py b/tests/integration/src/sqlbuild/compiler/analysis_session/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py new file mode 100644 index 0000000000..8c14cdd29c --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class GeneratedAnalysisParityTestCase: + """Seeded generated projects whose model analysis both engines must agree on.""" + + description: str + seed: int + count: int + model_count: int + dialects: tuple[str | None, ...] + expected_minimum_native: int + expected_minimum_expression_shapes: int + + +@dataclass(frozen=True) +class AnalysisFallbackTestCase: + """A project the native session hands back to Python whole, recording why.""" + + description: str + files: dict[str, str] + allow_compact_analysis: bool + keeps_catalog: bool + expected_kind: str + + +@dataclass(frozen=True) +class SessionFailureTestCase: + """A session that fails after Python answered deferrals, so Python analyses everything.""" + + description: str + seed: int + model_count: int + expected_kinds: dict[str, int] diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py new file mode 100644 index 0000000000..080c39fe00 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py @@ -0,0 +1,430 @@ +from __future__ import annotations + +import json +import random +from collections import Counter +from collections.abc import Callable +from contextlib import suppress +from dataclasses import dataclass, field, replace +from functools import partial +from itertools import chain +from pathlib import Path +from typing import Any, cast + +import pytest + +import sqlbuild._native as native_module +import sqlbuild.compiler.compile._helpers.assembly.project as project_assembly +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.compiler.analysis_session.main._analyze_native_model_sql import ( + analyze_native_model_sql, +) +from sqlbuild.compiler.analysis_session.main._infer_native_expression_source_shapes import ( + infer_native_expression_source_shapes, +) +from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest +from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( + get_expression_source_shapes, +) +from sqlbuild.compiler.compile.exceptions import CompileInputError +from sqlbuild.compiler.compile.main._build_compile_inputs import build_compile_inputs +from sqlbuild.compiler.compile.models import ( + CompactLineageFacts, + CompileAdapterContext, + CompileProjectInputs, + ModelSqlAnalysis, +) +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily +from sqlbuild.sql_values.types import CollectionRendering + +_PROJECT_TOML: str = ( + 'name = "orders_analysis"\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) + - name: raw_events + description: Untyped events feed. + columns: + - name: event_id + - name: order_id + - name: kind +""" +_INITIAL_RELATIONS: tuple[tuple[str, tuple[str, ...]], ...] = ( + ('__source("raw_orders")', ("order_id", "customer_id", "amount", "status")), + ('__source("raw_customers")', ("customer_id", "Customer_Name", "region")), + ('__source("raw_events")', ("event_id", "order_id", "kind")), +) +_ADAPTER_CONTEXT: CompileAdapterContext = CompileAdapterContext( + value_renderer=DuckDbAdapter(), + collection_rendering=CollectionRendering.VALUE_LIST, + python_functions_inherit_default_namespace=True, + sql_lexical_syntax=DuckDbAdapter.sql_lexical_syntax, +) +_PIVOT_MODEL: str = """MODEL ( + description "Order amounts pivoted by status", + contract enforced, + materialized table, + columns (customer_id (type INTEGER)), + dynamic_columns ( + status_amounts ( + pivot_column status, + value_column amount, + aggregate MAX, + type DOUBLE + ) + ), +); + +PIVOT __source("raw_orders") +ON status +USING MAX(amount) +GROUP BY customer_id +""" + +type _Relation = tuple[str, tuple[str, ...]] +type _Template = Callable[[random.Random, list[_Relation]], tuple[str, str, tuple[str, ...]]] + + +def _star(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tuple[str, ...]]: + relation, columns = inputs[0] + return "", f"SELECT * FROM {relation}", columns + + +def _qualified_star( + rng: random.Random, inputs: list[_Relation] +) -> tuple[str, 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[_Relation]) -> tuple[str, 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 _set_operation(rng: random.Random, inputs: list[_Relation]) -> tuple[str, 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 _expressions(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tuple[str, ...]]: + relation, columns = inputs[0] + first, last = columns[0], columns[-1] + sql: str = ( + f"SELECT CAST({first} AS VARCHAR) AS as_text, {first} IS NULL AS is_missing,\n" + f" COALESCE({first}, {last}) AS filled, {last} AS passthrough,\n" + f" ROW_NUMBER() OVER (PARTITION BY {last} ORDER BY {first}) AS position\n" + f"FROM {relation}\nWHERE {first} IS NOT NULL" + ) + return "", sql, ("as_text", "is_missing", "filled", "passthrough", "position") + + +def _aggregate(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tuple[str, ...]]: + relation, columns = inputs[0] + column: str = rng.choice(columns) + sql: str = ( + f"SELECT {column} AS group_key, COUNT(*) AS row_count, MAX({columns[-1]}) AS latest\n" + f"FROM {relation}\nGROUP BY 1" + ) + return "", sql, ("group_key", "row_count", "latest") + + +def _quoted(rng: random.Random, inputs: list[_Relation]) -> tuple[str, 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 _contract(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tuple[str, ...]]: + relation, columns = inputs[0] + sql: str = ( + f"SELECT CAST({columns[0]} AS INTEGER) AS contract_key,\n" + f" CAST({columns[-1]} AS VARCHAR) AS contract_label\nFROM {relation}" + ) + header: str = ( + " contract enforced,\n" + " columns (contract_key (type INTEGER), contract_label (type VARCHAR)),\n" + ) + return header, sql, ("contract_key", "contract_label") + + +def _unknown_column( + rng: random.Random, inputs: list[_Relation] +) -> tuple[str, str, tuple[str, ...]]: + relation, columns = inputs[0] + sql: str = f"SELECT {columns[0]}, missing_column FROM {relation}" + return "", sql, (columns[0], "missing_column") + + +def _subquery(rng: random.Random, inputs: list[_Relation]) -> tuple[str, 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, + _set_operation, + _expressions, + _aggregate, + _quoted, + _contract, + _unknown_column, + _subquery, +) + + +def generated_analysis_files(*, rng: random.Random, model_count: int) -> dict[str, str]: + """Sources (typed, inferred and untyped), a pivot and models drawing on every template.""" + + files: dict[str, str] = { + "sqlbuild_project.toml": _PROJECT_TOML, + "sources/raw.yml": _SOURCES, + "models/marts/status_amounts.sql": _PIVOT_MODEL, + } + relations: list[_Relation] = list(_INITIAL_RELATIONS) + for index in range(model_count): + template: _Template = _TEMPLATES[index % len(_TEMPLATES)] + header, sql, columns = template(rng, rng.sample(relations, k=2)) + name: str = f"orders_model_{index}" + files[f"models/layer_{index % 3}/{name}.sql"] = ( + f'MODEL (\n description "Generated orders model {index}",\n{header});\n\n{sql}\n' + ) + relations.append((f'__ref("{name}")', columns)) + return files + + +def compile_inputs(*, project_dir: Path, files: dict[str, str]) -> CompileProjectInputs: + """Write a project and attach its compile inputs with the DuckDB compile context.""" + + 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_compile_inputs( + discovered_inputs=discover_project_inputs(project_dir=project_dir), + adapter_context=_ADAPTER_CONTEXT, + run_id="integration_run", + ) + + +@dataclass +class AnalysisParity: + """Python's and the native session's views of every analysis seam call.""" + + names: list[object] = field(default_factory=list) + python: list[object] = field(default_factory=list) + native: list[object] = field(default_factory=list) + analysed_models: int = 0 + expression_shapes: int = 0 + + +def compare_analyses( + *, + inputs: CompileProjectInputs, + dialect: str | None, + parity: AnalysisParity, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Assemble `inputs`, analysing every seam call with both engines into `parity`.""" + + def models_by_both( + *, + python_analysis: partial[dict[str, ModelSqlAnalysis]], + dynamic_families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], + ) -> dict[str, ModelSqlAnalysis]: + keywords: dict[str, Any] = {**python_analysis.keywords, "analysis_cache": None} + profile: ExpressionInferenceProfile = keywords["inference_profile"] + native_catalog: Any = profile.binding_catalog.with_relations({}) + native: dict[str, ModelSqlAnalysis] | None = analyze_native_model_sql( + request=NativeModelAnalysisRequest( + **{ + **keywords, + "inference_profile": replace(profile, binding_catalog=native_catalog), + }, + dynamic_families_by_table=dynamic_families_by_table, + ) + ) + python: dict[str, ModelSqlAnalysis] = python_analysis.func(**keywords) + parity.analysed_models += len(python) + _append(parity, "model analyses", _analysis_views(python), _analysis_views(native)) + _append( + parity, + "catalog", + _catalog_view(profile.binding_catalog), + _catalog_view(native_catalog), + ) + return python + + def shapes_by_both( + *, expressions: tuple[str, ...], profile: ExpressionInferenceProfile + ) -> tuple[dict[str, str] | None, ...]: + catalog: Any = cast(Any, profile.binding_catalog) + native_catalog: Any = catalog.with_relations({}) + native: tuple[dict[str, str] | None, ...] | None = infer_native_expression_source_shapes( + expressions=expressions, profile=replace(profile, binding_catalog=native_catalog) + ) + python: tuple[dict[str, str] | None, ...] = get_expression_source_shapes( + expressions=expressions, profile=profile + ) + parity.expression_shapes += sum(shape is not None for shape in python) + _append( + parity, + "expression shapes", + (python, catalog.expression_shapes), + (native, native_catalog.expression_shapes), + ) + return python + + with monkeypatch.context() as patch: + patch.setattr(project_assembly, "analyze_model_sql_by_engine", models_by_both) + patch.setattr(project_assembly, "expression_source_shapes_by_engine", shapes_by_both) + with suppress(CompileInputError): + _ = project_assembly.assemble_compiled_project( + inputs=inputs, + inference_profile=ExpressionInferenceProfile(sql_analysis_dialect=dialect), + ) + + +def analysis_request( + *, inputs: CompileProjectInputs, monkeypatch: pytest.MonkeyPatch +) -> NativeModelAnalysisRequest: + """The native request the model analysis seam builds while assembling `inputs`.""" + + requests: list[NativeModelAnalysisRequest] = [] + + def captured( + *, + python_analysis: partial[dict[str, ModelSqlAnalysis]], + dynamic_families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], + ) -> dict[str, ModelSqlAnalysis]: + requests.append( + NativeModelAnalysisRequest( + **python_analysis.keywords, dynamic_families_by_table=dynamic_families_by_table + ) + ) + return python_analysis() + + with monkeypatch.context() as patch, suppress(CompileInputError): + patch.setattr(project_assembly, "analyze_model_sql_by_engine", captured) + _ = project_assembly.assemble_compiled_project( + inputs=inputs, + inference_profile=ExpressionInferenceProfile(sql_analysis_dialect="duckdb"), + ) + return requests[0] + + +def deferral_kinds(record_dir: Path) -> Counter[str]: + """Count the `site:kind` deferral records below `record_dir`.""" + + lines: list[str] = list( + chain.from_iterable( + path.read_text("utf-8").splitlines() + for path in sorted(record_dir.glob("analysis-deferrals-*.jsonl")) + ) + ) + records: list[dict[str, str]] = [json.loads(line) for line in lines] + return Counter(f"{record['site']}:{record['kind']}" for record in records) + + +def _append(parity: AnalysisParity, name: str, python: object, native: object) -> None: + parity.names.append(name) + parity.python.append(python) + parity.native.append(native) + + +def _analysis_views(analyses: dict[str, ModelSqlAnalysis] | None) -> object: + return analyses and { + name: ( + analysis.polyglot_analysis.analysis_succeeded, + analysis.polyglot_analysis.columns, + isinstance(analysis.polyglot_analysis.lineage_columns, CompactLineageFacts), + tuple(analysis.polyglot_analysis.lineage_columns), + analysis.polyglot_analysis.has_star, + analysis.polyglot_analysis.star_resolved, + analysis.polyglot_analysis.binding_diagnostics, + analysis.polyglot_analysis.binding_validated, + analysis.cleaned_sql, + analysis.placeholders, + analysis.dynamic_column_contract, + ) + for name, analysis in analyses.items() + } + + +def _catalog_view(catalog: Any) -> tuple[object, ...]: + return (dict(catalog.schemas), dict(catalog.analysis_shapes)) + + +class FailingProvideSession: + """A native session whose answers to deferrals are refused, as a failed session refuses.""" + + def __init__(self, session: Any) -> None: + self._session: Any = session + self.answered: int = 0 + self.failure: str = "injected session failure" + + def run(self) -> object: + return self._session.run() + + def provide(self, answers: list[object]) -> bool: + self.answered += len(answers) + return False + + def finish(self) -> object: + return self._session.finish() + + +def failing_provide_sessions(*, monkeypatch: pytest.MonkeyPatch) -> list[FailingProvideSession]: + """Make every native session refuse Python's answers; return the sessions started.""" + + started: list[FailingProvideSession] = [] + original: Callable[..., Any] = native_module.start_model_analysis_session + + def start(catalog: object, request: tuple[object, ...]) -> FailingProvideSession: + started.append(FailingProvideSession(original(catalog, request))) + return started[-1] + + monkeypatch.setattr(native_module, "start_model_analysis_session", start) + return started diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py new file mode 100644 index 0000000000..ee5c88edd7 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py @@ -0,0 +1,173 @@ +"""The native model analysis session equals Python's model analysis on generated projects.""" + +from __future__ import annotations + +import random +from collections import Counter +from dataclasses import replace +from pathlib import Path + +import pytest + +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.analysis_session.main._analyze_native_model_sql import ( + analyze_native_model_sql, +) +from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest +from sqlbuild.compiler.compile.models import CompileProjectInputs +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR +from tests.integration.src.sqlbuild.compiler.analysis_session._test_types import ( + AnalysisFallbackTestCase, + GeneratedAnalysisParityTestCase, + SessionFailureTestCase, +) +from tests.integration.src.sqlbuild.compiler.analysis_session.helpers import ( + AnalysisParity, + FailingProvideSession, + analysis_request, + compare_analyses, + compile_inputs, + deferral_kinds, + failing_provide_sessions, + generated_analysis_files, +) +from tests.integration.src.sqlbuild.compiler.helpers import mismatches + +_ORDERS_PROJECT: dict[str, str] = { + "sqlbuild_project.toml": 'name = "orders_cycle"\nadapter = "duckdb"\n', + "models/orders.sql": 'MODEL (description "Orders");\n\nSELECT order_id FROM __ref("returns")\n', + "models/returns.sql": 'MODEL (description "Returns");\n\nSELECT order_id FROM __ref("orders")\n', +} + + +@pytest.mark.parametrize( + "test_case", + [ + GeneratedAnalysisParityTestCase( + description="stars, CTEs, set operations, untyped inputs, contracts and a pivot", + seed=20261008, + count=6, + model_count=24, + dialects=("duckdb", "postgres", "snowflake", "bigquery"), + expected_minimum_native=500, + expected_minimum_expression_shapes=40, + ) + ], + ids=lambda case: case.description, +) +def test_given_generated_projects_when_analysing_natively_then_matches_python( + test_case: GeneratedAnalysisParityTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + record_dir: Path = tmp_path / "records" + monkeypatch.setenv(ANALYSIS_RECORD_DIR_ENV_VAR, str(record_dir)) + rng: random.Random = random.Random(test_case.seed) + parity: AnalysisParity = AnalysisParity() + for index in range(test_case.count): + inputs: CompileProjectInputs = compile_inputs( + project_dir=tmp_path / f"project_{index}", + files=generated_analysis_files(rng=rng, model_count=test_case.model_count), + ) + for dialect in test_case.dialects: + compare_analyses(inputs=inputs, dialect=dialect, parity=parity, monkeypatch=monkeypatch) + kinds: Counter[str] = deferral_kinds(record_dir) + + assert mismatches(inputs=parity.names, expected=parity.python, actual=parity.native) == [] + assert kinds["analysis_session:session"] == kinds["analysis_session:expression_shapes"] == 0 + assert ( + parity.analysed_models - kinds["analysis_session:legacy_analysis"] + >= test_case.expected_minimum_native + ) + assert parity.expression_shapes >= test_case.expected_minimum_expression_shapes + + +@pytest.mark.parametrize( + "test_case", + [ + AnalysisFallbackTestCase( + description="models that reference each other", + files=_ORDERS_PROJECT, + allow_compact_analysis=True, + keeps_catalog=True, + expected_kind="session", + ), + AnalysisFallbackTestCase( + description="lineage disabled, so compact analysis is off", + files=_ORDERS_PROJECT, + allow_compact_analysis=False, + keeps_catalog=True, + expected_kind="no_compact_analysis", + ), + AnalysisFallbackTestCase( + description="an inference profile without a binding catalog", + files=_ORDERS_PROJECT, + allow_compact_analysis=True, + keeps_catalog=False, + expected_kind="no_analysis_catalog", + ), + ], + ids=lambda case: case.description, +) +def test_given_unsupported_analysis_when_analysing_natively_then_python_analyses_and_records( + test_case: AnalysisFallbackTestCase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + record_dir: Path = tmp_path / "records" + request: NativeModelAnalysisRequest = analysis_request( + inputs=compile_inputs(project_dir=tmp_path / "project", files=test_case.files), + monkeypatch=monkeypatch, + ) + profile: ExpressionInferenceProfile = request.inference_profile + monkeypatch.setenv(ANALYSIS_RECORD_DIR_ENV_VAR, str(record_dir)) + + analyses: object = analyze_native_model_sql( + request=replace( + request, + allow_compact_analysis=test_case.allow_compact_analysis, + inference_profile=replace( + profile, binding_catalog=(None, profile.binding_catalog)[test_case.keeps_catalog] + ), + ) + ) + + assert analyses is None + assert deferral_kinds(record_dir) == Counter({f"analysis_session:{test_case.expected_kind}": 1}) + + +@pytest.mark.parametrize( + "test_case", + [ + SessionFailureTestCase( + description="untyped inputs defer enrichments before the session fails", + seed=20261009, + model_count=12, + expected_kinds={"analysis_session:session": 1}, + ) + ], + ids=lambda case: case.description, +) +def test_given_session_failure_after_deferrals_when_analysing_then_records_only_the_session( + test_case: SessionFailureTestCase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + record_dir: Path = tmp_path / "records" + request: NativeModelAnalysisRequest = analysis_request( + inputs=compile_inputs( + project_dir=tmp_path / "project", + files=generated_analysis_files( + rng=random.Random(test_case.seed), model_count=test_case.model_count + ), + ), + monkeypatch=monkeypatch, + ) + sessions: list[FailingProvideSession] = failing_provide_sessions(monkeypatch=monkeypatch) + monkeypatch.setenv(ANALYSIS_RECORD_DIR_ENV_VAR, str(record_dir)) + + analyses: object = analyze_native_model_sql(request=request) + + assert analyses is None + assert sum(session.answered for session in sessions) > 0 + assert deferral_kinds(record_dir) == Counter(test_case.expected_kinds) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/integration/src/sqlbuild/compiler/attachments/test_native_attachment_project.py b/tests/integration/src/sqlbuild/compiler/attachments/test_native_attachment_project.py index 41bc763176..1dd66a3c41 100644 --- a/tests/integration/src/sqlbuild/compiler/attachments/test_native_attachment_project.py +++ b/tests/integration/src/sqlbuild/compiler/attachments/test_native_attachment_project.py @@ -102,7 +102,7 @@ + '__expected__orders AS (SELECT id, 1.5 AS amount, label FROM __source("order_events")),\n' + "__assert__no_rows AS (SELECT 1 FROM __source(order_events))\n" }, - expected_preview_entries=frozenset({"SqlTestTargetCatalog"}), + expected_native_entries=frozenset({"SqlTestTargetCatalog"}), expected_outcome_fragment="error: SQL scenario file tests/scenarios/orders_scenario.sql CTE " "'__expected__orders' must not reference project source", ), @@ -114,7 +114,7 @@ + "__seed__channel_codes AS (SELECT 1 AS id, 'web' AS label),\n" + "__expected__orders AS (SELECT 1 AS id FROM __source(order_events))\n" }, - expected_preview_entries=frozenset({"SqlTestTargetCatalog"}), + expected_native_entries=frozenset({"SqlTestTargetCatalog"}), expected_outcome_fragment="error: SQL scenario file tests/scenarios/orders_scenario.sql " "references unknown source 'missing_events'", ), @@ -126,7 +126,7 @@ + "__seed__channel_codes AS (SELECT 1 AS id FROM __source(order_events)),\n" + "__expected__orders AS (SELECT 1 AS id, 1.5 AS amount, 'web' AS label)\n" }, - expected_preview_entries=frozenset({"SqlTestTargetCatalog"}), + expected_native_entries=frozenset({"SqlTestTargetCatalog"}), expected_outcome_fragment="error: SQL scenario file tests/scenarios/orders_scenario.sql " "references unknown source 'missing_events'", ), @@ -137,7 +137,7 @@ + "WITH\n__seed__ AS (SELECT 1 AS id, 'web' AS label),\n" + "__expected__orders AS (SELECT 1 AS id, 1.5 AS amount, 'web' AS label)\n" }, - expected_preview_entries=frozenset({"SqlTestTargetCatalog"}), + expected_native_entries=frozenset({"SqlTestTargetCatalog"}), expected_outcome_fragment="error: SQL scenario 'tests/scenarios/orders_scenario.sql' " "must use __seed__ to identify a target", ), @@ -149,7 +149,7 @@ + "__macro__tidy AS (SELECT 'x'),\n" + "__expected__orders AS (SELECT 1 AS id, 1.5 AS amount, 'web' AS label)\n" }, - expected_preview_entries=frozenset({"SqlTestTargetCatalog"}), + expected_native_entries=frozenset({"SqlTestTargetCatalog"}), expected_outcome_fragment="does not support macro mock CTE '__macro__tidy'", ), AttachmentProjectTestCase( @@ -161,7 +161,7 @@ + "__expected__orders AS (SELECT 1 AS id, 1.5 AS amount, 'web' AS label " + "FROM __assert__no_rows)\n" }, - expected_preview_entries=frozenset({"SqlTestTargetCatalog"}), + expected_native_entries=frozenset({"SqlTestTargetCatalog"}), expected_outcome_fragment="check CTE '__expected__orders' must not depend on " "'__assert__no_rows'", ),