diff --git a/Cargo.lock b/Cargo.lock index 0882c97e86..a5d4e8ec33 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1148,6 +1148,7 @@ dependencies = [ "polyglot-sql-function-catalogs", "rayon", "regex", + "rustc-hash", "serde", "serde_json", "sha2", @@ -1315,6 +1316,7 @@ dependencies = [ "polyglot-sql", "rayon", "regex", + "rustc-hash", "serde", "serde_json", "sha2", diff --git a/Cargo.toml b/Cargo.toml index 1eed898f30..bc0e40fc5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,6 +84,8 @@ ruff_python_ast = "=0.0.4" ruff_python_parser = "=0.0.4" ruff_text_size = "=0.0.4" # Unix only: `sqb` shares the compute log prune lock with Python's `filelock` through `flock`. +# Fast deterministic hasher for hot in-memory maps whose iteration order never reaches output. +rustc-hash = "2" rustix = { version = "1.1.4", default-features = false, features = ["fs", "std"] } saphyr-parser = "0.1" regex = "1" diff --git a/crates/sqlbuild-analysis/Cargo.toml b/crates/sqlbuild-analysis/Cargo.toml index d9f8c47650..034ba496ec 100644 --- a/crates/sqlbuild-analysis/Cargo.toml +++ b/crates/sqlbuild-analysis/Cargo.toml @@ -17,6 +17,7 @@ polyglot-sql.workspace = true polyglot-sql-function-catalogs.workspace = true rayon.workspace = true regex.workspace = true +rustc-hash.workspace = true serde.workspace = true serde_json.workspace = true sha2.workspace = true diff --git a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/extraction.rs b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/extraction.rs index 99a04ce4a9..ced6725397 100644 --- a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/extraction.rs +++ b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/extraction.rs @@ -39,6 +39,19 @@ const SCENARIO_CONTEXT: &str = "SQL scenario"; const DECLARATION_CALLS: [&str; 4] = ["enum", "const", "var", "param"]; /// Longest projection the alias help repeats verbatim. const ALIAS_HELP_EXPRESSION_LIMIT: usize = 80; +const FROM_KEYWORD: &str = "FROM"; +const SELECT_LIST_END_KEYWORDS: [&str; 10] = [ + FROM_KEYWORD, + "WHERE", + "GROUP", + "HAVING", + "QUALIFY", + "WINDOW", + "ORDER", + "LIMIT", + "OFFSET", + "FETCH", +]; const RESERVED_NAMES: [&str; 5] = [ "__actual", "__expected__typed", @@ -777,15 +790,7 @@ fn projection_names( } } } - let mut end = select_list_from(branch, select_end, syntax)?.unwrap_or(branch.len()); - for keyword in [ - "WHERE", "GROUP", "HAVING", "QUALIFY", "WINDOW", "ORDER", "LIMIT", "OFFSET", "FETCH", - ] { - if let Some(position) = find_top_level_clause_keyword(branch, select_end, keyword, syntax)? - { - end = end.min(position); - } - } + let end = select_list_end(branch, select_end, syntax)?.unwrap_or(branch.len()); let expressions = split_top_level(&branch[select_end..end], b',', syntax)?; if expressions.is_empty() { return Err( @@ -1379,20 +1384,32 @@ pub(crate) fn find_top_level_keyword( }) } -/// The top-level `FROM` that ends a select list, skipping the `FROM` of `IS [NOT] DISTINCT FROM`. -fn select_list_from( +/// The first top-level clause keyword ending a select list, skipping names and `IS DISTINCT FROM`. +fn select_list_end( sql: &str, start: usize, syntax: &LexicalSyntax, ) -> Result, String> { - let mut search_start = start; - while let Some(position) = find_top_level_clause_keyword(sql, search_start, "FROM", syntax)? { - if !trailing_distinct_operator(&sql[start..position], syntax)? { - return Ok(Some(position)); + scan_code(sql, start, syntax, |index, depth| { + let Some(first) = byte_at(sql, index).filter(|_| depth == 0) else { + return Ok(CodeStep::Advance); + }; + let first = first.to_ascii_uppercase(); + let Some(keyword) = SELECT_LIST_END_KEYWORDS + .iter() + .filter(|keyword| keyword.as_bytes()[0] == first) + .find(|keyword| consume_keyword(sql, index, keyword).is_some()) + else { + return Ok(CodeStep::Advance); + }; + if previous_code_byte(sql, index, syntax)? == Some(b'.') { + return Ok(CodeStep::Advance); } - search_start = position + "FROM".len(); - } - Ok(None) + if *keyword == FROM_KEYWORD && trailing_distinct_operator(&sql[start..index], syntax)? { + return Ok(CodeStep::Advance); + } + Ok(CodeStep::Stop(index)) + }) } /// Whether `prefix` ends with `IS DISTINCT` or `IS NOT DISTINCT`, ignoring comments. @@ -1437,22 +1454,6 @@ fn trailing_word<'a>( Ok(word_start.map(|start| (&code[..start], &code[start..]))) } -fn find_top_level_clause_keyword( - sql: &str, - start: usize, - keyword: &str, - syntax: &LexicalSyntax, -) -> Result, String> { - let mut search_start = start; - while let Some(position) = find_top_level_keyword(sql, search_start, keyword, syntax)? { - if previous_code_byte(sql, position, syntax)? != Some(b'.') { - return Ok(Some(position)); - } - search_start = position + keyword.len(); - } - Ok(None) -} - fn previous_code_byte(sql: &str, end: usize, syntax: &LexicalSyntax) -> Result, String> { let mut previous = None; let mut index = 0; diff --git a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/helper_scope.rs b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/helper_scope.rs index dc7af49e4f..eb628449de 100644 --- a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/helper_scope.rs +++ b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/helper_scope.rs @@ -5,6 +5,7 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use sqlbuild_core::graph_walk::main::depth_first_postorder::depth_first_postorder; use polyglot_sql::Dialect; +use rustc_hash::FxHashSet; use crate::compiler::_helpers::sql_tests::cte_rename::defined_cte_keys; use crate::compiler::_helpers::sql_tests::cte_slices::SliceDialect; @@ -549,7 +550,7 @@ fn is_generated_cte_name(name: &str) -> bool { /// Distinct lowercase identifiers outside strings and comments, in first-appearance order. fn identifier_tokens(sql: &str, patterns: &SqlTestPatterns) -> Vec { let protected = protected_ranges(&patterns.lexical, sql); - let mut seen: HashSet = HashSet::new(); + let mut seen: FxHashSet = FxHashSet::default(); let mut tokens: Vec = Vec::new(); let mut push = |token: &str| { let lowered = token.to_ascii_lowercase(); diff --git a/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_extraction_rules.rs b/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_extraction_rules.rs index 46a1a60aa5..b1125b3969 100644 --- a/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_extraction_rules.rs +++ b/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_extraction_rules.rs @@ -169,6 +169,22 @@ fn given_authored_sql_tests_when_extracting_then_each_rule_holds() { sql: "WITH __ref__orders AS (SELECT 1 AS id), __expected__orders AS (SELECT 'completed' IS DISTINCT FROM 'x' AS status, a IS NOT /* c */ DISTINCT FROM b AS same FROM __ref__orders) SELECT 1", expected_outcome: accepted_outcome(false), }, + ExtractionRuleTestCase { + description: "clause keywords after a dot or inside parentheses do not end the select list", + mode: "model", + raw: false, + syntax: "generic", + sql: "WITH __ref__orders AS (SELECT 1 AS id), __expected__orders AS (SELECT t.order, t.limit AS l, (SELECT max(id) FROM __ref__orders WHERE id > 0 ORDER BY 1) AS m FROM __ref__orders t WHERE t.id > 0 ORDER BY 1) SELECT 1", + expected_outcome: accepted_outcome(false), + }, + ExtractionRuleTestCase { + description: "the first top-level clause keyword ends a select list without FROM", + mode: "model", + raw: false, + syntax: "generic", + sql: "WITH __ref__orders AS (SELECT 1 AS id), __expected__orders AS (SELECT 1 AS id WHERE TRUE LIMIT 1) SELECT 1", + expected_outcome: accepted_outcome(false), + }, ExtractionRuleTestCase { description: "IS DISTINCT FROM without an alias is rejected", mode: "model", diff --git a/crates/sqlbuild-analysis/src/query_analysis/_helpers/borrowed_facts.rs b/crates/sqlbuild-analysis/src/query_analysis/_helpers/borrowed_facts.rs index 36bd0f9c60..d90013ddbf 100644 --- a/crates/sqlbuild-analysis/src/query_analysis/_helpers/borrowed_facts.rs +++ b/crates/sqlbuild-analysis/src/query_analysis/_helpers/borrowed_facts.rs @@ -1,13 +1,15 @@ //! Immutable CTE output graph over one borrowed syntax tree. use std::cell::Cell; -use std::collections::{BTreeSet, HashMap, HashSet}; +use std::collections::BTreeSet; use std::rc::Rc; use polyglot_sql::expressions::{Column, Join, JoinKind, Select, Star, Values, With}; use polyglot_sql::{ DialectType, Expression, ExpressionWalk, ProjectionNullability, TransformKind, ValidationSchema, }; +use rustc_hash::{FxHashMap, FxHashSet}; + const SQL_WILDCARD: &str = "*"; #[derive(Clone, Debug)] @@ -22,14 +24,14 @@ pub(super) struct OutputFact { #[derive(Clone)] struct Source { outputs: Rc>, - columns: Rc>>, + columns: Rc>>, physical: Option, null_extended: bool, } impl Source { fn new(outputs: Vec, physical: Option) -> Self { - let mut columns: HashMap> = HashMap::new(); + let mut columns: FxHashMap> = FxHashMap::default(); for (index, output) in outputs.iter().enumerate() { let key = output.name.to_lowercase(); let value = (!columns.contains_key(&key)).then_some(index); @@ -44,13 +46,13 @@ impl Source { } } -type Relations = HashMap; +type Relations = FxHashMap; #[derive(Default)] struct Sources { items: Vec<(String, Source)>, - using_columns: HashMap, + using_columns: FxHashMap, } -type Aliases = HashMap>; +type Aliases = FxHashMap>; struct BindingProof { enabled: bool, @@ -110,7 +112,7 @@ fn infer_with_binding_proof( | DialectType::Databricks ), }; - let mut relations = Relations::new(); + let mut relations = Relations::default(); if let Some(schema) = schema { for table in &schema.tables { let outputs: Vec = table @@ -147,7 +149,7 @@ fn query(expression: &Expression, inherited: &Relations, proof: &BindingProof) - )); let mut relations = inherited.clone(); if let Some(with) = query_with(expression) { - let mut names: HashSet = HashSet::new(); + let mut names: FxHashSet = FxHashSet::default(); for cte in &with.ctes { if proof.enabled { proof.require(!with.recursive && names.insert(cte.alias.name.to_lowercase())); @@ -195,11 +197,19 @@ fn set_outputs( right: Vec, by_name: bool, ) -> Vec { + let mut right_by_name: FxHashMap = FxHashMap::default(); + if by_name { + for (index, output) in right.iter().enumerate() { + right_by_name + .entry(output.name.to_ascii_lowercase()) + .or_insert(index); + } + } for (index, output) in left.iter_mut().enumerate() { let other = if by_name { - right - .iter() - .find(|other| other.name.eq_ignore_ascii_case(&output.name)) + right_by_name + .get(&output.name.to_ascii_lowercase()) + .map(|position| &right[*position]) } else { right.get(index) }; @@ -214,11 +224,12 @@ fn set_outputs( } } if by_name { + let mut left_names: FxHashSet = left + .iter() + .map(|output| output.name.to_ascii_lowercase()) + .collect(); for mut output in right { - if !left - .iter() - .any(|other| other.name.eq_ignore_ascii_case(&output.name)) - { + if left_names.insert(output.name.to_ascii_lowercase()) { output.nullability = ProjectionNullability::Nullable; left.push(output); } @@ -273,7 +284,7 @@ fn select_outputs(select: &Select, relations: &Relations, proof: &BindingProof) sources = joined_sources(sources, join, relations, proof); } let mut outputs: Vec = Vec::new(); - let mut aliases = Aliases::new(); + let mut aliases = Aliases::default(); for projection in &select.expressions { let projection = unwrapped(projection); if let Some(star) = projection_star(projection) { @@ -402,7 +413,7 @@ fn source( Expression::Lateral(lateral) if matches!(unwrapped(&lateral.this), Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN")) => { proof.require(false); - let fact = expression_fact(&lateral.this, preceding, &Aliases::new()); + let fact = expression_fact(&lateral.this, preceding, &Aliases::default()); let names: Vec = if lateral.column_aliases.is_empty() { ["SEQ", "KEY", "PATH", "INDEX", "VALUE", "THIS"] .into_iter() @@ -438,7 +449,7 @@ fn values_outputs(values: &Values) -> Vec { for row in &values.expressions { let mut facts: Vec = Vec::new(); for (index, expression) in row.expressions.iter().enumerate() { - let mut fact = expression_fact(expression, &Sources::default(), &Aliases::new()); + let mut fact = expression_fact(expression, &Sources::default(), &Aliases::default()); fact.name = values.column_aliases.get(index).map_or_else( || format!("column{}", index + 1), |alias| alias.name.clone(), @@ -483,7 +494,7 @@ fn joined_sources( items: vec![(name.clone(), right.clone())], ..Default::default() }; - let mut joined: HashMap = HashMap::new(); + let mut joined: FxHashMap = FxHashMap::default(); for identifier in &join.using { let column = Column { name: identifier.clone(), @@ -493,8 +504,8 @@ fn joined_sources( span: None, inferred_type: None, }; - let left = column_fact(&column, &sources, &Aliases::new()); - let right = column_fact(&column, &right_scope, &Aliases::new()); + let left = column_fact(&column, &sources, &Aliases::default()); + let right = column_fact(&column, &right_scope, &Aliases::default()); proof.require( left.as_ref().is_some_and(|fact| fact.resolved) && right.as_ref().is_some_and(|fact| fact.resolved), @@ -535,7 +546,7 @@ fn joined_sources( { proof.require( scalar_binding_expression(on) - && expression_fact(on, &sources, &Aliases::new()).resolved, + && expression_fact(on, &sources, &Aliases::default()).resolved, ); } sources @@ -708,7 +719,7 @@ fn nullability( fn star_outputs(star: &Star, sources: &Sources) -> Vec { let mut outputs: Vec = Vec::new(); - let mut merged: HashSet = HashSet::new(); + let mut merged: FxHashSet = FxHashSet::default(); for (name, source) in &sources.items { if star .table @@ -756,7 +767,7 @@ fn star_outputs(star: &Star, sources: &Sources) -> Vec { .find(|alias| alias.alias.name.eq_ignore_ascii_case(&output.name)) { let name = output.name.clone(); - output = expression_fact(&alias.this, sources, &Aliases::new()); + output = expression_fact(&alias.this, sources, &Aliases::default()); output.name = name; } if let Some((_, alias)) = star diff --git a/crates/sqlbuild-analysis/src/query_analysis/tests/helpers.rs b/crates/sqlbuild-analysis/src/query_analysis/tests/helpers.rs index 8f53fef98e..360e565241 100644 --- a/crates/sqlbuild-analysis/src/query_analysis/tests/helpers.rs +++ b/crates/sqlbuild-analysis/src/query_analysis/tests/helpers.rs @@ -542,6 +542,25 @@ pub(crate) fn borrowed_facts_preserve_union_dependencies_and_join_nullability() facts[0].nullability, polyglot_sql::ProjectionNullability::NonNull ); + let expression = polyglot_sql::parse_one( + "SELECT order_id FROM orders UNION ALL BY NAME SELECT customer_id AS extra_id, customer_id AS ORDER_ID FROM customers", + polyglot_sql::DialectType::Snowflake, + ).expect("valid SQL"); + let facts = crate::query_analysis::borrowed_facts::infer( + &expression, + Some(&schema), + polyglot_sql::DialectType::Snowflake, + ); + assert_eq!( + facts + .iter() + .map(|fact| (fact.name.as_str(), fact.nullability, fact.upstream.len())) + .collect::>(), + vec![ + ("order_id", polyglot_sql::ProjectionNullability::NonNull, 2), + ("extra_id", polyglot_sql::ProjectionNullability::Nullable, 1), + ] + ); let expression = polyglot_sql::parse_one( "SELECT o.order_id, c.customer_id FROM orders o LEFT JOIN customers c ON o.order_id = c.customer_id", polyglot_sql::DialectType::Snowflake, diff --git a/crates/sqlbuild-analysis/src/semantic_validation/_helpers/function_names.rs b/crates/sqlbuild-analysis/src/semantic_validation/_helpers/function_names.rs index b1c0b34a6a..c4c20a44d1 100644 --- a/crates/sqlbuild-analysis/src/semantic_validation/_helpers/function_names.rs +++ b/crates/sqlbuild-analysis/src/semantic_validation/_helpers/function_names.rs @@ -4,10 +4,10 @@ use crate::semantic_validation::models::FunctionProbes; use crate::semantic_validation::types::ProbeKey; use polyglot_sql::{Dialect, DialectType, Expression, ValidationError}; use polyglot_sql_function_catalogs::{CatalogSink, FunctionNameCase, FunctionSignature}; +use rustc_hash::{FxHashMap, FxHashSet}; use sqlbuild_sqltext::sql_scan::main::comment_end::comment_end; use sqlbuild_sqltext::sql_scan::main::non_code_end::non_code_end; use sqlbuild_sqltext::sql_scan::models::{QuotePolicy, Unclosed}; -use std::collections::{HashMap, HashSet}; use std::mem::Discriminant; use std::sync::{LazyLock, OnceLock, PoisonError}; @@ -37,20 +37,21 @@ const NOT_KEYWORD: &str = "NOT"; const MATERIALIZED_KEYWORD: &str = "MATERIALIZED"; const DISTINCT_KEYWORD: &str = "DISTINCT"; -type SpellingIndex = HashMap>; +type SpellingIndex = FxHashMap>; -static KEYWORDS: LazyLock> = +static KEYWORDS: LazyLock> = LazyLock::new(|| NON_CALL_KEYWORDS.split_whitespace().collect()); -static EXPRESSION_STARTS: LazyLock> = +static EXPRESSION_STARTS: LazyLock> = LazyLock::new(|| EXPRESSION_KEYWORDS.split_whitespace().collect()); -static CLAUSE_KEYWORDS: LazyLock> = +static CLAUSE_KEYWORDS: LazyLock> = LazyLock::new(|| CLAUSE_KEYWORD_LIST.split_whitespace().collect()); static ROOT_FRAME: Frame = Frame { select_seen: false, clause: Clause::Other, }; -static SNOWFLAKE: LazyLock> = LazyLock::new(|| accepted_names(SNOWFLAKE_FUNCTIONS)); -static DUCKDB: LazyLock> = LazyLock::new(|| { +static SNOWFLAKE: LazyLock> = + LazyLock::new(|| accepted_names(SNOWFLAKE_FUNCTIONS)); +static DUCKDB: LazyLock> = LazyLock::new(|| { let mut sink: CatalogNames = CatalogNames::default(); polyglot_sql_function_catalogs::register_enabled_catalogs(&mut sink); sink.names @@ -106,7 +107,7 @@ enum Signature { /// Upper-case DuckDB names from the SQL library's introspected function catalogue. #[derive(Default)] struct CatalogNames { - names: HashSet, + names: FxHashSet, } /// Incremental byte-to-character position, so locating every call stays linear in the SQL. @@ -185,7 +186,7 @@ impl FunctionProbes { } } -fn accepted_names(source: &str) -> HashSet { +fn accepted_names(source: &str) -> FxHashSet { source .lines() .map(str::trim) @@ -194,7 +195,7 @@ fn accepted_names(source: &str) -> HashSet { .collect() } -fn accepted(dialect: DialectType) -> Option<&'static HashSet> { +fn accepted(dialect: DialectType) -> Option<&'static FxHashSet> { match dialect { DialectType::Snowflake => Some(&SNOWFLAKE), DialectType::DuckDB => Some(&DUCKDB), @@ -483,7 +484,7 @@ fn spelling_index(dialect: DialectType, arity: usize) -> Option<&'static Spellin _ => return None, }; Some(indexes.get(arity)?.get_or_init(|| { - let mut index: SpellingIndex = HashMap::new(); + let mut index: SpellingIndex = FxHashMap::default(); for name in names { if let Some(signature) = parsed_signature(dialect, name, arity) { index.entry(signature).or_default().push(name.clone()); diff --git a/crates/sqlbuild-cache/src/hosted_shapes/constants.rs b/crates/sqlbuild-cache/src/hosted_shapes/constants.rs index 9bd44bb448..98262c5aab 100644 --- a/crates/sqlbuild-cache/src/hosted_shapes/constants.rs +++ b/crates/sqlbuild-cache/src/hosted_shapes/constants.rs @@ -1,7 +1,7 @@ //! The Python shape names of the natively hosted compiler shapes. /// The Python model type a natively hosted compiler shape is captured as. -pub const HOSTED_SHAPE_TYPES: [(&str, &str); 491] = [ +pub const HOSTED_SHAPE_TYPES: [(&str, &str); 492] = [ ( "sqlbuild._native:AlignedInterval", "sqlbuild.cursor_algebra.models:AlignedInterval", @@ -810,6 +810,10 @@ pub const HOSTED_SHAPE_TYPES: [(&str, &str); 491] = [ "sqlbuild._native:InferredColumn", "sqlbuild.compiler.compile.models:InferredColumn", ), + ( + "sqlbuild._native:InferredNullability", + "sqlbuild.compiler.lineage.types:InferredNullability", + ), ( "sqlbuild._native:InitialValidFrom", "sqlbuild.compiler.planner.types:InitialValidFrom", diff --git a/crates/sqlbuild-cache/src/project_reuse/_helpers/entry.rs b/crates/sqlbuild-cache/src/project_reuse/_helpers/entry.rs index e3603b78d5..fdd365f3e7 100644 --- a/crates/sqlbuild-cache/src/project_reuse/_helpers/entry.rs +++ b/crates/sqlbuild-cache/src/project_reuse/_helpers/entry.rs @@ -44,9 +44,6 @@ pub(crate) fn read_entry(path: &Path) -> Option<(StoredInputs, StoredOutput)> { return None; }; output.stdout = stdout; - if stdout_key(&output.stdout) != digest { - return None; - } Some((inputs, output)) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis/mod.rs index d13e5033ba..f7aa48b0a2 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod binding_locations; mod catalog; pub(crate) mod conversions; pub(crate) mod explicit_references; +pub(crate) mod model_reads; pub(crate) mod normalization; pub(crate) mod normalization_results; mod positions; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis/model_reads.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis/model_reads.rs new file mode 100644 index 0000000000..a15ec21969 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis/model_reads.rs @@ -0,0 +1,119 @@ +//! A compiled model's analysis outcomes as Rust data: views read their owner, records their fields. + +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult}; + +use crate::bindings::_helpers::compiled_project::compiled_model_view::{ + DeclaredRow, InferredRow, LineageInput, NativeCompiledModel, +}; + +/// Each inferred output column as `(name, type, nullability)`, `None` when none were inferred. +pub(crate) fn inferred_rows(model: &Bound<'_, PyAny>) -> PyResult>> { + if let Ok(view) = model.downcast::() { + return view.get().inferred_rows(model.py()); + } + let rows: Bound<'_, PyAny> = model.getattr("inferred_rows")?; + if rows.is_none() { + return Ok(None); + } + let mut inferred: Vec = Vec::new(); + for row in rows.try_iter()? { + let (name, column_type, nullability): (String, Option, Bound<'_, PyAny>) = + row?.extract()?; + inferred.push((name, column_type, nullability.str()?.extract()?)); + } + Ok(Some(inferred)) +} + +/// The names of the inferred output columns, `None` when the shape is unknown. +pub(crate) fn inferred_names(model: &Bound<'_, PyAny>) -> PyResult>> { + Ok(inferred_rows(model)?.map(column_names)) +} + +fn column_names(rows: Vec) -> Vec { + rows.into_iter().map(|(name, _, _)| name).collect() +} + +/// Each declared column of the attached schema as `(name, type)`, `None` without a schema. +pub(crate) fn declared_rows(model: &Bound<'_, PyAny>) -> PyResult>> { + if let Ok(view) = model.downcast::() { + return view.get().schema_column_rows(model.py()); + } + model.getattr("schema_column_rows")?.extract() +} + +/// Whether the attached schema declares dynamic column families. +pub(crate) fn dynamic_columns_declared(model: &Bound<'_, PyAny>) -> PyResult { + if let Ok(view) = model.downcast::() { + return view.get().dynamic_columns_declared(model.py()); + } + model.getattr("has_dynamic_columns")?.is_truthy() +} + +/// `(has star, star resolved)` of the model's analysed query. +pub(crate) fn star_flags(model: &Bound<'_, PyAny>) -> PyResult<(bool, bool)> { + if let Ok(view) = model.downcast::() { + let view: &NativeCompiledModel = view.get(); + return Ok(( + view.fast_lineage_has_star(model.py())?, + view.fast_lineage_star_resolved(model.py())?, + )); + } + Ok(( + model.getattr("fast_lineage_has_star")?.is_truthy()?, + model.getattr("fast_lineage_star_resolved")?.is_truthy()?, + )) +} + +/// Output columns whose types semantic recovery forgot, in any order. +pub(crate) fn unchecked_columns(model: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(view) = model.downcast::() { + return Ok(view.get().unchecked_columns().to_vec()); + } + model + .getattr("unchecked_output_columns")? + .try_iter()? + .map(|column| column?.extract()) + .collect() +} + +/// Each lineage output with the `(resource, column)` pairs it reads; empty without lineage. +pub(crate) fn lineage_inputs(model: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(view) = model.downcast::() { + return Ok(view.get().lineage_inputs(model.py())?.unwrap_or_default()); + } + let columns: Bound<'_, PyAny> = model.getattr("fast_lineage_columns")?; + if !columns.is_truthy()? { + return Ok(Vec::new()); + } + let mut lineage: Vec = Vec::new(); + for output in columns.try_iter()? { + let output: Bound<'_, PyAny> = output?; + let mut upstream: Vec<(String, String)> = Vec::new(); + for source in output.getattr("upstream_columns")?.try_iter()? { + let source: Bound<'_, PyAny> = source?; + upstream.push(( + source.getattr("resource_name")?.extract()?, + source.getattr("column_name")?.extract()?, + )); + } + lineage.push((output.getattr("output_column")?.extract()?, upstream)); + } + Ok(lineage) +} + +/// Whether the native session validated the model's bindings. +pub(crate) fn binding_validated(model: &Bound<'_, PyAny>) -> PyResult { + if let Ok(view) = model.downcast::() { + return view.get().binding_validated(model.py()); + } + model.getattr("binding_validated")?.is_truthy() +} + +/// Whether the model's dynamic pivot proof proved its output columns. +pub(crate) fn output_proven(model: &Bound<'_, PyAny>) -> PyResult { + if let Ok(view) = model.downcast::() { + return view.get().output_proven(model.py()); + } + let proof: Bound<'_, PyAny> = model.getattr("dynamic_column_contract")?; + Ok(!proof.is_none() && proof.getattr("output_proven")?.is_truthy()?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/fixture_columns.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/fixture_columns.rs index 79386119f1..18d65c3499 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/fixture_columns.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/fixture_columns.rs @@ -18,7 +18,6 @@ use crate::bindings::_helpers::records::ids::RecordId; use crate::bindings::_helpers::records::registry; use crate::bindings::types::CompilerDetach; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const PROJECTION_FACTS_CONTEXT: &str = "native fixture projection facts"; const LEGACY_MODE: &str = "legacy"; @@ -53,9 +52,7 @@ fn infer_fixture_column_facts<'py>( .map_err(|reason| compiler_error(native_failure(PROJECTION_FACTS_CONTEXT, &reason)))?; let simple_types: HashMap = facts.types.into_iter().collect(); let column_class: Bound<'py, PyAny> = registry::class(py, RecordId::InferredColumn)?; - let nullability_class: Bound<'py, PyAny> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")?; + let nullability_class: Bound<'py, PyAny> = registry::class(py, RecordId::InferredNullability)?; let mut built: Vec> = Vec::with_capacity(columns.len()); for (name, data_type, nullability) in columns { let data_type: Option = data_type diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs index 0253298b28..b6d43b56f8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs @@ -19,7 +19,6 @@ use crate::bindings::_helpers::records::ids::{NativeAttr, RecordId}; use crate::bindings::_helpers::records::registry; use crate::bindings::_helpers::semantic_shapes::shapes::{VALID_FROM_DEFAULT, VALID_TO_DEFAULT}; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const LOGGING_MODULE: &str = "logging"; const DEBUG_LOGGER: &str = "sqlbuild.compile"; const ANALYSIS_FAILURE_MESSAGE: &str = "native compact query analysis failed"; @@ -151,10 +150,8 @@ fn run_session<'py>( for failure in failures.try_iter()? { log_debug(py, ANALYSIS_FAILURE_MESSAGE, &failure?)?; } - let unknown: Bound<'py, PyAny> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")? - .getattr("UNKNOWN")?; + let unknown: Bound<'py, PyAny> = + registry::class(py, RecordId::InferredNullability)?.getattr("UNKNOWN")?; for item in publications.try_iter()? { let (name, shape): (Bound<'py, PyAny>, Bound<'py, PyAny>) = item?.extract()?; let published: Bound<'py, PyAny> = dict_of(&shape)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/outcome_objects.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/outcome_objects.rs index cd67a2362c..a93bbd7512 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/outcome_objects.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/outcome_objects.rs @@ -12,7 +12,6 @@ use crate::bindings::_helpers::record_behaviours::lineage_facts::{ use crate::bindings::_helpers::records::ids::RecordId; use crate::bindings::_helpers::records::registry; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const LINEAGE_NATIVE: &str = "native"; type Object<'py> = Bound<'py, PyAny>; @@ -140,9 +139,7 @@ fn inferred_columns<'py>(rows: &Bound<'py, PyAny>) -> PyResult return Ok(rows.clone()); } let column_class: Bound<'py, PyAny> = registry::class(py, RecordId::InferredColumn)?; - let nullability_class: Bound<'py, PyAny> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")?; + let nullability_class: Bound<'py, PyAny> = registry::class(py, RecordId::InferredNullability)?; let mut columns: Vec> = Vec::new(); for row in rows.try_iter()? { let (name, data_type, nullability): ( diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/profile_inputs.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/profile_inputs.rs index 2115c45c4f..175bf780ee 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/profile_inputs.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/profile_inputs.rs @@ -9,8 +9,9 @@ use crate::bindings::_helpers::analysis::binding_catalog::{ }; use crate::bindings::_helpers::native_host::adapter::AdapterHost; use crate::bindings::_helpers::native_host::host; +use crate::bindings::_helpers::records::ids::RecordId; +use crate::bindings::_helpers::records::registry; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const GENERIC_DIALECT: &str = "generic"; const ADAPTER_RULE_ID: &str = "python"; /// The nullability rules the native analysis knows, with their native rule ids. @@ -53,9 +54,7 @@ pub(crate) fn nullability_rules(profile: &Bound<'_, PyAny>) -> PyResult<(Pairs, if adapter.is_empty() { return Ok((rows, None)); } - let nullability: Bound<'_, PyAny> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")?; + let nullability: Bound<'_, PyAny> = registry::class(py, RecordId::InferredNullability)?; Ok((rows, Some((adapter.unbind(), nullability.unbind())))) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/stage_functions.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/stage_functions.rs index ba26c4d945..493e4c5d6d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/stage_functions.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/stage_functions.rs @@ -34,7 +34,6 @@ const LINEAGE_RESOURCE_TYPES: [(&str, &str); 4] = [ ("seed", "seed"), ("table_fn", "table_fn"), ]; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; /// Each query's `(succeeded, column rows, has star)`; a native internal failure raises. #[pyfunction] @@ -72,9 +71,7 @@ fn infer_native_query_columns<'py>( )?; let query_columns: Bound<'py, PyAny> = registry::class(py, RecordId::NativeQueryColumns)?; let column_class: Bound<'py, PyAny> = registry::class(py, RecordId::InferredColumn)?; - let nullability_class: Bound<'py, PyAny> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")?; + let nullability_class: Bound<'py, PyAny> = registry::class(py, RecordId::InferredNullability)?; let mut analyses: Vec> = Vec::with_capacity(rows.len()); for (succeeded, columns, has_star) in rows { let columns: Bound<'py, PyAny> = match columns { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs index 6b3728c6e3..b1c5cf17fd 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs @@ -17,7 +17,6 @@ use crate::bindings::_helpers::native_host::host; use crate::bindings::_helpers::records::ids::{NativeAttr, RecordId}; use crate::bindings::_helpers::records::registry; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const GENERIC_DIALECT: &str = "generic"; const UDF_PREFIX: &str = "__sqlbuild_udf_"; @@ -290,7 +289,7 @@ fn record_expression_shapes<'py>( ))); } let unknown: Bound<'py, PyAny> = - imported(py, LINEAGE_TYPES_MODULE, "InferredNullability")?.getattr("UNKNOWN")?; + registry::class(py, RecordId::InferredNullability)?.getattr("UNKNOWN")?; for (source_input, shape) in sources.iter().zip(shapes) { if shape.is_none() { continue; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs index 2e17eb2a39..3903924712 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs @@ -1201,11 +1201,7 @@ const R_INFERRED_COLUMN: RecordSpec = RecordSpec { FieldSpec::new("type", FieldDefault::None), FieldSpec::new( "nullability", - FieldDefault::Member( - "sqlbuild.compiler.lineage.types", - "InferredNullability", - "UNKNOWN", - ), + FieldDefault::NativeMember(RecordId::InferredNullability, "UNKNOWN"), ), ], }; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_analysis.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_analysis.rs index fc7c9d8483..af04b0ec9c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_analysis.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_analysis.rs @@ -175,6 +175,13 @@ fn located( } impl AnalysisRecord { + /// The session's lineage rows, or `None` when the model kept no lineage. + pub(crate) fn lineage_rows(&self) -> Option<&[LineageRow]> { + self.lineage.as_ref().map(|lineage| match lineage { + LineageFacts::Native(rows) | LineageFacts::NativeFacts(rows) => rows.as_slice(), + }) + } + /// The session's lineage as `CompactLineageFacts` or plain facts, or `None`. pub(crate) fn lineage_object( &self, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_view.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_view.rs index 1fabe3eb7c..ac38fc0f72 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_view.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/compiled_model_view.rs @@ -5,6 +5,7 @@ use pyo3::sync::PyOnceLock; use pyo3::types::{PyDict, PyDictMethods, PyFrozenSet, PyString, PyTuple, PyTupleMethods}; use pyo3::{FromPyObject, pyclass, pymethods}; +use sqlbuild_analysis::assembly::analysis_session::models::LineageRow; use sqlbuild_analysis::compiled_project::models::{DynamicProofFact, ModelAnalysisFact}; use crate::bindings::_helpers::compiled_project::compiled_model_analysis::{ @@ -38,6 +39,9 @@ pub(crate) type DeclaredRow = (String, Option); /// `(database, schema, name)` of a relation. pub(crate) type DestinationRow = (Option, Option, String); +/// `(output column, [(resource name, column name)])` of one lineage output. +pub(crate) type LineageInput = (String, Vec<(String, String)>); + /// The Python classes a compiled model's projected attributes are built with. #[derive(FromPyObject)] #[pyo3(from_item_all)] @@ -279,6 +283,32 @@ impl NativeCompiledModel { }) } + /// Each output column of the kept lineage with its `(resource name, column)` sources. + pub(crate) fn lineage_inputs(&self, py: Python<'_>) -> PyResult>> { + self.compiled(py, |compiled, _| { + Ok(compiled + .analysis + .lineage_rows() + .map(|rows| rows.iter().map(lineage_input).collect())) + }) + } + + /// Whether the model's dynamic pivot proof proved its output columns. + pub(crate) fn output_proven(&self, py: Python<'_>) -> PyResult { + self.compiled(py, |compiled, _| { + Ok(compiled + .analysis + .proof + .as_ref() + .is_some_and(|proof| proof.output_proven)) + }) + } + + /// Output columns whose types semantic recovery forgot. + pub(crate) fn unchecked_columns(&self) -> &[String] { + &self.overlay.poisoned + } + fn model_file_attribute(&self, py: Python<'_>, name: &str) -> PyResult> { self.read(py, |record, py| { Ok(record.model_file.bind(py).getattr(name)?.unbind()) @@ -286,6 +316,15 @@ impl NativeCompiledModel { } } +/// A lineage row's output column with its `(resource name, column)` sources. +fn lineage_input(row: &LineageRow) -> LineageInput { + let mut sources: Vec<(String, String)> = Vec::with_capacity(row.sources.len()); + for (_, resource, column) in &row.sources { + sources.push((resource.clone(), column.clone())); + } + (row.output_column.clone(), sources) +} + /// Each family's name casefolded as Python's, with its inferred type. fn casefolded_families(families: &[(String, Option)]) -> Vec<(String, Option)> { families @@ -338,7 +377,7 @@ impl NativeCompiledModel { } #[getter] - fn name(&self, py: Python<'_>) -> PyResult { + pub(crate) fn name(&self, py: Python<'_>) -> PyResult { self.read(py, |record, py| model_name(record, py)) } @@ -426,7 +465,7 @@ impl NativeCompiledModel { /// Each inferred column as `(name, type, nullability)`, or `None` when none were inferred. #[getter] - fn inferred_rows(&self, py: Python<'_>) -> PyResult>> { + pub(crate) fn inferred_rows(&self, py: Python<'_>) -> PyResult>> { self.compiled(py, |compiled, _| { let Some(rows) = &compiled.analysis.inferred else { return Ok(None); @@ -443,7 +482,7 @@ impl NativeCompiledModel { /// Each declared column of the attached schema as `(name, type)`, or `None` without one. #[getter] - fn schema_column_rows(&self, py: Python<'_>) -> PyResult>> { + pub(crate) fn schema_column_rows(&self, py: Python<'_>) -> PyResult>> { self.read(py, |record, py| { let Some(schema) = &record.schema else { return Ok(None); @@ -461,7 +500,7 @@ impl NativeCompiledModel { /// Whether the attached schema declares dynamic columns. #[getter(has_dynamic_columns)] - fn dynamic_columns_declared(&self, py: Python<'_>) -> PyResult { + pub(crate) fn dynamic_columns_declared(&self, py: Python<'_>) -> PyResult { self.read(py, |record, py| { Ok(record .schema @@ -530,12 +569,12 @@ impl NativeCompiledModel { } #[getter] - fn fast_lineage_has_star(&self, py: Python<'_>) -> PyResult { + pub(crate) fn fast_lineage_has_star(&self, py: Python<'_>) -> PyResult { self.compiled(py, |compiled, _| Ok(compiled.analysis.has_star)) } #[getter] - fn fast_lineage_star_resolved(&self, py: Python<'_>) -> PyResult { + pub(crate) fn fast_lineage_star_resolved(&self, py: Python<'_>) -> PyResult { self.compiled(py, |compiled, _| Ok(compiled.analysis.star_resolved)) } @@ -597,7 +636,7 @@ impl NativeCompiledModel { } #[getter] - fn binding_validated(&self, py: Python<'_>) -> PyResult { + pub(crate) fn binding_validated(&self, py: Python<'_>) -> PyResult { self.compiled(py, |compiled, _| Ok(compiled.analysis.binding_validated)) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/loop_classes.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/loop_classes.rs index b22859dafd..f7fe29395d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/loop_classes.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/loop_classes.rs @@ -10,7 +10,6 @@ use crate::bindings::_helpers::records::registry; const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; const AUDITING_TYPES_MODULE: &str = "sqlbuild.compiler.auditing.types"; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; /// The materializations a model's storage policies are inherited for, besides none. const TABLE_BACKED_MATERIALIZATIONS: [&str; 3] = ["TABLE", "INCREMENTAL", "SNAPSHOT"]; /// The declaration kinds whose project-owned records the loop names. @@ -241,7 +240,7 @@ fn input_view_classes<'py>(py: Python<'py>) -> PyResult> { ), ( "nullability", - class(py, LINEAGE_TYPES_MODULE, "InferredNullability")?, + registry::class(py, RecordId::InferredNullability)?, ), ( "compiler_diagnostic", diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/records.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/records.rs index 79c7d56990..c99a29a70c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/records.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/records.rs @@ -15,7 +15,6 @@ use crate::bindings::_helpers::compiled_project::project_view::NativeProjectView use crate::bindings::_helpers::records::ids::RecordId; use crate::bindings::_helpers::records::registry; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const DATACLASSES_MODULE: &str = "dataclasses"; const DECIMAL_KIND: &str = "decimal"; const COLLECTION_KINDS: [&str; 2] = ["list", "set"]; @@ -387,7 +386,7 @@ pub(crate) fn record_model_analyses<'py>( .collect::>()?; let mut transferred: usize = NativeCompiledProject::record_view_analyses(project, views)?; let nullable: Bound<'py, PyAny> = - imported(py, LINEAGE_TYPES_MODULE, "InferredNullability")?.getattr("NULLABLE")?; + registry::class(py, RecordId::InferredNullability)?.getattr("NULLABLE")?; for model in &models { if model.is_instance_of::() { continue; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiler_enums/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiler_enums/constants.rs index b9f26149d2..17cb132b3c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiler_enums/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiler_enums/constants.rs @@ -79,6 +79,15 @@ const E_COLUMN_LINEAGE_CONFIDENCE: EnumSpec = EnumSpec { ("UNKNOWN", "unknown"), ], }; +const E_INFERRED_NULLABILITY: EnumSpec = EnumSpec { + name: "InferredNullability", + doc: "Conservative nullability state inferred for an output column.", + members: &[ + ("NON_NULL", "non_null"), + ("NULLABLE", "nullable"), + ("UNKNOWN", "unknown"), + ], +}; const E_COLUMN_LINEAGE_MODE: EnumSpec = EnumSpec { name: "ColumnLineageMode", doc: "Column lineage analyzer mode.", @@ -715,6 +724,7 @@ pub(crate) const COMPILER_ENUMS: &[EnumSpec] = &[ E_COLUMN_TRANSFORM_KIND, E_COLUMN_LINEAGE_CONFIDENCE, E_COLUMN_LINEAGE_MODE, + E_INFERRED_NULLABILITY, E_POLYGLOT_ANALYSIS_DIALECT, E_MIGRATION_DISCOVERY, E_MIGRATION_COMPATIBILITY, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiler_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiler_records/constants.rs index 63489d0281..85bd0efb7a 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiler_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiler_records/constants.rs @@ -199,11 +199,7 @@ const R_COLUMN_LINEAGE: RecordSpec = RecordSpec { FieldSpec::required("upstream_columns"), FieldSpec::new( "nullability", - FieldDefault::Member( - "sqlbuild.compiler.lineage.types", - "InferredNullability", - "UNKNOWN", - ), + FieldDefault::NativeMember(RecordId::InferredNullability, "UNKNOWN"), ), FieldSpec::new( "confidence", diff --git a/crates/sqlbuild-python/src/bindings/_helpers/hand_built/compiled_model.rs b/crates/sqlbuild-python/src/bindings/_helpers/hand_built/compiled_model.rs index d3a44f5d8b..ccf852eb11 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/hand_built/compiled_model.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/hand_built/compiled_model.rs @@ -6,7 +6,6 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; use pyo3::types::{PyDict, PyList, PyString, PyTuple}; const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; type Object<'py> = Bound<'py, PyAny>; @@ -177,9 +176,7 @@ fn fallback_columns<'py>(rows: &Object<'py>) -> PyResult> { return Ok(rows.clone()); } let column_class: Object<'py> = registry::class(py, RecordId::InferredColumn)?; - let nullability: Object<'py> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")?; + let nullability: Object<'py> = registry::class(py, RecordId::InferredNullability)?; let mut columns: Vec> = Vec::new(); for row in rows.try_iter()? { let (name, data_type, value): (Object<'py>, Object<'py>, Object<'py>) = row?.extract()?; @@ -259,8 +256,7 @@ fn located<'py>( let types = py.import(COMPILE_TYPES_MODULE)?; let diagnostic_class: Object<'py> = registry::class(py, RecordId::CompilerDiagnostic)?; let location_class: Object<'py> = registry::class(py, RecordId::SourceLocation)?; - let phase: Object<'py> = - registry::class(types.py(), RecordId::DiagnosticPhase)?.getattr("COMPILE")?; + let phase: Object<'py> = registry::class(py, RecordId::DiagnosticPhase)?.getattr("COMPILE")?; let model: Object<'py> = types.getattr("CompiledResourceType")?.getattr("MODEL")?; let name: Object<'py> = model_file.getattr("file_path")?.getattr("stem")?; let mut built: Vec> = Vec::new(); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/lineage/lineage_objects.rs b/crates/sqlbuild-python/src/bindings/_helpers/lineage/lineage_objects.rs index 59832a8c6a..a66b3e90f2 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/lineage/lineage_objects.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/lineage/lineage_objects.rs @@ -4,11 +4,10 @@ use crate::bindings::_helpers::native_host::executor::ExecutorHost; use crate::bindings::_helpers::native_host::host; use crate::bindings::_helpers::records::ids::RecordId; use crate::bindings::_helpers::records::registry; -use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyResult, Python}; +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; use pyo3::types::{PyDict, PyTuple}; use sqlbuild_analysis::lineage::models::{LineageColumn, LineageSource}; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; const LOGGING_MODULE: &str = "logging"; const LINEAGE_LOGGER: &str = "sqlbuild.lineage"; @@ -35,11 +34,10 @@ pub(crate) struct LineageClasses<'py> { impl<'py> LineageClasses<'py> { pub(crate) fn new(py: Python<'py>) -> PyResult { - let types: Bound<'py, PyModule> = py.import(LINEAGE_TYPES_MODULE)?; let resource_type: Bound<'py, PyAny> = py .import(COMPILE_TYPES_MODULE)? .getattr("CompiledResourceType")?; - let nullability: Bound<'py, PyAny> = types.getattr("InferredNullability")?; + let nullability: Bound<'py, PyAny> = registry::class(py, RecordId::InferredNullability)?; Ok(Self { qualified_column: registry::class(py, RecordId::QualifiedLineageColumn)?, lineage_source: registry::class(py, RecordId::ColumnLineageSource)?, @@ -50,9 +48,9 @@ impl<'py> LineageClasses<'py> { column_fact: registry::class(py, RecordId::CompiledLineageColumnFact)?, source_fact: registry::class(py, RecordId::CompiledLineageSourceFact)?, compact_facts: registry::class(py, RecordId::CompactLineageFacts)?, - transform_kind: registry::class(types.py(), RecordId::ColumnTransformKind)?, - confidence: registry::class(types.py(), RecordId::ColumnLineageConfidence)?, - analysis_dialect: registry::class(types.py(), RecordId::PolyglotAnalysisDialect)?, + transform_kind: registry::class(py, RecordId::ColumnTransformKind)?, + confidence: registry::class(py, RecordId::ColumnLineageConfidence)?, + analysis_dialect: registry::class(py, RecordId::PolyglotAnalysisDialect)?, model_resource: resource_type.getattr("MODEL")?, unknown_nullability: nullability.getattr("UNKNOWN")?, nullability, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_lineage.rs b/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_lineage.rs index 8fdcbdc561..7439f16ab5 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_lineage.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_lineage.rs @@ -16,6 +16,7 @@ use sqlbuild_analysis::lineage_graph::main::analysis_dialect::analysis_dialect; use sqlbuild_analysis::lineage_graph::models::AnalysisDialect; use sqlbuild_analysis::semantic_validation::models as validation; +use crate::bindings::_helpers::analysis::model_reads::star_flags; use crate::bindings::_helpers::boundary::panics::compiler_error; use crate::bindings::_helpers::lineage::column_lineage_graph::ProjectColumnLineage; use crate::bindings::_helpers::lineage::lineage_objects::{LineageClasses, log_lineage_event}; @@ -137,9 +138,9 @@ fn selected_models<'py>( /// One model's FAST lineage work: its compact facts and whether its root star expands. struct FastModel<'py> { - model: Bound<'py, PyAny>, name: String, facts: Option>, + has_star: bool, expands_star: bool, } @@ -154,8 +155,8 @@ fn fast_lineage<'py>( let mut requests: Vec = Vec::new(); for model in models { let facts: Bound<'py, PyAny> = model.getattr("fast_lineage_columns")?; - let expands_star: bool = model.getattr("fast_lineage_has_star")?.is_truthy()? - && !model.getattr("fast_lineage_star_resolved")?.is_truthy()?; + let (has_star, star_resolved): (bool, bool) = star_flags(model)?; + let expands_star: bool = has_star && !star_resolved; if facts.is_none() { requests.push(FastLineageModel::Parse { query_sql: model.getattr("query_sql")?.extract()?, @@ -168,9 +169,9 @@ fn fast_lineage<'py>( }); } work.push(FastModel { - model: model.clone(), name: model.getattr("name")?.extract()?, facts: (!facts.is_none()).then_some(facts), + has_star, expands_star, }); } @@ -211,10 +212,7 @@ fn fast_lineage<'py>( } columns = PyTuple::new(py, expanded)?.into_any(); } - compact.set_item( - &model.name, - (columns, model.model.getattr("fast_lineage_has_star")?), - )?; + compact.set_item(&model.name, (columns, model.has_star))?; model_order.push(model.name); continue; } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_schema.rs b/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_schema.rs index 43085f139d..052ea85341 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_schema.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/lineage/project_schema.rs @@ -5,6 +5,8 @@ use sqlbuild_analysis::lineage::models::{ LineageResourceType, LineageSchemaResource, RichSchemaResource, }; +use crate::bindings::_helpers::analysis::model_reads::{declared_rows, inferred_rows}; + type TypedColumns = Vec<(String, Option)>; /// The project's models, then sources, then seeds, each with its typed and schema columns. @@ -90,23 +92,13 @@ fn project_columns(project: &Bound<'_, PyAny>) -> PyResult> } fn inferred_columns(model: &Bound<'_, PyAny>) -> PyResult { - let rows: Bound<'_, PyAny> = model.getattr("inferred_rows")?; - if !rows.is_truthy()? { - return Ok(Vec::new()); - } - rows.try_iter()? - .map(|row| { - let (name, column_type, _nullability): (String, Option, Bound<'_, PyAny>) = - row?.extract()?; - Ok((name, column_type)) - }) - .collect() + Ok(inferred_rows(model)? + .unwrap_or_default() + .into_iter() + .map(|(name, column_type, _)| (name, column_type)) + .collect()) } fn schema_columns(model: &Bound<'_, PyAny>) -> PyResult { - let rows: Bound<'_, PyAny> = model.getattr("schema_column_rows")?; - if !rows.is_truthy()? { - return Ok(Vec::new()); - } - rows.try_iter()?.map(|row| row?.extract()).collect() + Ok(declared_rows(model)?.unwrap_or_default()) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_changes/detect.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_changes/detect.rs index e7b772f5f6..274e121ff8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_changes/detect.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_changes/detect.rs @@ -1,5 +1,6 @@ //! Detect how each selected model and function changed since its recorded build. +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::planner_changes::results::native_enum; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; use pyo3::types::{PyDict, PyDictMethods, PyString, PyTuple}; @@ -427,14 +428,12 @@ pub(crate) fn model_change<'py>( /// The model's resolved `replay_on_change` policy. fn replay_policy(model: &Bound<'_, PyAny>) -> PyResult { let py: Python<'_> = model.py(); - let raw: Option = call_with( - ®istry::native(py, NativeAttr::GetConfigStr)?, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", PyString::new(py, REPLAY_CONFIG_KEY).into_any()), - ], + let raw: Option = authored_config_str( + &model.getattr("config")?.getattr("values")?, + REPLAY_CONFIG_KEY, )? - .extract()?; + .map(|value| value.extract()) + .transpose()?; resolve_replay_on_change(raw.as_deref()).map_err(|error| replay_policy_error(py, error)) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/display_plan.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/display_plan.rs index 889c8ea83b..a778e2173f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/display_plan.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/display_plan.rs @@ -1,5 +1,6 @@ //! Display-only plans for external plan surfaces when warehouse planning cannot resolve. +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::planner_entries::values::{ boolean, native, none, planner_model, planner_type, string, }; @@ -33,13 +34,9 @@ fn display_action( } fn config<'py>(model: &Bound<'py, PyAny>, key: &str) -> PyResult> { - native( - model.py(), - NativeAttr::GetConfigStr, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", string(model.py(), key)), - ], + Ok( + authored_config_str(&model.getattr("config")?.getattr("values")?, key)? + .unwrap_or_else(|| model.py().None().into_bound(model.py())), ) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/execution_plan.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/execution_plan.rs index 6a8888bc25..6fa68d01ba 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/execution_plan.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/execution_plan.rs @@ -5,95 +5,17 @@ use crate::bindings::_helpers::compile_driver::python_calls::{ }; use crate::bindings::_helpers::native_host::adapter::AdapterHost; use crate::bindings::_helpers::native_host::host; -use crate::bindings::_helpers::native_threads::exit_signal::signal_interpreter_exit; -use crate::bindings::_helpers::planner_entries::values::{ - COMPILE_TYPES, boolean, native, replaced, string, +use crate::bindings::_helpers::planner_entries::test_planning::{ + PlanningRequest, SqlTestPlanning, start_planning, }; +use crate::bindings::_helpers::planner_entries::values::{COMPILE_TYPES, native, replaced}; use crate::bindings::_helpers::records::ids::{NativeAttr, RecordId}; -use crate::bindings::_helpers::records::registry; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; use pyo3::types::{PyDict, PyTuple}; use pyo3::{pyfunction, wrap_pyfunction}; use std::time::Instant; const INSPECTION_CATALOG: AdapterHost = AdapterHost::OpenInspectionCatalog; -const TEST_PLANNING_THREAD: &str = "sqlbuild-test-planning"; -/// How long interpreter exit waits for each SQL test planning thread still running native code. -const PLANNING_EXIT_JOIN_SECONDS: f64 = 5.0; - -/// The planned SQL tests' future and, when planning runs in the background, its daemon thread. -type TestPlanning<'py> = (Bound<'py, PyAny>, Option>); - -/// `thread.join(timeout)` for a live thread; no timeout when `timeout` is `None`. -fn join_thread(thread: &Bound<'_, PyAny>, timeout: Option) -> PyResult<()> { - if thread.call_method0("is_alive")?.is_truthy()? { - let options = PyDict::new(thread.py()); - options.set_item("timeout", timeout)?; - thread.call_method("join", (), Some(&options))?; - } - Ok(()) -} - -/// Join SQL test planning threads at exit: finalization ending one inside native planning aborts. -#[pyfunction] -fn join_sql_test_planning_threads(py: Python<'_>) -> PyResult<()> { - signal_interpreter_exit(py)?; - for thread in imported(py, "threading", "enumerate")? - .call0()? - .try_iter()? - { - let thread: Bound<'_, PyAny> = thread?; - if thread.getattr("name")?.eq(TEST_PLANNING_THREAD)? { - join_thread(&thread, Some(PLANNING_EXIT_JOIN_SECONDS))?; - } - } - Ok(()) -} - -/// A future resolving to the planned SQL tests, planned on a daemon thread when enabled. -fn sql_test_planning<'py>( - (project, adapter): (&Bound<'py, PyAny>, &Bound<'py, PyAny>), - (selected_keys, sql_test_selection): (&Bound<'py, PyAny>, &Bound<'py, PyAny>), - enabled: bool, -) -> PyResult> { - let py = project.py(); - let future = imported(py, "concurrent.futures", "Future")?.call0()?; - if !enabled { - future.call_method1( - "set_result", - (native( - py, - RecordId::PlannedSqlTests, - &[ - ("selected_keys", selected_keys.clone()), - ("sql_test_selection", sql_test_selection.clone()), - ], - )?,), - )?; - return Ok((future, None)); - } - let kwargs = PyDict::new(py); - kwargs.set_item("future", &future)?; - kwargs.set_item("project", project)?; - kwargs.set_item("adapter", adapter)?; - kwargs.set_item("selected_keys", selected_keys)?; - kwargs.set_item("sql_test_selection", sql_test_selection)?; - let thread = call_with( - &imported(py, "threading", "Thread")?, - &[ - ( - "target", - registry::native(py, NativeAttr::PlanSelectedSqlTestsInto)?, - ), - ("kwargs", kwargs.into_any()), - ("name", string(py, TEST_PLANNING_THREAD)), - ("daemon", boolean(py, true)?), - ], - )?; - thread.call_method0("start")?; - Ok((future, Some(thread))) -} - /// Selected scopes, and the overrides with cursor overrides typed for the selected models. fn scopes_and_cursor_overrides<'py>( args: &KeywordArguments<'py>, @@ -192,7 +114,7 @@ fn planned_output<'py>( args: &KeywordArguments<'py>, runtime: &Bound<'py, PyAny>, (scopes, overrides): (&Bound<'py, PyAny>, &Bound<'py, PyAny>), - test_planning: &Bound<'py, PyAny>, + test_planning: &SqlTestPlanning, ) -> PyResult> { let project = args.value("project")?; let py = project.py(); @@ -295,7 +217,7 @@ fn planned_output<'py>( ("source_freshness", source_freshness.clone()), ], )?; - let planned_sql_tests = test_planning.call_method0("result")?; + let planned_sql_tests = test_planning.result(py)?.into_bound(py); let mut plan_output = native( py, NativeAttr::AssembleBasePlanOutput, @@ -398,12 +320,14 @@ fn build_execution_plan<'py>( )?; let (scopes, overrides) = scopes_and_cursor_overrides(&args)?; let inspection = scopes.getattr("inspection_scope")?; - let (test_planning, planning_thread) = sql_test_planning( - (&project, &adapter), - ( - &inspection.getattr("selected_keys")?, - &inspection.getattr("sql_test_selection")?, - ), + let test_planning = start_planning( + py, + PlanningRequest { + project: project.clone().unbind(), + adapter: adapter.clone().unbind(), + selected_keys: inspection.getattr("selected_keys")?.unbind(), + sql_test_selection: inspection.getattr("sql_test_selection")?.unbind(), + }, args.value("policies")? .getattr("plan_sql_tests")? .is_truthy()?, @@ -414,56 +338,16 @@ fn build_execution_plan<'py>( &[("adapter", adapter), ("connection", connection)], )?; within(&catalog, |_| { - planned_output(&args, &runtime, (&scopes, &overrides), &test_planning) + planned_output(&args, &runtime, (&scopes, &overrides), test_planning.get()) }) })(); - let joined = planning_thread - .as_ref() - .map_or(Ok(()), |thread| join_thread(thread, None)); + let joined = test_planning.get().result(py); let planned = planned?; joined?; Ok(planned) } -/// Start SQL test planning on a daemon thread, or resolve to no tests when disabled. -#[pyfunction] -#[pyo3(signature = (**kwargs))] -fn start_sql_test_planning<'py>( - py: Python<'py>, - kwargs: Option>, -) -> PyResult> { - let args = KeywordArguments::read( - py, - "start_sql_test_planning", - kwargs, - ( - &[ - "project", - "adapter", - "selected_keys", - "sql_test_selection", - "enabled", - ], - &[], - ), - )?; - let (future, _) = sql_test_planning( - (&args.value("project")?, &args.value("adapter")?), - ( - &args.value("selected_keys")?, - &args.value("sql_test_selection")?, - ), - args.flag("enabled")?, - )?; - Ok(future) -} - -/// Add the execution plan entry points and join SQL test planning threads at interpreter exit. +/// Add the execution plan entry point. pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(build_execution_plan, module)?)?; - module.add_function(wrap_pyfunction!(start_sql_test_planning, module)?)?; - let hook = wrap_pyfunction!(join_sql_test_planning_threads, module)?; - module.add_function(hook.clone())?; - imported(module.py(), "atexit", "register")?.call1((hook,))?; - Ok(()) + module.add_function(wrap_pyfunction!(build_execution_plan, module)?) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/mod.rs index 452f3c67a2..31de37adfb 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/mod.rs @@ -8,4 +8,5 @@ mod model_entries; mod plan_output; pub(crate) mod registration; mod relations; +mod test_planning; mod values; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/model_entries.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/model_entries.rs index 65a5a8a122..efef439137 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/model_entries.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/model_entries.rs @@ -1,15 +1,24 @@ //! Per-model plan entries from resolved change results and warehouse snapshot facts. -use crate::bindings::_helpers::compile_driver::python_calls::{KeywordArguments, imported}; +use crate::bindings::_helpers::compile_driver::python_calls::{ + KeywordArguments, imported, keywords, +}; use crate::bindings::_helpers::planner_entries::values::{ COMPILE_TYPES, MICROBATCH_CONCURRENCY_MESSAGE, boolean, error_of, native, none, planner_type, replaced, string, }; +use crate::bindings::_helpers::planner_phases::full_refresh::resolve_model_full_refresh; use crate::bindings::_helpers::planner_resolution::cursor_ranges::{ CursorWindowInputs, config_str, cursor_safety_configs, empty_cursor_inputs, microbatch_range, model_cursor_overrides, plan_cursor_bounds, }; use crate::bindings::_helpers::planner_resolution::cursor_warnings::with_cursor_safety_warnings; +use crate::bindings::_helpers::planner_resolution::orchestration::{ + resolve_model_hook_entries, resolve_model_sql, +}; +use crate::bindings::_helpers::planner_strategy::microbatch_limits::apply_configured_microbatch_limits; +use crate::bindings::_helpers::planner_strategy::model_entries::assemble_model_plan_entry; +use crate::bindings::_helpers::planner_strategy::strategy::resolve_model_plan_action; use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr, RecordId}; use crate::bindings::_helpers::records::registry; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; @@ -68,15 +77,7 @@ fn plan_model_from_change<'py>( ), ], )?; - let resolved = native( - py, - NativeAttr::ResolveModelPlanAction, - &[ - ("model", model.clone()), - ("change_result", request.change_result.clone()), - ("full_refresh", full_refresh.clone()), - ], - )?; + let resolved = resolve_model_plan_action(model, &request.change_result, request.full_refresh)?; let (action, reason) = (resolved.get_item(0)?, resolved.get_item(1)?); let materialization_type = native( py, @@ -123,37 +124,32 @@ fn plan_model_from_change<'py>( ("end_cursor_override", request.overrides.1.clone()), ], )?; - let resolved_sql = native( + let resolved_sql = resolve_model_sql( py, - NativeAttr::ResolveModelSql, - &[ - ("adapter", request.adapter.clone()), - ("model", model.clone()), - ("snapshot", request.snapshot.clone()), - ("context", context.clone()), - ("backfill", backfill.clone()), - ("full_refresh", full_refresh.clone()), - ("cursor_overrides", pair), - ( - "external_sql_reference_resolver", - request.external_resolver.clone(), - ), - ("replaces_relation", boolean(py, replaces_relation)?), - ("full_rebuild_cause", full_rebuild_cause), - ], + Some(keywords( + py, + &[ + ("adapter", request.adapter.clone()), + ("model", model.clone()), + ("snapshot", request.snapshot.clone()), + ("context", context.clone()), + ("backfill", backfill.clone()), + ("full_refresh", full_refresh.clone()), + ("cursor_overrides", pair), + ( + "external_sql_reference_resolver", + request.external_resolver.clone(), + ), + ("replaces_relation", boolean(py, replaces_relation)?), + ("full_rebuild_cause", full_rebuild_cause), + ], + )?), )?; - let hooks = native( - py, - NativeAttr::ResolveModelHookEntries, - &[ - ("model", model.clone()), - ("adapter", request.adapter.clone()), - ("context", context.clone()), - ( - "external_sql_reference_resolver", - request.external_resolver.clone(), - ), - ], + let hooks = resolve_model_hook_entries( + model, + request.adapter.clone(), + context.clone(), + request.external_resolver.clone(), )?; let (future_config, start_config) = cursor_safety_configs(model, context)?; native( @@ -206,9 +202,8 @@ fn plan_model_from_change<'py>( let cursor_bounds = plan_cursor_bounds(model, &window)?; let (empty_inputs, waiting_on_empty_inputs) = empty_cursor_inputs(model, &window)?; let range = microbatch_range(model, &window)?; - let assembled = native( + let assembled = assemble_model_plan_entry(Some(&keywords( py, - NativeAttr::AssembleModelPlanEntry, &[ ("model", model.clone()), ("adapter", request.adapter.clone()), @@ -238,7 +233,7 @@ fn plan_model_from_change<'py>( ("start_cursor_config", start_config), ("invocation_time", context.getattr("invocation_time")?), ], - )?; + )?))?; let bounds = if range.is_truthy()? { range } else { @@ -504,15 +499,7 @@ fn plan_entries<'py>(args: &KeywordArguments<'py>) -> PyResult &cursor_overrides, (start_override.clone(), end_override.clone()), )?; - let effective_full_refresh = native( - py, - NativeAttr::ResolveModelFullRefresh, - &[ - ("model", model.clone()), - ("cli_full_refresh", full_refresh.clone()), - ], - )? - .is_truthy()?; + let effective_full_refresh = resolve_model_full_refresh(&model, full_refresh.is_truthy()?)?; let (mut entry, entry_warnings) = plan_model_from_change(&ModelPlanRequest { model: model.clone(), snapshot: snapshot.clone(), @@ -560,11 +547,7 @@ fn plan_entries<'py>(args: &KeywordArguments<'py>) -> PyResult }; entry = replaced(&entry, &[("unaccounted_partition_policy", policy)])?; if entry.getattr("action")?.ne(&skip)? { - let limited = native( - py, - NativeAttr::ApplyConfiguredMicrobatchLimits, - &[("entry", entry), ("inputs", inputs.clone())], - )?; + let limited = apply_configured_microbatch_limits(&entry, &inputs)?; entry = limited.get_item(0)?; for warning in limited.get_item(1)?.try_iter()? { warnings.push(warning?); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/plan_output.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/plan_output.rs index a308ef4b52..7e21078603 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/plan_output.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/plan_output.rs @@ -761,45 +761,6 @@ fn plan_selected_sql_tests<'py>( planned_sql_tests(project, adapter, selected_keys, sql_test_selection) } -/// Plan selected SQL tests into `future`, which re-raises any failure on the joining thread. -#[pyfunction] -#[pyo3(signature = (**kwargs))] -fn plan_selected_sql_tests_into<'py>( - py: Python<'py>, - kwargs: Option>, -) -> PyResult<()> { - let args = KeywordArguments::read( - py, - "plan_selected_sql_tests_into", - kwargs, - ( - &[ - "future", - "project", - "adapter", - "selected_keys", - "sql_test_selection", - ], - &[], - ), - )?; - let future = args.value("future")?; - match native( - py, - NativeAttr::PlanSelectedSqlTests, - &[ - ("project", args.value("project")?), - ("adapter", args.value("adapter")?), - ("selected_keys", args.value("selected_keys")?), - ("sql_test_selection", args.value("sql_test_selection")?), - ], - ) { - Ok(planned) => future.call_method1("set_result", (planned,))?, - Err(error) => future.call_method1("set_exception", (error.value(future.py()),))?, - }; - Ok(()) -} - /// Return model closures resolved natively in one project batch, keyed by SQL test key. #[pyfunction] #[pyo3(signature = (*, project, tests))] @@ -863,7 +824,6 @@ pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(build_plan_output, module)?)?; module.add_function(wrap_pyfunction!(plan_function, module)?)?; module.add_function(wrap_pyfunction!(plan_selected_sql_tests, module)?)?; - module.add_function(wrap_pyfunction!(plan_selected_sql_tests_into, module)?)?; module.add_function(wrap_pyfunction!(sql_test_is_selected, module)?)?; module.add_function(wrap_pyfunction!(sql_test_model_chain_names_by_key, module)?)?; module.add_function(wrap_pyfunction!(build_selected_test_entries, module)?) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/registration.rs index 32e77e42bf..15f31ed296 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/registration.rs @@ -9,5 +9,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { super::execution_plan::register(module)?; super::model_entries::register(module)?; super::plan_output::register(module)?; - super::relations::register(module) + super::relations::register(module)?; + super::test_planning::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/relations.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/relations.rs index a81b416a48..723a14b74c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/relations.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/relations.rs @@ -2,6 +2,7 @@ use crate::bindings::_helpers::compile_driver::python_calls::KeywordArguments; use crate::bindings::_helpers::planner_entries::values::{native, none, planner_model}; +use crate::bindings::_helpers::planner_strategy::audit_entries::audit_is_selected; use crate::bindings::_helpers::records::ids::{NativeAttr, RecordId}; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; use pyo3::types::{PyDict, PyDictMethods, PyFrozenSet, PySet, PySetMethods, PyTuple}; @@ -16,16 +17,7 @@ fn selected_audit_relation_keys<'py>( let keys = PySet::empty(py)?; for audit in project.getattr("audits")?.try_iter()? { let audit = audit?; - if native( - py, - NativeAttr::AuditIsSelected, - &[ - ("audit", audit.clone()), - ("selected_keys", selected_keys.clone()), - ], - )? - .is_truthy()? - { + if audit_is_selected(&audit, selected_keys)? { for dependency in audit.getattr("scope_deps")?.try_iter()? { keys.add(dependency?)?; } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs new file mode 100644 index 0000000000..9c13e79504 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs @@ -0,0 +1,269 @@ +//! SQL test planning on a native thread, joined by the planner or, at interpreter exit, within a bound. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel}; +use std::sync::{Mutex, PoisonError}; +use std::time::{Duration, Instant}; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; +use pyo3::sync::PyOnceLock; +use pyo3::types::PyDict; +use pyo3::{pyclass, pyfunction, pymethods, wrap_pyfunction}; + +use crate::bindings::_helpers::compile_driver::python_calls::{ + KeywordArguments, call_with, imported, +}; +use crate::bindings::_helpers::native_threads::exit_signal::signal_interpreter_exit; +use crate::bindings::_helpers::records::ids::{NativeAttr, RecordId}; +use crate::bindings::_helpers::records::registry; + +const TEST_PLANNING_THREAD: &str = "sqlbuild-test-planning"; +/// How long interpreter exit waits for SQL test planning threads still running native code. +const PLANNING_EXIT_JOIN: Duration = Duration::from_secs(5); +/// How often a waiting planner wakes to notice interrupt signals. +const RESULT_POLL: Duration = Duration::from_millis(100); + +/// How often interpreter exit checks whether planning threads finished. +const EXIT_POLL: Duration = Duration::from_millis(10); + +/// The process-wide count of running planning threads. +static PLANNING_THREADS: PyOnceLock> = PyOnceLock::new(); + +/// Planning threads still running native code. +#[pyclass(module = "sqlbuild._native", frozen)] +struct PlanningThreads { + running: AtomicUsize, +} + +/// Marks one planning thread finished when dropped, including when the thread panics. +struct RunningPlanning(Py); + +impl Drop for RunningPlanning { + fn drop(&mut self) { + self.0.get().finished(); + } +} + +impl PlanningThreads { + fn started(&self) { + self.running.fetch_add(1, Ordering::AcqRel); + } + + fn finished(&self) { + self.running.fetch_sub(1, Ordering::AcqRel); + } + + /// Wait without the GIL until no planning thread runs or `limit` passes. + fn wait(&self, limit: Duration) { + let deadline: Instant = Instant::now() + limit; + while self.running.load(Ordering::Acquire) > 0 && Instant::now() < deadline { + std::thread::sleep(EXIT_POLL); + } + } +} + +/// The running planning threads, created on first use. +fn planning_threads(py: Python<'_>) -> PyResult<&Py> { + PLANNING_THREADS.get_or_try_init(py, || { + Py::new( + py, + PlanningThreads { + running: AtomicUsize::new(0), + }, + ) + }) +} + +/// The planned SQL tests, or the error planning raised. +type Planned = Result, PyErr>; + +/// The planned SQL tests of one plan, planned on a native thread when enabled. +#[pyclass(module = "sqlbuild._native", frozen)] +pub(crate) struct SqlTestPlanning { + pending: Mutex>>, + finished: Mutex>, +} + +impl SqlTestPlanning { + fn of(pending: Option>, finished: Option) -> Self { + Self { + pending: Mutex::new(pending), + finished: Mutex::new(finished), + } + } + + /// The outcome the planning thread sends, waited for without the GIL or any lock held. + fn received(&self, py: Python<'_>, mut receiver: Receiver) -> PyResult { + loop { + let (waited, next) = py.detach(move || { + let next = receiver.recv_timeout(RESULT_POLL); + (receiver, next) + }); + receiver = waited; + match next { + Ok(planned) => return Ok(planned), + Err(RecvTimeoutError::Timeout) => { + if let Err(signal) = py.check_signals() { + *self.pending.lock().unwrap_or_else(PoisonError::into_inner) = + Some(receiver); + return Err(signal); + } + } + Err(RecvTimeoutError::Disconnected) => { + return Ok(Err(PyRuntimeError::new_err( + "SQL test planning thread stopped", + ))); + } + } + } + } +} + +#[pymethods] +impl SqlTestPlanning { + /// Wait for planning to finish and return the planned tests, raising its error. + pub(crate) fn result(&self, py: Python<'_>) -> PyResult> { + let pending: Option> = self + .pending + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + if let Some(receiver) = pending { + let planned: Planned = self.received(py, receiver)?; + *self.finished.lock().unwrap_or_else(PoisonError::into_inner) = Some(planned); + } + match self + .finished + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + { + Some(Ok(planned)) => Ok(planned.clone_ref(py)), + Some(Err(error)) => Err(error.clone_ref(py)), + None => Err(PyRuntimeError::new_err( + "SQL test planning result is awaited elsewhere", + )), + } + } +} + +/// The inputs one planning run reads, owned so they can move to the planning thread. +pub(crate) struct PlanningRequest { + pub(crate) project: Py, + pub(crate) adapter: Py, + pub(crate) selected_keys: Py, + pub(crate) sql_test_selection: Py, +} + +/// `plan_selected_sql_tests(**request)`, looked up on the module so tests can replace it. +fn plan(py: Python<'_>, request: &PlanningRequest) -> Planned { + call_with( + ®istry::native(py, NativeAttr::PlanSelectedSqlTests)?, + &[ + ("project", request.project.bind(py).clone()), + ("adapter", request.adapter.bind(py).clone()), + ("selected_keys", request.selected_keys.bind(py).clone()), + ( + "sql_test_selection", + request.sql_test_selection.bind(py).clone(), + ), + ], + ) + .map(Bound::unbind) +} + +/// Plan on a native thread when `enabled`, or resolve at once to no planned tests. +pub(crate) fn start_planning( + py: Python<'_>, + request: PlanningRequest, + enabled: bool, +) -> PyResult> { + if !enabled { + let empty = call_with( + ®istry::class(py, RecordId::PlannedSqlTests)?, + &[ + ("selected_keys", request.selected_keys.bind(py).clone()), + ( + "sql_test_selection", + request.sql_test_selection.bind(py).clone(), + ), + ], + )?; + return Py::new(py, SqlTestPlanning::of(None, Some(Ok(empty.unbind())))); + } + let (sender, receiver): (Sender, Receiver) = channel(); + let threads: Py = planning_threads(py)?.clone_ref(py); + threads.get().started(); + let spawned = std::thread::Builder::new() + .name(TEST_PLANNING_THREAD.to_owned()) + .spawn(move || { + let _running = RunningPlanning(threads); + let planned: Planned = Python::attach(|py| { + let planned: Planned = plan(py, &request); + drop(request); + planned + }); + let _receiver_gone = sender.send(planned); + }); + if let Err(error) = spawned { + planning_threads(py)?.get().finished(); + return Err(PyRuntimeError::new_err(error.to_string())); + } + Py::new(py, SqlTestPlanning::of(Some(receiver), None)) +} + +/// Join SQL test planning threads at exit: finalization ending one inside native planning aborts. +#[pyfunction] +fn join_sql_test_planning_threads(py: Python<'_>) -> PyResult<()> { + signal_interpreter_exit(py)?; + let threads: &PlanningThreads = planning_threads(py)?.get(); + py.detach(|| threads.wait(PLANNING_EXIT_JOIN)); + Ok(()) +} + +/// Start SQL test planning on a native thread, or resolve to no tests when disabled. +#[pyfunction] +#[pyo3(signature = (**kwargs))] +fn start_sql_test_planning<'py>( + py: Python<'py>, + kwargs: Option>, +) -> PyResult> { + let args = KeywordArguments::read( + py, + "start_sql_test_planning", + kwargs, + ( + &[ + "project", + "adapter", + "selected_keys", + "sql_test_selection", + "enabled", + ], + &[], + ), + )?; + start_planning( + py, + PlanningRequest { + project: args.value("project")?.unbind(), + adapter: args.value("adapter")?.unbind(), + selected_keys: args.value("selected_keys")?.unbind(), + sql_test_selection: args.value("sql_test_selection")?.unbind(), + }, + args.flag("enabled")?, + ) +} + +/// Add the planning handle and entry points, and join planning threads at interpreter exit. +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_function(wrap_pyfunction!(start_sql_test_planning, module)?)?; + let hook = wrap_pyfunction!(join_sql_test_planning_threads, module)?; + module.add_function(hook.clone())?; + imported(module.py(), "atexit", "register")?.call1((hook,))?; + Ok(()) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_migrations/column_objects.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_migrations/column_objects.rs index 8408fc47cf..50b4e856d7 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_migrations/column_objects.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_migrations/column_objects.rs @@ -1,6 +1,7 @@ //! Boundary conversions for the column migration planner. use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::records::ids::{NativeAttr, NativeItem, RecordId}; use crate::bindings::_helpers::records::registry; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; @@ -42,13 +43,9 @@ pub(super) fn native<'py>( } pub(super) fn config<'py>(model: &Bound<'py, PyAny>, key: &str) -> PyResult> { - native( - model.py(), - NativeAttr::GetConfigStr, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", text(model.py(), key)), - ], + Ok( + authored_config_str(&model.getattr("config")?.getattr("values")?, key)? + .unwrap_or_else(|| model.py().None().into_bound(model.py())), ) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_phases/objects.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_phases/objects.rs index dece55e662..81a56d9ab4 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_phases/objects.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_phases/objects.rs @@ -1,9 +1,10 @@ //! Boundary helpers shared by the native planner phases. use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported, keywords}; +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::native_host::host; use crate::bindings::_helpers::native_host::models::HostId; -use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr, NativeItem, RecordId}; +use crate::bindings::_helpers::records::ids::{ErrorId, NativeItem, RecordId}; use crate::bindings::_helpers::records::registry; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; use pyo3::types::{PyFrozenSet, PyString, PyTuple}; @@ -154,15 +155,9 @@ pub(crate) fn by_name<'py>( /// A model config value as text, `None` when absent or not text. pub(crate) fn config_str(model: &Bound<'_, PyAny>, key: &str) -> PyResult> { - native( - model.py(), - NativeAttr::GetConfigStr, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", text(model.py(), key)), - ], - )? - .extract() + authored_config_str(&model.getattr("config")?.getattr("values")?, key)? + .map(|value| value.extract()) + .transpose() } /// The compiled object key of the model `name`. diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_model.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_model.rs index 45114be861..7a41207e0b 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_model.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_model.rs @@ -1,6 +1,9 @@ //! Model cursor orchestration over live warehouse snapshot facts. use crate::bindings::_helpers::compile_driver::python_calls::{KeywordArguments, call_with}; +use crate::bindings::_helpers::model_config::config_values::{ + config_cursor_bound, config_str as authored_config_str, +}; use crate::bindings::_helpers::native_host::adapter::AdapterHost; use crate::bindings::_helpers::native_host::executor::ExecutorHost; use crate::bindings::_helpers::native_host::host; @@ -16,20 +19,13 @@ use pyo3::types::PyDict; use pyo3::{IntoPyObject, PyErr, pyfunction, wrap_pyfunction}; fn config<'py>(model: &Bound<'py, PyAny>, key: &str) -> PyResult> { - call_with( - ®istry::native( - model.py(), - if key == CURSOR_START { - NativeAttr::GetConfigCursorBound - } else { - NativeAttr::GetConfigStr - }, - )?, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", string(model.py(), key)), - ], - ) + let values: Bound<'py, PyAny> = model.getattr("config")?.getattr("values")?; + let value: Option> = if key == CURSOR_START { + config_cursor_bound(&values, key)? + } else { + authored_config_str(&values, key)? + }; + Ok(value.unwrap_or_else(|| model.py().None().into_bound(model.py()))) } fn unbounded(py: Python<'_>) -> PyResult> { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_overrides.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_overrides.rs index ca59b2a270..33ba90bdd1 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_overrides.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_overrides.rs @@ -1,8 +1,9 @@ //! Selected-model inspection and untyped override conversion. use crate::bindings::_helpers::compile_driver::python_calls::{KeywordArguments, call_with}; +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::planner_resolution::values::{method, string}; -use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr, RecordId}; +use crate::bindings::_helpers::records::ids::{ErrorId, RecordId}; use crate::bindings::_helpers::records::registry; use pyo3::prelude::{ Bound, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, @@ -27,7 +28,6 @@ fn type_cursor_overrides<'py>( if overrides.is_none() || !overrides.getattr("has_untyped")?.is_truthy()? { return Ok(overrides); } - let get_config = registry::native(py, NativeAttr::GetConfigStr)?; let mut models: Vec<(String, String)> = Vec::new(); let selected = if args.passed("selected_models")? { args.value("selected_models")? @@ -37,11 +37,8 @@ fn type_cursor_overrides<'py>( for model in selected.try_iter()? { let model = model?; let values = model.getattr("config")?.getattr("values")?; - let config = |key: &str| { - call_with( - &get_config, - &[("values", values.clone()), ("key", string(py, key))], - ) + let config = |key: &str| -> PyResult> { + Ok(authored_config_str(&values, key)?.unwrap_or_else(|| py.None().into_bound(py))) }; if !config("materialized")?.eq("incremental")? || config("cursor")?.is_none() { continue; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_ranges.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_ranges.rs index 97421226ba..c13ec8d03f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_ranges.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/cursor_ranges.rs @@ -6,6 +6,7 @@ use crate::bindings::_helpers::compile_driver::python_calls::{ use crate::bindings::_helpers::cursor::conversion::CursorSession; use crate::bindings::_helpers::cursor::scalars::{parse_value, sentinel_to_token}; use crate::bindings::_helpers::cursor::values::{local_date_time, replaced}; +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::native_host::executor::ExecutorHost; use crate::bindings::_helpers::native_host::host; use crate::bindings::_helpers::planner_resolution::cursor_bounds::duration; @@ -73,13 +74,9 @@ fn native<'py>( /// One string-valued model configuration entry. pub(crate) fn config_str<'py>(model: &Bound<'py, PyAny>, key: &str) -> PyResult> { - native( - model.py(), - NativeAttr::GetConfigStr, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", string(model.py(), key)), - ], + Ok( + authored_config_str(&model.getattr("config")?.getattr("values")?, key)? + .unwrap_or_else(|| model.py().None().into_bound(model.py())), ) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/mod.rs index 6fa15a9dae..df9dfd5f7e 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/mod.rs @@ -12,7 +12,7 @@ mod cursor_windows; mod future_safety; mod locations; mod maximum_safety; -mod orchestration; +pub(crate) mod orchestration; mod refs; pub(crate) mod registration; mod sources; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/orchestration.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/orchestration.rs index 3496d94b0c..862b1f7b5e 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/orchestration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_resolution/orchestration.rs @@ -4,6 +4,7 @@ use crate::bindings::_helpers::compile_driver::python_calls::{ KeywordArguments, call_with, imported, }; use crate::bindings::_helpers::cursor::scalars::sentinel_to_token; +use crate::bindings::_helpers::model_config::config_values::config_str as authored_config_str; use crate::bindings::_helpers::planner_resolution::values::string; use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr, RecordId}; use crate::bindings::_helpers::records::registry; @@ -135,7 +136,7 @@ impl<'py> Resolution<'py> { #[pyfunction] #[pyo3(signature = (**kwargs))] -fn resolve_model_sql<'py>( +pub(crate) fn resolve_model_sql<'py>( py: Python<'py>, kwargs: Option>, ) -> PyResult> { @@ -166,13 +167,9 @@ fn resolve_model_sql<'py>( context: args.value("context")?, external: args.value("external_sql_reference_resolver")?, }; - let cursor_type = call_with( - ®istry::native(py, NativeAttr::GetConfigStr)?, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", string(py, "cursor_type")), - ], - )?; + let cursor_type = + authored_config_str(&model.getattr("config")?.getattr("values")?, "cursor_type")? + .unwrap_or_else(|| py.None().into_bound(py)); let overrides = args.value("cursor_overrides")?; let bounds = call_with( ®istry::native(py, NativeAttr::ComputeModelCursorBounds)?, @@ -242,13 +239,11 @@ fn resolve_model_sql<'py>( /// Effective append cursor lower-bound inclusivity. fn append_cursor_inclusive<'py>(model: &Bound<'py, PyAny>) -> PyResult> { let py = model.py(); - let strategy = call_with( - ®istry::native(py, NativeAttr::GetConfigStr)?, - &[ - ("values", model.getattr("config")?.getattr("values")?), - ("key", string(py, "incremental_strategy")), - ], - )?; + let strategy = authored_config_str( + &model.getattr("config")?.getattr("values")?, + "incremental_strategy", + )? + .unwrap_or_else(|| py.None().into_bound(py)); let raw = model .getattr("config")? .getattr("values")? @@ -352,7 +347,7 @@ fn resolve_function_sql<'py>( #[pyfunction] #[pyo3(signature = (*, model, adapter, context, external_sql_reference_resolver))] -fn resolve_model_hook_entries<'py>( +pub(crate) fn resolve_model_hook_entries<'py>( model: &Bound<'py, PyAny>, adapter: Bound<'py, PyAny>, context: Bound<'py, PyAny>, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs index 793f79e5e3..6c190d691c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs @@ -1,9 +1,12 @@ //! Assemble audit plan dataclasses using native rendering and scheduling. +use crate::bindings::_helpers::planner_strategy::audits::{ + reads_outside_target_lineage, resolve_attachment_kind, resolve_effective_run_scope, +}; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult}; use pyo3::pyfunction; use pyo3::types::{ - PyDict, PyDictMethods, PyList, PyListMethods, PyString, PyTuple, PyTupleMethods, + PyBool, PyDict, PyDictMethods, PyList, PyListMethods, PyString, PyTuple, PyTupleMethods, }; use crate::bindings::_helpers::boundary::keyword_request::Request; @@ -81,14 +84,8 @@ pub(crate) fn plan_audit<'py>(request: Option<&Bound<'py, PyDict>>) -> PyResult< &format!("audit '{name}' planned evidence SQL"), )? }; - let attachment = call_with( - ®istry::native(py, NativeAttr::ResolveAttachmentKind)?, - &[ - ("audit", audit.clone()), - ("downstream_deps", request.get("lineage_downstream_deps")?), - ], - )?; - let attachment = attachment.cast::()?; + let attachment: Bound<'_, PyTuple> = + resolve_attachment_kind(&audit, &request.get("lineage_downstream_deps")?)?; let kind = attachment.get_item(0)?; let target = attachment.get_item(1)?; let kind_text: String = kind.extract()?; @@ -101,13 +98,7 @@ pub(crate) fn plan_audit<'py>(request: Option<&Bound<'py, PyDict>>) -> PyResult< }; let requested = audit_enum(&audit, "run_scope", "AuditRunScope", "final")?; let severity = audit_enum(&audit, "severity", "AuditSeverity", "warn")?; - let effective = call_with( - ®istry::native(py, NativeAttr::ResolveEffectiveRunScope)?, - &[ - ("requested_run_scope", requested.clone()), - ("attached_model_materialization", materialization), - ], - )?; + let effective = resolve_effective_run_scope(&requested, &materialization)?; let mut target_kind = audit.getattr("attached_target_kind")?; if target_kind.is_none() && (kind_text == MODEL_ATTACHMENT || kind_text == SOURCE_ATTACHMENT) { target_kind = @@ -153,13 +144,12 @@ pub(crate) fn plan_audit<'py>(request: Option<&Bound<'py, PyDict>>) -> PyResult< ("attached_target_name", target), ( "reads_outside_target_lineage", - call_with( - ®istry::native(py, NativeAttr::ReadsOutsideTargetLineage)?, - &[ - ("audit", audit), - ("upstream_deps", request.get("lineage_upstream_deps")?), - ], - )?, + PyBool::new( + py, + reads_outside_target_lineage(&audit, &request.get("lineage_upstream_deps")?)?, + ) + .to_owned() + .into_any(), ), ] { values.set_item(field, value)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audits.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audits.rs index ba45519683..6b8649f776 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audits.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audits.rs @@ -58,7 +58,7 @@ fn attachment_value(kind: AuditAttachmentKind) -> &'static str { /// Audit attachment kind and attached target name; attached audits gate their target. #[pyfunction] #[pyo3(signature = (*, audit, downstream_deps))] -fn resolve_attachment_kind<'py>( +pub(crate) fn resolve_attachment_kind<'py>( audit: &Bound<'py, PyAny>, downstream_deps: &Bound<'py, PyAny>, ) -> PyResult> { @@ -110,7 +110,7 @@ fn resolve_attachment_kind<'py>( /// Whether an attached audit reads anything besides its target and the target's upstream. #[pyfunction] #[pyo3(signature = (*, audit, upstream_deps))] -fn reads_outside_target_lineage( +pub(crate) fn reads_outside_target_lineage( audit: &Bound<'_, PyAny>, upstream_deps: &Bound<'_, PyAny>, ) -> PyResult { @@ -135,7 +135,7 @@ fn reads_outside_target_lineage( /// The run scope after degradation: delta-and-final only for incremental and snapshot targets. #[pyfunction] #[pyo3(signature = (*, requested_run_scope, attached_model_materialization))] -fn resolve_effective_run_scope<'py>( +pub(crate) fn resolve_effective_run_scope<'py>( requested_run_scope: &Bound<'py, PyAny>, attached_model_materialization: &Bound<'py, PyAny>, ) -> PyResult> { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/strategy.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/strategy.rs index 389a5f139b..cf13bb91c8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/strategy.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/strategy.rs @@ -72,7 +72,7 @@ fn materialization_of<'py>(model: &Bound<'py, PyAny>) -> PyResult( +pub(crate) fn resolve_model_plan_action<'py>( model: &Bound<'py, PyAny>, change_result: &Bound<'py, PyAny>, full_refresh: bool, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/project_assembly/declared_relations.rs b/crates/sqlbuild-python/src/bindings/_helpers/project_assembly/declared_relations.rs index 1644ce7838..37cc938dd0 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/project_assembly/declared_relations.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/project_assembly/declared_relations.rs @@ -4,13 +4,14 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyRes use pyo3::types::{PyBool, PyDict, PyDictMethods}; use pyo3::{pyfunction, wrap_pyfunction}; +use crate::bindings::_helpers::records::ids::RecordId; +use crate::bindings::_helpers::records::registry; + const UNKNOWN_TYPE: &str = "UNKNOWN"; const NOT_NULL_AUDIT_NAME: &str = "not_null"; const ENFORCED_CONTRACT: &str = "enforced"; const CONTRACT_KEY: &str = "contract"; const TABLE_FUNCTION_PREFIX: &str = "__sqlbuild_table_function_"; -const NULLABILITY_MODULE: &str = "sqlbuild.compiler.lineage.types"; -const NULLABILITY_CLASS: &str = "InferredNullability"; const NON_NULL_MEMBER: &str = "NON_NULL"; const UNKNOWN_MEMBER: &str = "UNKNOWN"; @@ -65,8 +66,7 @@ fn table_function_name(function: &Bound<'_, PyAny>) -> PyResult { #[pyfunction] fn declared_relations<'py>(inputs: Bound<'py, PyAny>) -> PyResult> { let py: Python<'py> = inputs.py(); - let nullability: Bound<'py, PyAny> = - py.import(NULLABILITY_MODULE)?.getattr(NULLABILITY_CLASS)?; + let nullability: Bound<'py, PyAny> = registry::class(py, RecordId::InferredNullability)?; let types: Bound<'py, PyDict> = PyDict::new(py); let nulls: Bound<'py, PyDict> = PyDict::new(py); let complete: Bound<'py, PyDict> = PyDict::new(py); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs b/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs index 9a8610eb09..78b2df5d8d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs @@ -275,6 +275,7 @@ native_ids! { IncrementalMode => "IncrementalMode", IncrementalStrategy => "IncrementalStrategy", InferredColumn => "InferredColumn", + InferredNullability => "InferredNullability", InitialValidFrom => "InitialValidFrom", IntegerValue => "IntegerValue", IntegrationLoaderConfig => "IntegrationLoaderConfig", @@ -1037,7 +1038,6 @@ native_ids! { PlanRuleFixes => "plan_rule_fixes", PlanRuleRun => "plan_rule_run", PlanSelectedSqlTests => "plan_selected_sql_tests", - PlanSelectedSqlTestsInto => "plan_selected_sql_tests_into", PlanSqlTests => "plan_sql_tests", PlanSqlTestsNatively => "plan_sql_tests_natively", PlanStaticSqlTests => "plan_static_sql_tests", diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs index 033e2c85ad..0018608488 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs @@ -9,7 +9,7 @@ use std::time::Duration; use pyo3::exceptions::PyKeyboardInterrupt; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; -use pyo3::types::{PyBytes, PyBytesMethods, PyDict, PyDictMethods, PyTuple}; +use pyo3::types::{PyDict, PyDictMethods, PyTuple}; use serde_json::{Map, Value, json}; use sqlbuild_rules::custom_host_pool::constants::{ CGROUP_ROOT, HOST_TIMEOUT_MILLIS, MIN_INVOCATIONS_PER_HOST, PROC_CGROUP, @@ -375,23 +375,36 @@ fn host_outcome(spec: &str) -> HostOutcome { /// Publish the pickled project and config atomically; hosts started earlier wait for it. fn write_inputs(py: Python<'_>, path: &Path, request: &HostRequest<'_, '_>) -> PyResult<()> { let payload: Bound<'_, PyAny> = host_payload_project(request.project)?; - let pickled: Bound<'_, PyAny> = imported(py, "pickle", "dumps")? - .call1((PyTuple::new(py, [payload, request.config.clone()])?,))?; - let bytes: &[u8] = pickled.cast::()?.as_bytes(); let parent: &Path = path.parent().unwrap_or(path); let partial: PathBuf = parent.join(format!( "writing-{}-{}.pickle", std::process::id(), INPUT_COUNTER.fetch_add(1, Ordering::Relaxed) )); - if let Err(error) = std::fs::write(&partial, bytes) { + let written: PyResult<()> = pickle_into( + &partial, + &PyTuple::new(py, [payload, request.config.clone()])?, + ); + if let Err(error) = written { remove_file(&partial); - return Err(error.into()); + return Err(error); } std::fs::rename(&partial, path)?; Ok(()) } +/// Pickle `value` straight into a new file at `path`, never holding the whole pickle in memory. +fn pickle_into(path: &Path, value: &Bound<'_, PyTuple>) -> PyResult<()> { + let py: Python<'_> = value.py(); + let handle: Bound<'_, PyAny> = imported(py, "builtins", "open")?.call1((text(path), "wb"))?; + let dumped: PyResult> = + imported(py, "pickle", "dump")?.call1((value, &handle)); + let closed: PyResult> = handle.call_method0("close"); + dumped?; + closed?; + Ok(()) +} + /// The project hosts decode, without model state no rule fact reads. pub(crate) fn host_payload_project<'py>( project: &Bound<'py, PyAny>, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/sql_rules.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/sql_rules.rs index 6d9ca62c97..c7de972fb2 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/sql_rules.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/sql_rules.rs @@ -10,6 +10,7 @@ use pyo3::types::{ }; use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; +use crate::bindings::_helpers::analysis::model_reads::output_proven; use crate::bindings::_helpers::boundary::keyword_request::Request; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported, keywords}; use crate::bindings::_helpers::records::ids::{ErrorId, NativeAttr, RecordId}; @@ -388,8 +389,7 @@ fn evaluated_sql_rules<'py>( } else { let mut dynamic: Vec> = Vec::new(); for (_, model) in &models { - let contract: Bound<'py, PyAny> = model.getattr("dynamic_column_contract")?; - if !contract.is_none() && contract.getattr("output_proven")?.is_truthy()? { + if output_proven(model)? { dynamic.push( project_dir .div(model.getattr("relative_path")?)? diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/completion.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/completion.rs index a853c62be7..66565fa232 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/completion.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/completion.rs @@ -8,7 +8,9 @@ use sqlbuild_analysis::semantic_checks::models::{ CompletedDiagnostic, CompletionModel, CompletionRequest, SemanticDiagnostic, SemanticLocation, }; +use crate::bindings::_helpers::analysis::model_reads::{inferred_names, lineage_inputs}; use crate::bindings::_helpers::boundary::panics::compiler_error; +use crate::bindings::_helpers::compiled_project::compiled_model_view::LineageInput; use crate::bindings::_helpers::semantic_checks::failures::semantic_error; use crate::bindings::_helpers::semantic_checks::metadata::{ShapeRows, shape_rows}; use crate::bindings::_helpers::semantic_checks::session_facts::{ @@ -16,8 +18,7 @@ use crate::bindings::_helpers::semantic_checks::session_facts::{ }; use crate::bindings::_helpers::semantic_checks::stage_objects::StageClasses; use crate::bindings::_helpers::semantic_checks::type_recovery::{ - LineageInput, StageShared, diagnostic_ids, id_of, inferred_names, items, lineage_outputs, - lineage_payload, + StageShared, diagnostic_ids, id_of, items, lineage_outputs, }; use crate::bindings::_helpers::semantic_shapes::shapes::project_shapes; use crate::bindings::types::CompilerDetach; @@ -182,7 +183,7 @@ fn completion_model( .lineage .clone() } else { - lineage_payload(model)? + lineage_inputs(model)? }; let opt_out: Bound<'_, PyAny> = model.getattr("rejected_sql_analysis_opt_out")?; let mut binding_codes: Vec = Vec::new(); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/metadata.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/metadata.rs index b8446a8998..8cfb954000 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/metadata.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/metadata.rs @@ -12,6 +12,7 @@ use sqlbuild_analysis::semantic_checks::models::{ }; use sqlbuild_analysis::semantic_validation::models as validation; +use crate::bindings::_helpers::analysis::model_reads::binding_validated; use crate::bindings::_helpers::boundary::panics::compiler_error; use crate::bindings::_helpers::semantic_checks::failures::semantic_error; use crate::bindings::_helpers::semantic_checks::inputs::{ @@ -406,7 +407,7 @@ fn checked(classes: &StageClasses<'_>, model: &Bound<'_, PyAny>) -> PyResult) -> PyResult> { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/stage.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/stage.rs index 9cd3ca1d7b..52d5c5f15a 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/stage.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/stage.rs @@ -6,6 +6,7 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyRef use pyo3::types::{PyDict, PyTuple}; use pyo3::{pyfunction, wrap_pyfunction}; +use crate::bindings::_helpers::analysis::model_reads::output_proven; use crate::bindings::_helpers::analysis_session::session::NativeModelAnalysisSession; use crate::bindings::_helpers::boundary::keyword_request::Request; use crate::bindings::_helpers::boundary::panics::NativeCompilerError; @@ -106,8 +107,7 @@ fn session_fact_models( if !kept.contains(&name) { continue; } - let proof: Bound<'_, PyAny> = model.getattr("dynamic_column_contract")?; - if !proof.is_none() && proof.getattr("output_proven")?.is_truthy()? { + if output_proven(&model)? { continue; } models.insert(name); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/type_recovery.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/type_recovery.rs index 1509c5c736..2b722c58c2 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/type_recovery.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_checks/type_recovery.rs @@ -12,8 +12,10 @@ use sqlbuild_analysis::semantic_checks::models::{ use sqlbuild_analysis::semantic_checks::types::RevisedBinding; use sqlbuild_analysis::semantic_validation::models as validation; +use crate::bindings::_helpers::analysis::model_reads::{inferred_names, lineage_inputs}; use crate::bindings::_helpers::analysis_session::session::NativeModelAnalysisSession; use crate::bindings::_helpers::boundary::panics::compiler_error; +use crate::bindings::_helpers::compiled_project::compiled_model_view::LineageInput; use crate::bindings::_helpers::semantic_checks::failures::{internal_error, semantic_error}; use crate::bindings::_helpers::semantic_checks::inputs::revalidation_request; use crate::bindings::_helpers::semantic_checks::session_facts::{ @@ -26,9 +28,6 @@ use crate::bindings::types::CompilerDetach; const UNCHANGED: &str = "unchanged"; const UNKNOWN_RECOVERED_TYPE: &str = "UNKNOWN"; -/// `(output_column, [(resource_name, column_name)])`. -pub(crate) type LineageInput = (String, Vec<(String, String)>); - /// The analysis state every semantic completion step shares. pub(crate) struct StageShared<'a> { pub(crate) catalog: &'a validation::ProjectCatalog, @@ -164,28 +163,6 @@ pub(crate) fn id_of(ids: &Bound<'_, PyDict>, item: &Bound<'_, PyAny>) -> PyResul ids.as_any().get_item(item)?.extract() } -/// Each output column with the `(resource, column)` pairs it reads. -pub(crate) fn lineage_payload(model: &Bound<'_, PyAny>) -> PyResult> { - let columns: Bound<'_, PyAny> = model.getattr("fast_lineage_columns")?; - if !columns.is_truthy()? { - return Ok(Vec::new()); - } - let mut lineage: Vec = Vec::new(); - for output in columns.try_iter()? { - let output: Bound<'_, PyAny> = output?; - let mut upstream: Vec<(String, String)> = Vec::new(); - for source in output.getattr("upstream_columns")?.try_iter()? { - let source: Bound<'_, PyAny> = source?; - upstream.push(( - source.getattr("resource_name")?.extract()?, - source.getattr("column_name")?.extract()?, - )); - } - lineage.push((output.getattr("output_column")?.extract()?, upstream)); - } - Ok(lineage) -} - pub(crate) fn lineage_outputs(lineage: Vec) -> Vec { lineage .into_iter() @@ -196,19 +173,6 @@ pub(crate) fn lineage_outputs(lineage: Vec) -> Vec .collect() } -/// The names of a model's inferred columns, or `None` when its shape is unknown. -pub(crate) fn inferred_names(model: &Bound<'_, PyAny>) -> PyResult>> { - let columns: Bound<'_, PyAny> = model.getattr("inferred_columns")?; - if columns.is_none() { - return Ok(None); - } - columns - .try_iter()? - .map(|column| column?.getattr("name")?.extract()) - .collect::>() - .map(Some) -} - fn blocking_binding_found(models: &[Bound<'_, PyAny>]) -> PyResult { for model in models { for item in model.getattr("binding_diagnostics")?.try_iter()? { @@ -257,7 +221,7 @@ fn recovery_model( let facts = session_model_facts(shared.session, &name).map_err(compiler_error)?; (column_names(facts), facts.lineage.clone()) } else { - (inferred_names(model)?, lineage_payload(model)?) + (inferred_names(model)?, lineage_inputs(model)?) }; let mut references: Vec = Vec::new(); for reference in model.getattr("references")?.try_iter()? { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/coverage.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/coverage.rs index 749b42b474..a274ad1f7b 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/coverage.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/coverage.rs @@ -6,6 +6,8 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyRes use pyo3::types::{PyBool, PyDict, PyDictMethods, PyTuple}; use pyo3::{pyfunction, wrap_pyfunction}; +use crate::bindings::_helpers::analysis::model_reads::{inferred_rows, star_flags}; +use crate::bindings::_helpers::compiled_project::compiled_model_view::InferredRow; use crate::bindings::_helpers::semantic_shapes::shapes::project_shapes; const TYPE_CHECKED_DIALECTS: [&str; 4] = ["duckdb", "snowflake", "postgresql", "bigquery"]; @@ -99,19 +101,17 @@ pub(crate) fn semantic_coverage<'py>(project: &Bound<'py, PyAny>) -> PyResult = model.getattr("inferred_rows")?; - if !inferred.is_truthy()? { + let inferred: Vec = inferred_rows(model)?.unwrap_or_default(); + if inferred.is_empty() { partial.insert("output shape could not be inferred".to_owned()); - } else if model.getattr("fast_lineage_has_star")?.is_truthy()? && !shapes.contains(name)? { + } else if star_flags(model)?.0 && !shapes.contains(name)? { partial.insert("unresolved star over open inputs".to_owned()); } - if inferred.is_truthy()? { - for row in inferred.try_iter()? { - if row?.get_item(1)?.is_none() { - partial.insert("some output types are unknown".to_owned()); - break; - } - } + if inferred + .iter() + .any(|(_, column_type, _)| column_type.is_none()) + { + partial.insert("some output types are unknown".to_owned()); } if !TYPE_CHECKED_DIALECTS.contains(&dialect.as_str()) { partial.insert("type checking awaits dialect coercion support".to_owned()); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/shapes.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/shapes.rs index 469b5c98eb..4c51f02c5f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/shapes.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/shapes.rs @@ -7,7 +7,11 @@ use pyo3::{pyfunction, wrap_pyfunction}; use crate::bindings::_helpers::analysis::binding_catalog::{ BindingCatalog, CASE_SENSITIVE_BINDING_DIALECTS, }; +use crate::bindings::_helpers::analysis::model_reads::{ + declared_rows, dynamic_columns_declared, inferred_rows, star_flags, unchecked_columns, +}; use crate::bindings::_helpers::boundary::keyword_request::Request; +use crate::bindings::_helpers::compiled_project::compiled_model_view::{DeclaredRow, InferredRow}; use crate::bindings::_helpers::native_host::adapter::AdapterHost; use crate::bindings::_helpers::native_host::host; use crate::bindings::_helpers::records::ids::{NativeAttr, RecordId}; @@ -152,15 +156,19 @@ fn or_generic<'py>(dialect: &Bound<'py, PyAny>) -> PyResult> { } /// `{name: type or "UNKNOWN"}` over `(name, type, ...)` rows. -fn typed_columns<'py>(py: Python<'py>, rows: &Bound<'py, PyAny>) -> PyResult> { +fn typed_columns<'py>(py: Python<'py>, rows: &[InferredRow]) -> PyResult> { let columns: Bound<'py, PyDict> = PyDict::new(py); - for row in rows.try_iter()? { - let row: Bound<'py, PyAny> = row?; - columns.set_item(row.get_item(0)?, or_unknown(&row.get_item(1)?)?)?; + for (name, column_type, _) in rows { + columns.set_item(name, text_or_unknown(column_type.as_deref()))?; } Ok(columns) } +/// `value or "UNKNOWN"` for an optional type name. +fn text_or_unknown(value: Option<&str>) -> &str { + value.filter(|text| !text.is_empty()).unwrap_or(UNKNOWN) +} + fn or_unknown<'py>(value: &Bound<'py, PyAny>) -> PyResult> { if value.is_truthy()? { return Ok(value.clone()); @@ -287,8 +295,8 @@ fn model_shapes<'py>( while !pending.is_empty() { let mut remaining: Vec> = Vec::new(); for model in &pending { - let declared: Bound<'py, PyAny> = model.getattr("schema_column_rows")?; - let inferred: Bound<'py, PyAny> = model.getattr("inferred_rows")?; + let declared: Option> = declared_rows(model)?; + let inferred: Option> = inferred_rows(model)?; let rows: Bound<'py, PyAny> = model.getattr("reference_rows")?; let relations: Vec = relation_row_names(&rows)?; let contract: Bound<'py, PyAny> = model @@ -296,32 +304,29 @@ fn model_shapes<'py>( .getattr("values")? .call_method1("get", ("contract",))?; if contract.eq(enforced)? - && !declared.is_none() - && !model.getattr("has_dynamic_columns")?.is_truthy()? + && let Some(declared) = &declared + && !dynamic_columns_declared(model)? { - let unchecked: Bound<'py, PyAny> = model.getattr("unchecked_output_columns")?; + let unchecked: Vec = unchecked_columns(model)?; let shape: Bound<'py, PyDict> = PyDict::new(py); - for row in declared.try_iter()? { - let row: Bound<'py, PyAny> = row?; - let name: Bound<'py, PyAny> = row.get_item(0)?; - let column_type: Bound<'py, PyAny> = if unchecked.contains(&name)? { - PyString::new(py, UNKNOWN).into_any() + for (name, column_type) in declared { + let column_type: &str = if unchecked.contains(name) { + UNKNOWN } else { - or_unknown(&row.get_item(1)?)? + text_or_unknown(column_type.as_deref()) }; shape.set_item(name, column_type)?; } shapes.set_item(model.getattr("name")?, shape)?; continue; } - let closes: bool = inferred.is_truthy()? - && (!model.getattr("fast_lineage_has_star")?.is_truthy()? - || (model.getattr("fast_lineage_star_resolved")?.is_truthy()? - && rows.is_truthy()? - && all_known(&relations, shapes)?)); - if closes { + let (has_star, star_resolved): (bool, bool) = star_flags(model)?; + let closes: bool = inferred.as_ref().is_some_and(|rows| !rows.is_empty()) + && (!has_star + || (star_resolved && rows.is_truthy()? && all_known(&relations, shapes)?)); + if let (true, Some(inferred)) = (closes, &inferred) { let inputs: Bound<'py, PyDict> = model_inputs(model, shapes)?; - let columns: Bound<'py, PyDict> = typed_columns(py, &inferred)?; + let columns: Bound<'py, PyDict> = typed_columns(py, inferred)?; shapes.set_item( model.getattr("name")?, published_shape(model, &columns, &inputs, profile)?, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/source_bindings.rs b/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/source_bindings.rs index 976de6687d..ca768df886 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/source_bindings.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/semantic_shapes/source_bindings.rs @@ -23,7 +23,6 @@ use crate::bindings::_helpers::semantic_shapes::shapes::{ const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const BINDING_UNKNOWN_TABLE_INTERNAL_CODE: &str = "B000"; const DATACLASSES_MODULE: &str = "dataclasses"; const GRAPHLIB_MODULE: &str = "graphlib"; @@ -156,10 +155,8 @@ fn rebound_models<'py>( .call1((graph,))? .call_method0("static_order")?; let query_class: Object<'py> = registry::class(py, RecordId::NativeColumnQuery)?; - let unknown: Object<'py> = py - .import(LINEAGE_TYPES_MODULE)? - .getattr("InferredNullability")? - .getattr("UNKNOWN")?; + let unknown: Object<'py> = + registry::class(py, RecordId::InferredNullability)?.getattr("UNKNOWN")?; let infer: Object<'py> = registry::native(py, NativeAttr::InferNativeQueryColumns)?; let mut affected: HashSet = HashSet::new(); let mut models: Vec> = Vec::new(); @@ -303,9 +300,8 @@ fn validated_diagnostics<'py>( let skipped: Object<'py> = PyString::new(py, BINDING_UNKNOWN_TABLE_INTERNAL_CODE).into_any(); let types = py.import(COMPILE_TYPES_MODULE)?; let diagnostic_class: Object<'py> = registry::class(py, RecordId::CompilerDiagnostic)?; - let severity_class: Object<'py> = registry::class(types.py(), RecordId::DiagnosticSeverity)?; - let phase: Object<'py> = - registry::class(types.py(), RecordId::DiagnosticPhase)?.getattr("COMPILE")?; + let severity_class: Object<'py> = registry::class(py, RecordId::DiagnosticSeverity)?; + let phase: Object<'py> = registry::class(py, RecordId::DiagnosticPhase)?.getattr("COMPILE")?; let model_type: Object<'py> = types.getattr("CompiledResourceType")?.getattr("MODEL")?; let mut diagnostics: Vec> = Vec::new(); for ((model, request), result) in models.iter().zip(&requests).zip(results.try_iter()?) { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/type_system/nullability.rs b/crates/sqlbuild-python/src/bindings/_helpers/type_system/nullability.rs index 5ca8f2ccff..7a48d910e9 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/type_system/nullability.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/type_system/nullability.rs @@ -6,15 +6,14 @@ use pyo3::{pyfunction, wrap_pyfunction}; use sqlbuild_analysis::nullability_rules::main::conditional_result_nullability::conditional_result_nullability; use sqlbuild_analysis::nullability_rules::main::safe_cast_nullability::safe_cast_nullability; -use crate::bindings::_helpers::compile_driver::python_calls::imported; +use crate::bindings::_helpers::records::ids::RecordId; +use crate::bindings::_helpers::records::registry; -const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; -const NULLABILITY_CLASS: &str = "InferredNullability"; const UNKNOWN_MEMBER: &str = "UNKNOWN"; /// The `InferredNullability` member whose value is `value`. fn member<'py>(args: &Bound<'py, PyAny>, value: &str) -> PyResult> { - imported(args.py(), LINEAGE_TYPES_MODULE, NULLABILITY_CLASS)?.call1((value,)) + registry::class(args.py(), RecordId::InferredNullability)?.call1((value,)) } /// The text of one argument; an argument that is not a string matches no state. @@ -36,8 +35,7 @@ fn argument_values(args: &Bound<'_, PyAny>, count: usize) -> PyResult(args: &Bound<'py, PyAny>) -> PyResult> { if !args.is_truthy()? { - return imported(args.py(), LINEAGE_TYPES_MODULE, NULLABILITY_CLASS)? - .getattr(UNKNOWN_MEMBER); + return registry::class(args.py(), RecordId::InferredNullability)?.getattr(UNKNOWN_MEMBER); } args.get_item(0) } diff --git a/crates/sqlbuild-rules/Cargo.toml b/crates/sqlbuild-rules/Cargo.toml index 16312b86c4..0679b5241b 100644 --- a/crates/sqlbuild-rules/Cargo.toml +++ b/crates/sqlbuild-rules/Cargo.toml @@ -13,6 +13,7 @@ globset.workspace = true polyglot-sql.workspace = true rayon.workspace = true regex.workspace = true +rustc-hash.workspace = true serde.workspace = true serde_json.workspace = true sha2.workspace = true diff --git a/crates/sqlbuild-rules/src/lint_text/_helpers/sites.rs b/crates/sqlbuild-rules/src/lint_text/_helpers/sites.rs index 9b9cc86bea..ef48017747 100644 --- a/crates/sqlbuild-rules/src/lint_text/_helpers/sites.rs +++ b/crates/sqlbuild-rules/src/lint_text/_helpers/sites.rs @@ -10,7 +10,7 @@ use sqlbuild_sqltext::sql_scan::main::non_code_end::non_code_end; use sqlbuild_sqltext::sql_scan::models::QuotePolicy; use crate::lint_text::constants::{ - FORMAT_SITE_PATTERN, FUNCTION_NAMES, LINT_SITE_PATTERN, SENTINEL_SUFFIX, + FORMAT_SITE_PATTERN, FUNCTION_NAMES, LINT_SITE_PATTERN, SENTINEL_SUFFIX, SITE_START_BYTES, }; use crate::lint_text::errors::LintTextError; use crate::lint_text::models::{InterpolationSite, Neutralized}; @@ -63,7 +63,11 @@ pub(crate) fn site_starts(text: &str, pattern: &Regex, policy: QuotePolicy) -> V Err(_) => break, Ok(None) => {} } - if let Some(found) = text.get(index..).and_then(|rest| pattern.find(rest)) { + let found = SITE_START_BYTES + .contains(&bytes[index]) + .then(|| text.get(index..).and_then(|rest| pattern.find(rest))) + .flatten(); + if let Some(found) = found { starts.push(index); index += found.end(); } else { diff --git a/crates/sqlbuild-rules/src/lint_text/constants.rs b/crates/sqlbuild-rules/src/lint_text/constants.rs index fbc6ff9d4b..fd41451e50 100644 --- a/crates/sqlbuild-rules/src/lint_text/constants.rs +++ b/crates/sqlbuild-rules/src/lint_text/constants.rs @@ -20,6 +20,8 @@ pub(crate) const FUNCTION_NAMES: [&str; 9] = [ "__udf", ]; +/// Bytes every lint and format site pattern match starts with. +pub(crate) const SITE_START_BYTES: [u8; 3] = *b"@$_"; pub(crate) const LINT_SITE_PATTERN: &str = r"^(?:@@|\$\{|@|(?:__dbt_ref|__ref|__seed|__source|__table_fn|__udf)\()"; pub(crate) const FORMAT_SITE_PATTERN: &str = r"(?i)^(?:@@|\$\{|@|(?:__cursor_end|__cursor_start|__dbt_ref|__empty_fixture|__ref|__seed|__source|__table_fn|__udf)\s*\()"; diff --git a/crates/sqlbuild-rules/src/lint_text/tests/test_lint_text.rs b/crates/sqlbuild-rules/src/lint_text/tests/test_lint_text.rs index 1a887c5bcc..d5466360bc 100644 --- a/crates/sqlbuild-rules/src/lint_text/tests/test_lint_text.rs +++ b/crates/sqlbuild-rules/src/lint_text/tests/test_lint_text.rs @@ -1,6 +1,7 @@ +use crate::lint_text::_helpers::sites::site_pattern; use crate::lint_text::constants::{ FORBIDDEN_FIX_REASON, FORBIDDEN_MESSAGE, FORBIDDEN_REMEDIATION, INVALID_MESSAGE, - INVALID_REMEDIATION, REASON_REQUIRED_FIX_REASON, UNUSED_REMEDIATION, + INVALID_REMEDIATION, REASON_REQUIRED_FIX_REASON, SITE_START_BYTES, UNUSED_REMEDIATION, }; use crate::lint_text::errors::LintTextError; use crate::lint_text::main::apply_suppressions::apply_suppressions; @@ -17,11 +18,45 @@ use crate::lint_text::models::{ }; use crate::lint_text::tests::test_types::{ Finding, FormatBodiesTestCase, LintBodiesTestCase, NeutralizeTestCase, ParameterTestCase, - RestoreTestCase, SuppressionTestCase, TokenAtTestCase, + RestoreTestCase, SitePatternTestCase, SuppressionTestCase, TokenAtTestCase, }; const SUPPRESSED_FILE: &str = "-- sqb: ignore SQBRSQL005 because kept\n\n-- note\nselect 1\n-- sqb: ignore X because y\n-- sqb: ignore nope\nselect 2\n"; +#[test] +fn given_site_patterns_when_text_opens_with_another_byte_then_no_site_starts_there() +-> Result<(), LintTextError> { + let test_cases = [ + SitePatternTestCase { + description: "lint sites", + formatting: false, + expected_other_start_bytes: Vec::new(), + }, + SitePatternTestCase { + description: "format sites", + formatting: true, + expected_other_start_bytes: Vec::new(), + }, + ]; + for test_case in test_cases { + let pattern = site_pattern(test_case.formatting)?; + let other_start_bytes: Vec = (0..=127u8) + .filter(|byte| !SITE_START_BYTES.contains(byte)) + .filter(|byte| { + ["@@x", "@x", "${x}", "__ref(", "__REF (", "__udf("] + .iter() + .any(|site| pattern.is_match(&format!("{}{site}", char::from(*byte)))) + }) + .collect(); + assert_eq!( + other_start_bytes, test_case.expected_other_start_bytes, + "{}", + test_case.description + ); + } + Ok(()) +} + #[test] fn given_authored_sql_when_neutralizing_then_sites_use_character_offsets() -> Result<(), LintTextError> { diff --git a/crates/sqlbuild-rules/src/lint_text/tests/test_types.rs b/crates/sqlbuild-rules/src/lint_text/tests/test_types.rs index 264314566b..3df6eae669 100644 --- a/crates/sqlbuild-rules/src/lint_text/tests/test_types.rs +++ b/crates/sqlbuild-rules/src/lint_text/tests/test_types.rs @@ -86,3 +86,9 @@ pub(super) struct FixtureBlockTestCase { pub model_mode: bool, pub expected_sql: Option<&'static str>, } + +pub(crate) struct SitePatternTestCase { + pub(crate) description: &'static str, + pub(crate) formatting: bool, + pub(crate) expected_other_start_bytes: Vec, +} diff --git a/crates/sqlbuild-rules/src/sql_lint/_helpers/additional.rs b/crates/sqlbuild-rules/src/sql_lint/_helpers/additional.rs index caa570c1b4..a23e97824f 100644 --- a/crates/sqlbuild-rules/src/sql_lint/_helpers/additional.rs +++ b/crates/sqlbuild-rules/src/sql_lint/_helpers/additional.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use polyglot_sql::tokens::{Span, Token, TokenType}; @@ -9,6 +9,7 @@ use crate::sql_lint::_helpers::engine::{ }; use crate::sql_lint::constants::RELATION_MODIFIER_KEYWORDS; use crate::sql_lint::models::{AdditionalFactOptions, AdditionalQueryFacts}; +use rustc_hash::{FxHashMap, FxHashSet}; use sqlbuild_core::constants::{CLOSE_PAREN, OPEN_PAREN, SQL_WILDCARD}; const COUNT_ONE_LITERAL: &str = "1"; @@ -32,7 +33,7 @@ struct QuerySlice<'a> { struct ReferenceFactContext<'a> { query: &'a QuerySlice<'a>, - outer_relations: &'a HashSet, + outer_relations: &'a FxHashSet, external_identifiers: &'a HashSet, query_start: usize, query_end: usize, @@ -42,7 +43,7 @@ struct ReferenceCandidateContext<'a> { tokens: &'a [Token], references: &'a [usize], external_identifiers: &'a HashSet, - projection_aliases: &'a HashMap, + projection_aliases: &'a FxHashMap, } struct SelectFactContext<'a> { @@ -751,8 +752,8 @@ fn unqualified_reference( .then_some(index) } -fn relation_names(tokens: &[Token], relation_clause: &[usize]) -> HashSet { - let mut names: HashSet = HashSet::new(); +fn relation_names(tokens: &[Token], relation_clause: &[usize]) -> FxHashSet { + let mut names: FxHashSet = FxHashSet::default(); for (position, &index) in relation_clause.iter().enumerate() { let starts_relation = position == 0 || matches!( @@ -801,8 +802,8 @@ fn outer_relation_names( depths: &[usize], select_index: usize, depth: usize, -) -> HashSet { - let mut names: HashSet = HashSet::new(); +) -> FxHashSet { + let mut names: FxHashSet = FxHashSet::default(); for outer_select in (0..select_index).filter(|&index| { tokens[index].token_type == TokenType::Select && depths[index] < depth @@ -1022,8 +1023,8 @@ fn is_context_value(token: &Token) -> bool { ) } -fn projection_alias_positions(tokens: &[Token], projection: &[usize]) -> HashMap { - let mut aliases: HashMap = HashMap::new(); +fn projection_alias_positions(tokens: &[Token], projection: &[usize]) -> FxHashMap { + let mut aliases: FxHashMap = FxHashMap::default(); for window in projection.windows(2) { if tokens[window[0]].token_type == TokenType::As && is_identifier(&tokens[window[1]]) { aliases.insert(tokens[window[1]].text.to_ascii_lowercase(), window[1]); @@ -1519,7 +1520,7 @@ fn ambiguous_order_direction_spans(tokens: &[Token], direct: &[usize]) -> Vec Vec { - let mut seen: HashSet = HashSet::new(); + let mut seen: FxHashSet = FxHashSet::default(); let mut duplicates: Vec = Vec::new(); for window in indices.windows(2) { if tokens[window[0]].token_type != TokenType::As || !is_identifier(&tokens[window[1]]) { diff --git a/crates/sqlbuild-rules/src/sql_quality/_helpers/cte_columns.rs b/crates/sqlbuild-rules/src/sql_quality/_helpers/cte_columns.rs index 8b5e6b2b3c..c60f04a74a 100644 --- a/crates/sqlbuild-rules/src/sql_quality/_helpers/cte_columns.rs +++ b/crates/sqlbuild-rules/src/sql_quality/_helpers/cte_columns.rs @@ -1,7 +1,6 @@ //! SQBRSQL042: output columns of non-final CTEs that no later step reads. use std::borrow::Cow; -use std::collections::{HashMap, HashSet}; use polyglot_sql::expressions::{Cte, JoinKind, Select}; use polyglot_sql::scope::walk_in_scope; @@ -18,6 +17,7 @@ use crate::sql_quality::syntax::{ Projection, columns, cte_select_tokens, lower, projection_removal, projections, qualifier, sources, }; +use rustc_hash::{FxHashMap, FxHashSet}; use sqlbuild_analysis::sql_tokens::main::projection_spans::projection_spans; const HARNESS_PREFIX: &str = "__"; @@ -27,7 +27,7 @@ const LAYOUT_FIX_REFUSAL: &str = "The projection list could not be mapped to aut /// What one later step reads from a CTE. #[derive(Default)] struct Usage { - names: HashSet, + names: FxHashSet, everything: bool, passes_to: Vec, read: bool, @@ -66,8 +66,8 @@ pub(crate) fn diagnostics(request: &QualityRequest<'_>) -> Vec { terminal.with = None; let terminal: Expression = Expression::Select(Box::new(terminal)); let names: Vec = with.ctes.iter().map(|cte| lower(&cte.alias.name)).collect(); - let used: HashMap = reads(&with.ctes, &names, &terminal); - let select_tokens: HashMap = cte_select_tokens(request.tokens); + let used: FxHashMap = reads(&with.ctes, &names, &terminal); + let select_tokens: FxHashMap = cte_select_tokens(request.tokens); for (index, cte) in with .ctes .iter() @@ -94,15 +94,15 @@ pub(crate) fn diagnostics(request: &QualityRequest<'_>) -> Vec { } /// Names each CTE's later readers use, following `SELECT *` pass-through to the reader. -fn reads(ctes: &[Cte], names: &[String], terminal: &Expression) -> HashMap { +fn reads(ctes: &[Cte], names: &[String], terminal: &Expression) -> FxHashMap { let bodies: Vec<(Option, &Expression)> = ctes .iter() .map(|cte| (Some(lower(&cte.alias.name)), &cte.this)) .chain(std::iter::once((None, terminal))) .collect(); - let known: HashSet<&str> = names.iter().map(String::as_str).collect(); - let outputs: Vec>> = ctes.iter().map(output_names).collect(); - let mut readers: HashMap> = HashMap::new(); + let known: FxHashSet<&str> = names.iter().map(String::as_str).collect(); + let outputs: Vec>> = ctes.iter().map(output_names).collect(); + let mut readers: FxHashMap> = FxHashMap::default(); for (index, (_, body)) in bodies.iter().enumerate() { for name in read_relations(body) { if known.contains(name.as_str()) { @@ -110,7 +110,7 @@ fn reads(ctes: &[Cte], names: &[String], terminal: &Expression) -> HashMap = HashMap::new(); + let mut usage: FxHashMap = FxHashMap::default(); for (index, name) in names.iter().enumerate() { let mut current: Usage = Usage::default(); if index + 1 == names.len() { @@ -146,7 +146,7 @@ fn reads(ctes: &[Cte], names: &[String], terminal: &Expression) -> HashMap Option> { +fn output_names(cte: &Cte) -> Option> { let Expression::Select(select) = &cte.this else { return None; }; @@ -159,7 +159,7 @@ fn output_names(cte: &Cte) -> Option> { .collect() } -fn read_relations(body: &Expression) -> HashSet { +fn read_relations(body: &Expression) -> FxHashSet { body.dfs() .filter_map(|node| match node { Expression::Table(table) => Some(lower(&table.name.name)), @@ -171,13 +171,13 @@ fn read_relations(body: &Expression) -> HashSet { /// The CTE whose readers are being inspected. struct Target<'a> { name: &'a str, - outputs: Option<&'a HashSet>, + outputs: Option<&'a FxHashSet>, } fn reader_usage(target: &Target<'_>, reader: Option, body: &Expression) -> Usage { let name: &str = target.name; let mut usage: Usage = Usage::default(); - let mut aliases: HashSet = HashSet::new(); + let mut aliases: FxHashSet = FxHashSet::default(); let mut top_level_star: bool = false; for node in body.dfs() { let Expression::Select(select) = node else { @@ -254,7 +254,7 @@ fn reader_usage(target: &Target<'_>, reader: Option, body: &Expression) } /// Whether a reader consumes whole rows: a nested star, `COLUMNS(...)`, or the alias itself. -fn reads_as_row(select: &Select, local: &[String], outputs: Option<&HashSet>) -> bool { +fn reads_as_row(select: &Select, local: &[String], outputs: Option<&FxHashSet>) -> bool { for clause in clause_roots(select) { for node in walk_in_scope(clause, false) { if reads_row_at(select, node, local, outputs) { @@ -269,7 +269,7 @@ fn reads_row_at( select: &Select, node: &Expression, local: &[String], - outputs: Option<&HashSet>, + outputs: Option<&FxHashSet>, ) -> bool { match node { Expression::Star(star) => { @@ -308,12 +308,12 @@ fn clause_roots(select: &Select) -> Vec<&Expression> { } /// Lateral alias reuse: a later projection or a clause of the CTE reading an earlier output. -fn self_reads(body: &Expression) -> HashSet { - let mut reads: HashSet = HashSet::new(); +fn self_reads(body: &Expression) -> FxHashSet { + let mut reads: FxHashSet = FxHashSet::default(); let Expression::Select(select) = body else { return reads; }; - let mut earlier: HashSet = HashSet::new(); + let mut earlier: FxHashSet = FxHashSet::default(); for item in projections(select) { if let Projection::Named { name, expression } = item { for read in unqualified_reads(expression) { @@ -357,7 +357,7 @@ fn unqualified_reads(expression: &Expression) -> Vec { fn unused_outputs( request: &QualityRequest<'_>, - select_tokens: &HashMap, + select_tokens: &FxHashMap, outputs: CteOutputs<'_>, ) -> Vec { let CteOutputs { cte, name, usage } = outputs; @@ -381,7 +381,7 @@ fn unused_outputs( .get(name) .map(|&select_token| projection_spans(request.tokens, select_token)) .unwrap_or_default(); - let grouped: HashSet = select + let grouped: FxHashSet = select .group_by .iter() .flat_map(|group| group.expressions.iter()) diff --git a/crates/sqlbuild-rules/src/sql_quality/_helpers/syntax.rs b/crates/sqlbuild-rules/src/sql_quality/_helpers/syntax.rs index bb52d54da7..7f79990f0f 100644 --- a/crates/sqlbuild-rules/src/sql_quality/_helpers/syntax.rs +++ b/crates/sqlbuild-rules/src/sql_quality/_helpers/syntax.rs @@ -1,12 +1,11 @@ //! Shared syntax helpers for SQL quality Rules. -use std::collections::HashMap; - use polyglot_sql::expressions::{Column, Select, Star}; use polyglot_sql::tokens::{Token, TokenType}; use polyglot_sql::{Expression, ExpressionWalk}; use crate::sql_quality::constants::MIN_REMOVABLE_OUTPUTS; +use rustc_hash::FxHashMap; /// One projection of a select list. pub(super) enum Projection<'a> { @@ -83,11 +82,11 @@ pub(super) fn sources(select: &Select) -> Vec<(String, Option, &Expressi } /// Token index of the SELECT opening each `name AS (SELECT` CTE body, by lower-cased name. -pub(super) fn cte_select_tokens(tokens: &[Token]) -> HashMap { +pub(super) fn cte_select_tokens(tokens: &[Token]) -> FxHashMap { let significant: Vec = (0..tokens.len()) .filter(|&index| !is_trivia(&tokens[index])) .collect(); - let mut found: HashMap = HashMap::new(); + let mut found: FxHashMap = FxHashMap::default(); for window in significant.windows(4) { let [name_index, as_index, paren_index, select_index] = window else { continue; diff --git a/crates/sqlbuild-sqltext/src/sql_scan/_helpers/dialect.rs b/crates/sqlbuild-sqltext/src/sql_scan/_helpers/dialect.rs index 4d4f6ee5bb..1a70ecdec5 100644 --- a/crates/sqlbuild-sqltext/src/sql_scan/_helpers/dialect.rs +++ b/crates/sqlbuild-sqltext/src/sql_scan/_helpers/dialect.rs @@ -13,6 +13,12 @@ pub(crate) fn non_code_end( index: usize, syntax: &LexicalSyntax, ) -> Result, Unclosed> { + if sql + .get(index) + .is_some_and(|first| !may_open_non_code(*first, syntax)) + { + return Ok(None); + } if sql[index..].starts_with(b"/*") { return block_comment_end(sql, index, syntax).map(Some); } @@ -35,6 +41,15 @@ pub(crate) fn non_code_end( } } +/// Whether a comment or quoted text may open with `byte`; most code bytes are rejected here cheaply. +fn may_open_non_code(byte: u8, syntax: &LexicalSyntax) -> bool { + matches!(byte, b'/' | b'$' | b'\'' | b'"' | b'`') + || syntax + .line_comment_prefixes + .iter() + .any(|prefix| prefix.as_bytes().first().is_none_or(|start| *start == byte)) +} + /// Return every comment and quoted-text range in `sql`; an unclosed one extends to the end. pub(crate) fn non_code_ranges(sql: &str, syntax: &LexicalSyntax) -> Vec<(usize, usize)> { let bytes = sql.as_bytes(); diff --git a/crates/sqlbuild-sqltext/src/sql_scan/tests/test_dialect_scan.rs b/crates/sqlbuild-sqltext/src/sql_scan/tests/test_dialect_scan.rs index 64291fbdd4..b23fd27f99 100644 --- a/crates/sqlbuild-sqltext/src/sql_scan/tests/test_dialect_scan.rs +++ b/crates/sqlbuild-sqltext/src/sql_scan/tests/test_dialect_scan.rs @@ -1,11 +1,93 @@ +use crate::sql_scan::main::dialect_non_code_end::dialect_non_code_end; use crate::sql_scan::main::dialect_non_code_ranges::dialect_non_code_ranges; +use crate::sql_scan::models::{LexicalSyntax, Unclosed}; use crate::sql_scan::tests::helpers::{ escape_strings, generic, raw_strings, single_quote_backslashes, triple_quotes, }; use crate::sql_scan::tests::test_types::{ - DialectFirstNonCodeRangeTestCase, DialectNonCodeRangesTestCase, + DialectFirstNonCodeRangeTestCase, DialectNonCodeEndTestCase, DialectNonCodeRangesTestCase, }; +#[test] +fn given_an_offset_when_finding_dialect_non_code_end_then_only_opening_bytes_start_non_code() { + let test_cases = [ + DialectNonCodeEndTestCase { + description: "a code byte opens nothing", + sql: "select 1", + index: 0, + syntax: generic(), + expected_end: Ok(None), + }, + DialectNonCodeEndTestCase { + description: "a hash is code when it is not a line comment prefix", + sql: "# note\nx", + index: 0, + syntax: generic(), + expected_end: Ok(None), + }, + DialectNonCodeEndTestCase { + description: "a hash opens a line comment when the dialect says so", + sql: "x # note\ny", + index: 2, + syntax: triple_quotes(), + expected_end: Ok(Some(9)), + }, + DialectNonCodeEndTestCase { + description: "a lone slash is code", + sql: "a / b", + index: 2, + syntax: generic(), + expected_end: Ok(None), + }, + DialectNonCodeEndTestCase { + description: "a dollar tag opens a dollar-quoted string", + sql: "$tag$ x $tag$ y", + index: 0, + syntax: generic(), + expected_end: Ok(Some(13)), + }, + DialectNonCodeEndTestCase { + description: "a dollar parameter is code", + sql: "$1 + 2", + index: 0, + syntax: generic(), + expected_end: Ok(None), + }, + DialectNonCodeEndTestCase { + description: "an unclosed quote is reported", + sql: "x 'open", + index: 2, + syntax: generic(), + expected_end: Err(Unclosed::Quote), + }, + DialectNonCodeEndTestCase { + description: "the end of the text opens nothing", + sql: "x", + index: 1, + syntax: generic(), + expected_end: Ok(None), + }, + DialectNonCodeEndTestCase { + description: "an empty line comment prefix matches at every offset", + sql: "ab\ncd", + index: 1, + syntax: LexicalSyntax { + line_comment_prefixes: vec![String::new()], + ..LexicalSyntax::default() + }, + expected_end: Ok(Some(3)), + }, + ]; + for test_case in test_cases { + assert_eq!( + dialect_non_code_end(test_case.sql.as_bytes(), test_case.index, &test_case.syntax), + test_case.expected_end, + "{}", + test_case.description + ); + } +} + #[test] fn given_dialect_lexical_syntax_when_scanning_non_code_then_dialect_decides_boundaries() { let test_cases = [ diff --git a/crates/sqlbuild-sqltext/src/sql_scan/tests/test_types.rs b/crates/sqlbuild-sqltext/src/sql_scan/tests/test_types.rs index 86883ba217..4507180b90 100644 --- a/crates/sqlbuild-sqltext/src/sql_scan/tests/test_types.rs +++ b/crates/sqlbuild-sqltext/src/sql_scan/tests/test_types.rs @@ -15,6 +15,14 @@ pub(crate) struct NonCodeEndTestCase { pub(crate) expected_end: Result, Unclosed>, } +pub(crate) struct DialectNonCodeEndTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) index: usize, + pub(crate) syntax: LexicalSyntax, + pub(crate) expected_end: Result, Unclosed>, +} + pub(crate) struct DialectFirstNonCodeRangeTestCase { pub(crate) description: &'static str, pub(crate) sql: &'static str, diff --git a/python-boundary.toml b/python-boundary.toml index 1f018ef03d..a2b95848b8 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -712,11 +712,6 @@ path = "cost/types.py" rule = 1 note = "Shape rule: data enums and aliases only (CostCapability, CostStatus), no logic; consumed by adapters/snowflake/classes/snowflake_adapter.py." -[[stays]] -path = "compiler/lineage/types.py" -rule = 1 -note = "Shape rule: data enums and aliases only (InferredNullability), no logic; consumed by adapter/contract/types.py, adapter/type_system/main/conditional_result_nullability.py, adapter/type_system/main/first_arg_nullability.py and adapter/type_system/main/safe_cast_nullability.py, and read natively to build nullability results." - [[stays]] path = "compiler/macro_bridge/models.py" rule = 2 diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 378b4a7620..54f6ce2037 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -65,9 +65,6 @@ from sqlbuild.compiler.compile.types import ( SqlTestMode, TypedSqlValueRenderer, ) -from sqlbuild.compiler.lineage.types import ( - InferredNullability, -) from sqlbuild.compiler.python_nodes.types import PythonNodeKind from sqlbuild.compiler.references.models import ( ProjectRelation, @@ -6778,14 +6775,6 @@ def plan_selected_sql_tests( selected_keys: frozenset[CompiledObjectKey], sql_test_selection: SqlTestSelection, ) -> PlannedSqlTests: ... -def plan_selected_sql_tests_into( - *, - future: Future[PlannedSqlTests], - project: CompiledProjectView, - adapter: BaseAdapter, - selected_keys: frozenset[CompiledObjectKey], - sql_test_selection: SqlTestSelection, -) -> None: ... def sql_test_is_selected( *, test: CompiledSqlTest, @@ -6891,6 +6880,12 @@ def build_execution_plan( project_config: ProjectConfig | None = None, local_config: LocalConfig | None = None, ) -> PlanOutput: ... + +class SqlTestPlanning: + """The planned SQL tests of one plan, planned on a native thread when enabled.""" + + def result(self) -> PlannedSqlTests: ... + def start_sql_test_planning( *, project: CompiledProjectView, @@ -6898,7 +6893,7 @@ def start_sql_test_planning( selected_keys: frozenset[CompiledObjectKey], sql_test_selection: SqlTestSelection, enabled: bool, -) -> Future[PlannedSqlTests]: ... +) -> SqlTestPlanning: ... def join_sql_test_planning_threads() -> None: ... def build_display_only_sqlbuild_plan( *, project: CompiledProjectView, selected_model_names: tuple[str, ...], full_refresh: bool @@ -11174,6 +11169,13 @@ class DiagnosticSeverity(StrEnum): WARNING = "warning" INFO = "info" +class InferredNullability(StrEnum): + """Conservative nullability state inferred for an output column.""" + + NON_NULL = "non_null" + NULLABLE = "nullable" + UNKNOWN = "unknown" + class ColumnTransformKind(StrEnum): """High-level transform classification for one output column.""" diff --git a/src/sqlbuild/adapter/contract/types.py b/src/sqlbuild/adapter/contract/types.py index 65a2965d67..b1c018e69e 100644 --- a/src/sqlbuild/adapter/contract/types.py +++ b/src/sqlbuild/adapter/contract/types.py @@ -6,7 +6,7 @@ from enum import StrEnum from typing import Literal, Protocol -from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild._native import InferredNullability type FunctionNullabilityRule = Callable[[tuple[InferredNullability, ...]], InferredNullability] type StatementSizeLimit = tuple[int, Literal["bytes", "characters"]] diff --git a/src/sqlbuild/adapter/type_system/main/conditional_result_nullability.py b/src/sqlbuild/adapter/type_system/main/conditional_result_nullability.py index 5e18a0780e..69ac3fdfc7 100644 --- a/src/sqlbuild/adapter/type_system/main/conditional_result_nullability.py +++ b/src/sqlbuild/adapter/type_system/main/conditional_result_nullability.py @@ -1,7 +1,6 @@ """Public adapter conditional nullability inference operation.""" -from sqlbuild._native import infer_conditional_result_nullability -from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild._native import InferredNullability, infer_conditional_result_nullability def conditional_result_nullability( diff --git a/src/sqlbuild/adapter/type_system/main/first_arg_nullability.py b/src/sqlbuild/adapter/type_system/main/first_arg_nullability.py index 1e4e38ceee..e2a289218c 100644 --- a/src/sqlbuild/adapter/type_system/main/first_arg_nullability.py +++ b/src/sqlbuild/adapter/type_system/main/first_arg_nullability.py @@ -1,7 +1,6 @@ """Public adapter unary nullability inference operation.""" -from sqlbuild._native import infer_first_arg_nullability -from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild._native import InferredNullability, infer_first_arg_nullability def first_arg_nullability(args: tuple[InferredNullability, ...]) -> InferredNullability: diff --git a/src/sqlbuild/adapter/type_system/main/safe_cast_nullability.py b/src/sqlbuild/adapter/type_system/main/safe_cast_nullability.py index 03829b08e0..6dfb2a5fa0 100644 --- a/src/sqlbuild/adapter/type_system/main/safe_cast_nullability.py +++ b/src/sqlbuild/adapter/type_system/main/safe_cast_nullability.py @@ -1,7 +1,6 @@ """Public adapter safe-cast nullability inference operation.""" -from sqlbuild._native import infer_safe_cast_nullability -from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild._native import InferredNullability, infer_safe_cast_nullability def safe_cast_nullability(args: tuple[InferredNullability, ...]) -> InferredNullability: diff --git a/src/sqlbuild/compiler/lineage/types.py b/src/sqlbuild/compiler/lineage/types.py deleted file mode 100644 index 5c1151ce5a..0000000000 --- a/src/sqlbuild/compiler/lineage/types.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Column lineage domain types.""" - -from __future__ import annotations - -from enum import StrEnum - - -class InferredNullability(StrEnum): - """Conservative nullability state inferred for an output column.""" - - NON_NULL = "non_null" - NULLABLE = "nullable" - UNKNOWN = "unknown" diff --git a/tests/integration/src/sqlbuild/adapters/bigquery/_test_types.py b/tests/integration/src/sqlbuild/adapters/bigquery/_test_types.py index 88320a2a3e..7e3ac1bc18 100644 --- a/tests/integration/src/sqlbuild/adapters/bigquery/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/bigquery/_test_types.py @@ -1,6 +1,7 @@ from dataclasses import dataclass, field from datetime import datetime +from sqlbuild._native import InferredNullability from sqlbuild.adapter.contract.models import ( ColumnInfo, QueryResult, @@ -8,7 +9,6 @@ RowDiffSampleRow, SchemaDiffResult, ) -from sqlbuild.compiler.lineage.types import InferredNullability @dataclass(frozen=True) diff --git a/tests/integration/src/sqlbuild/adapters/bigquery/test_client.py b/tests/integration/src/sqlbuild/adapters/bigquery/test_client.py index 593f0693a7..0bc8c859f7 100644 --- a/tests/integration/src/sqlbuild/adapters/bigquery/test_client.py +++ b/tests/integration/src/sqlbuild/adapters/bigquery/test_client.py @@ -11,6 +11,7 @@ FINGERPRINT_TABLE_NAME, Fingerprint, FingerprintSet, + InferredNullability, SourceFreshnessValueKind, normalize_source_freshness_data_version, read_latest_fingerprints, @@ -31,7 +32,6 @@ ) from sqlbuild.adapter.contract.types import FunctionNullabilityRule from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.adapters.bigquery._test_types import ( BigQueryBuildFlowTestCase, BigQueryDeleteInsertCursorTestCase, diff --git a/tests/integration/src/sqlbuild/adapters/databricks/_test_types.py b/tests/integration/src/sqlbuild/adapters/databricks/_test_types.py index fe08d50555..78030f257d 100644 --- a/tests/integration/src/sqlbuild/adapters/databricks/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/databricks/_test_types.py @@ -1,5 +1,6 @@ from dataclasses import dataclass, field +from sqlbuild._native import InferredNullability from sqlbuild.adapter.contract.models import ( ColumnInfo, QueryResult, @@ -7,7 +8,6 @@ RowDiffSampleRow, SchemaDiffResult, ) -from sqlbuild.compiler.lineage.types import InferredNullability @dataclass(frozen=True) diff --git a/tests/integration/src/sqlbuild/adapters/databricks/test_client.py b/tests/integration/src/sqlbuild/adapters/databricks/test_client.py index 415ed47ebe..af1f8ae885 100644 --- a/tests/integration/src/sqlbuild/adapters/databricks/test_client.py +++ b/tests/integration/src/sqlbuild/adapters/databricks/test_client.py @@ -10,6 +10,7 @@ FINGERPRINT_TABLE_NAME, Fingerprint, FingerprintSet, + InferredNullability, SourceFreshnessValueKind, normalize_source_freshness_data_version, read_latest_fingerprints, @@ -30,7 +31,6 @@ ) from sqlbuild.adapter.contract.types import FunctionNullabilityRule from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.adapters.databricks._test_types import ( DatabricksBuildFlowTestCase, DatabricksExpressionNullabilityRuleTestCase, diff --git a/tests/integration/src/sqlbuild/adapters/duckdb/_test_types.py b/tests/integration/src/sqlbuild/adapters/duckdb/_test_types.py index 38119cc4da..09db89e6db 100644 --- a/tests/integration/src/sqlbuild/adapters/duckdb/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/duckdb/_test_types.py @@ -1,6 +1,6 @@ from dataclasses import dataclass, field -from sqlbuild._native import SeedCsvSettings +from sqlbuild._native import InferredNullability, SeedCsvSettings from sqlbuild.adapter.contract.models import ( ColumnInfo, CursorValue, @@ -9,7 +9,6 @@ RowDiffTolerances, SchemaDiffResult, ) -from sqlbuild.compiler.lineage.types import InferredNullability @dataclass(frozen=True) diff --git a/tests/integration/src/sqlbuild/adapters/duckdb/test_client.py b/tests/integration/src/sqlbuild/adapters/duckdb/test_client.py index 606a918d2c..60e9669bbb 100644 --- a/tests/integration/src/sqlbuild/adapters/duckdb/test_client.py +++ b/tests/integration/src/sqlbuild/adapters/duckdb/test_client.py @@ -7,7 +7,7 @@ import duckdb import pytest -from sqlbuild._native import SeedCsvSettings +from sqlbuild._native import InferredNullability, SeedCsvSettings from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.models import ( ColumnInfo, @@ -23,7 +23,6 @@ ) from sqlbuild.adapter.contract.types import FunctionNullabilityRule from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.adapters.duckdb._test_types import ( CompositeRowDiffSamplingTestCase, ConnectSettingsTestCase, diff --git a/tests/integration/src/sqlbuild/adapters/snowflake/_test_types.py b/tests/integration/src/sqlbuild/adapters/snowflake/_test_types.py index 24169d45bf..4201cd89dc 100644 --- a/tests/integration/src/sqlbuild/adapters/snowflake/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/snowflake/_test_types.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field from datetime import datetime -from sqlbuild._native import TableType +from sqlbuild._native import InferredNullability, TableType from sqlbuild.adapter.contract.models import ( ColumnInfo, QueryResult, @@ -9,7 +9,6 @@ RowDiffSampleRow, SchemaDiffResult, ) -from sqlbuild.compiler.lineage.types import InferredNullability from sqlbuild.cost.types import CostStatus diff --git a/tests/integration/src/sqlbuild/adapters/snowflake/test_client.py b/tests/integration/src/sqlbuild/adapters/snowflake/test_client.py index 5d5c601f5d..75b0f3fb6b 100644 --- a/tests/integration/src/sqlbuild/adapters/snowflake/test_client.py +++ b/tests/integration/src/sqlbuild/adapters/snowflake/test_client.py @@ -8,6 +8,7 @@ import pytest from sqlbuild._native import ( + InferredNullability, SourceFreshnessValueKind, normalize_source_freshness_data_version, source_freshness_data_version_hash, @@ -26,7 +27,6 @@ ) from sqlbuild.adapter.contract.types import FunctionNullabilityRule from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.adapters.snowflake._test_types import ( SnowflakeBuildFlowTestCase, SnowflakeExpressionNullabilityRuleTestCase, diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py index a9d2585e20..f6f471fc2f 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py @@ -22,6 +22,7 @@ CompileAdapterContext, DynamicColumnContractProof, InferredColumn, + InferredNullability, ModelSqlAnalysis, NativeModelAnalyses, NativeModelAnalysisRequest, @@ -44,7 +45,6 @@ from sqlbuild.cli.entry.main.entry import main from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.compiler.analysis_session._test_types import ( AdapterRuleCallbackTestCase, ) diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py index 73d1870fc9..26028f3a36 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py @@ -13,6 +13,7 @@ import sqlbuild._native as native_module from sqlbuild._native import ( ColumnLineageMode, + InferredNullability, NativeModelAnalysisRequest, analyze_native_model_sql, assemble_project, @@ -21,7 +22,6 @@ from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.compiler.analysis_session._test_types import ( AdapterRuleCallbackTestCase, CyclicAnalysisTestCase, diff --git a/tests/integration/src/sqlbuild/compiler/contracts/helpers.py b/tests/integration/src/sqlbuild/compiler/contracts/helpers.py index eef1bc48dd..1f6c0b4873 100644 --- a/tests/integration/src/sqlbuild/compiler/contracts/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/contracts/helpers.py @@ -21,6 +21,7 @@ DynamicColumnFamilyProof, IncrementalMode, InferredColumn, + InferredNullability, MaterializationType, SchemaColumn, SchemaDynamicColumnFamily, @@ -34,7 +35,6 @@ from sqlbuild.adapter.contract.types import TypeDialect from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView -from sqlbuild.compiler.lineage.types import InferredNullability _PROJECT_TOML: str = ( 'name = "orders_contracts"\nadapter = "duckdb"\n\n[connection]\ndatabase = "orders.duckdb"\n' diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/test_main.py b/tests/integration/src/sqlbuild/compiler/pipeline/test_main.py index b8e8be854b..32590ee704 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/test_main.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/test_main.py @@ -11,6 +11,7 @@ CompilePipelineResult, DiscoveredProjectInputs, InferredColumn, + InferredNullability, ModelPlanEntry, RulesError, build_compiled_project, @@ -21,7 +22,6 @@ from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.compiler.lineage.types import InferredNullability from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( AppendCursorPipelineIntegrationTestCase, CompileProgressIntegrationTestCase, diff --git a/tests/unit/src/sqlbuild/adapter/type_system/_test_types.py b/tests/unit/src/sqlbuild/adapter/type_system/_test_types.py index 9f110cd9d6..6c9d3a7617 100644 --- a/tests/unit/src/sqlbuild/adapter/type_system/_test_types.py +++ b/tests/unit/src/sqlbuild/adapter/type_system/_test_types.py @@ -1,7 +1,7 @@ from dataclasses import dataclass +from sqlbuild._native import InferredNullability from sqlbuild.adapter.contract.models import NormalizedType -from sqlbuild.compiler.lineage.types import InferredNullability @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/adapter/type_system/test_nullability_rules.py b/tests/unit/src/sqlbuild/adapter/type_system/test_nullability_rules.py index 17c57e0ce8..1c6304a6aa 100644 --- a/tests/unit/src/sqlbuild/adapter/type_system/test_nullability_rules.py +++ b/tests/unit/src/sqlbuild/adapter/type_system/test_nullability_rules.py @@ -1,11 +1,11 @@ import pytest +from sqlbuild._native import InferredNullability from sqlbuild.adapter.type_system.main.conditional_result_nullability import ( conditional_result_nullability, ) from sqlbuild.adapter.type_system.main.first_arg_nullability import first_arg_nullability from sqlbuild.adapter.type_system.main.safe_cast_nullability import safe_cast_nullability -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.adapter.type_system._test_types import ( FirstArgIdentityTestCase, NullabilityRuleTestCase, diff --git a/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py b/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py index 49945ac720..d63ff26d46 100644 --- a/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py @@ -1,13 +1,13 @@ from dataclasses import dataclass from datetime import datetime +from sqlbuild._native import InferredNullability from sqlbuild.adapter.contract.models import ( ColumnInfo, RowDiffResult, RowDiffSampleRow, SchemaDiffResult, ) -from sqlbuild.compiler.lineage.types import InferredNullability @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/adapters/bigquery/test_client.py b/tests/unit/src/sqlbuild/adapters/bigquery/test_client.py index 32f5e1245e..19cccaa5c1 100644 --- a/tests/unit/src/sqlbuild/adapters/bigquery/test_client.py +++ b/tests/unit/src/sqlbuild/adapters/bigquery/test_client.py @@ -4,7 +4,7 @@ import pytest -from sqlbuild._native import FunctionArgument, FunctionReturnColumn +from sqlbuild._native import FunctionArgument, FunctionReturnColumn, InferredNullability from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.models import ( ColumnInfo, @@ -28,7 +28,6 @@ from sqlbuild.adapters.bigquery.classes.bigquery_connection import _BigQueryConnection from sqlbuild.adapters.bigquery.classes.bigquery_cursor import _BigQueryCursor from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.lineage.types import InferredNullability from sqlbuild.observability import ( EventDispatcher, LifecycleEvent, diff --git a/tests/unit/src/sqlbuild/adapters/databricks/_test_types.py b/tests/unit/src/sqlbuild/adapters/databricks/_test_types.py index 747801d996..f4974ee6f2 100644 --- a/tests/unit/src/sqlbuild/adapters/databricks/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/databricks/_test_types.py @@ -1,8 +1,8 @@ from dataclasses import dataclass from datetime import datetime +from sqlbuild._native import InferredNullability from sqlbuild.adapter.contract.models import RelationInfo -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.adapters.databricks.helpers import FakeDatabricksMetadataCursor diff --git a/tests/unit/src/sqlbuild/adapters/databricks/test_client.py b/tests/unit/src/sqlbuild/adapters/databricks/test_client.py index 9618419132..046e6ad499 100644 --- a/tests/unit/src/sqlbuild/adapters/databricks/test_client.py +++ b/tests/unit/src/sqlbuild/adapters/databricks/test_client.py @@ -5,7 +5,7 @@ import pytest -from sqlbuild._native import FunctionArgument, FunctionReturnColumn +from sqlbuild._native import FunctionArgument, FunctionReturnColumn, InferredNullability from sqlbuild.adapter.contract.models import ( ExpressionInferenceProfile, RelationInfo, @@ -15,7 +15,6 @@ from sqlbuild.adapter.contract.types import FunctionNullabilityRule, LoaderLogicalType from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.adapters.databricks._test_types import ( DatabricksDeltaDetailAgeTestCase, DatabricksExpressionInferenceProfileTestCase, diff --git a/tests/unit/src/sqlbuild/adapters/duckdb/_test_types.py b/tests/unit/src/sqlbuild/adapters/duckdb/_test_types.py index 7317b8a384..8500d1ab73 100644 --- a/tests/unit/src/sqlbuild/adapters/duckdb/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/duckdb/_test_types.py @@ -1,6 +1,6 @@ from dataclasses import dataclass -from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild._native import InferredNullability @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/adapters/duckdb/test_client.py b/tests/unit/src/sqlbuild/adapters/duckdb/test_client.py index e3b4bb573e..bfa13727c1 100644 --- a/tests/unit/src/sqlbuild/adapters/duckdb/test_client.py +++ b/tests/unit/src/sqlbuild/adapters/duckdb/test_client.py @@ -2,11 +2,10 @@ import pytest -from sqlbuild._native import FunctionArgument, FunctionReturnColumn +from sqlbuild._native import FunctionArgument, FunctionReturnColumn, InferredNullability from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.adapter.contract.types import CursorKind, FunctionNullabilityRule from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.adapter.contract.classes.base_adapter.helpers import ( RecordingConnection, ) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py index 427104900e..4637b11380 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py @@ -1,10 +1,10 @@ from dataclasses import dataclass from datetime import datetime +from sqlbuild._native import InferredNullability from sqlbuild.adapter.contract.models import FunctionInfo, SchemaDiffResult from sqlbuild.adapter.contract.types import TableFreshnessStatus from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter -from sqlbuild.compiler.lineage.types import InferredNullability from sqlbuild.cost.types import CostStatus diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py b/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py index 4daf453e63..80b645cce2 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py @@ -7,7 +7,7 @@ import pytest -from sqlbuild._native import FunctionArgument, FunctionReturnColumn +from sqlbuild._native import FunctionArgument, FunctionReturnColumn, InferredNullability from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.models import ( @@ -25,7 +25,6 @@ ) from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.adapters.snowflake._test_types import ( SnowflakeConnectConfigTestCase, SnowflakeExpressionInferenceProfileTestCase, diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py index d3d7488e73..4ae65df22c 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py @@ -11,6 +11,7 @@ CompileSqlReference, ExpansionSpan, InferredColumn, + InferredNullability, ResolvedTimeTravelRetention, SourceLocation, SqlLexicalSyntax, @@ -20,9 +21,6 @@ ) from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.compiler.compile.types import AttachedAuditTargetKind -from sqlbuild.compiler.lineage.types import ( - InferredNullability, -) if TYPE_CHECKING: from sqlbuild._native import CompileInputsView diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/helpers.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/helpers.py index 904478c5df..5f7b1a7378 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/helpers.py @@ -30,6 +30,7 @@ DeclarationScopeResolver, DiscoveredMacroFile, DiscoveredProjectInputs, + InferredNullability, LoadedMacro, MacroMetadata, NativeColumnQuery, @@ -68,7 +69,6 @@ CompiledProjectView, CompiledResourceType, ) -from sqlbuild.compiler.lineage.types import InferredNullability if TYPE_CHECKING: from sqlbuild._native import CompileInputsView, ModelInput diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/render/test_cursor_intrinsics.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/render/test_cursor_intrinsics.py index 49e2ce9725..aedb960f59 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/render/test_cursor_intrinsics.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/render/test_cursor_intrinsics.py @@ -4,6 +4,7 @@ from sqlbuild._native import ( InferredColumn, + InferredNullability, NativeColumnQuery, cursor_intrinsics_analysis_sql, infer_native_query_columns, @@ -12,7 +13,6 @@ ) from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.compiler.compile.exceptions import CompileInputError -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.compiler.compile._helpers.render._test_types import ( CursorIntrinsicAnalysisTestCase, CursorIntrinsicErrorTestCase, diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py index 4ef9a72518..4d4345dc32 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_analysis_columns.py @@ -9,6 +9,7 @@ CompiledLineageSourceFact, CompileSqlReference, InferredColumn, + InferredNullability, NativeColumnQuery, NativeQueryColumns, SqlReferenceKind, @@ -17,7 +18,6 @@ ) from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.compiler.compile.types import CompiledResourceType -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( InferColumnsTestCase, PolyglotAnalysisTestCase, diff --git a/tests/unit/src/sqlbuild/compiler/contracts/_test_types.py b/tests/unit/src/sqlbuild/compiler/contracts/_test_types.py index ee757c5903..8f5b20ce8c 100644 --- a/tests/unit/src/sqlbuild/compiler/contracts/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/contracts/_test_types.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from pathlib import Path -from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild._native import InferredNullability @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/compiler/contracts/helpers.py b/tests/unit/src/sqlbuild/compiler/contracts/helpers.py index 61471d09ce..0275c794c9 100644 --- a/tests/unit/src/sqlbuild/compiler/contracts/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/contracts/helpers.py @@ -12,6 +12,7 @@ CompileModelConfig, DynamicColumnContractProof, InferredColumn, + InferredNullability, SchemaAuditInstance, SchemaColumn, SchemaDynamicColumnFamily, @@ -20,7 +21,6 @@ SourceLocation, ) from sqlbuild.compiler.compile.types import CompiledProjectView, CompiledResourceType -from sqlbuild.compiler.lineage.types import InferredNullability def make_contract_project( diff --git a/tests/unit/src/sqlbuild/compiler/contracts/test_validate.py b/tests/unit/src/sqlbuild/compiler/contracts/test_validate.py index 6bb37dd1d1..02e284d25b 100644 --- a/tests/unit/src/sqlbuild/compiler/contracts/test_validate.py +++ b/tests/unit/src/sqlbuild/compiler/contracts/test_validate.py @@ -8,12 +8,12 @@ ContractValidationResult, DynamicColumnContractProof, DynamicColumnFamilyProof, + InferredNullability, SchemaDynamicColumnFamily, SourceLocation, evaluate_model_contracts, ) from sqlbuild.adapter.contract.types import TypeDialect -from sqlbuild.compiler.lineage.types import InferredNullability from tests.unit.src.sqlbuild.compiler.contracts._test_types import ( ContractLocationTestCase, ContractMissingDeclarationsTestCase, diff --git a/tests/unit/src/sqlbuild/compiler/planner/classes/test_sql_test_planning.py b/tests/unit/src/sqlbuild/compiler/planner/classes/test_sql_test_planning.py index e58cae2ad0..91f8d27399 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/classes/test_sql_test_planning.py +++ b/tests/unit/src/sqlbuild/compiler/planner/classes/test_sql_test_planning.py @@ -1,7 +1,6 @@ from __future__ import annotations import threading -from concurrent.futures import Future import pytest @@ -11,6 +10,7 @@ PlannedSqlTests, PlannerInputError, PlanWarning, + SqlTestPlanning, SqlTestSelection, WarningSeverity, start_sql_test_planning, @@ -60,7 +60,7 @@ def plan( monkeypatch.setattr(native_module, "plan_selected_sql_tests", plan) - planning: Future[PlannedSqlTests] = start_sql_test_planning( + planning: SqlTestPlanning = start_sql_test_planning( sql_test_selection=SqlTestSelection(), project=object(), # ty: ignore[invalid-argument-type] adapter=object(), # ty: ignore[invalid-argument-type] @@ -102,7 +102,7 @@ def plan( monkeypatch.setattr(native_module, "plan_selected_sql_tests", plan) - planning: Future[PlannedSqlTests] = start_sql_test_planning( + planning: SqlTestPlanning = start_sql_test_planning( sql_test_selection=SqlTestSelection(), project=object(), # ty: ignore[invalid-argument-type] adapter=object(), # ty: ignore[invalid-argument-type] @@ -142,7 +142,7 @@ def plan( monkeypatch.setattr(native_module, "plan_selected_sql_tests", plan) - planning: Future[PlannedSqlTests] = start_sql_test_planning( + planning: SqlTestPlanning = start_sql_test_planning( sql_test_selection=SqlTestSelection(), project=object(), # ty: ignore[invalid-argument-type] adapter=object(), # ty: ignore[invalid-argument-type] @@ -157,14 +157,14 @@ def plan( "test_case", [ BackgroundSqlTestPlanningTestCase( - description="running planning never keeps the process alive", + description="running planning leaves no Python thread to keep the process alive", enabled=True, expected_planner_calls=1, ), ], ids=lambda case: case.description, ) -def test_given_running_planning_when_planner_fails_then_worker_is_daemon( +def test_given_running_planning_when_listing_threads_then_no_python_worker_blocks_exit( test_case: BackgroundSqlTestPlanningTestCase, monkeypatch: pytest.MonkeyPatch ) -> None: started: threading.Event = threading.Event() @@ -184,7 +184,7 @@ def plan( monkeypatch.setattr(native_module, "plan_selected_sql_tests", plan) - planning: Future[PlannedSqlTests] = start_sql_test_planning( + planning: SqlTestPlanning = start_sql_test_planning( sql_test_selection=SqlTestSelection(), project=object(), # ty: ignore[invalid-argument-type] adapter=object(), # ty: ignore[invalid-argument-type] @@ -198,5 +198,5 @@ def plan( release.set() planning.result() - assert daemon_by_name["sqlbuild-test-planning"] is True + assert daemon_by_name.get("sqlbuild-test-planning", True) is True assert test_case.expected_planner_calls == 1