diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs index fd4aabb6c4..267114a012 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs @@ -38,18 +38,24 @@ fn normalize_dialect_sql(py: Python<'_>, sql: &str, dialect: &str) -> PyResult, dialect: &str, - requests: Vec<(String, std::collections::HashMap)>, -) -> PyResult> { - py.compiler_detach(|| { - crate::semantic_validation::main::normalize_batch::normalize_analysis_sqls( - dialect, requests, + requests: Vec, +) -> PyResult>> { + let results = py.compiler_detach(|| { + Ok( + crate::semantic_validation::main::normalize_batch::normalize_analysis_sqls( + dialect, requests, None, + ), ) - }) - .map_err(value_error) + }); + crate::bindings::main::normalization_results::normalization_results( + py, + results.map_err(value_error)?, + ) } #[pyfunction] diff --git a/crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs b/crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs new file mode 100644 index 0000000000..f87fc5b8ec --- /dev/null +++ b/crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs @@ -0,0 +1,5 @@ +pub(crate) fn catch_compiler_panic( + operation: impl FnOnce() -> Result, +) -> Result { + crate::bindings::_helpers::panics::catch_compiler_panic(operation) +} diff --git a/crates/sqlbuild-rules-native/src/bindings/main/mod.rs b/crates/sqlbuild-rules-native/src/bindings/main/mod.rs index 3b8691d5be..9f6415af70 100644 --- a/crates/sqlbuild-rules-native/src/bindings/main/mod.rs +++ b/crates/sqlbuild-rules-native/src/bindings/main/mod.rs @@ -1,3 +1,5 @@ +pub(crate) mod catch_compiler_panic; pub(crate) mod compiler_error; pub(crate) mod compiler_guard; +pub(crate) mod normalization_results; pub mod register; diff --git a/crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs b/crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs new file mode 100644 index 0000000000..88cb1a632a --- /dev/null +++ b/crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs @@ -0,0 +1,17 @@ +use pyo3::prelude::{IntoPyObject, Py, PyAny, PyResult, Python}; + +/// Return normalized SQL in order, carrying each failure as the exception a single call raises. +pub(crate) fn normalization_results( + py: Python<'_>, + results: Vec>, +) -> PyResult>> { + results + .into_iter() + .map(|result| match result { + Ok(sql) => Ok(sql.into_pyobject(py)?.into_any().unbind()), + Err(error) => Ok(crate::bindings::_helpers::panics::compiler_error(error) + .into_value(py) + .into_any()), + }) + .collect() +} diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs b/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs index 37ff225df0..a44bde62ff 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs @@ -1,14 +1,19 @@ //! Compile-owned schema catalog and native, batched binding requests. use crate::bindings::main::compiler_error::compiler_error; +use crate::bindings::main::normalization_results::normalization_results; use crate::bindings::types::CompilerDetach; +use crate::semantic_validation::main::normalize_batch; use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use crate::semantic_validation::models::{ CatalogInput, Columns, CompactAnalysisJob, FunctionProbes, ProjectCatalog, }; -use crate::semantic_validation::types::{BindingRequest, DiagnosticRow, Relations}; +use crate::semantic_validation::types::{ + BindingRequest, DiagnosticRow, NormalizationRequest, Relations, +}; use crate::semantic_validation::{ _helpers::{diagnostics, identifiers}, main as validation, @@ -19,7 +24,7 @@ use polyglot_sql::{ ValidationSchema, }; use pyo3::exceptions::PyValueError; -use pyo3::prelude::{Bound, FromPyObject, PyAny, PyAnyMethods, PyResult, Python}; +use pyo3::prelude::{Bound, FromPyObject, Py, PyAny, PyAnyMethods, PyResult, Python}; use pyo3::pymethods; use pyo3::types::{PyBytes, PyDict, PyDictMethods}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; @@ -137,6 +142,7 @@ impl ProjectCatalog { analysis_tables: HashMap::new(), function_probes: Arc::new(FunctionProbes::default()), analysis_pool: Arc::default(), + running_analyses: Arc::default(), }; catalog.update_relations(relations); Ok(catalog) @@ -158,6 +164,7 @@ impl ProjectCatalog { analysis_tables: self.analysis_tables.clone(), function_probes: Arc::new(FunctionProbes::default()), analysis_pool: Arc::clone(&self.analysis_pool), + running_analyses: Arc::clone(&self.running_analyses), }; catalog.update_relations(relations); catalog @@ -234,6 +241,7 @@ impl ProjectCatalog { ) -> PyResult>> { py.compiler_detach(|| { let pool = self.analysis_pool()?; + let _running = RunningAnalysis::start(&self.running_analyses); pool.install(|| { requests .into_par_iter() @@ -253,6 +261,44 @@ impl ProjectCatalog { }) .map_err(compiler_error) } + + /// Normalize a preparation batch on the idle analysis pool, else in turn, in one detached call. + fn normalize_analysis_sqls( + &self, + py: Python<'_>, + dialect: &str, + requests: Vec, + ) -> PyResult>> { + let results = py + .compiler_detach(|| { + let pool = (self.running_analyses.load(Ordering::Acquire) == 0) + .then(|| self.analysis_pool()) + .transpose()?; + Ok(normalize_batch::normalize_analysis_sqls( + dialect, + requests, + pool.as_deref(), + )) + }) + .map_err(compiler_error)?; + normalization_results(py, results) + } +} + +/// Count one analysis batch as running on the catalog's pool until dropped. +struct RunningAnalysis<'a>(&'a AtomicUsize); + +impl<'a> RunningAnalysis<'a> { + fn start(running: &'a AtomicUsize) -> Self { + running.fetch_add(1, Ordering::AcqRel); + Self(running) + } +} + +impl Drop for RunningAnalysis<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } } #[pymethods] @@ -265,7 +311,10 @@ impl CompactAnalysisJob { .take() .ok_or_else(|| PyValueError::new_err("compact analysis job already ran"))?; let result = py - .compiler_detach(|| analysis(&self.catalog)) + .compiler_detach(|| { + let _running = RunningAnalysis::start(&self.catalog.running_analyses); + analysis(&self.catalog) + }) .map_err(compiler_error)?; Ok(PyBytes::new(py, result.as_bytes())) } @@ -283,6 +332,7 @@ impl ProjectCatalog { analysis_tables: HashMap::new(), function_probes: Arc::clone(&self.function_probes), analysis_pool: Arc::clone(&self.analysis_pool), + running_analyses: Arc::clone(&self.running_analyses), } } @@ -321,6 +371,7 @@ impl ProjectCatalog { analysis_tables: HashMap::new(), function_probes: Arc::new(FunctionProbes::default()), analysis_pool: Arc::default(), + running_analyses: Arc::default(), } } pub(crate) fn analysis_schema( diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs b/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs index 5337ae4090..0ad8408745 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs @@ -1,19 +1,27 @@ -use std::collections::HashMap; +use rayon::ThreadPool; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use crate::bindings::main::catch_compiler_panic::catch_compiler_panic; +use crate::semantic_validation::types::NormalizationRequest; + +/// Normalize every request, keeping each failure in place so callers raise the first in order. pub(crate) fn normalize_analysis_sqls( dialect: &str, - requests: Vec<(String, HashMap)>, -) -> Result, String> { - let mut results: Vec = Vec::with_capacity(requests.len()); - for (sql, placeholders) in requests { - results.push( + requests: Vec, + pool: Option<&ThreadPool>, +) -> Vec> { + let normalize = |(sql, stubs, placeholders): NormalizationRequest| { + catch_compiler_panic(|| { crate::semantic_validation::_helpers::normalization::normalize_analysis_sql( &sql, dialect, - HashMap::new(), + stubs, placeholders, - )?, - ); + ) + }) + }; + match pool { + Some(pool) => pool.install(|| requests.into_par_iter().map(normalize).collect()), + None => requests.into_iter().map(normalize).collect(), } - Ok(results) } diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/models.rs b/crates/sqlbuild-rules-native/src/semantic_validation/models.rs index 46afc66c34..8e94c6ec28 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/models.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/models.rs @@ -6,6 +6,7 @@ use polyglot_sql::{DialectType, SchemaValidationOptions}; use pyo3::{FromPyObject, pyclass}; use rayon::ThreadPool; use std::collections::HashMap; +use std::sync::atomic::AtomicUsize; use std::sync::{Arc, Mutex}; #[derive(Default, Debug)] @@ -57,6 +58,8 @@ pub(crate) struct ProjectCatalog { pub(crate) function_probes: Arc, /// One analysis pool per compile catalog, built on first use and shared by its views. pub(super) analysis_pool: Arc>>>, + /// Analysis batches running on the pool; preparation only borrows the pool while none run. + pub(super) running_analyses: Arc, } /// One resolved compact analysis batch that runs without borrowing its project catalog. diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs b/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs index bd969772d8..418d1ace8a 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs @@ -1,5 +1,7 @@ #[path = "tests/test_function_names.rs"] mod function_names; +#[path = "tests/test_normalize_batch.rs"] +mod normalize_batch; #[path = "tests/test_positions.rs"] mod positions; #[path = "tests/test_semantic_validation.rs"] diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs new file mode 100644 index 0000000000..3cdbaceca7 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs @@ -0,0 +1,89 @@ +use std::collections::HashMap; + +use crate::semantic_validation::main::normalize_batch::normalize_analysis_sqls; +use crate::semantic_validation::tests::test_types::NormalizationBatchTestCase; +use crate::semantic_validation::types::NormalizationRequest; + +#[test] +fn given_batch_when_normalizing_on_pool_or_in_turn_then_matches_single_calls_in_order() { + let test_cases = [ + NormalizationBatchTestCase { + description: "stubbed references and placeholders", + dialect: "snowflake", + requests: &[ + ( + r#"SELECT * FROM __ref("orders") JOIN __ref("customers") USING (id)"#, + &[("orders", "__sqb_rel_0"), ("customers", "__sqb_rel_1")], + ), + ("SELECT payload:item[0] AS item FROM inventory", &[]), + ("SELECT 1 LIMIT @@@limit", &[]), + ], + expected_failures: &[], + }, + NormalizationBatchTestCase { + description: "failures stay at their own positions", + dialect: "snowflake", + requests: &[ + ("SELECT 'unterminated FROM orders", &[]), + ("SELECT id FROM products", &[]), + ("SELECT /* unterminated FROM customers", &[]), + ], + expected_failures: &[0, 2], + }, + ]; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(3) + .build() + .expect("test pool must build"); + + for test_case in test_cases { + let requests = || -> Vec { + test_case + .requests + .iter() + .map(|(sql, stubs)| { + ( + (*sql).to_owned(), + stubs + .iter() + .map(|(name, stub)| ((*name).to_owned(), (*stub).to_owned())) + .collect(), + HashMap::from([("limit".to_owned(), "10".to_owned())]), + ) + }) + .collect() + }; + let expected: Vec> = requests() + .into_iter() + .map(|(sql, stubs, placeholders)| { + crate::semantic_validation::_helpers::normalization::normalize_analysis_sql( + &sql, + test_case.dialect, + stubs, + placeholders, + ) + }) + .collect(); + let failures: Vec = (0..expected.len()) + .filter(|index| expected[*index].is_err()) + .collect(); + + assert_eq!( + failures, test_case.expected_failures, + "{}", + test_case.description + ); + assert_eq!( + normalize_analysis_sqls(test_case.dialect, requests(), Some(&pool)), + expected, + "{}", + test_case.description + ); + assert_eq!( + normalize_analysis_sqls(test_case.dialect, requests(), None), + expected, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs index 030a8221de..521de7be8e 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs @@ -45,3 +45,10 @@ pub(crate) struct TypeLocationTestCase { pub(crate) findings: &'static [(&'static str, &'static str)], pub(crate) expected_starts: &'static [&'static str], } + +pub(crate) struct NormalizationBatchTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) requests: &'static [(&'static str, &'static [(&'static str, &'static str)])], + pub(crate) expected_failures: &'static [usize], +} diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/types.rs b/crates/sqlbuild-rules-native/src/semantic_validation/types.rs index 1e171ebe37..87ebd9353e 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/types.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/types.rs @@ -3,6 +3,8 @@ use std::collections::HashMap; pub(crate) type Relations = HashMap; pub(crate) type BindingRequest = (String, Vec<(String, bool)>, Relations); +/// One analysis SQL with its relation stubs and placeholder defaults. +pub(crate) type NormalizationRequest = (String, HashMap, HashMap); pub(crate) type DiagnosticRow = ( String, String, diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py index 145dd4de11..3833027435 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py @@ -33,6 +33,7 @@ from sqlbuild.compiler.compile._helpers.sharing.binding import ( binding_query_fields, lineage_reference_map, + normalized_analysis_sql, prepare_binding_queries, relation_stubs_collide, remembered_shared_results, @@ -171,7 +172,9 @@ from sqlbuild.compiler.sql_analysis.main._binding_catalog import create_binding_catalog from sqlbuild.compiler.sql_analysis.main._decode_schema_validation import decode_schema_validation from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql -from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import normalize_analysis_sqls +from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( + normalize_analysis_sql_results, +) from sqlbuild.compiler.sql_analysis.main._schema_validation import get_schema_validations from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql from sqlbuild.compiler.sql_analysis.models import ( @@ -666,11 +669,18 @@ def _compact_batch_context( if name in inputs.column_nullability_by_table }, ) - normalized_sqls: list[str] = normalize_analysis_sqls( - sqls=inputs.query_sqls, + requests: list[tuple[str, None, dict[str, str] | None]] = [ + (query_sql, None, query_placeholders) + for query_sql, query_placeholders in zip( + inputs.query_sqls, inputs.placeholders, strict=True + ) + ] + results: list[str | Exception] = normalize_analysis_sql_results( dialect=dialect, - placeholders=inputs.placeholders, + requests=requests, + catalog=binding_catalog.native if binding_catalog is not None else None, ) + normalized_sqls: list[str] = [normalized_analysis_sql(result) for result in results] binding_queries: tuple[PreparedBindingQuery | None, ...] = ( prepare_binding_queries( binding_catalog=binding_catalog, @@ -697,7 +707,6 @@ def _prepare_compact_analysis_batch( ) -> CompactBatchPreparation: query_sqls: tuple[str, ...] = inputs.query_sqls references: tuple[tuple[CompileSqlReference, ...], ...] = inputs.references - placeholders: tuple[dict[str, str] | None, ...] = inputs.placeholders column_nullability_by_table: dict[str, dict[str, InferredNullability]] = ( inputs.column_nullability_by_table ) @@ -735,63 +744,40 @@ def _prepare_compact_analysis_batch( ) binding_queries: tuple[PreparedBindingQuery | None, ...] = context.binding_queries shared_query_indexes: set[int] = set() - for ( - query_sql, + members: list[tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]] + stubbed_sqls: dict[int, str | Exception] + members, stubbed_sqls = _member_analysis_stubs( + inputs=inputs, + context=context, + binding_schemas=member_binding_schemas, + ) + for index, ( cleaned_sql, query_references, - query_placeholders, query_recover_cte_facts, binding_schema, binding_query, - ) in zip( - query_sqls, - normalized_sqls, - references, - placeholders, - recover_cte_facts, - member_binding_schemas, - binding_queries, - strict=True, - ): - lineage_references: dict[str, tuple[CompiledResourceType, str]] = lineage_reference_map( - query_references + (lineage_references, canonical_stubs), + ) in enumerate( + zip( + normalized_sqls, + references, + recover_cte_facts, + member_binding_schemas, + binding_queries, + members, + strict=True, ) + ): shared: SharedBindingQuery | None = ( binding_query.shared if binding_query is not None else None ) - stubbed_reference_names: frozenset[str] = ( - frozenset() - if binding_schema is not None - or relation_stubs_collide(cleaned_sql=cleaned_sql, names=lineage_references) - else frozenset(lineage_references) - - _qualified_reference_names( - query_sql=cleaned_sql, - reference_names=lineage_references.keys(), - ) - ) - canonical_stubs: dict[str, str] = ( - shared.stubs - if shared is not None - else { - name: ( - f"{COMPACT_RELATION_STUB_PREFIX}{index}" - if name in stubbed_reference_names - else name - ) - for index, name in enumerate(lineage_references) - } - ) analysis_sql: str = ( shared.sql if shared is not None else cleaned_sql if binding_schema is not None - else _cleaned_analysis_sql( - query_sql=query_sql, - placeholders=query_placeholders, - dialect=dialect, - relation_stubs=canonical_stubs, - ) + else normalized_analysis_sql(stubbed_sqls[index]) ) query: dict[str, object] = { "sql": analysis_sql, @@ -920,6 +906,86 @@ def _prepare_compact_analysis_batch( ) +def _member_analysis_stubs( + *, + inputs: CompactBatchInputs, + context: CompactBatchContext, + binding_schemas: tuple[dict[str, dict[str, str]] | None, ...], +) -> tuple[ + list[tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]], + dict[int, str | Exception], +]: + """Return each member's relation stubs and, in one native batch, its stubbed SQL.""" + + members: list[tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]] = [] + stubbed_requests: dict[int, tuple[str, dict[str, str], dict[str, str] | None]] = {} + for index, ( + query_sql, + cleaned_sql, + query_references, + query_placeholders, + member_schema, + member_query, + ) in enumerate( + zip( + inputs.query_sqls, + context.normalized_sqls, + inputs.references, + inputs.placeholders, + binding_schemas, + context.binding_queries, + strict=True, + ) + ): + member: tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]] = ( + _member_relation_stubs( + cleaned_sql=cleaned_sql, + references=query_references, + binding_schema=member_schema, + binding_query=member_query, + ) + ) + members.append(member) + if member_schema is None and (member_query is None or member_query.shared is None): + stubbed_requests[index] = (query_sql, member[1], query_placeholders) + stubbed_sqls: list[str | Exception] = normalize_analysis_sql_results( + dialect=inputs.inference_profile.sql_analysis_dialect, + requests=tuple(stubbed_requests.values()), + catalog=(context.binding_catalog.native if context.binding_catalog is not None else None), + ) + return members, dict(zip(stubbed_requests, stubbed_sqls, strict=True)) + + +def _member_relation_stubs( + *, + cleaned_sql: str, + references: tuple[CompileSqlReference, ...], + binding_schema: dict[str, dict[str, str]] | None, + binding_query: PreparedBindingQuery | None, +) -> tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]: + """Return a member's lineage references and the relation names its analysis SQL uses.""" + + lineage_references: dict[str, tuple[CompiledResourceType, str]] = lineage_reference_map( + references + ) + if binding_query is not None and binding_query.shared is not None: + return lineage_references, binding_query.shared.stubs + stubbed_reference_names: frozenset[str] = ( + frozenset() + if binding_schema is not None + or relation_stubs_collide(cleaned_sql=cleaned_sql, names=lineage_references) + else frozenset(lineage_references) + - _qualified_reference_names( + query_sql=cleaned_sql, + reference_names=lineage_references.keys(), + ) + ) + return lineage_references, { + name: f"{COMPACT_RELATION_STUB_PREFIX}{index}" if name in stubbed_reference_names else name + for index, name in enumerate(lineage_references) + } + + def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> object: """Analyse every project size with the full worker pool; the rules phase sets peak memory.""" diff --git a/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py b/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py index 21ba613ae5..709e95c39a 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py +++ b/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py @@ -30,7 +30,9 @@ ) from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.compiler.sql_analysis.constants import SQL_QUOTED_IDENTIFIER_DELIMITER -from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql +from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( + normalize_analysis_sql_results, +) from sqlbuild.compiler.sql_analysis.models import SqlSchemaValidationRequest @@ -108,6 +110,35 @@ def prepare_binding_queries( if prekey is not None and count >= MIN_SHARED_BINDING_QUERY_MEMBERS ) ) + shared_stubs: dict[int, dict[str, str]] = {} + for index, binding_schema in enumerate(binding_schemas): + prekey: str | None = prekeys[index] + if binding_schema is None or prekey is None or prekey not in candidates: + continue + names: list[str] | None = _shared_reference_names( + cleaned_sql=cleaned_sqls[index], + references=inputs.references[index], + binding_schema=binding_schema, + ) + if names is not None: + shared_stubs[index] = { + name: f"{COMPACT_RELATION_STUB_PREFIX}{position}" + for position, name in enumerate(names) + } + shared_sqls: dict[int, str | Exception] = dict( + zip( + shared_stubs, + normalize_analysis_sql_results( + dialect=dialect, + requests=[ + (inputs.query_sqls[index], stubs, inputs.placeholders[index]) + for index, stubs in shared_stubs.items() + ], + catalog=binding_catalog.native, + ), + strict=True, + ) + ) prepared: list[PreparedBindingQuery | None] = [] for index, binding_schema in enumerate(binding_schemas): if binding_schema is None: @@ -120,19 +151,18 @@ def prepare_binding_queries( ) ] )[0] - prekey: str | None = prekeys[index] - shared: SharedBindingQuery | None = None - if prekey is not None and prekey in candidates: - shared = _shared_binding_query( + shared: SharedBindingQuery | None = ( + _shared_binding_query( binding_catalog=binding_catalog, - query=SqlSchemaValidationRequest( - sql=inputs.query_sqls[index], dialect=dialect, schema=binding_schema - ), - cleaned_sql=cleaned_sqls[index], - references=inputs.references[index], - placeholders=inputs.placeholders[index], + sql=normalized_analysis_sql(shared_sqls[index]), + dialect=dialect, + schema=binding_schema, + stubs=shared_stubs[index], overrides=overrides, ) + if index in shared_stubs + else None + ) prepared.append( PreparedBindingQuery(references=binding_references, overrides=overrides, shared=shared) ) @@ -190,6 +220,14 @@ def shared_result_is_exact(*, validation: object, template: object) -> bool: return isinstance(validation, dict) and cast(dict[str, object], validation).get("errors") == [] +def normalized_analysis_sql(result: str | Exception) -> str: + """Return one normalized batch member, raising the error its own normalization raised.""" + + if isinstance(result, Exception): + raise result + return result + + def relation_stubs_collide(*, cleaned_sql: str, names: Iterable[str]) -> bool: """Return whether authored SQL or relation names already use the relation stub prefix.""" @@ -226,29 +264,20 @@ def _shared_reference_names( def _shared_binding_query( *, binding_catalog: Any, - query: SqlSchemaValidationRequest, - cleaned_sql: str, - references: tuple[CompileSqlReference, ...], - placeholders: dict[str, str] | None, + sql: str, + dialect: str | None, + schema: Mapping[str, Mapping[str, str]], + stubs: dict[str, str], overrides: Mapping[str, Mapping[str, str]], -) -> SharedBindingQuery | None: - names: list[str] | None = _shared_reference_names( - cleaned_sql=cleaned_sql, references=references, binding_schema=query.schema - ) - if names is None: - return None - stubs: dict[str, str] = { - name: f"{COMPACT_RELATION_STUB_PREFIX}{index}" for index, name in enumerate(names) - } - sql: str = normalize_analysis_sql( - sql=query.sql, dialect=query.dialect, stubs=stubs, placeholders=placeholders - ) +) -> SharedBindingQuery: + """Identify one query from its SQL normalized with relation `stubs`.""" + analysis_shapes: list[tuple[object, ...]] = [] - for name in names: + for name in stubs: types, nullability = binding_catalog.analysis_shapes.get(name, ({}, {})) analysis_shapes.append((stubs[name], tuple(types.items()), tuple(nullability.items()))) binding_shapes: list[tuple[object, ...]] = [] - for relation, columns in query.schema.items(): + for relation, columns in schema.items(): effective: Mapping[str, str] = ( overrides[relation] if relation in overrides @@ -264,7 +293,7 @@ def _shared_binding_query( stubs=stubs, key=( sql, - query.dialect or "generic", + dialect or "generic", tuple(analysis_shapes), tuple(sorted(binding_shapes)), ), diff --git a/src/sqlbuild/compiler/sql_analysis/constants.py b/src/sqlbuild/compiler/sql_analysis/constants.py index cfeba4649c..87feca62c3 100644 --- a/src/sqlbuild/compiler/sql_analysis/constants.py +++ b/src/sqlbuild/compiler/sql_analysis/constants.py @@ -181,3 +181,4 @@ "motherduck": "duckdb", "sqlserver": "tsql", } +ANALYSIS_NORMALIZATION_CHUNK_SIZE: int = 256 diff --git a/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py b/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py index c6ebd32ad1..a8dc96e619 100644 --- a/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py +++ b/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py @@ -1,21 +1,37 @@ -"""Normalize a compile batch under one native GIL-release boundary.""" +"""Normalize analysis SQL batches under one native GIL-release boundary.""" +from collections.abc import Callable, Mapping, Sequence from typing import cast import sqlbuild._native as _native -from sqlbuild.compiler.sql_analysis.types import NativePositionsModule +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_NORMALIZATION_CHUNK_SIZE +from sqlbuild.compiler.sql_analysis.types import ( + NativeNormalizationRequest, + NativePositionsModule, + NativeProjectCatalog, +) -def normalize_analysis_sqls( +def normalize_analysis_sql_results( *, - sqls: tuple[str, ...], dialect: str | None, - placeholders: tuple[dict[str, str] | None, ...], -) -> list[str]: - requests: list[tuple[str, dict[str, str]]] = [ - (sql, defaults or {}) for sql, defaults in zip(sqls, placeholders, strict=True) - ] - return cast(NativePositionsModule, _native).normalize_analysis_sqls( - dialect=dialect or "generic", - requests=requests, + requests: Sequence[tuple[str, Mapping[str, str] | None, Mapping[str, str] | None]], + catalog: NativeProjectCatalog | None, +) -> list[str | Exception]: + """Normalize `(sql, stubs, placeholders)` in order, keeping each failure in its place.""" + + normalize: Callable[..., list[str | Exception]] = ( + catalog.normalize_analysis_sqls + if catalog is not None + else cast(NativePositionsModule, _native).normalize_analysis_sqls ) + results: list[str | Exception] = [] + for start in range(0, len(requests), ANALYSIS_NORMALIZATION_CHUNK_SIZE): + chunk: list[NativeNormalizationRequest] = [ + (sql, dict(stubs or {}), dict(placeholders or {})) + for sql, stubs, placeholders in requests[ + start : start + ANALYSIS_NORMALIZATION_CHUNK_SIZE + ] + ] + results.extend(normalize(dialect=dialect or "generic", requests=chunk)) + return results diff --git a/src/sqlbuild/compiler/sql_analysis/types.py b/src/sqlbuild/compiler/sql_analysis/types.py index f490eb9c05..7e69e03f36 100644 --- a/src/sqlbuild/compiler/sql_analysis/types.py +++ b/src/sqlbuild/compiler/sql_analysis/types.py @@ -7,6 +7,7 @@ type NativeDiagnosticRow = tuple[str, str, int | None, int | None, int | None, int | None, str] type NativeBindingRequest = tuple[str, list[tuple[str, bool]], dict[str, Mapping[str, str]]] +type NativeNormalizationRequest = tuple[str, dict[str, str], dict[str, str]] class NativeCompactAnalysisJob(Protocol): @@ -30,6 +31,9 @@ def update_relations(self, relations: Mapping[str, Mapping[str, str]]) -> None: def binding_results( self, requests: list[NativeBindingRequest] ) -> list[list[NativeDiagnosticRow]]: ... + def normalize_analysis_sqls( + self, *, dialect: str, requests: list[NativeNormalizationRequest] + ) -> list[str | Exception]: ... class NativeCatalogModule(Protocol): @@ -47,8 +51,8 @@ def position( class NativePositionsModule(Protocol): def normalize_analysis_sqls( - self, *, dialect: str, requests: list[tuple[str, dict[str, str]]] - ) -> list[str]: ... + self, *, dialect: str, requests: list[NativeNormalizationRequest] + ) -> list[str | Exception]: ... def binding_diagnostics( self, *, sql: str, dialect: str, rows: list[NativeDiagnosticRow] ) -> list[NativeDiagnosticRow]: ... diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py index a4f6dea82a..b399e498dd 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py @@ -6,6 +6,8 @@ from dataclasses import dataclass, field from pathlib import Path +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import PreparedCompile + @dataclass(frozen=True) class SourceSemanticBindingCase: @@ -759,3 +761,13 @@ class RenderLoadNoticeTestCase: description: str notice_bytes: int expected_notice: bool + + +@dataclass(frozen=True) +class DenseBatchedPreparationTestCase: + description: str + model_count: int + expected_exit_code: int + + +type DensePreparedCompile = PreparedCompile diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py new file mode 100644 index 0000000000..fffac8f29e --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py @@ -0,0 +1,61 @@ +"""Dense batched analysis preparation must match per-model preparation byte for byte.""" + +from pathlib import Path + +import pytest + +from scripts.cold_compile_performance._helpers.dense_project import write_dense_compile_project +from sqlbuild.compiler.compile.constants import COMPACT_RELATION_STUB_PREFIX +from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import ( + DenseBatchedPreparationTestCase, + DensePreparedCompile, +) +from tests.integration.src.sqlbuild.compiler.pipeline.helpers import ( + keep_batched_normalization, + prepared_compile, + use_per_model_normalization, +) + + +@pytest.mark.performance +@pytest.mark.parametrize( + "test_case", + (DenseBatchedPreparationTestCase("dense_models_3000", 3000, 0),), + ids=lambda case: case.description, +) +def test_given_dense_project_when_compiling_cold_then_batched_preparation_matches_per_model( + test_case: DenseBatchedPreparationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + write_dense_compile_project(project_dir=batched_dir, model_count=test_case.model_count) + write_dense_compile_project(project_dir=reference_dir, model_count=test_case.model_count) + + expected: DensePreparedCompile = prepared_compile( + project_dir=reference_dir, + args=(), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: DensePreparedCompile = prepared_compile( + project_dir=batched_dir, + args=(), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=keep_batched_normalization, + ) + + assert actual[3] == expected[3] + assert actual[1] == expected[1] + assert actual[2] == expected[2] + assert actual[0] == expected[0] + assert actual[0][0] == test_case.expected_exit_code + assert any(COMPACT_RELATION_STUB_PREFIX.encode() in payload for payload in actual[3]) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"])) diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py b/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py index 2642c01a53..81a1f14d2d 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py @@ -2,6 +2,7 @@ from collections.abc import Callable from dataclasses import dataclass, field +from pathlib import Path from scripts.cold_compile_performance.models import RandomDagProject from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -393,6 +394,9 @@ class FunctionArgumentTypeCase: type CompileOutcome = tuple[int, dict[str, object], dict[str, bytes]] +type PreparedCompile = tuple[ + CompileOutcome, tuple[str, ...], dict[str, object | None], tuple[bytes, ...] +] @dataclass(frozen=True) @@ -477,3 +481,18 @@ class DataflowStartFailureCase: expected_error: str expected_live_workers: tuple[str, ...] expected_notices: int + + +@dataclass(frozen=True) +class BatchedPreparationCase: + description: str + write_project: Callable[[Path], tuple[str, ...]] + reshaped: tuple[int, ...] + expected_exit_code: int + + +@dataclass(frozen=True) +class PerturbedPreparationCase: + description: str + project: RandomDagProject + expected_payloads_match: bool diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py b/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py index e0b685df6a..66f7ada616 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py @@ -3,16 +3,23 @@ import itertools import json import random +import shutil import signal import threading import time -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from contextlib import suppress from pathlib import Path -from typing import Any +from types import ModuleType +from typing import Any, cast +import orjson import pytest +import sqlbuild._native as _native +from scripts.cold_compile_performance._helpers.dense_project import write_dense_compile_project +from scripts.cold_compile_performance._helpers.random_dag_project import write_random_dag_project +from scripts.cold_compile_performance.models import RandomDagProject from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.cli.commands.main.entrypoint.entry import main @@ -21,6 +28,7 @@ from sqlbuild.compiler.compile._helpers.assembly import binding_waves from sqlbuild.compiler.compile._helpers.assembly import project as project_assembly from sqlbuild.compiler.compile._helpers.assembly.binding_waves import analyze_binding_waves +from sqlbuild.compiler.compile._helpers.sharing import binding as binding_sharing from sqlbuild.compiler.compile.classes.binding_dataflow import BindingDataflow from sqlbuild.compiler.compile.main._build_compile_inputs import build_compile_inputs from sqlbuild.compiler.compile.models import ( @@ -37,14 +45,25 @@ from sqlbuild.compiler.pipeline.main.compile import run_compile_pipeline from sqlbuild.compiler.pipeline.main.project import compile_project from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult +from sqlbuild.compiler.sql_analysis.main import _normalize_analysis_batch as normalize_batch +from sqlbuild.compiler.sql_analysis.types import NativePositionsModule, NativeProjectCatalog from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.sql_values.types import CollectionRendering from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( CompileOutcome, DataflowScheduleCase, + PreparedCompile, SharedBindingQueryCase, ) +_REPOSITORY_ROOT: Path = Path(__file__).resolve().parents[6] +_NORMALIZATION_MODULES: tuple[ModuleType, ...] = (normalize_batch, binding_sharing, compact) +_ANALYSIS_CACHE_METRICS: tuple[str, ...] = ( + "analysis_batch_cache_hits", + "analysis_entry_cache_hits", + "analysis_cache_misses", + "analysis_cache_bypasses", +) _SCHEMA_FIXTURE_PATH: Path = ( Path(__file__).resolve().parents[5] / "fixtures" / "dbt_manifest_v12_schema.json" ) @@ -263,11 +282,19 @@ def compile_outcome( ) -> CompileOutcome: """Compile through the CLI; return the exit code, timing-free JSON, and artifacts.""" + return timed_compile_outcome(project_dir=project_dir, args=args, capsys=capsys)[0] + + +def timed_compile_outcome( + *, project_dir: Path, args: tuple[str, ...], capsys: pytest.CaptureFixture[str] +) -> tuple[CompileOutcome, dict[str, object]]: + """Compile through the CLI; return the timing-free outcome and its compile timings.""" + exit_code: int = main(["--project-dir", str(project_dir), "compile", "--json", *args]) payload: dict[str, object] = json.loads( capsys.readouterr().out.replace(str(project_dir), "") ) - _ = payload.pop("compile_timings", None) + timings: dict[str, object] = cast(dict[str, object], payload.pop("compile_timings", {})) compiled: Path = project_dir / "target" / "compiled" return ( exit_code, @@ -276,7 +303,7 @@ def compile_outcome( path.relative_to(compiled).as_posix(): path.read_bytes() for path in compiled.rglob("*.sql") }, - ) + ), timings def reshape_models(*, project_dir: Path, names: tuple[str, ...], indexes: tuple[int, ...]) -> None: @@ -513,3 +540,121 @@ def live_model_analysis_threads() -> list[str]: thread.name ) return sorted(names_by_kind.get(True, [])) + + +def random_dag_writer(project: RandomDagProject) -> Callable[[Path], tuple[str, ...]]: + return lambda project_dir: write_random_dag_project(project_dir=project_dir, project=project) + + +def fixture_writer(relative: str) -> Callable[[Path], tuple[str, ...]]: + def write(project_dir: Path) -> tuple[str, ...]: + _ = shutil.copytree( + _REPOSITORY_ROOT / relative, + project_dir, + ignore=shutil.ignore_patterns("target", "*.duckdb"), + ) + return () + + return write + + +def dense_writer(model_count: int) -> Callable[[Path], tuple[str, ...]]: + def write(project_dir: Path) -> tuple[str, ...]: + write_dense_compile_project(project_dir=project_dir, model_count=model_count) + return () + + return write + + +def use_per_model_normalization(monkeypatch: pytest.MonkeyPatch) -> None: + """Normalize each batch member with its own native call, the reference preparation.""" + + for module in _NORMALIZATION_MODULES: + monkeypatch.setattr(module, "normalize_analysis_sql_results", _per_model_normalization) + + +def keep_batched_normalization(monkeypatch: pytest.MonkeyPatch) -> None: + """Leave batched normalization in place, the preparation under test.""" + + _ = monkeypatch + + +def perturb_batched_normalization(monkeypatch: pytest.MonkeyPatch) -> None: + """Change the last member of every batched normalization, which comparisons must catch.""" + + batched: Callable[..., list[str | Exception]] = normalize_batch.normalize_analysis_sql_results + + def perturbed(**arguments: Any) -> list[str | Exception]: + results: list[str | Exception] = batched(**arguments) + return [*results[:-1], *(f"{result} " for result in results[-1:])] + + for module in _NORMALIZATION_MODULES: + monkeypatch.setattr(module, "normalize_analysis_sql_results", perturbed) + + +def prepared_compile( + *, + project_dir: Path, + args: tuple[str, ...], + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, + normalization: Callable[[pytest.MonkeyPatch], None], +) -> PreparedCompile: + """Compile once; return the outcome, cache keys, cache metrics, and native batch payloads.""" + + with monkeypatch.context() as patch: + normalization(patch) + preparations: list[CompactBatchPreparation] = trace_native_compact_batches(patch) + cache_keys: list[str] = _trace_analysis_cache_keys(patch) + outcome: CompileOutcome + timings: dict[str, object] + outcome, timings = timed_compile_outcome(project_dir=project_dir, args=args, capsys=capsys) + return ( + outcome, + tuple(sorted(cache_keys)), + {name: timings.get(name) for name in _ANALYSIS_CACHE_METRICS}, + tuple(_preparation_payload(preparation) for preparation in preparations), + ) + + +def _trace_analysis_cache_keys(monkeypatch: pytest.MonkeyPatch) -> list[str]: + original: Callable[..., str] = project_assembly.model_analysis_cache_key + keys: list[str] = [] + + def traced(**arguments: Any) -> str: + key: str = original(**arguments) + keys.append(key) + return key + + monkeypatch.setattr(project_assembly, "model_analysis_cache_key", traced) + return keys + + +def _per_model_normalization( + *, + dialect: str | None, + requests: Sequence[tuple[str, Mapping[str, str] | None, Mapping[str, str] | None]], + catalog: NativeProjectCatalog | None, +) -> list[str | Exception]: + _ = catalog + native: NativePositionsModule = cast(NativePositionsModule, _native) + return [ + native.normalize_analysis_sqls( + dialect=dialect or "generic", + requests=[(sql, dict(stubs or {}), dict(placeholders or {}))], + )[0] + for sql, stubs, placeholders in requests + ] + + +def _preparation_payload(preparation: CompactBatchPreparation) -> bytes: + return orjson.dumps( + { + "cleaned_sql": preparation.cleaned_sql, + "queries": preparation.queries, + "templates": preparation.templates, + "projections": preparation.projections, + "shared_query_indexes": sorted(preparation.shared_query_indexes), + }, + option=orjson.OPT_SORT_KEYS, + ) diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py b/tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py new file mode 100644 index 0000000000..c60d1408c3 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py @@ -0,0 +1,199 @@ +"""Batched analysis preparation must match per-model preparation byte for byte.""" + +from pathlib import Path + +import pytest + +from scripts.cold_compile_performance.models import RandomDagProject +from sqlbuild.compiler.compile.main import _build_compile_inputs +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( + BatchedPreparationCase, + PerturbedPreparationCase, + PreparedCompile, +) +from tests.integration.src.sqlbuild.compiler.pipeline.helpers import ( + dense_writer, + fixture_writer, + keep_batched_normalization, + perturb_batched_normalization, + prepared_compile, + random_dag_writer, + reshape_models, + use_per_model_normalization, + use_wave_analysis, +) + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationCase( + "inferred shapes", + random_dag_writer(RandomDagProject(seed=11, model_count=48)), + (16,), + 0, + ), + BatchedPreparationCase( + "binding errors mid-graph", + random_dag_writer(RandomDagProject(seed=23, model_count=48, errors=True)), + (24,), + 1, + ), + BatchedPreparationCase( + "run ids and analysis opt-outs", + random_dag_writer( + RandomDagProject(seed=37, model_count=40, run_ids=True, analysis_opt_outs=True) + ), + (20,), + 0, + ), + BatchedPreparationCase( + "reference cycle", + random_dag_writer(RandomDagProject(seed=67, model_count=24, cycle=True)), + (8,), + 0, + ), + BatchedPreparationCase( + "waffle shop fixture", fixture_writer("tests/e2e/fixtures/waffle_shop"), (), 0 + ), + BatchedPreparationCase( + "tidy shop example", fixture_writer("website/examples/tidy-shop"), (), 0 + ), + BatchedPreparationCase("dense single batch", dense_writer(120), (), 0), + ], + ids=lambda case: case.description, +) +def test_given_project_when_compiling_in_waves_then_batched_preparation_matches_per_model( + test_case: BatchedPreparationCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + names: tuple[str, ...] = test_case.write_project(batched_dir) + _ = test_case.write_project(reference_dir) + monkeypatch.setattr( + _build_compile_inputs, + "resolve_run_id", + lambda *, selected_run_id: selected_run_id or "run-0", + ) + use_wave_analysis(monkeypatch) + for step, reshaped, args in ( + ("cold", (), ()), + ("warm", (), ()), + ("edit", test_case.reshaped, ()), + ("uncached", (), ("--no-cache",)), + ): + reshape_models(project_dir=batched_dir, names=names, indexes=reshaped) + reshape_models(project_dir=reference_dir, names=names, indexes=reshaped) + expected: PreparedCompile = prepared_compile( + project_dir=reference_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: PreparedCompile = prepared_compile( + project_dir=batched_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=keep_batched_normalization, + ) + + assert actual[3] == expected[3], f"{step} native batch payloads differ" + assert actual[1] == expected[1], f"{step} analysis cache keys differ" + assert actual[2] == expected[2], f"{step} analysis cache hits differ" + assert actual[0] == expected[0], f"{step} compile differs from per-model preparation" + assert actual[0][0] == test_case.expected_exit_code + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationCase( + "inferred shapes", random_dag_writer(RandomDagProject(seed=13, model_count=48)), (), 0 + ), + BatchedPreparationCase( + "binding errors mid-graph", + random_dag_writer(RandomDagProject(seed=29, model_count=48, errors=True)), + (), + 1, + ), + ], + ids=lambda case: case.description, +) +def test_given_duckdb_project_when_compiling_dataflow_with_and_without_cache_then_matches_per_model( + test_case: BatchedPreparationCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + _ = test_case.write_project(batched_dir) + _ = test_case.write_project(reference_dir) + monkeypatch.setattr( + _build_compile_inputs, + "resolve_run_id", + lambda *, selected_run_id: selected_run_id or "run-0", + ) + for args in ((), (), ("--no-cache",)): + expected: PreparedCompile = prepared_compile( + project_dir=reference_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: PreparedCompile = prepared_compile( + project_dir=batched_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=keep_batched_normalization, + ) + + assert actual[:3] == expected[:3] + assert actual[0][0] == test_case.expected_exit_code + + +@pytest.mark.parametrize( + "test_case", + [PerturbedPreparationCase("last member", RandomDagProject(seed=11, model_count=24), False)], + ids=lambda case: case.description, +) +def test_given_perturbed_batch_result_when_compiling_then_payloads_differ_from_per_model( + test_case: PerturbedPreparationCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + _ = random_dag_writer(test_case.project)(batched_dir) + _ = random_dag_writer(test_case.project)(reference_dir) + use_wave_analysis(monkeypatch) + + expected: PreparedCompile = prepared_compile( + project_dir=reference_dir, + args=("--no-cache",), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: PreparedCompile = prepared_compile( + project_dir=batched_dir, + args=("--no-cache",), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=perturb_batched_normalization, + ) + + assert len(actual[3]) == len(expected[3]) > 0 + assert (actual[3] == expected[3]) is test_case.expected_payloads_match + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py index 159458e9b0..d3b072c465 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py @@ -1,7 +1,10 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass +from sqlbuild.compiler.sql_analysis.types import NativeProjectCatalog + @dataclass(frozen=True) class SemanticDiagnosticMappingCase: @@ -47,3 +50,33 @@ class PolyglotExplicitGuardTestCase: function_depth: int maximum_function_depth: int expected_error_pattern: str + + +@dataclass(frozen=True) +class NormalizationBatchTestCase: + description: str + dialect: str + requests: tuple[tuple[str, dict[str, str], dict[str, str]], ...] + catalog: Callable[[str], NativeProjectCatalog | None] + expected_stubbed_sql: str + + +@dataclass(frozen=True) +class NormalizationBatchFailureTestCase: + description: str + dialect: str + before: tuple[str, ...] + failing: str + after: tuple[str, ...] + catalog: Callable[[str], NativeProjectCatalog | None] + expected_error_type: type[Exception] + expected_result_types: tuple[str, ...] + + +@dataclass(frozen=True) +class NormalizationChunkTestCase: + description: str + valid_count: int + failing_index: int + catalog: Callable[[str], NativeProjectCatalog | None] + expected_error_type: type[Exception] diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py index c73c266f3a..efd7dcee2a 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py @@ -1,8 +1,29 @@ """Helpers for SQL analysis public-entry tests.""" +from sqlbuild.compiler.sql_analysis.classes.binding_catalog import BindingCatalog +from sqlbuild.compiler.sql_analysis.types import NativeProjectCatalog + def nested_coalesce_sql(*, function_depth: int) -> str: expression: str = "value" for _ in range(function_depth): expression = f"COALESCE({expression}, 0)" return f"SELECT {expression} AS resolved_value FROM records" + + +def pool_catalog(dialect: str) -> NativeProjectCatalog: + """Return a compile catalog, which normalizes batches on its analysis pool.""" + + return BindingCatalog( + dialect=dialect, + quoted_ignore_case=False, + known_functions=(), + known_types=(), + relations={}, + ).native + + +def without_catalog(dialect: str) -> None: + """Return no catalog, which normalizes batches in turn.""" + + _ = dialect diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py new file mode 100644 index 0000000000..98e6448a9d --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_NORMALIZATION_CHUNK_SIZE +from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql +from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( + normalize_analysis_sql_results, +) +from tests.unit.src.sqlbuild.compiler.sql_analysis.main._test_types import ( + NormalizationBatchFailureTestCase, + NormalizationBatchTestCase, + NormalizationChunkTestCase, +) +from tests.unit.src.sqlbuild.compiler.sql_analysis.main.helpers import ( + pool_catalog, + without_catalog, +) + + +@pytest.mark.parametrize( + "test_case", + ( + NormalizationBatchTestCase( + description="stubbed relations and placeholders in turn", + dialect="snowflake", + requests=( + ( + 'SELECT * FROM __ref("orders") JOIN __ref("customers") USING (customer_id)', + {"orders": "__sqb_rel_0", "customers": "__sqb_rel_1"}, + {}, + ), + ("SELECT payload:item[0] AS item FROM inventory", {}, {}), + ("SELECT 1 AS quantity LIMIT @@@limit", {}, {"limit": "10"}), + ), + catalog=without_catalog, + expected_stubbed_sql="SELECT * FROM __sqb_rel_0 JOIN __sqb_rel_1 USING (customer_id)", + ), + NormalizationBatchTestCase( + description="stubbed relations and placeholders on the catalog pool", + dialect="snowflake", + requests=( + ( + 'SELECT * FROM __ref("orders") JOIN __ref("customers") USING (customer_id)', + {"orders": "__sqb_rel_0", "customers": "__sqb_rel_1"}, + {}, + ), + ("SELECT payload:item[0] AS item FROM inventory", {}, {}), + ("SELECT 1 AS quantity LIMIT @@@limit", {}, {"limit": "10"}), + ), + catalog=pool_catalog, + expected_stubbed_sql="SELECT * FROM __sqb_rel_0 JOIN __sqb_rel_1 USING (customer_id)", + ), + ), + ids=lambda case: case.description, +) +def test_given_batch_when_normalizing_then_matches_single_normalizations_in_order( + test_case: NormalizationBatchTestCase, +) -> None: + expected: list[str | Exception] = [ + normalize_analysis_sql( + sql=sql, dialect=test_case.dialect, stubs=stubs, placeholders=placeholders + ) + for sql, stubs, placeholders in test_case.requests + ] + + results: list[str | Exception] = normalize_analysis_sql_results( + dialect=test_case.dialect, + requests=test_case.requests, + catalog=test_case.catalog(test_case.dialect), + ) + + assert results == expected + assert results[0] == test_case.expected_stubbed_sql + + +@pytest.mark.parametrize( + "test_case", + ( + NormalizationBatchFailureTestCase( + description="unterminated quote in turn", + dialect="snowflake", + before=("SELECT id FROM products",), + failing="SELECT 'unterminated FROM orders", + after=("SELECT /* unterminated FROM customers", "SELECT id FROM customers"), + catalog=without_catalog, + expected_error_type=ValueError, + expected_result_types=("str", "ValueError", "ValueError", "str"), + ), + NormalizationBatchFailureTestCase( + description="unterminated quote on the catalog pool", + dialect="snowflake", + before=("SELECT id FROM products",), + failing="SELECT 'unterminated FROM orders", + after=("SELECT /* unterminated FROM customers", "SELECT id FROM customers"), + catalog=pool_catalog, + expected_error_type=ValueError, + expected_result_types=("str", "ValueError", "ValueError", "str"), + ), + ), + ids=lambda case: case.description, +) +def test_given_failing_member_when_normalizing_batch_then_keeps_its_single_call_error_in_place( + test_case: NormalizationBatchFailureTestCase, +) -> None: + with pytest.raises(test_case.expected_error_type) as raised: + _ = normalize_analysis_sql(sql=test_case.failing, dialect=test_case.dialect) + + results: list[str | Exception] = normalize_analysis_sql_results( + dialect=test_case.dialect, + requests=[ + (sql, None, None) for sql in (*test_case.before, test_case.failing, *test_case.after) + ], + catalog=test_case.catalog(test_case.dialect), + ) + first_error: str | Exception = results[len(test_case.before)] + + assert results[: len(test_case.before)] == [ + normalize_analysis_sql(sql=sql, dialect=test_case.dialect) for sql in test_case.before + ] + assert type(first_error) is type(raised.value) + assert str(first_error) == str(raised.value) + assert tuple(type(result).__name__ for result in results) == test_case.expected_result_types + + +@pytest.mark.parametrize( + "test_case", + ( + NormalizationChunkTestCase( + description="failure on a chunk boundary in turn", + valid_count=ANALYSIS_NORMALIZATION_CHUNK_SIZE * 2 + 3, + failing_index=ANALYSIS_NORMALIZATION_CHUNK_SIZE, + catalog=without_catalog, + expected_error_type=ValueError, + ), + NormalizationChunkTestCase( + description="failure on a chunk boundary on the catalog pool", + valid_count=ANALYSIS_NORMALIZATION_CHUNK_SIZE * 2 + 3, + failing_index=ANALYSIS_NORMALIZATION_CHUNK_SIZE, + catalog=pool_catalog, + expected_error_type=ValueError, + ), + ), + ids=lambda case: case.description, +) +def test_given_batch_beyond_one_chunk_when_normalizing_then_keeps_every_member_in_order( + test_case: NormalizationChunkTestCase, +) -> None: + valid: list[str] = [ + f"SELECT {index} AS order_id FROM orders" for index in range(test_case.valid_count) + ] + sqls: list[str] = [ + *valid[: test_case.failing_index], + "SELECT 'unterminated FROM orders", + *valid[test_case.failing_index :], + ] + + results: list[str | Exception] = normalize_analysis_sql_results( + dialect="snowflake", + requests=[(sql, None, None) for sql in sqls], + catalog=test_case.catalog("snowflake"), + ) + + assert len(results) == len(sqls) + assert isinstance(results[test_case.failing_index], test_case.expected_error_type) + assert [*results[: test_case.failing_index], *results[test_case.failing_index + 1 :]] == valid + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"])