From ac984eb586f30ef234bc5ae45a80565e1eeb6611 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:33:26 +0100 Subject: [PATCH 01/13] refactor(analysis): host InferredNullability natively InferredNullability becomes a native StrEnum registered with the compiler enums, so Rust reads it through the typed registry instead of importing sqlbuild.compiler.lineage.types in thirteen places. Python callers import it from sqlbuild._native and the shape-only module is deleted. Stage captures keep the original captured shape name. --- .../sqlbuild-cache/src/hosted_shapes/constants.rs | 6 +++++- .../_helpers/analysis_session/fixture_columns.rs | 5 +---- .../_helpers/analysis_session/model_analyses.rs | 7 ++----- .../_helpers/analysis_session/outcome_objects.rs | 5 +---- .../_helpers/analysis_session/profile_inputs.rs | 7 +++---- .../_helpers/analysis_session/stage_functions.rs | 5 +---- .../_helpers/compile_driver/assembly_analysis.rs | 3 +-- .../bindings/_helpers/compile_records/constants.rs | 6 +----- .../_helpers/compiled_project/loop_classes.rs | 3 +-- .../bindings/_helpers/compiled_project/records.rs | 3 +-- .../bindings/_helpers/compiler_enums/constants.rs | 10 ++++++++++ .../bindings/_helpers/compiler_records/constants.rs | 6 +----- .../bindings/_helpers/hand_built/compiled_model.rs | 8 ++------ .../bindings/_helpers/lineage/lineage_objects.rs | 12 +++++------- .../_helpers/project_assembly/declared_relations.rs | 8 ++++---- .../src/bindings/_helpers/records/ids.rs | 1 + .../_helpers/semantic_shapes/source_bindings.rs | 12 ++++-------- .../bindings/_helpers/type_system/nullability.rs | 10 ++++------ python-boundary.toml | 5 ----- src/sqlbuild/_native.pyi | 10 +++++++--- src/sqlbuild/adapter/contract/types.py | 2 +- .../main/conditional_result_nullability.py | 3 +-- .../type_system/main/first_arg_nullability.py | 3 +-- .../type_system/main/safe_cast_nullability.py | 3 +-- src/sqlbuild/compiler/lineage/types.py | 13 ------------- .../src/sqlbuild/adapters/bigquery/_test_types.py | 2 +- .../src/sqlbuild/adapters/bigquery/test_client.py | 2 +- .../src/sqlbuild/adapters/databricks/_test_types.py | 2 +- .../src/sqlbuild/adapters/databricks/test_client.py | 2 +- .../src/sqlbuild/adapters/duckdb/_test_types.py | 3 +-- .../src/sqlbuild/adapters/duckdb/test_client.py | 3 +-- .../src/sqlbuild/adapters/snowflake/_test_types.py | 3 +-- .../src/sqlbuild/adapters/snowflake/test_client.py | 2 +- .../sqlbuild/compiler/analysis_session/helpers.py | 2 +- .../analysis_session/test_native_model_analysis.py | 2 +- .../src/sqlbuild/compiler/contracts/helpers.py | 2 +- .../src/sqlbuild/compiler/pipeline/test_main.py | 2 +- .../src/sqlbuild/adapter/type_system/_test_types.py | 2 +- .../adapter/type_system/test_nullability_rules.py | 2 +- .../src/sqlbuild/adapters/bigquery/_test_types.py | 2 +- .../src/sqlbuild/adapters/bigquery/test_client.py | 3 +-- .../src/sqlbuild/adapters/databricks/_test_types.py | 2 +- .../src/sqlbuild/adapters/databricks/test_client.py | 3 +-- .../src/sqlbuild/adapters/duckdb/_test_types.py | 2 +- .../src/sqlbuild/adapters/duckdb/test_client.py | 3 +-- .../src/sqlbuild/adapters/snowflake/_test_types.py | 2 +- .../src/sqlbuild/adapters/snowflake/test_client.py | 3 +-- .../compiler/compile/_helpers/_test_types.py | 4 +--- .../sqlbuild/compiler/compile/_helpers/helpers.py | 2 +- .../_helpers/render/test_cursor_intrinsics.py | 2 +- .../compile/_helpers/test_sql_analysis_columns.py | 2 +- .../src/sqlbuild/compiler/contracts/_test_types.py | 2 +- .../unit/src/sqlbuild/compiler/contracts/helpers.py | 2 +- .../sqlbuild/compiler/contracts/test_validate.py | 2 +- 54 files changed, 88 insertions(+), 135 deletions(-) delete mode 100644 src/sqlbuild/compiler/lineage/types.py 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-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/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/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..8235026556 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", 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/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..31668b1004 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, @@ -11174,6 +11171,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, From 91ea22c0f00766ed556cb36630d6cf950ea732d0 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:04:48 +0100 Subject: [PATCH 02/13] refactor(analysis): read compiled models' analysis as native data Semantic shapes, coverage, type recovery, completion, metadata and lineage read each compiled model's inferred columns, declared columns, star flags, unchecked outputs, lineage, binding validation and dynamic proof through typed accessors. Native model views answer from their owner's Rust facts without building Python projections; hand-built records are converted at the read. --- .../src/bindings/_helpers/analysis/mod.rs | 1 + .../bindings/_helpers/analysis/model_reads.rs | 119 ++++++++++++++++++ .../compiled_model_analysis.rs | 7 ++ .../compiled_project/compiled_model_view.rs | 53 ++++++-- .../_helpers/lineage/project_lineage.rs | 14 +-- .../_helpers/lineage/project_schema.rs | 24 ++-- .../_helpers/semantic_checks/completion.rs | 7 +- .../_helpers/semantic_checks/metadata.rs | 3 +- .../_helpers/semantic_checks/stage.rs | 4 +- .../_helpers/semantic_checks/type_recovery.rs | 42 +------ .../_helpers/semantic_shapes/coverage.rs | 20 +-- .../_helpers/semantic_shapes/shapes.rs | 49 ++++---- 12 files changed, 235 insertions(+), 108 deletions(-) create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/analysis/model_reads.rs 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/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/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/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)?, From e1a20aca93cc283fb4507ddb44ef81393507f8ce Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:22:01 +0100 Subject: [PATCH 03/13] refactor(rules): read proven dynamic contracts natively for SQL rule paths --- .../sqlbuild-python/src/bindings/_helpers/rules/sql_rules.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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")?)? From 217c58c3e3d5d10161592235f38c1e800490033e Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:04:55 +0100 Subject: [PATCH 04/13] refactor: plan SQL tests on a native thread instead of a Python thread and future SQL test planning ran on a Python threading.Thread whose target was the native planning function, reporting through a concurrent.futures.Future. It now runs on a Rust thread that attaches to the interpreter and sends its outcome over a channel; the planner and `start_sql_test_planning` read it through the native `SqlTestPlanning` handle. The atexit hook waits, without the GIL and for at most five seconds, until no planning thread runs, so finalization never ends one inside native code. The planning entry is still looked up on the module, so tests can replace it. --- .../planner_entries/execution_plan.rs | 150 ++-------- .../bindings/_helpers/planner_entries/mod.rs | 1 + .../_helpers/planner_entries/plan_output.rs | 40 --- .../_helpers/planner_entries/registration.rs | 3 +- .../_helpers/planner_entries/test_planning.rs | 260 ++++++++++++++++++ .../src/bindings/_helpers/records/ids.rs | 1 - src/sqlbuild/_native.pyi | 16 +- .../planner/classes/test_sql_test_planning.py | 16 +- 8 files changed, 295 insertions(+), 192 deletions(-) create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs 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/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/test_planning.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs new file mode 100644 index 0000000000..84582b5084 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs @@ -0,0 +1,260 @@ +//! 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, +} + +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 planned: Planned = Python::attach(|py| { + let planned: Planned = plan(py, &request); + drop(request); + planned + }); + let _receiver_gone = sender.send(planned); + threads.get().finished(); + }); + 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/records/ids.rs b/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs index 8235026556..78b2df5d8d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/records/ids.rs @@ -1038,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/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 31668b1004..54f6ce2037 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -6775,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, @@ -6888,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, @@ -6895,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 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 From 98f9c51899571a59ca862693ef059b32b031c021 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:25:54 +0100 Subject: [PATCH 05/13] refactor: call per-model planner helpers directly instead of through the module Model plan entries, audit entries, relation contexts and cursor resolution looked up native helpers on `sqlbuild._native` and called them with keyword dictionaries once or more per model. They now call the Rust functions directly: authored config strings and cursor bounds, the model plan action, SQL, hooks, full-refresh and microbatch limit resolution, model plan entry assembly, audit selection, attachment kind, effective run scope and lineage reads. Results are unchanged. --- .../_helpers/planner_changes/detect.rs | 13 ++- .../_helpers/planner_entries/display_plan.rs | 11 +-- .../_helpers/planner_entries/model_entries.rs | 95 ++++++++----------- .../_helpers/planner_entries/relations.rs | 12 +-- .../planner_migrations/column_objects.rs | 11 +-- .../_helpers/planner_phases/objects.rs | 15 +-- .../planner_resolution/cursor_model.rs | 24 ++--- .../planner_resolution/cursor_overrides.rs | 11 +-- .../planner_resolution/cursor_ranges.rs | 11 +-- .../_helpers/planner_resolution/mod.rs | 2 +- .../planner_resolution/orchestration.rs | 27 +++--- .../planner_strategy/audit_entries.rs | 36 +++---- .../_helpers/planner_strategy/audits.rs | 6 +- .../_helpers/planner_strategy/strategy.rs | 2 +- 14 files changed, 107 insertions(+), 169 deletions(-) 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/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/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_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, From ac68b74a700e761c01b00e84c689cdb0860ae8b4 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:02:58 +0100 Subject: [PATCH 06/13] perf(sqltext): reject code bytes before dialect comment and quote checks Dialect non-code scanning is called once per byte by many SQL test and analysis scanners. Return early when the byte cannot open a block comment, line comment, dollar quote or quoted text, instead of running prefix comparisons at every offset. Results are unchanged. --- .../src/sql_scan/_helpers/dialect.rs | 15 ++++ .../src/sql_scan/tests/test_dialect_scan.rs | 84 ++++++++++++++++++- .../src/sql_scan/tests/test_types.rs | 8 ++ 3 files changed, 106 insertions(+), 1 deletion(-) 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, From e9eff5fe0a5d02db303cd725a36fb478d72d3cc2 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:02:58 +0100 Subject: [PATCH 07/13] perf(rules): run the interpolation-site pattern only at possible site starts Lint and format site patterns can only match at '@', '$' or '_'; skip the anchored regex search at every other code byte. --- .../src/lint_text/_helpers/sites.rs | 8 +++- .../sqlbuild-rules/src/lint_text/constants.rs | 2 + .../src/lint_text/tests/test_lint_text.rs | 39 ++++++++++++++++++- .../src/lint_text/tests/test_types.rs | 6 +++ 4 files changed, 51 insertions(+), 4 deletions(-) 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, +} From 69a7bef63821ccdf86469772d22e83ec1c899178 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:02:58 +0100 Subject: [PATCH 08/13] perf(analysis): find the end of a SQL test select list in one scan Projection-name validation scanned each branch once per clause keyword. Scan once and stop at the first top-level FROM or clause keyword with the same name and IS [NOT] DISTINCT FROM rules. --- .../compiler/_helpers/sql_tests/extraction.rs | 69 ++++++++++--------- .../tests/test_sql_test_extraction_rules.rs | 16 +++++ 2 files changed, 51 insertions(+), 34 deletions(-) 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/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", From d34db431d2924b9f6abc604a670410afbff90cda Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:02:58 +0100 Subject: [PATCH 09/13] perf(cache): stop re-hashing replay stdout when reading a stored compile The store's whole-body checksum already verifies the stdout record stored under its digest, so reading a project reuse entry no longer hashes the full compile output a second time. Persisted bytes are unchanged. --- crates/sqlbuild-cache/src/project_reuse/_helpers/entry.rs | 3 --- 1 file changed, 3 deletions(-) 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)) } From 13bebe593ac53a9b0495cb10710d6df2698311b5 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:02:58 +0100 Subject: [PATCH 10/13] perf: use FxHash indexes for hot in-memory lookups Add rustc-hash (already in the lock file) as the workspace fast hasher. Index hosted shape names once instead of scanning every shape per lookup, match UNION BY NAME outputs through a name index instead of quadratic scans, and use FxHash for hot sets and maps in function-name validation, CTE column rules, lint additional facts and SQL test identifier tokens. None of these maps' iteration order reaches output. --- Cargo.lock | 3 + Cargo.toml | 2 + crates/sqlbuild-analysis/Cargo.toml | 1 + .../_helpers/sql_tests/helper_scope.rs | 3 +- .../query_analysis/_helpers/borrowed_facts.rs | 59 +++++++++++-------- .../src/query_analysis/tests/helpers.rs | 19 ++++++ .../_helpers/function_names.rs | 23 ++++---- crates/sqlbuild-cache/Cargo.toml | 1 + .../src/hosted_shapes/_helpers/mod.rs | 1 + .../hosted_shapes/_helpers/native_names.rs | 43 ++++++++++++++ crates/sqlbuild-rules/Cargo.toml | 1 + .../src/sql_lint/_helpers/additional.rs | 21 +++---- .../src/sql_quality/_helpers/cte_columns.rs | 40 ++++++------- .../src/sql_quality/_helpers/syntax.rs | 7 +-- 14 files changed, 154 insertions(+), 70 deletions(-) create mode 100644 crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs create mode 100644 crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs diff --git a/Cargo.lock b/Cargo.lock index 0882c97e86..98dc6d83bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1148,6 +1148,7 @@ dependencies = [ "polyglot-sql-function-catalogs", "rayon", "regex", + "rustc-hash", "serde", "serde_json", "sha2", @@ -1175,6 +1176,7 @@ version = "0.156.1" dependencies = [ "blake3", "rayon", + "rustc-hash", "serde_json", "sha2", "sqlbuild-core", @@ -1315,6 +1317,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/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/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/Cargo.toml b/crates/sqlbuild-cache/Cargo.toml index 617208de4e..89b1f64b69 100644 --- a/crates/sqlbuild-cache/Cargo.toml +++ b/crates/sqlbuild-cache/Cargo.toml @@ -10,6 +10,7 @@ rust-version.workspace = true [dependencies] blake3.workspace = true rayon.workspace = true +rustc-hash.workspace = true serde_json.workspace = true sha2.workspace = true sqlbuild-core.workspace = true diff --git a/crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs b/crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs new file mode 100644 index 0000000000..2b116399fe --- /dev/null +++ b/crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs @@ -0,0 +1 @@ +pub(crate) mod native_names; diff --git a/crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs b/crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs new file mode 100644 index 0000000000..c9b6c20faa --- /dev/null +++ b/crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs @@ -0,0 +1,43 @@ +//! Every hosted Python shape name with the one native class hosting it, indexed once. + +use std::sync::LazyLock; + +use rustc_hash::FxHashMap; + +use crate::hosted_shapes::constants::{HOSTED_SHAPE_TYPES, NATIVE_MODULE}; +use crate::hosted_shapes::errors::AmbiguousShapeName; + +static NATIVE_NAMES: LazyLock>> = + LazyLock::new(|| { + let mut index = FxHashMap::default(); + for (native, python) in HOSTED_SHAPE_TYPES { + let Some(shape) = python.split_once(':') else { + continue; + }; + let Some(native_name) = native + .strip_prefix(NATIVE_MODULE) + .and_then(|rest| rest.strip_prefix(':')) + else { + continue; + }; + index + .entry(shape) + .and_modify(|found: &mut Option<&'static str>| { + if found.is_some_and(|existing| existing != native_name) { + *found = None; + } + }) + .or_insert(Some(native_name)); + } + index + }); + +/// The native class hosting `module:class`, an ambiguity, or `None` when no native class hosts it. +pub(crate) fn native_name( + module: &str, + class: &str, +) -> Option> { + NATIVE_NAMES + .get(&(module, class)) + .map(|found| found.ok_or(AmbiguousShapeName)) +} 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/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; From b0e9cf44be8800a54e1a550bc8925a93f402adc8 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:41:02 +0100 Subject: [PATCH 11/13] perf: stream the custom-rule host payload pickle to disk The parent pickled the whole host payload into one bytes object before writing it, holding a second copy of the payload (about 285 MB for a 10,000-model project) alive during the rules phase. Pickle straight into the partial file instead; the atomic rename and the payload the hosts read are unchanged. --- .../bindings/_helpers/rules/custom_hosts.rs | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) 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>, From 65bcdbcc7515be8fc3d775590049d3522c3e1eb7 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:35:17 +0100 Subject: [PATCH 12/13] chore: drop the unused hosted shape index after the registry replaced reverse lookups --- .../src/hosted_shapes/_helpers/mod.rs | 1 - .../hosted_shapes/_helpers/native_names.rs | 43 ------------------- 2 files changed, 44 deletions(-) delete mode 100644 crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs delete mode 100644 crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs diff --git a/crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs b/crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs deleted file mode 100644 index 2b116399fe..0000000000 --- a/crates/sqlbuild-cache/src/hosted_shapes/_helpers/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) mod native_names; diff --git a/crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs b/crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs deleted file mode 100644 index c9b6c20faa..0000000000 --- a/crates/sqlbuild-cache/src/hosted_shapes/_helpers/native_names.rs +++ /dev/null @@ -1,43 +0,0 @@ -//! Every hosted Python shape name with the one native class hosting it, indexed once. - -use std::sync::LazyLock; - -use rustc_hash::FxHashMap; - -use crate::hosted_shapes::constants::{HOSTED_SHAPE_TYPES, NATIVE_MODULE}; -use crate::hosted_shapes::errors::AmbiguousShapeName; - -static NATIVE_NAMES: LazyLock>> = - LazyLock::new(|| { - let mut index = FxHashMap::default(); - for (native, python) in HOSTED_SHAPE_TYPES { - let Some(shape) = python.split_once(':') else { - continue; - }; - let Some(native_name) = native - .strip_prefix(NATIVE_MODULE) - .and_then(|rest| rest.strip_prefix(':')) - else { - continue; - }; - index - .entry(shape) - .and_modify(|found: &mut Option<&'static str>| { - if found.is_some_and(|existing| existing != native_name) { - *found = None; - } - }) - .or_insert(Some(native_name)); - } - index - }); - -/// The native class hosting `module:class`, an ambiguity, or `None` when no native class hosts it. -pub(crate) fn native_name( - module: &str, - class: &str, -) -> Option> { - NATIVE_NAMES - .get(&(module, class)) - .map(|found| found.ok_or(AmbiguousShapeName)) -} From 5b36f278cb49f2dff480f244f5e0626ccfa6217f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:53:54 +0100 Subject: [PATCH 13/13] fix: mark SQL test planning finished on panic and drop an unused cache dependency --- Cargo.lock | 1 - crates/sqlbuild-cache/Cargo.toml | 1 - .../_helpers/planner_entries/test_planning.rs | 11 ++++++++++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 98dc6d83bc..a5d4e8ec33 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1176,7 +1176,6 @@ version = "0.156.1" dependencies = [ "blake3", "rayon", - "rustc-hash", "serde_json", "sha2", "sqlbuild-core", diff --git a/crates/sqlbuild-cache/Cargo.toml b/crates/sqlbuild-cache/Cargo.toml index 89b1f64b69..617208de4e 100644 --- a/crates/sqlbuild-cache/Cargo.toml +++ b/crates/sqlbuild-cache/Cargo.toml @@ -10,7 +10,6 @@ rust-version.workspace = true [dependencies] blake3.workspace = true rayon.workspace = true -rustc-hash.workspace = true serde_json.workspace = true sha2.workspace = true sqlbuild-core.workspace = true 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 index 84582b5084..9c13e79504 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_entries/test_planning.rs @@ -38,6 +38,15 @@ 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); @@ -191,13 +200,13 @@ pub(crate) fn start_planning( 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); - threads.get().finished(); }); if let Err(error) = spawned { planning_threads(py)?.get().finished();