diff --git a/.github/workflows/verify.yml b/.github/workflows/verify.yml index b918513f98..8827d87d94 100644 --- a/.github/workflows/verify.yml +++ b/.github/workflows/verify.yml @@ -73,6 +73,7 @@ jobs: ;; src/sqlbuild/rule_engine/*|src/sqlbuild/lint/*|scripts/rules_benchmark/*|scripts/benchmark_rules_ci.py|tests/performance/*) broad=true + performance=true format_performance=true rules_performance=true ;; @@ -324,6 +325,53 @@ jobs: make test-e2e-cold-compile-performance \ SQLBUILD_BENCHMARK_MODELS="${BENCHMARK_MODELS}" + dense_compile_performance: + name: Dense all-rules compile (${{ matrix.models }} models) + needs: [changes, lockfile] + if: needs.changes.outputs.performance == 'true' + runs-on: ubuntu-24.04 + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + models: [1000, 3000, 5000, 10000] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v5 + with: + version: "0.12.3" + enable-cache: true + cache-dependency-glob: | + uv.lock + Cargo.lock + - uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + - name: Install Python + run: uv python install 3.12 + - name: Install dependencies + run: uv sync --locked --all-groups --all-extras + - name: Run dense all-rules guard under 4 GiB + env: + BENCHMARK_MODELS: ${{ matrix.models }} + MEMORY_LIMIT_BYTES: "4294967296" + run: | + sudo systemd-run \ + --unit="sqlbuild-dense-${BENCHMARK_MODELS}-${GITHUB_RUN_ID}" \ + --uid="$(id -u)" \ + --gid="$(id -g)" \ + --same-dir \ + --setenv="HOME=${HOME}" \ + --setenv="PATH=${PATH}" \ + --setenv="SQLBUILD_REQUIRE_MEMORY_LIMIT_BYTES=${MEMORY_LIMIT_BYTES}" \ + --property="MemoryMax=${MEMORY_LIMIT_BYTES}" \ + --property="MemorySwapMax=0" \ + --wait \ + --collect \ + --pipe \ + make test-e2e-dense-compile-performance \ + SQLBUILD_BENCHMARK_MODELS="${BENCHMARK_MODELS}" + cache_compile_performance: name: Fresh-process compile cache (${{ matrix.models }} models) needs: [changes, lockfile] @@ -464,6 +512,7 @@ jobs: format_performance, performance, cold_compile_performance, + dense_compile_performance, cache_compile_performance, rules_performance, postgres, @@ -479,6 +528,7 @@ jobs: FORMAT_PERFORMANCE_RESULT: ${{ needs.format_performance.result }} PERFORMANCE_RESULT: ${{ needs.performance.result }} COLD_COMPILE_PERFORMANCE_RESULT: ${{ needs.cold_compile_performance.result }} + DENSE_COMPILE_PERFORMANCE_RESULT: ${{ needs.dense_compile_performance.result }} CACHE_COMPILE_PERFORMANCE_RESULT: ${{ needs.cache_compile_performance.result }} RULES_PERFORMANCE_RESULT: ${{ needs.rules_performance.result }} POSTGRES_RESULT: ${{ needs.postgres.result }} @@ -494,6 +544,7 @@ jobs: "${FORMAT_PERFORMANCE_RESULT}" \ "${PERFORMANCE_RESULT}" \ "${COLD_COMPILE_PERFORMANCE_RESULT}" \ + "${DENSE_COMPILE_PERFORMANCE_RESULT}" \ "${CACHE_COMPILE_PERFORMANCE_RESULT}" \ "${RULES_PERFORMANCE_RESULT}" \ "${POSTGRES_RESULT}"; do diff --git a/Cargo.lock b/Cargo.lock index eef298b0cf..1127e3396b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -311,7 +311,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "polyglot-sql" version = "0.12.1" -source = "git+https://github.com/chio-labs/polyglot?rev=190b1f8265b6135a9a6a100356599dd99487c1b4#190b1f8265b6135a9a6a100356599dd99487c1b4" +source = "git+https://github.com/chio-labs/polyglot.git?rev=cd599007a1ba6c4cc1f2d99938fc0e6d3b376a9b#cd599007a1ba6c4cc1f2d99938fc0e6d3b376a9b" dependencies = [ "polyglot-sql-ast-derive", "serde", @@ -324,7 +324,7 @@ dependencies = [ [[package]] name = "polyglot-sql-ast-derive" version = "0.12.1" -source = "git+https://github.com/chio-labs/polyglot?rev=190b1f8265b6135a9a6a100356599dd99487c1b4#190b1f8265b6135a9a6a100356599dd99487c1b4" +source = "git+https://github.com/chio-labs/polyglot.git?rev=cd599007a1ba6c4cc1f2d99938fc0e6d3b376a9b#cd599007a1ba6c4cc1f2d99938fc0e6d3b376a9b" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 7ae47ac5fe..5a1970bbee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ fensu-policy = "=0.22.0" globset = "0.4" proc-macro2 = { version = "1", features = ["span-locations"] } pyo3 = { version = "0.26", features = ["abi3-py312"] } -polyglot-sql = { git = "https://github.com/chio-labs/polyglot", rev = "190b1f8265b6135a9a6a100356599dd99487c1b4", default-features = false, features = [ +polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "cd599007a1ba6c4cc1f2d99938fc0e6d3b376a9b", default-features = false, features = [ "ast-tools", "builder", "dialect-bigquery", diff --git a/Makefile b/Makefile index 11c675682d..052054a51c 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,8 @@ SHELL := /bin/bash test-e2e-duckdb-build-incremental test-e2e-duckdb-build-virtual \ test-e2e-duckdb-cli-data test-e2e-duckdb-cli test-e2e-duckdb-virtual \ test-e2e-duckdb-integrations test-e2e-performance \ - test-e2e-cold-compile-performance test-e2e-cache-compile-performance + test-e2e-cold-compile-performance test-e2e-cache-compile-performance \ + test-e2e-dense-compile-performance format: uv run ruff format . @@ -238,6 +239,14 @@ test-e2e-cold-compile-performance: -m cold_compile_performance -k "models_$(SQLBUILD_BENCHMARK_MODELS)" \ -vv --log-cli-level=INFO --color=yes +test-e2e-dense-compile-performance: + test -n "$(SQLBUILD_BENCHMARK_MODELS)" + env PYTHONUNBUFFERED=1 SQLBUILD_CONCURRENCY=4 uv run pytest \ + tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py \ + -n auto --dist loadfile \ + -m cold_compile_performance -k "models_$(SQLBUILD_BENCHMARK_MODELS)_" \ + -vv -rP --log-level=INFO --log-cli-level=INFO --color=yes + test-e2e-cache-compile-performance: test -n "$(SQLBUILD_BENCHMARK_MODELS)" env PYTHONUNBUFFERED=1 SQLBUILD_CONCURRENCY=4 uv run pytest \ diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs index e5761efd93..817bc8f596 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs @@ -21,8 +21,9 @@ fn value_error(error: impl std::fmt::Display) -> PyErr { } #[pyfunction] -fn evaluate_json(request_json: &str) -> PyResult { - evaluate::evaluate_json(request_json).map_err(value_error) +fn evaluate_json(py: Python<'_>, request_json: &str) -> PyResult { + py.detach(|| evaluate::evaluate_json(request_json)) + .map_err(value_error) } #[pyfunction] @@ -31,6 +32,31 @@ fn lint_sql_json(py: Python<'_>, request_json: &str) -> PyResult { .map_err(value_error) } +#[pyfunction] +fn finalize_rule_findings_json(py: Python<'_>, request_json: &str) -> PyResult { + py.detach(|| crate::engine::main::finalize::finalize_findings_json(request_json)) + .map_err(value_error) +} + +#[pyfunction] +fn prepare_lint_sql( + py: Python<'_>, + expanded: &str, + before_expansion: &str, + prior_sites: Vec, +) -> PyResult> { + py.detach(|| { + crate::sql_lint::main::preparation::prepare(expanded, before_expansion, &prior_sites) + }) + .map_err(value_error) +} + +#[pyfunction] +fn lint_sql_batch_json(py: Python<'_>, request_json: &str) -> PyResult { + py.detach(|| crate::sql_lint::main::batch_engine::lint_batch_json(request_json)) + .map_err(pyo3::exceptions::PyValueError::new_err) +} + #[pyfunction] fn format_sql_json(py: Python<'_>, request_json: &str) -> PyResult { py.detach(|| crate::sql_lint::main::formatter::format_json(request_json)) @@ -63,20 +89,24 @@ fn analyze_sql_uses_json(py: Python<'_>, request_json: &str) -> PyResult #[pyfunction] fn analyze_queries_json(py: Python<'_>, request_json: &str) -> PyResult { - py.detach(|| crate::query_analysis::main::analyze_json(request_json)) + py.detach(|| crate::query_analysis::main::analyze::analyze_json(request_json)) .map_err(value_error) } #[pyfunction] fn analyze_project_queries_json(py: Python<'_>, request_json: &str) -> PyResult { - py.detach(|| crate::query_analysis::main::analyze_project_json(request_json)) + py.detach(|| crate::query_analysis::main::analyze_project::analyze_project_json(request_json)) .map_err(value_error) } #[pyfunction] fn analyze_project_queries_compact_json(py: Python<'_>, request_json: &str) -> PyResult { - py.detach(|| crate::query_analysis::main::analyze_project_compact_json(request_json)) - .map_err(value_error) + py.detach(|| { + crate::query_analysis::main::analyze_project_compact::analyze_project_compact_json( + request_json, + ) + }) + .map_err(value_error) } #[pyfunction] @@ -215,6 +245,21 @@ fn tokenize_model_header(header: &str) -> PyResult> { crate::compiler::main::model_header_tokenizing::tokenize_one(header).map_err(value_error) } +#[pyfunction] +fn substitute_static_project_vars( + sqls: Vec, + variables: Vec<(String, String)>, +) -> Vec<(u8, Option)> { + crate::compiler::main::sql_interpolation::substitute_batch(&sqls, &variables) +} + +#[pyfunction] +fn extract_static_sql_references( + sql: &str, +) -> Option> { + crate::compiler::main::sql_references::extract(sql) +} + #[pyfunction] fn load_config_json(project_dir: &str) -> PyResult { load::load_config_json(std::path::Path::new(project_dir)).map_err(value_error) @@ -259,7 +304,10 @@ fn skill_freshness(content: Option<&str>, input_fingerprint: &str) -> String { pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(evaluate_json, module)?)?; + module.add_function(wrap_pyfunction!(finalize_rule_findings_json, module)?)?; module.add_function(wrap_pyfunction!(lint_sql_json, module)?)?; + module.add_function(wrap_pyfunction!(prepare_lint_sql, module)?)?; + module.add_function(wrap_pyfunction!(lint_sql_batch_json, module)?)?; module.add_function(wrap_pyfunction!(format_sql_json, module)?)?; module.add_function(wrap_pyfunction!(format_sql_batch_json, module)?)?; module.add_function(wrap_pyfunction!(schema_validation_json, module)?)?; @@ -276,6 +324,8 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(extract_sql_tests_json, module)?)?; module.add_function(wrap_pyfunction!(parse_model_headers, module)?)?; module.add_function(wrap_pyfunction!(tokenize_model_header, module)?)?; + module.add_function(wrap_pyfunction!(substitute_static_project_vars, module)?)?; + module.add_function(wrap_pyfunction!(extract_static_sql_references, module)?)?; module.add_function(wrap_pyfunction!(load_config_json, module)?)?; module.add_function(wrap_pyfunction!(catalogue_json, module)?)?; module.add_function(wrap_pyfunction!(selected_codes_json, module)?)?; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/mod.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/mod.rs index 5a6f7720e6..d4ab4655d5 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/mod.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/mod.rs @@ -1,2 +1,4 @@ pub(crate) mod model_headers; +pub(crate) mod sql_interpolation; +pub(crate) mod sql_references; pub(crate) mod sql_tests; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_interpolation/mod.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_interpolation/mod.rs new file mode 100644 index 0000000000..4d943c2ae2 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_interpolation/mod.rs @@ -0,0 +1 @@ +pub(crate) mod substitution; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_interpolation/substitution.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_interpolation/substitution.rs new file mode 100644 index 0000000000..8392ecf81c --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_interpolation/substitution.rs @@ -0,0 +1,164 @@ +//! Conservative native fast path for compile-wide scalar SQL variables. + +use std::collections::HashMap; + +pub(crate) const UNCHANGED: u8 = 0; +pub(crate) const SUBSTITUTED: u8 = 1; +pub(crate) const FALLBACK: u8 = 2; + +struct SubstitutionState { + output: Option, + copy_start: usize, +} + +pub(crate) fn substitute_batch( + sqls: &[String], + variables: &[(String, String)], +) -> Vec<(u8, Option)> { + let variables: HashMap<&str, &str> = variables + .iter() + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect(); + sqls.iter() + .map(|sql| substitute_one(sql, &variables)) + .collect() +} + +fn substitute_one(sql: &str, variables: &HashMap<&str, &str>) -> (u8, Option) { + if !sql.contains("@@") { + return (UNCHANGED, None); + } + + let bytes = sql.as_bytes(); + let mut state = SubstitutionState { + output: None, + copy_start: 0, + }; + let mut index = 0; + let mut quote: Option = None; + while index < bytes.len() { + if let Some(quote_byte) = quote { + if bytes[index] == quote_byte { + if matches!(quote_byte, b'\'' | b'"') + && bytes.get(index + 1).copied() == Some(quote_byte) + { + index += 2; + continue; + } + quote = None; + index += 1; + continue; + } + if bytes[index..].starts_with(b"@@") { + let Some((end, next_state)) = substitute_token(sql, index, variables, state) else { + return (FALLBACK, None); + }; + state = next_state; + index = end; + continue; + } + index += 1; + continue; + } + + if bytes[index..].starts_with(b"--") { + index = sql[index + 2..] + .find('\n') + .map_or(bytes.len(), |offset| index + 3 + offset); + continue; + } + if bytes[index..].starts_with(b"/*") { + let Some(offset) = sql[index + 2..].find("*/") else { + return (FALLBACK, None); + }; + index += offset + 4; + continue; + } + if matches!(bytes[index], b'\'' | b'"' | b'`') { + quote = Some(bytes[index]); + index += 1; + continue; + } + if bytes[index..].starts_with(b"@@") { + let Some((end, next_state)) = substitute_token(sql, index, variables, state) else { + return (FALLBACK, None); + }; + state = next_state; + index = end; + continue; + } + index += 1; + } + if quote.is_some() { + return (FALLBACK, None); + } + match state.output { + Some(mut rendered) => { + rendered.push_str(&sql[state.copy_start..]); + (SUBSTITUTED, Some(rendered)) + } + None => (UNCHANGED, None), + } +} + +fn substitute_token( + sql: &str, + start: usize, + variables: &HashMap<&str, &str>, + mut state: SubstitutionState, +) -> Option<(usize, SubstitutionState)> { + let bytes = sql.as_bytes(); + if bytes[start..].starts_with(b"@@@") { + let name_start = start + 3; + let end = if bytes + .get(name_start) + .copied() + .is_some_and(is_identifier_start) + { + consume_ascii_identifier(bytes, name_start)? + } else { + name_start + }; + return Some((end, state)); + } + let name_start = start + 2; + if bytes[name_start..].starts_with(b"ENV:") || bytes[name_start..].starts_with(b"CTX:") { + return None; + } + if !bytes + .get(name_start) + .copied() + .is_some_and(is_identifier_start) + { + return None; + } + let name_end = consume_ascii_identifier(bytes, name_start)?; + let name = &sql[name_start..name_end]; + let value = variables.get(name)?; + let rendered = state + .output + .get_or_insert_with(|| String::with_capacity(sql.len())); + rendered.push_str(&sql[state.copy_start..start]); + rendered.push_str(value); + state.copy_start = name_end; + Some((name_end, state)) +} + +fn consume_ascii_identifier(bytes: &[u8], start: usize) -> Option { + let mut end = start + 1; + while bytes.get(end).copied().is_some_and(is_identifier_continue) { + end += 1; + } + if bytes.get(end).is_some_and(|value| !value.is_ascii()) { + return None; + } + Some(end) +} + +fn is_identifier_start(value: u8) -> bool { + value.is_ascii_alphabetic() || value == b'_' +} + +fn is_identifier_continue(value: u8) -> bool { + value.is_ascii_alphanumeric() || value == b'_' +} diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs new file mode 100644 index 0000000000..c7b0a8fd77 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs @@ -0,0 +1,123 @@ +//! Conservative native fast path for logical SQL reference extraction. + +use crate::constants::{DBT_REFERENCE_KIND, TABLE_FUNCTION_REFERENCE_KIND}; + +const PREFIXES: [(&str, &str); 6] = [ + ("__dbt_ref(", "dbt_ref"), + ("__table_fn(", "table_function"), + ("__source(", "source"), + ("__seed(", "seed"), + ("__udf(", "udf"), + ("__ref(", "ref"), +]; + +pub(crate) type StaticReference = (String, String, Option, Option); + +pub(crate) fn extract(sql: &str) -> Option> { + let bytes = sql.as_bytes(); + let mut references: Vec = Vec::new(); + let mut index = 0; + while index < bytes.len() { + if bytes[index..].starts_with(b"--") { + index = sql[index + 2..] + .find('\n') + .map_or(bytes.len(), |offset| index + offset + 3); + continue; + } + if bytes[index..].starts_with(b"/*") { + let offset = sql[index + 2..].find("*/")?; + index += offset + 4; + continue; + } + if matches!(bytes[index], b'\'' | b'"' | b'`') { + index = quoted_end(bytes, index)?; + continue; + } + let Some((prefix, kind)) = PREFIXES + .iter() + .find(|(prefix, _)| bytes[index..].starts_with(prefix.as_bytes())) + else { + index += 1; + continue; + }; + if *kind == TABLE_FUNCTION_REFERENCE_KIND { + return None; + } + let (first_name, first_end) = reference_name(sql, index + prefix.len())?; + let mut cursor = ascii_whitespace_end(bytes, first_end); + let mut package = None; + let name; + if *kind == DBT_REFERENCE_KIND && bytes.get(cursor).copied() == Some(b',') { + cursor = ascii_whitespace_end(bytes, cursor + 1); + let (second_name, second_end) = reference_name(sql, cursor)?; + cursor = ascii_whitespace_end(bytes, second_end); + package = Some(first_name); + name = second_name; + } else { + name = first_name; + } + if bytes.get(cursor).copied() != Some(b')') { + return None; + } + references.push(((*kind).to_owned(), name, package, None)); + index = cursor + 1; + } + Some(references) +} + +fn reference_name(sql: &str, start: usize) -> Option<(String, usize)> { + let bytes = sql.as_bytes(); + let start = ascii_whitespace_end(bytes, start); + let first = bytes.get(start).copied()?; + if matches!(first, b'\'' | b'"') { + let end = quoted_end(bytes, start)?; + return Some((sql[start + 1..end - 1].to_owned(), end)); + } + if !is_identifier_start(first) { + return None; + } + let mut end = start + 1; + while bytes.get(end).copied().is_some_and(is_identifier_continue) { + end += 1; + } + if bytes.get(end).is_some_and(|value| !value.is_ascii()) { + return None; + } + Some((sql[start..end].to_owned(), end)) +} + +fn quoted_end(bytes: &[u8], start: usize) -> Option { + let quote = bytes[start]; + let mut index = start + 1; + while index < bytes.len() { + if bytes[index] != quote { + index += 1; + continue; + } + if matches!(quote, b'\'' | b'"') && bytes.get(index + 1).copied() == Some(quote) { + index += 2; + continue; + } + return Some(index + 1); + } + None +} + +fn ascii_whitespace_end(bytes: &[u8], start: usize) -> usize { + let mut end = start; + while bytes + .get(end) + .is_some_and(|value| value.is_ascii_whitespace()) + { + end += 1; + } + end +} + +fn is_identifier_start(value: u8) -> bool { + value.is_ascii_alphabetic() +} + +fn is_identifier_continue(value: u8) -> bool { + value.is_ascii_alphanumeric() || value == b'_' +} diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/mod.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/mod.rs new file mode 100644 index 0000000000..a19c34f070 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/mod.rs @@ -0,0 +1 @@ +pub(crate) mod extraction; diff --git a/crates/sqlbuild-rules-native/src/compiler/main/mod.rs b/crates/sqlbuild-rules-native/src/compiler/main/mod.rs index fe5931f2a7..a38f94aff9 100644 --- a/crates/sqlbuild-rules-native/src/compiler/main/mod.rs +++ b/crates/sqlbuild-rules-native/src/compiler/main/mod.rs @@ -1,5 +1,7 @@ pub(crate) mod model_header_parsing; pub(crate) mod model_header_tokenizing; +pub(crate) mod sql_interpolation; +pub(crate) mod sql_references; pub(crate) mod sql_test_extraction; pub(crate) mod sql_test_planning; pub(crate) mod sql_test_rendering; diff --git a/crates/sqlbuild-rules-native/src/compiler/main/sql_interpolation.rs b/crates/sqlbuild-rules-native/src/compiler/main/sql_interpolation.rs new file mode 100644 index 0000000000..471d2be4b8 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/main/sql_interpolation.rs @@ -0,0 +1,8 @@ +//! Public compiler entry point for conservative scalar SQL interpolation. + +pub(crate) fn substitute_batch( + sqls: &[String], + variables: &[(String, String)], +) -> Vec<(u8, Option)> { + crate::compiler::_helpers::sql_interpolation::substitution::substitute_batch(sqls, variables) +} diff --git a/crates/sqlbuild-rules-native/src/compiler/main/sql_references.rs b/crates/sqlbuild-rules-native/src/compiler/main/sql_references.rs new file mode 100644 index 0000000000..e9110f3fed --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/main/sql_references.rs @@ -0,0 +1,7 @@ +//! Public compiler entry point for conservative logical SQL references. + +use crate::compiler::_helpers::sql_references::extraction::StaticReference; + +pub(crate) fn extract(sql: &str) -> Option> { + crate::compiler::_helpers::sql_references::extraction::extract(sql) +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests.rs b/crates/sqlbuild-rules-native/src/compiler/tests.rs index 81acd88c77..0f31c82325 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests.rs @@ -2,6 +2,10 @@ mod helpers; #[path = "tests/test_model_header_tokenization.rs"] mod model_header_tokenization; +#[path = "tests/test_sql_interpolation.rs"] +mod sql_interpolation; +#[path = "tests/test_sql_references.rs"] +mod sql_references; #[path = "tests/test_sql_test_extraction.rs"] mod sql_test_extraction; #[path = "tests/test_sql_test_planning.rs"] diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs index 79abf7a186..840b7174f2 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs @@ -7,9 +7,79 @@ use serde_json::{Value, json}; use crate::compiler::_helpers::model_headers::tokenization::{ MAX_TOKENIZER_WORKERS, TOKENIZER_WORKER_STACK_BYTES, build_tokenizer_pool, parse_batch, }; +use crate::compiler::_helpers::sql_interpolation::substitution::{ + FALLBACK, SUBSTITUTED, UNCHANGED, substitute_batch, +}; +use crate::compiler::_helpers::sql_references::extraction::extract; use crate::compiler::main::sql_test_extraction::extract_batch_json; use crate::compiler::models::AuthoredValue; +pub(crate) fn scalar_variables_preserve_lexical_boundaries() -> bool { + let sqls = vec![ + "SELECT @@revision, '@@status', @@@window_start".to_owned(), + "-- @@revision\nSELECT /* @@status */ 1".to_owned(), + "SELECT '@@revision''s'".to_owned(), + ]; + substitute_batch( + &sqls, + &[ + ("revision".to_owned(), "7".to_owned()), + ("status".to_owned(), "ready".to_owned()), + ], + ) == vec![ + ( + SUBSTITUTED, + Some("SELECT 7, 'ready', @@@window_start".to_owned()), + ), + (UNCHANGED, None), + (SUBSTITUTED, Some("SELECT '7''s'".to_owned())), + ] +} + +pub(crate) fn dynamic_or_malformed_sql_requests_fallback() -> bool { + let sqls = vec![ + "SELECT @@ENV:USER".to_owned(), + "SELECT @@missing".to_owned(), + "SELECT @@revision, 'unterminated".to_owned(), + "SELECT @@revision /* unterminated".to_owned(), + "SELECT @@révision".to_owned(), + ]; + substitute_batch(&sqls, &[("revision".to_owned(), "7".to_owned())]) + == vec![(FALLBACK, None); sqls.len()] +} + +pub(crate) fn simple_references_preserve_authored_order() -> bool { + extract( + "SELECT * FROM __source('orders') UNION ALL SELECT * FROM __dbt_ref(\"shop\", \"customers\")", + ) == Some(vec![ + ("source".to_owned(), "orders".to_owned(), None, None), + ( + "dbt_ref".to_owned(), + "customers".to_owned(), + Some("shop".to_owned()), + None, + ), + ]) +} + +pub(crate) fn comments_and_quoted_text_hide_references() -> bool { + extract("-- __ref(\"ignored\")\nSELECT '__seed(\"also_ignored\")' FROM __ref(orders)") + == Some(vec![("ref".to_owned(), "orders".to_owned(), None, None)]) +} + +pub(crate) fn complex_or_malformed_sql_requests_fallback() -> bool { + [ + "SELECT * FROM __table_fn(\"orders\")(1)", + "SELECT * FROM __ref(concat('ord', 'ers'))", + "SELECT * FROM __ref(\"orders\"", + "SELECT * FROM __ref(\"orders\") /* unterminated", + "SELECT * FROM __ref(\"orders\") WHERE note = 'unterminated", + "SELECT * FROM __ref(örders)", + ] + .into_iter() + .all(|sql| extract(sql).is_none()) +} + pub(crate) fn nested_authored_headers_preserve_values_and_offsets() -> bool { let headers = vec![ " columns (café (type DECIMAL(10,2))), enabled true".to_owned(), diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_interpolation.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_interpolation.rs new file mode 100644 index 0000000000..fc13455259 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_interpolation.rs @@ -0,0 +1,29 @@ +use crate::compiler::tests::helpers::{ + dynamic_or_malformed_sql_requests_fallback, scalar_variables_preserve_lexical_boundaries, +}; +use crate::compiler::tests::test_types::StaticSqlOperationTestCase; + +#[test] +fn given_sql_interpolation_cases_when_substituting_then_expected_behavior_holds() { + let test_cases = [ + StaticSqlOperationTestCase { + description: "scalar variables preserve lexical boundaries", + run: scalar_variables_preserve_lexical_boundaries, + expected_success: true, + }, + StaticSqlOperationTestCase { + description: "dynamic and malformed SQL requests Python fallback", + run: dynamic_or_malformed_sql_requests_fallback, + expected_success: true, + }, + ]; + + for test_case in test_cases { + let actual_success = (test_case.run)(); + assert_eq!( + actual_success, test_case.expected_success, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_references.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_references.rs new file mode 100644 index 0000000000..cda5c3de96 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_references.rs @@ -0,0 +1,35 @@ +use crate::compiler::tests::helpers::{ + comments_and_quoted_text_hide_references, complex_or_malformed_sql_requests_fallback, + simple_references_preserve_authored_order, +}; +use crate::compiler::tests::test_types::StaticSqlOperationTestCase; + +#[test] +fn given_sql_reference_cases_when_extracting_then_expected_behavior_holds() { + let test_cases = [ + StaticSqlOperationTestCase { + description: "simple references preserve authored order", + run: simple_references_preserve_authored_order, + expected_success: true, + }, + StaticSqlOperationTestCase { + description: "comments and quoted text hide embedded references", + run: comments_and_quoted_text_hide_references, + expected_success: true, + }, + StaticSqlOperationTestCase { + description: "complex and malformed SQL requests Python fallback", + run: complex_or_malformed_sql_requests_fallback, + expected_success: true, + }, + ]; + + for test_case in test_cases { + let actual_success = (test_case.run)(); + assert_eq!( + actual_success, test_case.expected_success, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs index bd30bbf418..33c12780a5 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs @@ -4,6 +4,12 @@ pub(crate) struct ModelHeaderTokenizationTestCase { pub(crate) expected_success: bool, } +pub(crate) struct StaticSqlOperationTestCase { + pub(crate) description: &'static str, + pub(crate) run: fn() -> bool, + pub(crate) expected_success: bool, +} + pub(crate) struct SqlTestExtractionTestCase { pub(crate) description: &'static str, pub(crate) run: fn() -> bool, diff --git a/crates/sqlbuild-rules-native/src/constants.rs b/crates/sqlbuild-rules-native/src/constants.rs index 924ff6e398..1552e3d84c 100644 --- a/crates/sqlbuild-rules-native/src/constants.rs +++ b/crates/sqlbuild-rules-native/src/constants.rs @@ -11,6 +11,8 @@ pub(crate) const CUSTOM_RULE_NAMESPACE: &str = "XSQBR"; pub(crate) const TARGET_DIRECTORY: &str = "target"; pub(crate) const REFERENCE_KIND: &str = "ref"; pub(crate) const SOURCE_REFERENCE_KIND: &str = "source"; +pub(crate) const DBT_REFERENCE_KIND: &str = "dbt_ref"; +pub(crate) const TABLE_FUNCTION_REFERENCE_KIND: &str = "table_function"; pub(crate) const VIEW_MATERIALIZATION: &str = "view"; pub(crate) const ENFORCED_CONTRACT: &str = "enforced"; pub(crate) const SELECT_STAR_MODEL_RULE_CODE: &str = "SQBRMODEL102"; diff --git a/crates/sqlbuild-rules-native/src/engine/_helpers/evaluation.rs b/crates/sqlbuild-rules-native/src/engine/_helpers/evaluation.rs index 4df790994b..40d43b3f9b 100644 --- a/crates/sqlbuild-rules-native/src/engine/_helpers/evaluation.rs +++ b/crates/sqlbuild-rules-native/src/engine/_helpers/evaluation.rs @@ -4,12 +4,16 @@ use crate::constants::{ RULES_DIRECTORY, TARGET_DIRECTORY, }; use crate::engine::_helpers::cache::{Cache, RuleCacheBucket, RuleCacheEntry}; -use crate::models::{EvaluateRequest, EvaluateResponse, Fault, RuleMetadata, RulesCodeGrammar}; +use crate::models::{ + EvaluateRequest, EvaluateResponse, Fault, Model, RuleMetadata, RulesCodeGrammar, +}; use crate::rules::main::{ assemble_catalogue, evaluate as rules, evaluate_project, fingerprint, resolve_threshold_overrides, select, }; -use crate::rules::models::{ModelEvaluationRequest, ProjectEvaluationRequest}; +use crate::rules::models::{ + ModelEvaluationRequest, ProjectEvaluationRequest, ResolvedThresholdOverride, +}; use fensu_policy::lifecycle::constants::ANALYSIS_BATCH_SCHEMA_VERSION; use fensu_policy::lifecycle::errors::LifecycleError; use fensu_policy::lifecycle::models::{ @@ -18,11 +22,166 @@ use fensu_policy::lifecycle::models::{ RuntimeIdentity, ScopedIgnore, }; use fensu_policy::{apply_suppressions, evaluate_batch, run_custom_host}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, BTreeSet}; use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; +const NATIVE_RULE_WORKERS: usize = 4; +const NATIVE_RULE_STACK_BYTES: usize = 16 * 1024 * 1024; + +struct PendingModelRule<'a> { + index: usize, + model: &'a Model, + identity: Option, +} + +struct CompletedModelRule { + index: usize, + identity: Option, + faults: Result, String>, +} + +struct ModelRuleBatchContext<'a> { + request: &'a EvaluateRequest, + selected: &'a BTreeMap, + threshold_overrides: &'a [ResolvedThresholdOverride], +} + +struct NativeModelRuleRequest<'a> { + context: ModelRuleBatchContext<'a>, + cache: Option<&'a Cache>, + ruleset_fingerprint: &'a str, + project_fingerprint: Option<&'a str>, +} + +#[derive(Default)] +struct NativeModelEvaluation { + faults: Vec, + hits: usize, + misses: usize, +} + +fn evaluate_native_models( + input: NativeModelRuleRequest<'_>, +) -> Result { + let NativeModelRuleRequest { + context, + cache, + ruleset_fingerprint, + project_fingerprint, + } = input; + if !context + .selected + .values() + .any(|rule| !rule.custom && !rule.code.starts_with("SQBRSQL")) + { + return Ok(NativeModelEvaluation::default()); + } + let mut models: Vec<_> = context.request.models.iter().collect(); + models.sort_by(|left, right| left.relative_path.cmp(&right.relative_path)); + let all_models_required = context.selected.values().any(|rule| { + !rule.custom + && !rule.code.starts_with("SQBRSQL") + && !matches!( + rule.code.as_str(), + "SQBRDECLARATION201" | "SQBRDECLARATION301" | "SQBRTEST301" + ) + }); + let model_limit = if all_models_required { models.len() } else { 1 }; + let mut bucket = cache + .map(|store| store.native_rules_bucket(ruleset_fingerprint)) + .transpose()? + .unwrap_or_default(); + let mut bucket_changed = false; + let mut result = NativeModelEvaluation::default(); + let mut model_results: Vec>> = vec![None; model_limit]; + let mut pending_models: Vec> = Vec::new(); + for (index, model) in models.iter().take(model_limit).enumerate() { + let identity = cache + .map(|_| { + model_cache_identity( + model, + ruleset_fingerprint, + if index == 0 { + project_fingerprint + } else { + None + }, + ) + }) + .transpose()?; + if let Some(identity) = &identity + && let Some(entry) = bucket + .entries + .get(&model.relative_path) + .filter(|entry| entry.identity == *identity) + { + model_results[index] = Some(entry.faults.clone()); + result.hits += 1; + continue; + } + result.misses += 1; + pending_models.push(PendingModelRule { + index, + model, + identity, + }); + } + let completed = evaluate_pending_models(pending_models, context)?; + for completed_model in completed { + let faults = completed_model.faults?; + if let Some(identity) = completed_model.identity { + let _ = bucket.entries.insert( + models[completed_model.index].relative_path.clone(), + RuleCacheEntry { + identity, + faults: faults.clone(), + }, + ); + bucket_changed = true; + } + model_results[completed_model.index] = Some(faults); + } + result.faults = model_results.into_iter().flatten().flatten().collect(); + if bucket_changed && let Some(cache) = cache { + cache.put_native_rules_bucket(ruleset_fingerprint, &bucket)?; + } + Ok(result) +} + +fn evaluate_pending_models( + pending_models: Vec>, + context: ModelRuleBatchContext<'_>, +) -> Result, String> { + let evaluate_pending = |pending: PendingModelRule<'_>| CompletedModelRule { + index: pending.index, + identity: pending.identity, + faults: rules::evaluate_model(ModelEvaluationRequest { + model: pending.model, + config: &context.request.config, + selected: context.selected, + request: context.request, + threshold_overrides: context.threshold_overrides, + }), + }; + if pending_models.len() <= 1 { + return Ok(pending_models.into_iter().map(evaluate_pending).collect()); + } + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(NATIVE_RULE_WORKERS.min(pending_models.len())) + .stack_size(NATIVE_RULE_STACK_BYTES) + .build() + .map_err(|error| error.to_string())?; + Ok(pool.install(|| { + pending_models + .into_par_iter() + .map(evaluate_pending) + .collect() + })) +} + pub(crate) fn evaluate_json(request_json: &str) -> Result { let request: EvaluateRequest = serde_json::from_str(request_json) .map_err(|error| format!("invalid rules request: {error}"))?; @@ -46,7 +205,9 @@ pub(crate) fn evaluate_json(request_json: &str) -> Result { resolve_threshold_overrides::resolve_threshold_overrides(&request.config)?; if selected.is_empty() { validate_exception_paths(&request)?; - let _ = suppress_faults(&request, &selected, Vec::new())?; + if !request.defer_suppressions { + let _ = suppress_faults(&request, &selected, Vec::new())?; + } return serde_json::to_string(&EvaluateResponse { version: API_VERSION, faults: Vec::new(), @@ -78,110 +239,89 @@ pub(crate) fn evaluate_json(request_json: &str) -> Result { let cache = cache_enabled .then(|| Cache::open(Path::new(&request.project_dir))) .transpose()?; - let mut models: Vec<_> = request.models.iter().collect(); - models.sort_by(|left, right| left.relative_path.cmp(&right.relative_path)); - let built_in_started = Instant::now(); - let mut raw_faults = request.initial_findings.clone(); - raw_faults.extend(evaluate_project::evaluate_project( - ProjectEvaluationRequest { - selected: &selected_by_code, - request: &request, - }, - )?); - let mut cache_hits = 0; - let mut cache_misses = 0; - let native_model_selected = selected - .iter() - .any(|rule| !rule.custom && !rule.code.starts_with("SQBRSQL")); - if native_model_selected { - let all_models_required = selected.iter().any(|rule| { - !rule.custom - && !rule.code.starts_with("SQBRSQL") - && !matches!( - rule.code.as_str(), - "SQBRDECLARATION201" | "SQBRDECLARATION301" | "SQBRTEST301" - ) - }); - let native_model_limit = if all_models_required { models.len() } else { 1 }; - let mut bucket = cache - .as_ref() - .map(|store| store.native_rules_bucket(&ruleset_fingerprint)) - .transpose()? - .unwrap_or_default(); - let mut bucket_changed = false; - for (model_index, model) in models.iter().take(native_model_limit).enumerate() { - let fingerprint = cache - .as_ref() - .map(|_| { - model_cache_identity( - model, - &ruleset_fingerprint, - if model_index == 0 { - project_fingerprint.as_deref() - } else { - None - }, - ) - }) - .transpose()?; - if let Some(identity) = &fingerprint - && let Some(entry) = bucket - .entries - .get(&model.relative_path) - .filter(|entry| entry.identity == *identity) - { - raw_faults.extend(entry.faults.clone()); - cache_hits += 1; - continue; - } - cache_misses += 1; - let model_faults = rules::evaluate_model(ModelEvaluationRequest { - model, - config: &request.config, - selected: &selected_by_code, - request: &request, - threshold_overrides: &threshold_overrides, - })?; - if let Some(identity) = fingerprint { - let _ = bucket.entries.insert( - model.relative_path.clone(), - RuleCacheEntry { - identity, - faults: model_faults.clone(), - }, - ); - bucket_changed = true; - } - raw_faults.extend(model_faults); - } - if bucket_changed && let Some(store) = &cache { - store.put_native_rules_bucket(&ruleset_fingerprint, &bucket)?; - } - } - let built_in_ms = built_in_started.elapsed().as_millis() as u64; - let custom_started = Instant::now(); - let custom = match evaluate_custom_rules_cached(CustomRulesCacheRequest { + let (cached_custom, custom_probe_ms) = probe_custom_rules_cache(CustomRulesCacheRequest { request: &request, selected: &selected_by_code, cache: cache.as_ref(), project_fingerprint: project_fingerprint.as_deref(), - }) { - Ok(value) => value, - Err(error) if error == CUSTOM_HOST_REQUIRED_ERROR => { - return Err(format!("{error} [native_cache_misses={cache_misses}]")); - } - Err(error) => return Err(error), - }; - cache_hits += custom.hits; - cache_misses += custom.misses; - raw_faults.extend(custom.faults); - let custom_ms = custom_started.elapsed().as_millis() as u64; + })?; + let (raw_faults, cache_hits, cache_misses, built_in_ms, custom_ms) = + std::thread::scope(|scope| { + let custom_parallel_started = Instant::now(); + let custom_task = cached_custom.is_none().then(|| { + scope.spawn(|| { + evaluate_custom_rules_cached(CustomRulesCacheRequest { + request: &request, + selected: &selected_by_code, + cache: cache.as_ref(), + project_fingerprint: project_fingerprint.as_deref(), + }) + }) + }); + let built_in_started = Instant::now(); + let mut raw_faults = request.initial_findings.clone(); + raw_faults.extend(evaluate_project::evaluate_project( + ProjectEvaluationRequest { + selected: &selected_by_code, + request: &request, + }, + )?); + let native = evaluate_native_models(NativeModelRuleRequest { + context: ModelRuleBatchContext { + request: &request, + selected: &selected_by_code, + threshold_overrides: &threshold_overrides, + }, + cache: cache.as_ref(), + ruleset_fingerprint: &ruleset_fingerprint, + project_fingerprint: project_fingerprint.as_deref(), + })?; + let mut cache_hits = native.hits; + let mut cache_misses = native.misses; + raw_faults.extend(native.faults); + let built_in_ms = built_in_started.elapsed().as_millis() as u64; + let custom_started = if custom_task.is_some() { + custom_parallel_started + } else { + Instant::now() + }; + let custom_result = if let Some(task) = custom_task { + task.join() + .map_err(|_| "custom rule evaluation worker panicked".to_string())? + } else { + match cached_custom { + Some(custom) => Ok(custom), + None => evaluate_custom_rules_cached(CustomRulesCacheRequest { + request: &request, + selected: &selected_by_code, + cache: cache.as_ref(), + project_fingerprint: project_fingerprint.as_deref(), + }), + } + }; + let custom = match custom_result { + Ok(value) => value, + Err(error) if error == CUSTOM_HOST_REQUIRED_ERROR => { + return Err(format!("{error} [native_cache_misses={cache_misses}]")); + } + Err(error) => return Err(error), + }; + cache_hits += custom.hits; + cache_misses += custom.misses; + raw_faults.extend(custom.faults); + let custom_ms = custom_probe_ms + custom_started.elapsed().as_millis() as u64; + Ok::<_, String>((raw_faults, cache_hits, cache_misses, built_in_ms, custom_ms)) + })?; validate_exception_paths(&request)?; - let faults = suppress_faults(&request, &selected, raw_faults)?; + let faults = if request.defer_suppressions { + raw_faults + } else { + suppress_faults(&request, &selected, raw_faults)? + }; serde_json::to_string(&EvaluateResponse { version: API_VERSION, faults, - evaluated_models: models.len(), + evaluated_models: request.models.len(), cache_hits, cache_misses, ruleset_fingerprint, @@ -204,6 +344,18 @@ struct CustomEvaluation { misses: usize, } +fn probe_custom_rules_cache( + request: CustomRulesCacheRequest<'_>, +) -> Result<(Option, u64), String> { + let started = Instant::now(); + let cached = if request.request.custom_host.is_none() { + Some(evaluate_custom_rules_cached(request)?) + } else { + None + }; + Ok((cached, started.elapsed().as_millis() as u64)) +} + type LocalRuleMiss<'a> = ( &'a RuleMetadata, String, @@ -639,14 +791,42 @@ fn suppress_faults( selected: &[&RuleMetadata], faults: Vec, ) -> Result, String> { - let findings = faults - .into_iter() - .map(fault_to_finding) - .collect::, _>>()?; let evaluated_codes = selected .iter() .map(|rule| rule.code.clone()) .collect::>(); + apply_fault_policy(request, &evaluated_codes, faults) +} + +pub(crate) fn finalize_findings_json(request_json: &str) -> Result { + let input: crate::models::FinalizeFindingsRequest = serde_json::from_str(request_json) + .map_err(|error| format!("invalid findings request: {error}"))?; + if input.version != API_VERSION { + return Err(format!( + "unsupported rules native API version {}; expected {API_VERSION}", + input.version + )); + } + config::validate(&input.config)?; + let request = EvaluateRequest { + project_dir: input.project_dir, + config: input.config, + ..Default::default() + }; + validate_exception_paths(&request)?; + let faults = apply_fault_policy(&request, &input.evaluated_codes, input.findings)?; + serde_json::to_string(&faults).map_err(|error| error.to_string()) +} + +fn apply_fault_policy( + request: &EvaluateRequest, + evaluated_codes: &[String], + faults: Vec, +) -> Result, String> { + let findings = faults + .into_iter() + .map(fault_to_finding) + .collect::, _>>()?; let suppressions = request .config .rule_exceptions @@ -671,7 +851,7 @@ fn suppress_faults( let grammar = RulesCodeGrammar; apply_suppressions(ApplySuppressionsRequest { findings, - evaluated_codes: &evaluated_codes, + evaluated_codes, suppressions: &suppressions, scoped_ignores: &scoped_ignores, grammar: &grammar, diff --git a/crates/sqlbuild-rules-native/src/engine/main/finalize.rs b/crates/sqlbuild-rules-native/src/engine/main/finalize.rs new file mode 100644 index 0000000000..e28abf1c9c --- /dev/null +++ b/crates/sqlbuild-rules-native/src/engine/main/finalize.rs @@ -0,0 +1,3 @@ +pub(crate) fn finalize_findings_json(request_json: &str) -> Result { + crate::engine::_helpers::evaluation::finalize_findings_json(request_json) +} diff --git a/crates/sqlbuild-rules-native/src/engine/main/mod.rs b/crates/sqlbuild-rules-native/src/engine/main/mod.rs index b898ffbd13..0dea504272 100644 --- a/crates/sqlbuild-rules-native/src/engine/main/mod.rs +++ b/crates/sqlbuild-rules-native/src/engine/main/mod.rs @@ -1 +1,2 @@ pub(crate) mod evaluate; +pub(crate) mod finalize; diff --git a/crates/sqlbuild-rules-native/src/models.rs b/crates/sqlbuild-rules-native/src/models.rs index f9585e14f3..29ba2d0c4d 100644 --- a/crates/sqlbuild-rules-native/src/models.rs +++ b/crates/sqlbuild-rules-native/src/models.rs @@ -659,6 +659,7 @@ pub(crate) struct EvaluateRequest { pub public_constants: Vec, pub scope_index: ScopeIndexFacts, pub initial_findings: Vec, + pub defer_suppressions: bool, pub custom_rules: Vec, pub custom_host: Option, pub project_fingerprint: Option, @@ -696,6 +697,7 @@ impl Default for EvaluateRequest { public_constants: vec![], scope_index: ScopeIndexFacts::default(), initial_findings: vec![], + defer_suppressions: false, custom_rules: vec![], custom_host: None, project_fingerprint: None, @@ -703,6 +705,16 @@ impl Default for EvaluateRequest { } } +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct FinalizeFindingsRequest { + pub version: u32, + pub project_dir: String, + pub config: RulesConfig, + pub evaluated_codes: Vec, + pub findings: Vec, +} + #[derive(Debug, Serialize)] pub(crate) struct EvaluateResponse { pub version: u32, diff --git a/crates/sqlbuild-rules-native/src/query_analysis/_helpers/borrowed_facts.rs b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/borrowed_facts.rs new file mode 100644 index 0000000000..f7ae2b0367 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/borrowed_facts.rs @@ -0,0 +1,659 @@ +//! Immutable CTE output graph over one borrowed syntax tree. + +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::rc::Rc; + +use polyglot_sql::expressions::{Column, Join, JoinKind, Select, Star, Values, With}; +use polyglot_sql::{ + DialectType, Expression, ProjectionNullability, TransformKind, ValidationSchema, +}; +const SQL_WILDCARD: &str = "*"; + +#[derive(Clone, Debug)] +pub(super) struct OutputFact { + pub name: String, + pub upstream: Rc>, + pub nullability: ProjectionNullability, + pub resolved: bool, + pub transform: TransformKind, +} + +#[derive(Clone)] +struct Source { + outputs: Rc>, + columns: Rc>>, + physical: Option, + null_extended: bool, +} + +impl Source { + fn new(outputs: Vec, physical: Option) -> Self { + let mut columns: HashMap> = HashMap::new(); + for (index, output) in outputs.iter().enumerate() { + let key = output.name.to_lowercase(); + let value = (!columns.contains_key(&key)).then_some(index); + columns.insert(key, value); + } + Self { + outputs: Rc::new(outputs), + columns: Rc::new(columns), + physical, + null_extended: false, + } + } +} + +type Relations = HashMap; +#[derive(Default)] +struct Sources { + items: Vec<(String, Source)>, + using_columns: HashMap, +} +type Aliases = HashMap>; + +pub(super) fn infer( + expression: &Expression, + schema: Option<&ValidationSchema>, + dialect: DialectType, +) -> Vec { + let mut relations = Relations::new(); + if let Some(schema) = schema { + for table in &schema.tables { + let outputs: Vec = table + .columns + .iter() + .map(|column| OutputFact { + name: column.name.clone(), + upstream: Rc::new(BTreeSet::from([(table.name.clone(), column.name.clone())])), + nullability: match column.nullable { + Some(false) => ProjectionNullability::NonNull, + Some(true) => ProjectionNullability::Nullable, + None => ProjectionNullability::Unknown, + }, + resolved: true, + transform: TransformKind::Direct, + }) + .collect(); + relations.insert( + table.name.to_lowercase(), + Source::new(outputs, Some(table.name.clone())), + ); + } + } + query( + expression, + &relations, + matches!( + dialect, + DialectType::Snowflake + | DialectType::DuckDB + | DialectType::Redshift + | DialectType::Spark + | DialectType::Databricks + ), + ) +} + +fn query(expression: &Expression, inherited: &Relations, lateral_aliases: bool) -> Vec { + let expression = unwrapped(expression); + let mut relations = inherited.clone(); + if let Some(with) = query_with(expression) { + for cte in &with.ctes { + let mut outputs = query(&cte.this, &relations, lateral_aliases); + for (output, alias) in outputs.iter_mut().zip(&cte.columns) { + output.name = alias.name.clone(); + } + relations.insert(cte.alias.name.to_lowercase(), Source::new(outputs, None)); + } + } + match expression { + Expression::Select(select) => select_outputs(select, &relations, lateral_aliases), + Expression::Values(values) => values_outputs(values), + Expression::Union(set) => set_outputs( + query(&set.left, &relations, lateral_aliases), + query(&set.right, &relations, lateral_aliases), + set.by_name, + ), + Expression::Intersect(set) => set_outputs( + query(&set.left, &relations, lateral_aliases), + query(&set.right, &relations, lateral_aliases), + set.by_name, + ), + Expression::Except(set) => set_outputs( + query(&set.left, &relations, lateral_aliases), + query(&set.right, &relations, lateral_aliases), + set.by_name, + ), + Expression::Subquery(subquery) => query(&subquery.this, &relations, lateral_aliases), + _ => Vec::new(), + } +} + +fn set_outputs( + mut left: Vec, + right: Vec, + by_name: bool, +) -> Vec { + 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)) + } else { + right.get(index) + }; + if let Some(other) = other { + if !Rc::ptr_eq(&output.upstream, &other.upstream) { + Rc::make_mut(&mut output.upstream).extend(other.upstream.iter().cloned()); + } + output.resolved &= other.resolved; + output.nullability = merge_nullability(output.nullability, other.nullability); + } else { + output.nullability = ProjectionNullability::Nullable; + } + } + if by_name { + for mut output in right { + if !left + .iter() + .any(|other| other.name.eq_ignore_ascii_case(&output.name)) + { + output.nullability = ProjectionNullability::Nullable; + left.push(output); + } + } + } + left +} + +fn select_outputs( + select: &Select, + relations: &Relations, + lateral_aliases: bool, +) -> Vec { + let mut sources = Sources::default(); + for expression in select.from.iter().flat_map(|from| &from.expressions) { + if let Some(source) = source(expression, relations, &sources, lateral_aliases) { + sources.items.push(source); + } + } + for join in &select.joins { + sources = joined_sources(sources, join, relations, lateral_aliases); + } + let mut outputs: Vec = Vec::new(); + let mut aliases = Aliases::new(); + for projection in &select.expressions { + let projection = unwrapped(projection); + if let Some(star) = projection_star(projection) { + outputs.extend(star_outputs(&star, &sources)); + continue; + } + let name = projection.get_output_name(); + if name == SQL_WILDCARD { + continue; + } + let expression = match projection { + Expression::Alias(alias) => &alias.this, + _ => projection, + }; + let mut fact = expression_fact(expression, &sources, &aliases); + fact.name = name.to_string(); + if let Expression::Column(column) = unwrapped(expression) + && select + .where_clause + .as_ref() + .is_some_and(|clause| filtered_non_null(&clause.this, column)) + { + fact.nullability = ProjectionNullability::NonNull; + } + if lateral_aliases && matches!(projection, Expression::Alias(_)) { + let key = name.to_lowercase(); + let value = (!aliases.contains_key(&key)).then(|| fact.clone()); + aliases.insert(key, value); + } + outputs.push(fact); + } + outputs +} + +fn source( + expression: &Expression, + relations: &Relations, + preceding: &Sources, + lateral_aliases: bool, +) -> Option<(String, Source)> { + match unwrapped(expression) { + Expression::Table(table) => { + let name = table.alias.as_ref().unwrap_or(&table.name).name.clone(); + let mut source = relations + .get(&table.name.name.to_lowercase()) + .cloned() + .unwrap_or_else(|| Source::new(Vec::new(), Some(table.name.name.clone()))); + if !table.column_aliases.is_empty() { + let mut outputs = source.outputs.as_ref().clone(); + for (output, alias) in outputs.iter_mut().zip(&table.column_aliases) { + output.name = alias.name.clone(); + } + source = Source::new(outputs, source.physical); + } + Some((name, source)) + } + Expression::Subquery(subquery) => { + let mut outputs = query(&subquery.this, relations, lateral_aliases); + for (output, alias) in outputs.iter_mut().zip(&subquery.column_aliases) { + output.name = alias.name.clone(); + } + Some(( + subquery.alias.as_ref()?.name.clone(), + Source::new(outputs, None), + )) + } + Expression::Values(values) => Some(( + values + .alias + .as_ref() + .map_or(String::new(), |alias| alias.name.clone()), + Source::new(values_outputs(values), None), + )), + Expression::Lateral(lateral) if matches!(unwrapped(&lateral.this), Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN")) => + { + let fact = expression_fact(&lateral.this, preceding, &Aliases::new()); + let names: Vec = if lateral.column_aliases.is_empty() { + ["SEQ", "KEY", "PATH", "INDEX", "VALUE", "THIS"] + .into_iter() + .map(str::to_string) + .collect() + } else { + lateral.column_aliases.clone() + }; + Some(( + lateral.alias.clone()?, + Source::new( + names + .into_iter() + .map(|name| OutputFact { + name, + nullability: ProjectionNullability::Unknown, + ..fact.clone() + }) + .collect(), + None, + ), + )) + } + _ => None, + } +} + +fn values_outputs(values: &Values) -> Vec { + let mut outputs: Vec = Vec::new(); + 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()); + fact.name = values.column_aliases.get(index).map_or_else( + || format!("column{}", index + 1), + |alias| alias.name.clone(), + ); + facts.push(fact); + } + outputs = if outputs.is_empty() { + facts + } else { + set_outputs(outputs, facts, false) + }; + } + outputs +} + +fn joined_sources( + mut sources: Sources, + join: &Join, + relations: &Relations, + lateral_aliases: bool, +) -> Sources { + let Some((name, mut right)) = source(&join.this, relations, &sources, lateral_aliases) else { + return sources; + }; + let right_scope = Sources { + items: vec![(name.clone(), right.clone())], + ..Default::default() + }; + let mut joined: HashMap = HashMap::new(); + for identifier in &join.using { + let column = Column { + name: identifier.clone(), + table: None, + join_mark: false, + trailing_comments: Vec::new(), + span: None, + inferred_type: None, + }; + let left = column_fact(&column, &sources, &Aliases::new()); + let right = column_fact(&column, &right_scope, &Aliases::new()); + let output = match (left, right) { + (Some(mut left), Some(right)) if join.kind == JoinKind::Full => { + Rc::make_mut(&mut left.upstream).extend(right.upstream.iter().cloned()); + left.nullability = merge_nullability(left.nullability, right.nullability); + left.resolved &= right.resolved; + Some(left) + } + (_, right) if matches!(join.kind, JoinKind::Right | JoinKind::AsOfRight) => right, + (left, _) => left, + }; + if let Some(output) = output { + joined.insert(identifier.name.to_lowercase(), output); + } + } + if matches!( + join.kind, + JoinKind::Right | JoinKind::Full | JoinKind::AsOfRight + ) { + for (_, source) in &mut sources.items { + source.null_extended = true; + } + for output in sources.using_columns.values_mut() { + output.nullability = ProjectionNullability::Nullable; + } + } + right.null_extended = matches!( + join.kind, + JoinKind::Left | JoinKind::Full | JoinKind::AsOfLeft | JoinKind::LeftLateral + ); + sources.items.push((name, right)); + sources.using_columns.extend(joined); + sources +} + +fn column_fact(column: &Column, sources: &Sources, aliases: &Aliases) -> Option { + let mut matched = None; + let column_key = column.name.name.to_lowercase(); + if column.table.is_none() + && let Some(output) = sources.using_columns.get(&column_key) + { + return Some(output.clone()); + } + for (name, source) in &sources.items { + if column + .table + .as_ref() + .is_some_and(|table| !table.name.eq_ignore_ascii_case(name)) + { + continue; + } + let candidate = match source.columns.get(&column_key) { + Some(Some(index)) => source.outputs.get(*index), + Some(None) => return None, + None => None, + }; + let mut candidate = match (candidate, &source.physical) { + (Some(output), _) => output.clone(), + (None, Some(physical)) if source.outputs.is_empty() => OutputFact { + name: column.name.name.clone(), + upstream: Rc::new(BTreeSet::from([( + physical.clone(), + column.name.name.clone(), + )])), + nullability: ProjectionNullability::Unknown, + resolved: false, + transform: TransformKind::Direct, + }, + _ => continue, + }; + if source.null_extended { + candidate.nullability = ProjectionNullability::Nullable; + } + if matched.is_some() { + return None; + } + matched = Some(candidate); + } + matched.or_else(|| { + (column.table.is_none() + && sources + .items + .iter() + .all(|(_, source)| !source.outputs.is_empty())) + .then(|| { + aliases + .get(&column.name.name.to_lowercase()) + .cloned() + .flatten() + }) + .flatten() + }) +} + +fn expression_fact(expression: &Expression, sources: &Sources, aliases: &Aliases) -> OutputFact { + let expression = unwrapped(expression); + if let Expression::Column(column) = expression + && let Some(mut fact) = column_fact(column, sources, aliases) + { + fact.transform = TransformKind::Direct; + return fact; + } + if let Expression::Cast(cast) | Expression::TryCast(cast) | Expression::SafeCast(cast) = + expression + { + let mut fact = expression_fact(&cast.this, sources, aliases); + fact.transform = TransformKind::Cast; + fact.nullability = nullability(expression, sources, aliases); + return fact; + } + let mut upstream: BTreeSet<(String, String)> = BTreeSet::new(); + let mut resolved = true; + for node in polyglot_sql::scope::walk_in_scope(expression, false) { + if let Some(star) = projection_star(node) { + let outputs = star_outputs(&star, sources); + resolved &= !outputs.is_empty(); + for fact in outputs { + upstream.extend(fact.upstream.iter().cloned()); + resolved &= fact.resolved; + } + continue; + } + if let Expression::Column(column) = node { + if let Some(fact) = column_fact(column, sources, aliases) { + upstream.extend(fact.upstream.iter().cloned()); + resolved &= fact.resolved; + } else { + resolved = false; + } + } + } + let transform = match expression { + Expression::Column(_) | Expression::Identifier(_) => TransformKind::Direct, + Expression::Cast(_) | Expression::TryCast(_) | Expression::SafeCast(_) => { + TransformKind::Cast + } + Expression::Literal(_) | Expression::Boolean(_) | Expression::Null(_) => { + TransformKind::Constant + } + _ if polyglot_sql::traversal::contains_aggregate(expression) => TransformKind::Aggregation, + _ => TransformKind::Expression, + }; + OutputFact { + name: String::new(), + upstream: Rc::new(upstream), + nullability: nullability(expression, sources, aliases), + resolved, + transform, + } +} + +fn nullability( + expression: &Expression, + sources: &Sources, + aliases: &Aliases, +) -> ProjectionNullability { + match unwrapped(expression) { + Expression::Column(column) => column_fact(column, sources, aliases) + .map_or(ProjectionNullability::Unknown, |fact| fact.nullability), + Expression::Null(_) => ProjectionNullability::Nullable, + Expression::Literal(_) + | Expression::Boolean(_) + | Expression::Count(_) + | Expression::IsNull(_) => ProjectionNullability::NonNull, + Expression::Cast(cast) => nullability(&cast.this, sources, aliases), + Expression::Lower(function) | Expression::Upper(function) => { + nullability(&function.this, sources, aliases) + } + Expression::IfFunc(function) => { + let left = nullability(&function.true_value, sources, aliases); + let right = function + .false_value + .as_ref() + .map_or(ProjectionNullability::Nullable, |value| { + nullability(value, sources, aliases) + }); + merge_nullability(left, right) + } + Expression::Coalesce(function) => { + let values: Vec<_> = function + .expressions + .iter() + .map(|value| nullability(value, sources, aliases)) + .collect(); + if values.contains(&ProjectionNullability::NonNull) { + ProjectionNullability::NonNull + } else if !values.is_empty() + && values + .iter() + .all(|value| *value == ProjectionNullability::Nullable) + { + ProjectionNullability::Nullable + } else { + ProjectionNullability::Unknown + } + } + _ => ProjectionNullability::Unknown, + } +} + +fn star_outputs(star: &Star, sources: &Sources) -> Vec { + let mut outputs: Vec = Vec::new(); + let mut merged: HashSet = HashSet::new(); + for (name, source) in &sources.items { + if star + .table + .as_ref() + .is_some_and(|table| !table.name.eq_ignore_ascii_case(name)) + { + continue; + } + if source.outputs.is_empty() { + outputs.push(OutputFact { + name: SQL_WILDCARD.to_string(), + upstream: Rc::new(BTreeSet::new()), + nullability: ProjectionNullability::Unknown, + resolved: false, + transform: TransformKind::Star, + }); + continue; + } + for output in source.outputs.iter() { + if star + .except + .iter() + .flatten() + .any(|name| name.name.eq_ignore_ascii_case(&output.name)) + { + continue; + } + let key = output.name.to_lowercase(); + let using_output = star + .table + .is_none() + .then(|| sources.using_columns.get(&key)) + .flatten(); + if using_output.is_some() && !merged.insert(key) { + continue; + } + let mut output = using_output.unwrap_or(output).clone(); + if source.null_extended && using_output.is_none() { + output.nullability = ProjectionNullability::Nullable; + } + if let Some(alias) = star + .replace + .iter() + .flatten() + .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.name = name; + } + if let Some((_, alias)) = star + .rename + .iter() + .flatten() + .find(|(name, _)| name.name.eq_ignore_ascii_case(&output.name)) + { + output.name = alias.name.clone(); + } + outputs.push(output); + } + } + outputs +} + +pub(super) fn projection_star(expression: &Expression) -> Option { + match unwrapped(expression) { + Expression::Star(star) => Some(star.clone()), + Expression::Column(column) if column.name.name == SQL_WILDCARD => Some(Star { + table: column.table.clone(), + except: None, + replace: None, + rename: None, + trailing_comments: Vec::new(), + span: column.span, + }), + _ => None, + } +} + +fn filtered_non_null(expression: &Expression, column: &Column) -> bool { + match unwrapped(expression) { + Expression::And(binary) => { + filtered_non_null(&binary.left, column) || filtered_non_null(&binary.right, column) + } + Expression::IsNull(predicate) if predicate.not => { + matches!(unwrapped(&predicate.this), Expression::Column(filtered) + if filtered.name.name.eq_ignore_ascii_case(&column.name.name) && filtered.table.as_ref().map(|table| table.name.to_lowercase()) == column.table.as_ref().map(|table| table.name.to_lowercase())) + } + _ => false, + } +} + +fn merge_nullability( + left: ProjectionNullability, + right: ProjectionNullability, +) -> ProjectionNullability { + if left == ProjectionNullability::Nullable || right == ProjectionNullability::Nullable { + ProjectionNullability::Nullable + } else if left == ProjectionNullability::NonNull && right == ProjectionNullability::NonNull { + ProjectionNullability::NonNull + } else { + ProjectionNullability::Unknown + } +} + +fn unwrapped(mut expression: &Expression) -> &Expression { + loop { + expression = match expression { + Expression::Annotated(annotation) => &annotation.this, + Expression::Paren(paren) => &paren.this, + _ => return expression, + }; + } +} + +fn query_with(expression: &Expression) -> Option<&With> { + match expression { + Expression::Select(query) => query.with.as_ref(), + Expression::Union(query) => query.with.as_ref(), + Expression::Intersect(query) => query.with.as_ref(), + Expression::Except(query) => query.with.as_ref(), + _ => None, + } +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/_helpers/compatibility_types.rs b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/compatibility_types.rs new file mode 100644 index 0000000000..b254e14938 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/compatibility_types.rs @@ -0,0 +1,464 @@ +//! Borrowed-tree type evaluation with immutable lexical CTE outputs. + +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; + +use polyglot_sql::expressions::{DataType, Select, Star, Values, With}; +use polyglot_sql::{Dialect, DialectType, Expression, ValidationSchema}; + +type Outputs = Vec<(String, Option)>; +type Relations = HashMap>; +#[derive(Default)] +struct Sources { + items: Vec<(String, Rc)>, + using_columns: HashSet, +} +pub(super) const NULL_TYPE: &str = "\0NULL"; +const SQL_WILDCARD: &str = "*"; + +#[cfg(test)] +pub(super) fn infer( + expression: &Expression, + schema: Option<&ValidationSchema>, + functions: &HashMap, + dialect: DialectType, +) -> HashMap { + let mut types: HashMap = HashMap::new(); + for (name, data_type) in infer_outputs(expression, schema, functions, dialect) { + if let Some(value) = data_type.filter(|value| value != NULL_TYPE) { + types.insert(name, value); + } + } + types +} + +pub(super) fn infer_outputs( + expression: &Expression, + schema: Option<&ValidationSchema>, + functions: &HashMap, + dialect: DialectType, +) -> Outputs { + let mut relations = Relations::new(); + if let Some(schema) = schema { + for table in &schema.tables { + relations.insert( + table.name.to_lowercase(), + Rc::new( + table + .columns + .iter() + .map(|column| (column.name.clone(), Some(column.data_type.clone()))) + .collect(), + ), + ); + } + } + let context = TypeContext { functions, dialect }; + context.query(expression, &relations) +} + +struct TypeContext<'a> { + functions: &'a HashMap, + dialect: DialectType, +} + +impl TypeContext<'_> { + fn query(&self, expression: &Expression, inherited: &Relations) -> Outputs { + let expression = unwrapped(expression); + let mut relations = inherited.clone(); + if let Some(with) = query_with(expression) { + for cte in &with.ctes { + let mut outputs = self.query(&cte.this, &relations); + for ((name, _), alias) in outputs.iter_mut().zip(&cte.columns) { + *name = alias.name.clone(); + } + relations.insert(cte.alias.name.to_lowercase(), Rc::new(outputs)); + } + } + match expression { + Expression::Select(select) => self.select(select, &relations), + Expression::Values(values) => self.values_outputs(values), + Expression::Union(set) => self.set(&set.left, &set.right, set.by_name, &relations), + Expression::Intersect(set) => self.set(&set.left, &set.right, set.by_name, &relations), + Expression::Except(set) => self.set(&set.left, &set.right, set.by_name, &relations), + Expression::Subquery(subquery) => self.query(&subquery.this, &relations), + _ => Vec::new(), + } + } + + fn set( + &self, + left: &Expression, + right: &Expression, + by_name: bool, + relations: &Relations, + ) -> Outputs { + let left = self.query(left, relations); + let right = self.query(right, relations); + if by_name { + let mut result: Outputs = Vec::new(); + for (name, value) in &left { + let other = lookup(&right, name).unwrap_or(Some(NULL_TYPE)); + result.push((name.clone(), common_type(value.as_deref(), other))); + } + for (name, value) in right { + if lookup(&left, &name).is_none() { + result.push((name, value)); + } + } + return result; + } + if left.len() != right.len() { + return Vec::new(); + } + left.into_iter() + .zip(right) + .map(|((name, left), (_, right))| { + (name, common_type(left.as_deref(), right.as_deref())) + }) + .collect() + } + + fn select(&self, select: &Select, relations: &Relations) -> Outputs { + let mut aliases = Sources { + using_columns: select + .joins + .iter() + .flat_map(|join| &join.using) + .map(|column| column.name.to_lowercase()) + .collect(), + ..Default::default() + }; + let tables = select + .from + .iter() + .flat_map(|from| &from.expressions) + .chain(select.joins.iter().map(|join| &join.this)); + for source in tables { + if let Expression::Values(values) = unwrapped(source) { + aliases.items.push(( + values + .alias + .as_ref() + .map_or(String::new(), |alias| alias.name.to_lowercase()), + Rc::new(self.values_outputs(values)), + )); + } + if let Expression::Subquery(subquery) = unwrapped(source) + && let Some(alias) = &subquery.alias + { + let mut outputs = self.query(&subquery.this, relations); + for ((name, _), alias) in outputs.iter_mut().zip(&subquery.column_aliases) { + *name = alias.name.clone(); + } + aliases + .items + .push((alias.name.to_lowercase(), Rc::new(outputs))); + } + if let Expression::Table(table) = unwrapped(source) + && let Some(outputs) = relations.get(&table.name.name.to_lowercase()) + { + let outputs = if table.column_aliases.is_empty() { + outputs.clone() + } else { + let mut renamed = outputs.as_ref().clone(); + for ((name, _), alias) in renamed.iter_mut().zip(&table.column_aliases) { + *name = alias.name.clone(); + } + Rc::new(renamed) + }; + aliases.items.push(( + table + .alias + .as_ref() + .unwrap_or(&table.name) + .name + .to_lowercase(), + outputs, + )); + } + } + let mut result: Outputs = Vec::new(); + for projection in &select.expressions { + let projection = unwrapped(projection); + if let Some(star) = super::borrowed_facts::projection_star(projection) { + result.extend(star_outputs(&star, &aliases, self)); + } else { + let name = projection.get_output_name(); + if name != SQL_WILDCARD { + let inner = match projection { + Expression::Alias(alias) => &alias.this, + _ => projection, + }; + let value = self.projection_type(inner, &aliases); + result.push((name.to_string(), value)); + } + } + } + result + } + + fn values_outputs(&self, values: &Values) -> Outputs { + let mut outputs: Outputs = Vec::new(); + for row in &values.expressions { + for (index, expression) in row.expressions.iter().enumerate() { + let value = self.projection_type(expression, &Sources::default()); + if let Some((_, prior)) = outputs.get_mut(index) { + *prior = common_type(prior.as_deref(), value.as_deref()); + } else { + let name = values.column_aliases.get(index).map_or_else( + || format!("column{}", index + 1), + |alias| alias.name.clone(), + ); + outputs.push((name, value)); + } + } + } + outputs + } + + fn projection_type(&self, expression: &Expression, aliases: &Sources) -> Option { + if matches!(unwrapped(expression), Expression::Null(_)) { + Some(NULL_TYPE.to_string()) + } else { + self.expression(expression, aliases).or_else(|| { + expression + .inferred_type() + .filter(|data_type| **data_type != DataType::Unknown) + .and_then(|data_type| self.cast_type(data_type)) + }) + } + } + + fn expression(&self, expression: &Expression, aliases: &Sources) -> Option { + let expression = unwrapped(expression); + let function_name = match expression { + Expression::Function(function) => function.name.as_str(), + Expression::AggregateFunction(function) => function.name.as_str(), + _ => expression.variant_name(), + }; + if !matches!(expression, Expression::Column(_)) + && let Some(value) = self.functions.get(&function_name.to_uppercase()) + { + return Some(value.clone()); + } + match expression { + Expression::Alias(alias) => self.expression(&alias.this, aliases), + Expression::Cast(cast) | Expression::TryCast(cast) => self.cast_type(&cast.to), + Expression::Column(column) => { + if let Some(table) = &column.table { + let key = table.name.to_lowercase(); + let (_, source) = aliases.items.iter().find(|(name, _)| name == &key)?; + return lookup(source, &column.name.name) + .flatten() + .map(str::to_string); + } + let mut values = aliases + .items + .iter() + .filter_map(|(_, source)| lookup(source, &column.name.name)); + let value = values.next()??; + if aliases + .using_columns + .contains(&column.name.name.to_lowercase()) + { + for other in values { + if canonical(other?) != canonical(value) { + return None; + } + } + Some(value.to_string()) + } else { + values.next().is_none().then(|| value.to_string()) + } + } + Expression::Boolean(_) + | Expression::And(_) + | Expression::Between(_) + | Expression::Eq(_) + | Expression::Exists(_) + | Expression::Gt(_) + | Expression::Gte(_) + | Expression::ILike(_) + | Expression::In(_) + | Expression::Is(_) + | Expression::IsNull(_) + | Expression::Like(_) + | Expression::Lt(_) + | Expression::Lte(_) + | Expression::Neq(_) + | Expression::Not(_) + | Expression::Or(_) + | Expression::RegexpLike(_) => Some("BOOLEAN".to_string()), + Expression::Min(function) | Expression::Max(function) => { + self.expression(&function.this, aliases) + } + Expression::WindowFunction(function) => self.expression(&function.this, aliases), + Expression::WithinGroup(function) => self.expression(&function.this, aliases), + Expression::Case(case) => self.results( + case.whens + .iter() + .map(|(_, result)| result) + .chain(case.else_.iter()), + aliases, + ), + Expression::IfFunc(function) => self.results( + std::iter::once(&function.true_value).chain(function.false_value.iter()), + aliases, + ), + Expression::Coalesce(function) => self.results(function.expressions.iter(), aliases), + Expression::Concat(binary) => { + let left = self.string_operand(&binary.left, aliases)?; + let right = self.string_operand(&binary.right, aliases)?; + (is_text(&left) && is_text(&right)).then(|| "TEXT".to_string()) + } + _ => None, + } + } + + fn string_operand(&self, expression: &Expression, aliases: &Sources) -> Option { + if matches!(unwrapped(expression), Expression::Literal(literal) if matches!(literal.as_ref(), polyglot_sql::expressions::Literal::String(_))) + { + Some("TEXT".to_string()) + } else { + self.expression(expression, aliases) + } + } + + fn results<'a>( + &self, + values: impl Iterator, + aliases: &Sources, + ) -> Option { + let mut types = values + .filter(|value| !matches!(unwrapped(value), Expression::Null(_))) + .map(|value| self.expression(value, aliases)); + let first = types.next()??; + types + .all(|value| { + value + .as_deref() + .is_some_and(|value| canonical(value) == canonical(&first)) + }) + .then_some(first) + } + + fn cast_type(&self, data_type: &DataType) -> Option { + if let DataType::VarChar { + length: Some(length), + .. + } + | DataType::String { + length: Some(length), + } + | DataType::TextWithLength { length } = data_type + { + return Some(format!("VARCHAR({length})")); + } + match Dialect::get(self.dialect).generate(&Expression::DataType(data_type.clone())) { + Ok(rendered) => Some(rendered.replace(", ", ",")), + Err(_) => None, + } + } +} + +fn query_with(expression: &Expression) -> Option<&With> { + match expression { + Expression::Select(query) => query.with.as_ref(), + Expression::Union(query) => query.with.as_ref(), + Expression::Intersect(query) => query.with.as_ref(), + Expression::Except(query) => query.with.as_ref(), + _ => None, + } +} + +fn unwrapped(mut expression: &Expression) -> &Expression { + loop { + expression = match expression { + Expression::Annotated(annotation) => &annotation.this, + Expression::Paren(paren) => &paren.this, + _ => return expression, + }; + } +} + +fn lookup<'a>(outputs: &'a Outputs, name: &str) -> Option> { + outputs + .iter() + .find(|(key, _)| key == name) + .or_else(|| { + outputs + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + }) + .map(|(_, value)| value.as_deref()) +} + +fn star_outputs(star: &Star, aliases: &Sources, context: &TypeContext<'_>) -> Outputs { + let excluded: HashSet<_> = star + .except + .iter() + .flatten() + .map(|name| name.name.to_lowercase()) + .collect(); + let mut result: Outputs = Vec::new(); + for (alias, source) in &aliases.items { + if star + .table + .as_ref() + .is_some_and(|table| !table.name.eq_ignore_ascii_case(alias)) + { + continue; + } + for (name, value) in source.iter() { + let key = name.to_lowercase(); + if excluded.contains(&key) { + continue; + } + let replacement = star + .replace + .iter() + .flatten() + .find(|replacement| replacement.alias.name.eq_ignore_ascii_case(name)); + let value = replacement.map_or_else( + || value.clone(), + |replacement| context.projection_type(&replacement.this, aliases), + ); + let renamed = star + .rename + .iter() + .flatten() + .find(|(original, _)| original.name.eq_ignore_ascii_case(name)) + .map_or(name.as_str(), |(_, renamed)| renamed.name.as_str()); + result.push((renamed.to_string(), value)); + } + } + result +} + +fn canonical(value: &str) -> String { + let compact: String = value + .chars() + .filter(|c| !c.is_ascii_whitespace()) + .flat_map(char::to_uppercase) + .collect(); + match compact.as_str() { + "INTEGER" => "INT".to_string(), + "VARCHAR" | "STRING" => "TEXT".to_string(), + _ => compact, + } +} + +fn is_text(value: &str) -> bool { + matches!(value, "TEXT" | "STRING") || value.starts_with("VARCHAR") +} + +fn common_type(left: Option<&str>, right: Option<&str>) -> Option { + match (left, right) { + (Some("\0NULL"), Some("\0NULL")) => None, + (Some("\0NULL"), right) => right.map(str::to_string), + (left, Some("\0NULL")) => left.map(str::to_string), + (Some(left), Some(right)) if canonical(left) == canonical(right) => Some(left.to_string()), + _ => None, + } +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs new file mode 100644 index 0000000000..b7332a3de4 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs @@ -0,0 +1,1531 @@ +//! Coarse, bounded native batch boundary for compact SQL query analysis. + +use polyglot_sql::ExpressionWalk; +use polyglot_sql::{ + AnalyzeQueryOptions, DialectType, ProjectionNullability, QueryAnalysis, QueryShape, + ReferenceConfidence, SchemaValidationOptions, TransformKind, ValidationResult, + ValidationSchema, analyze_query, analyze_query_for_project_projections, +}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; + +use crate::constants::{SQL_WILDCARD, UNKNOWN_SQL_TYPE, VARCHAR_SQL_TYPE}; + +const DEFAULT_WORKERS: usize = 4; +const MAX_WORKERS: usize = 4; +const ANALYSIS_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; +const COMPACT_ANALYSIS_BATCH_SIZE: usize = 64; +const LARGE_COMPACT_SQL_BYTES: usize = 32 * 1024 * 1024; + +#[derive(Debug, Deserialize)] +struct AnalysisBatchRequest { + requests: Vec, + #[serde(default = "default_workers")] + workers: usize, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct AnalysisRequest { + sql: String, + #[serde(default = "default_dialect")] + dialect: String, + #[serde(default)] + schema: Option, + #[serde(default)] + binding_schema: Option, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum AnalysisResponse { + Success { analysis: QueryAnalysis }, + Failure { error: String }, +} + +#[derive(Debug, Deserialize)] +struct ProjectAnalysisBatchRequest { + requests: Vec, + #[serde(default = "default_workers")] + workers: usize, +} + +#[derive(Debug, Deserialize)] +struct CompactProjectAnalysisBatchRequest { + queries: Vec, + templates: Vec, + projections: Vec, + #[serde(default = "default_workers")] + workers: usize, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct CompactProjectAnalysisTemplateRequest { + query_index: usize, + #[serde(default)] + references: HashMap, + #[serde(default)] + function_return_types: HashMap, + #[serde(default)] + declared_column_order: Vec, + #[serde(default)] + recover_cte_facts: bool, + #[serde(default = "default_true")] + rich_type_inference: bool, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct CompactProjectAnalysisProjectionRequest { + template_index: usize, + #[serde(default)] + resource_names: HashMap, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct ProjectAnalysisRequest { + #[serde(flatten)] + query: AnalysisRequest, + #[serde(default)] + references: HashMap, + #[serde(default)] + function_return_types: HashMap, + #[serde(default)] + declared_column_order: Vec, + #[serde(default)] + recover_cte_facts: bool, + #[serde(default = "default_true")] + rich_type_inference: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct LineageResource { + resource_type: String, + resource_name: String, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum ProjectAnalysisResponse { + Success { analysis: ProjectAnalysis }, + Failure { error: String }, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct ProjectAnalysis { + columns: Vec, + lineage_columns: Vec, + has_star: bool, + #[serde(skip)] + requires_legacy_fallback: bool, +} + +struct ProjectAnalysisInputs<'a> { + analysis: &'a QueryAnalysis, + references: &'a HashMap, + function_return_types: &'a HashMap, + declared_column_order: &'a [String], + recover_cte_facts: bool, + rich_type_inference: bool, +} + +struct ProjectAnalysisProjection { + references: HashMap, + function_return_types: HashMap, + declared_column_order: Vec, + recover_cte_facts: bool, + rich_type_inference: bool, +} + +struct CompactQueryWork { + query: AnalysisRequest, + project_projections: bool, + projections: Vec<(usize, ProjectAnalysisProjection)>, +} + +struct CompiledQueryWorkResult { + projections: Vec<(usize, Result)>, + validation: Option>, +} + +struct ProjectAnalysisProjectionResult { + projection_index: usize, + resource_names: Vec<(String, String)>, +} + +#[derive(Debug, Serialize)] +struct ProjectColumn { + name: String, + #[serde(rename = "type")] + data_type: Option, + nullability: ProjectNullability, +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +enum ProjectNullability { + NonNull, + Nullable, + Unknown, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct ProjectLineageColumn { + output_column: String, + upstream_columns: Vec, + transform_kind: ProjectTransformKind, + confidence: ProjectConfidence, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct ProjectLineageSource { + resource_type: String, + resource_name: String, + column_name: String, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +enum ProjectTransformKind { + Direct, + Cast, + Expression, + Aggregation, + Star, + Constant, +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +enum ProjectConfidence { + High, + Medium, + Unknown, +} + +fn default_workers() -> usize { + DEFAULT_WORKERS +} + +fn default_dialect() -> String { + "generic".to_string() +} + +fn default_true() -> bool { + true +} + +pub(crate) fn analyze_json(request_json: &str) -> Result { + let request: AnalysisBatchRequest = + serde_json::from_str(request_json).map_err(|error| error.to_string())?; + let workers = request.workers.clamp(1, MAX_WORKERS); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers.min(request.requests.len().max(1))) + .stack_size(ANALYSIS_WORKER_STACK_BYTES) + .build() + .map_err(|error| error.to_string())?; + let responses: Vec = pool.install(|| { + request + .requests + .into_par_iter() + .map(analyze_request) + .collect() + }); + serde_json::to_string(&responses).map_err(|error| error.to_string()) +} + +pub(crate) fn analyze_project_json(request_json: &str) -> Result { + let request: ProjectAnalysisBatchRequest = + serde_json::from_str(request_json).map_err(|error| error.to_string())?; + let workers = request.workers.clamp(1, MAX_WORKERS); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers.min(request.requests.len().max(1))) + .stack_size(ANALYSIS_WORKER_STACK_BYTES) + .build() + .map_err(|error| error.to_string())?; + let responses: Vec = pool.install(|| { + request + .requests + .into_par_iter() + .map(analyze_project_request) + .collect() + }); + serde_json::to_string(&responses).map_err(|error| error.to_string()) +} + +pub(crate) fn analyze_project_compact_json(request_json: &str) -> Result { + let request: CompactProjectAnalysisBatchRequest = + serde_json::from_str(request_json).map_err(|error| error.to_string())?; + let workers = request.workers.clamp(1, MAX_WORKERS); + if request + .templates + .iter() + .any(|template| template.query_index >= request.queries.len()) + { + return Err("compact project analysis received an invalid query index".to_string()); + } + if request + .projections + .iter() + .any(|projection| projection.template_index >= request.templates.len()) + { + return Err("compact project analysis received an invalid template index".to_string()); + } + let unique_query_count = request.queries.len(); + let unique_projection_count = request.templates.len(); + let mut project_projection_queries = vec![false; unique_query_count]; + for template in &request.templates { + if template.recover_cte_facts { + project_projection_queries[template.query_index] = true; + } + } + let mut projections_by_query: Vec> = + (0..unique_query_count).map(|_| Vec::new()).collect(); + for (projection_index, template) in request.templates.into_iter().enumerate() { + projections_by_query[template.query_index].push(( + projection_index, + ProjectAnalysisProjection { + references: template.references, + function_return_types: template.function_return_types, + declared_column_order: template.declared_column_order, + recover_cte_facts: template.recover_cte_facts, + rich_type_inference: template.rich_type_inference, + }, + )); + } + let projection_results: Vec = request + .projections + .into_iter() + .map(|projection| { + let mut resource_names: Vec<(String, String)> = + projection.resource_names.into_iter().collect(); + resource_names.sort(); + ProjectAnalysisProjectionResult { + projection_index: projection.template_index, + resource_names, + } + }) + .collect(); + let sql_bytes: usize = request.queries.iter().map(|query| query.sql.len()).sum(); + let analysis_workers = workers.min(unique_query_count.max(1)); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(analysis_workers) + .stack_size(ANALYSIS_WORKER_STACK_BYTES) + .build() + .map_err(|error| error.to_string())?; + let query_work: Vec = request + .queries + .into_iter() + .zip(project_projection_queries) + .zip(projections_by_query) + .map( + |((query, project_projections), projections)| CompactQueryWork { + query, + project_projections, + projections, + }, + ) + .collect(); + let mut accumulator = CompactProjectAccumulator::new(unique_projection_count); + let mut validations = Vec::with_capacity(unique_query_count); + let batch_size = if sql_bytes >= LARGE_COMPACT_SQL_BYTES { + COMPACT_ANALYSIS_BATCH_SIZE + } else { + unique_query_count.max(1) + }; + let mut remaining = query_work.into_iter(); + loop { + let batch: Vec<_> = remaining.by_ref().take(batch_size).collect(); + if batch.is_empty() { + break; + } + let analysis_groups: Vec = pool.install(|| { + batch + .into_par_iter() + .map(analyze_compact_query_work) + .collect() + }); + for group in analysis_groups { + validations.push(group.validation.transpose()?); + for (projection_index, analysis) in group.projections { + accumulator.compact_analysis(projection_index, analysis)?; + } + } + } + let mut response = accumulator.finish( + projection_results, + unique_query_count, + unique_projection_count, + )?; + if validations.iter().any(Option::is_some) { + response.validations = Some(validations); + } + serde_json::to_string(&response).map_err(|error| error.to_string()) +} + +fn analyze_compact_query_work(work: CompactQueryWork) -> CompiledQueryWorkResult { + let work = match try_borrowed_query(work) { + Ok(result) => return result, + Err(work) => *work, + }; + let (query_result, validation, _) = compile_query(work.query, work.project_projections); + CompiledQueryWorkResult { + projections: project_query_templates(&query_result, work.projections), + validation, + } +} + +/// Fold fully bound lexical query graphs; unsupported shapes retain the existing resolver. +fn try_borrowed_query( + work: CompactQueryWork, +) -> Result> { + if work.query.schema.is_none() + || (work.query.binding_schema.is_none() + && !work + .query + .schema + .as_ref() + .is_some_and(|schema| schema.strict == Some(true))) + || work + .projections + .iter() + .any(|(_, projection)| !projection.rich_type_inference && !projection.recover_cte_facts) + { + return Err(Box::new(work)); + } + let Ok(dialect) = work.query.dialect.parse::() else { + return Err(Box::new(work)); + }; + if !matches!(dialect, DialectType::Snowflake | DialectType::DuckDB) { + return Err(Box::new(work)); + } + let Ok(mut statements) = polyglot_sql::parse_with_options( + &work.query.sql, + dialect, + &polyglot_sql::ParseOptions { + complexity_guard: Some(polyglot_sql::ComplexityGuardOptions { + max_function_call_depth: Some(128), + ..Default::default() + }), + }, + ) else { + return Err(Box::new(work)); + }; + if statements.len() != 1 { + return Err(Box::new(work)); + } + let Some(mut expression) = statements.pop() else { + return Err(Box::new(work)); + }; + if dialect == DialectType::Snowflake && expression.dfs().any(has_case_sensitive_binding) { + return Err(Box::new(work)); + } + if matches!(&expression, polyglot_sql::Expression::Select(select) + if select.with.is_none() && select.expressions.iter().any(|projection| super::borrowed_facts::projection_star(projection).is_some())) + { + return Err(Box::new(work)); + } + if !matches!( + expression, + polyglot_sql::Expression::Select(_) + | polyglot_sql::Expression::Union(_) + | polyglot_sql::Expression::Intersect(_) + | polyglot_sql::Expression::Except(_) + ) { + return Err(Box::new(work)); + } + if expression.dfs().any(|node| { + matches!( + node, + polyglot_sql::Expression::Pivot(_) | polyglot_sql::Expression::Unpivot(_) + | polyglot_sql::Expression::Intersect(_) | polyglot_sql::Expression::Except(_) + ) + || matches!(node, polyglot_sql::Expression::Table(table) if table.schema.is_some() || table.catalog.is_some()) + || matches!(node, polyglot_sql::Expression::Select(select) if select.with.as_ref().is_some_and(|with| with.recursive)) + || matches!(node, polyglot_sql::Expression::Union(union) if union.with.as_ref().is_some_and(|with| with.recursive)) + || matches!(node, polyglot_sql::Expression::Select(select) if !borrowed_sources_supported(select, dialect)) + }) { + return Err(Box::new(work)); + } + if expression.dfs().any(|node| { + matches!(node, polyglot_sql::Expression::Select(select) + if select.expressions.iter().any(|projection| projection.dfs().any(|child| matches!(child, + polyglot_sql::Expression::Select(_) | polyglot_sql::Expression::Union(_) + | polyglot_sql::Expression::Intersect(_) | polyglot_sql::Expression::Except(_))))) + }) { + return Err(Box::new(work)); + } + let validation = work + .query + .binding_schema + .as_ref() + .or(work.query.schema.as_ref()) + .map(|schema| { + let result = polyglot_sql::validation::validate_parsed_with_schema( + vec![expression.clone()], + dialect, + schema, + &SchemaValidationOptions { + check_types: false, + check_references: true, + strict: Some(true), + semantic: false, + strict_syntax: false, + ..Default::default() + }, + ); + crate::semantic_validation::main::complete_parsed_validation( + crate::semantic_validation::main::ParsedValidationRequest { + sql: &work.query.sql, + dialect, + schema, + result, + expression: Some(&expression), + }, + ) + }); + if !matches!(validation, Some(Ok(ref result)) if result.valid) { + return Err(Box::new(work)); + } + { + let schema = work.query.schema.as_ref().map(|schema| { + polyglot_sql::validation::mapping_schema_from_validation_schema_with_dialect( + schema, dialect, + ) + }); + let schema = schema + .as_ref() + .map(|schema| schema as &dyn polyglot_sql::schema::Schema); + polyglot_sql::optimizer::annotate_types::annotate_types( + &mut expression, + schema, + Some(dialect), + ); + } + let facts = super::borrowed_facts::infer(&expression, work.query.schema.as_ref(), dialect); + if facts.is_empty() + || facts + .iter() + .any(|fact| !fact.resolved || fact.name.is_empty() || fact.name == SQL_WILDCARD) + { + return Err(Box::new(work)); + } + let mut projections: Vec<(usize, Result)> = Vec::new(); + let mut facts_by_name: HashMap<&str, &super::borrowed_facts::OutputFact> = HashMap::new(); + for fact in &facts { + if facts_by_name.insert(&fact.name, fact).is_some() { + return Err(Box::new(work)); + } + } + for (index, projection) in &work.projections { + let mut outputs = super::compatibility_types::infer_outputs( + &expression, + work.query.schema.as_ref(), + &projection.function_return_types, + dialect, + ); + { + let types: HashMap<_, _> = outputs.into_iter().collect(); + outputs = facts + .iter() + .map(|fact| (fact.name.clone(), types.get(&fact.name).cloned().flatten())) + .collect(); + } + let columns: Vec<_> = outputs + .into_iter() + .map(|(name, data_type)| ProjectColumn { + nullability: facts_by_name + .get(name.as_str()) + .map_or(ProjectNullability::Unknown, |fact| { + project_nullability(fact.nullability) + }), + name, + data_type: known_compatibility_type(data_type, projection.rich_type_inference), + }) + .collect(); + let mut lineage_columns: Vec = Vec::new(); + for column in &columns { + let fact = facts_by_name.get(column.name.as_str()); + let mut upstream_columns: Vec = Vec::new(); + if let Some(fact) = fact { + for (table, source_column) in fact.upstream.iter() { + if let Some(resource) = projection.references.get(table) { + upstream_columns.push(ProjectLineageSource { + resource_type: resource.resource_type.clone(), + resource_name: resource.resource_name.clone(), + column_name: source_column.clone(), + }); + } + } + } + let upstream_columns = order_project_sources(upstream_columns); + lineage_columns.push(ProjectLineageColumn { + output_column: column.name.clone(), + transform_kind: fact.map_or(ProjectTransformKind::Direct, |fact| { + project_transform_kind(fact.transform, !upstream_columns.is_empty()) + }), + confidence: fact.map_or(ProjectConfidence::Unknown, |fact| { + if fact.resolved { + ProjectConfidence::High + } else { + ProjectConfidence::Unknown + } + }), + upstream_columns, + }); + } + projections.push(( + *index, + Ok(ProjectAnalysis { + columns, + lineage_columns, + has_star: false, + requires_legacy_fallback: false, + }), + )); + } + if projections.iter().any(|(_, result)| matches!(result, Ok(analysis) if analysis.columns.iter().any(|column| column.data_type.is_none()))) { + return Err(Box::new(work)); + } + Ok(CompiledQueryWorkResult { + projections, + validation: if work.query.binding_schema.is_some() { + validation + } else { + None + }, + }) +} + +fn has_case_sensitive_binding(expression: &polyglot_sql::Expression) -> bool { + use polyglot_sql::Expression; + match expression { + Expression::Column(column) => { + case_sensitive_identifier(&column.name) + || column.table.as_ref().is_some_and(case_sensitive_identifier) + } + Expression::Identifier(identifier) => case_sensitive_identifier(identifier), + Expression::Alias(alias) => { + case_sensitive_identifier(&alias.alias) + || alias.column_aliases.iter().any(case_sensitive_identifier) + } + Expression::Table(table) => { + case_sensitive_identifier(&table.name) + || table.alias.as_ref().is_some_and(case_sensitive_identifier) + || table.column_aliases.iter().any(case_sensitive_identifier) + } + Expression::Cte(cte) => { + case_sensitive_identifier(&cte.alias) + || cte.columns.iter().any(case_sensitive_identifier) + } + Expression::Subquery(query) => { + query.alias.as_ref().is_some_and(case_sensitive_identifier) + || query.column_aliases.iter().any(case_sensitive_identifier) + } + Expression::Select(select) => select + .with + .as_ref() + .is_some_and(case_sensitive_cte_bindings), + Expression::Union(union) => union.with.as_ref().is_some_and(case_sensitive_cte_bindings), + _ => false, + } +} + +fn case_sensitive_identifier(identifier: &polyglot_sql::expressions::Identifier) -> bool { + identifier.quoted && identifier.name != identifier.name.to_uppercase() +} + +fn case_sensitive_cte_bindings(with: &polyglot_sql::expressions::With) -> bool { + with.ctes.iter().any(|cte| { + case_sensitive_identifier(&cte.alias) || cte.columns.iter().any(case_sensitive_identifier) + }) +} + +fn borrowed_sources_supported( + select: &polyglot_sql::expressions::Select, + dialect: DialectType, +) -> bool { + use polyglot_sql::expressions::JoinKind; + if select.joins.iter().any(|join| !join.using.is_empty()) + && select + .expressions + .iter() + .any(|projection| super::borrowed_facts::projection_star(projection).is_some()) + { + return false; + } + select + .from + .iter() + .flat_map(|from| &from.expressions) + .all(|source| borrowed_source_supported(source, dialect)) + && select.joins.iter().all(|join| { + matches!( + join.kind, + JoinKind::Inner + | JoinKind::Left + | JoinKind::Right + | JoinKind::Full + | JoinKind::Cross + | JoinKind::Implicit + | JoinKind::AsOf + | JoinKind::AsOfLeft + | JoinKind::AsOfRight + | JoinKind::Lateral + | JoinKind::LeftLateral + ) && borrowed_source_supported(&join.this, dialect) + }) +} + +fn borrowed_source_supported(expression: &polyglot_sql::Expression, dialect: DialectType) -> bool { + use polyglot_sql::Expression; + match expression { + Expression::Table(_) => true, + Expression::Values(values) => { + dialect == DialectType::Snowflake || !values.column_aliases.is_empty() + } + Expression::Subquery(query) => query.alias.is_some(), + Expression::Lateral(lateral) => { + lateral.alias.is_some() + && matches!(lateral.this.as_ref(), Expression::Function(function) if !function.quoted && function.name.eq_ignore_ascii_case("FLATTEN")) + } + Expression::Paren(paren) => borrowed_source_supported(&paren.this, dialect), + Expression::Annotated(annotation) => borrowed_source_supported(&annotation.this, dialect), + _ => false, + } +} + +fn known_compatibility_type(data_type: Option, normalize: bool) -> Option { + data_type + .filter(|value| { + !value.is_empty() + && value != UNKNOWN_SQL_TYPE + && value != super::compatibility_types::NULL_TYPE + }) + .map(|value| { + if normalize { + normalize_project_type(&value) + } else { + value + } + }) +} + +type CompiledQueryResult = ( + Result, + Option>, + Option, +); + +fn compile_query(mut request: AnalysisRequest, project_projections: bool) -> CompiledQueryResult { + let Some(schema) = request.binding_schema.take() else { + return (query_analysis(request, project_projections), None, None); + }; + let dialect: DialectType = match request.dialect.parse() { + Ok(dialect) => dialect, + Err(_) => return (query_analysis(request, project_projections), None, None), + }; + let compiled = polyglot_sql::compile_required_query_analysis( + &request.sql, + AnalyzeQueryOptions { + complexity_guard: None, + dialect, + schema: request.schema, + }, + &schema, + &SchemaValidationOptions { + check_types: false, + check_references: true, + strict: Some(true), + semantic: false, + strict_syntax: false, + ..Default::default() + }, + project_projections, + ); + let validation = crate::semantic_validation::main::complete_parsed_validation( + crate::semantic_validation::main::ParsedValidationRequest { + sql: &request.sql, + dialect, + schema: &schema, + result: compiled.validation, + expression: compiled.expression.as_ref(), + }, + ); + ( + compiled.analysis.map_err(|error| error.to_string()), + Some(validation), + compiled.expression, + ) +} + +fn project_query_templates( + query_result: &Result, + projections: Vec<(usize, ProjectAnalysisProjection)>, +) -> Vec<(usize, Result)> { + let mut results: Vec<(usize, Result)> = + Vec::with_capacity(projections.len()); + for (projection_index, projection) in projections { + let projected = match query_result { + Ok(analysis) => Ok(project_analysis(ProjectAnalysisInputs { + analysis, + references: &projection.references, + function_return_types: &projection.function_return_types, + declared_column_order: &projection.declared_column_order, + recover_cte_facts: projection.recover_cte_facts, + rich_type_inference: projection.rich_type_inference, + })), + Err(error) => Err(error.clone()), + }; + results.push((projection_index, projected)); + } + results +} + +fn analyze_request(request: AnalysisRequest) -> AnalysisResponse { + match query_analysis(request, false) { + Ok(analysis) => AnalysisResponse::Success { analysis }, + Err(error) => AnalysisResponse::Failure { error }, + } +} + +fn query_analysis( + request: AnalysisRequest, + project_projections: bool, +) -> Result { + let dialect: DialectType = match request.dialect.parse() { + Ok(dialect) => dialect, + Err(error) => { + return Err(format!( + "unsupported Polyglot dialect '{}': {error}", + request.dialect + )); + } + }; + let options = AnalyzeQueryOptions { + complexity_guard: None, + dialect, + schema: request.schema, + }; + if project_projections { + analyze_query_for_project_projections(&request.sql, options) + } else { + analyze_query(&request.sql, options) + } + .map_err(|error| error.to_string()) +} + +fn analyze_project_request(request: ProjectAnalysisRequest) -> ProjectAnalysisResponse { + match analyze_project_result(request) { + Ok(analysis) => ProjectAnalysisResponse::Success { analysis }, + Err(error) => ProjectAnalysisResponse::Failure { error }, + } +} + +fn analyze_project_result(request: ProjectAnalysisRequest) -> Result { + let analysis = query_analysis(request.query, true)?; + Ok(project_analysis(ProjectAnalysisInputs { + analysis: &analysis, + references: &request.references, + function_return_types: &request.function_return_types, + declared_column_order: &request.declared_column_order, + recover_cte_facts: request.recover_cte_facts, + rich_type_inference: request.rich_type_inference, + })) +} + +type CompactProjectColumn = (usize, Option, u8, u8, u8, Vec<(usize, usize, usize)>); + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct CompactProjectBatch { + #[serde(skip_serializing_if = "Option::is_none")] + validations: Option>>, + strings: Vec, + facts: Vec, + templates: Vec, + analyses: Vec, + unique_query_count: usize, + unique_projection_count: usize, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum CompactProjectResponse { + Success((Vec, bool)), + Failure(String), +} + +type CompactProjectProjection = (usize, Vec<(usize, usize)>); + +#[derive(Default)] +struct StringInterner { + strings: Vec, + indexes: HashMap, +} + +struct CompactProjectAccumulator { + interner: StringInterner, + facts: Vec, + fact_indexes: HashMap, + templates: Vec>, +} + +impl StringInterner { + fn intern(&mut self, value: String) -> usize { + if let Some(index) = self.indexes.get(&value) { + return *index; + } + let index = self.strings.len(); + self.indexes.insert(value.clone(), index); + self.strings.push(value); + index + } +} + +impl CompactProjectAccumulator { + fn new(template_count: usize) -> Self { + Self { + interner: StringInterner::default(), + facts: Vec::new(), + fact_indexes: HashMap::new(), + templates: (0..template_count).map(|_| None).collect(), + } + } + + fn compact_analysis( + &mut self, + projection_index: usize, + result: Result, + ) -> Result<(), String> { + let analysis = match result { + Ok(analysis) => analysis, + Err(error) => { + self.set_template(projection_index, CompactProjectResponse::Failure(error))?; + return Ok(()); + } + }; + if analysis.requires_legacy_fallback { + self.set_template( + projection_index, + CompactProjectResponse::Failure( + "native project type recovery requires legacy fallback".to_string(), + ), + )?; + return Ok(()); + } + let mut columns = Vec::with_capacity(analysis.columns.len()); + for (column, lineage) in analysis.columns.into_iter().zip(analysis.lineage_columns) { + let upstream: Vec<(usize, usize, usize)> = lineage + .upstream_columns + .into_iter() + .map(|source| { + ( + self.interner.intern(source.resource_type), + self.interner.intern(source.resource_name), + self.interner.intern(source.column_name), + ) + }) + .collect(); + let fact = ( + self.interner.intern(column.name), + column.data_type.map(|value| self.interner.intern(value)), + nullability_code(column.nullability), + transform_code(lineage.transform_kind), + confidence_code(lineage.confidence), + upstream, + ); + let next_index = self.facts.len(); + let fact_index = *self.fact_indexes.entry(fact.clone()).or_insert_with(|| { + self.facts.push(fact); + next_index + }); + columns.push(fact_index); + } + self.set_template( + projection_index, + CompactProjectResponse::Success((columns, analysis.has_star)), + ) + } + + fn set_template( + &mut self, + projection_index: usize, + response: CompactProjectResponse, + ) -> Result<(), String> { + let Some(slot) = self.templates.get_mut(projection_index) else { + return Err("compact project analysis received an invalid projection index".to_owned()); + }; + if slot.is_some() { + return Err( + "compact project analysis received a duplicate projection index".to_owned(), + ); + } + *slot = Some(response); + Ok(()) + } + + fn finish( + mut self, + projections: Vec, + unique_query_count: usize, + unique_projection_count: usize, + ) -> Result { + let mut compact_projections: Vec = + Vec::with_capacity(projections.len()); + for projection in projections { + let mut resource_names: Vec<(usize, usize)> = + Vec::with_capacity(projection.resource_names.len()); + for (canonical_name, resource_name) in projection.resource_names { + resource_names.push(( + self.interner.intern(canonical_name), + self.interner.intern(resource_name), + )); + } + compact_projections.push((projection.projection_index, resource_names)); + } + let templates: Option> = self.templates.into_iter().collect(); + let Some(templates) = templates else { + return Err("compact project analysis omitted a projection".to_owned()); + }; + Ok(CompactProjectBatch { + validations: None, + strings: self.interner.strings, + facts: self.facts, + templates, + analyses: compact_projections, + unique_query_count, + unique_projection_count, + }) + } +} + +fn nullability_code(value: ProjectNullability) -> u8 { + match value { + ProjectNullability::Unknown => 0, + ProjectNullability::NonNull => 1, + ProjectNullability::Nullable => 2, + } +} + +fn transform_code(value: ProjectTransformKind) -> u8 { + match value { + ProjectTransformKind::Direct => 0, + ProjectTransformKind::Cast => 1, + ProjectTransformKind::Expression => 2, + ProjectTransformKind::Aggregation => 3, + ProjectTransformKind::Star => 4, + ProjectTransformKind::Constant => 5, + } +} + +fn confidence_code(value: ProjectConfidence) -> u8 { + match value { + ProjectConfidence::Unknown => 0, + ProjectConfidence::High => 1, + ProjectConfidence::Medium => 2, + } +} + +fn project_analysis(inputs: ProjectAnalysisInputs<'_>) -> ProjectAnalysis { + let cte_facts: CteColumnFacts = if inputs.recover_cte_facts { + recovered_cte_facts( + inputs.analysis, + inputs.function_return_types, + inputs.references, + inputs.rich_type_inference, + ) + } else { + HashMap::new() + }; + let infer_nullability = inputs.analysis.shape != QueryShape::SetOperation; + let mut columns: Vec = Vec::new(); + let mut lineage_columns: Vec = Vec::new(); + let mut requires_legacy_fallback = false; + for projection in inputs + .analysis + .projections + .iter() + .filter(|projection| !projection.is_star) + { + let Some(output_column) = projection + .name + .as_ref() + .filter(|name| !name.is_empty() && name.as_str() != SQL_WILDCARD) + else { + continue; + }; + let (upstream_columns, confidence) = if inputs.analysis.shape == QueryShape::SetOperation { + project_upstream_columns(projection, inputs.references) + } else { + recovered_projection_lineage(projection, inputs.references, &cte_facts) + }; + let has_upstream = !upstream_columns.is_empty(); + let transform_kind = project_transform_kind(projection.transform_kind, has_upstream); + requires_legacy_fallback |= + direct_projection_requires_legacy_fallback(projection, &inputs, &cte_facts); + columns.push(ProjectColumn { + name: output_column.clone(), + data_type: recovered_projection_type( + projection, + inputs.function_return_types, + &cte_facts, + ProjectionTypeOptions { + allow_direct_type_hint: inputs.rich_type_inference || inputs.recover_cte_facts, + rich_type_inference: inputs.rich_type_inference, + }, + ), + nullability: if infer_nullability { + recovered_projection_nullability(projection, &cte_facts) + } else { + ProjectNullability::Unknown + }, + }); + lineage_columns.push(ProjectLineageColumn { + output_column: output_column.clone(), + upstream_columns, + transform_kind, + confidence: if matches!(transform_kind, ProjectTransformKind::Constant) || has_upstream + { + confidence + } else { + ProjectConfidence::Unknown + }, + }); + } + if !inputs.analysis.star_projections.is_empty() && inputs.references.len() == 1 { + let order: HashMap<&str, usize> = inputs + .declared_column_order + .iter() + .enumerate() + .map(|(index, name)| (name.as_str(), index)) + .collect(); + let fallback = order.len(); + columns.sort_by_key(|column| order.get(column.name.as_str()).copied().unwrap_or(fallback)); + lineage_columns.sort_by_key(|column| { + order + .get(column.output_column.as_str()) + .copied() + .unwrap_or(fallback) + }); + } + ProjectAnalysis { + columns, + lineage_columns, + has_star: inputs.analysis.has_root_star, + requires_legacy_fallback, + } +} + +#[derive(Clone, Debug)] +struct CteColumnFact { + data_type: Option, + authoritative_type: bool, + nullability: ProjectNullability, + upstream_columns: Vec, + confidence: ProjectConfidence, +} + +#[derive(Clone, Copy)] +struct ProjectionTypeOptions { + allow_direct_type_hint: bool, + rich_type_inference: bool, +} + +type CteColumnFacts = HashMap<(String, String), CteColumnFact>; + +fn recovered_cte_facts( + analysis: &QueryAnalysis, + function_return_types: &HashMap, + references: &HashMap, + rich_type_inference: bool, +) -> CteColumnFacts { + let mut facts: CteColumnFacts = HashMap::new(); + for cte in &analysis.cte_facts { + for projection in &cte.projections { + let Some(name) = projection.name.as_ref() else { + continue; + }; + let data_type = recovered_projection_type( + projection, + function_return_types, + &facts, + ProjectionTypeOptions { + allow_direct_type_hint: true, + rich_type_inference, + }, + ); + let authoritative_type = + projection_type_is_authoritative(projection, function_return_types, &facts); + let nullability = if cte.shape == Some(QueryShape::SetOperation) { + project_nullability(projection.nullability) + } else { + recovered_projection_nullability(projection, &facts) + }; + let (upstream_columns, confidence) = if cte.shape == Some(QueryShape::SetOperation) { + project_upstream_columns(projection, references) + } else { + recovered_projection_lineage(projection, references, &facts) + }; + facts.insert( + (cte.name.to_lowercase(), name.to_lowercase()), + CteColumnFact { + data_type, + authoritative_type, + nullability, + upstream_columns, + confidence, + }, + ); + } + } + facts +} + +fn recovered_projection_type( + projection: &polyglot_sql::ProjectionFact, + function_return_types: &HashMap, + cte_facts: &CteColumnFacts, + options: ProjectionTypeOptions, +) -> Option { + if projection.transform_kind == TransformKind::Cast + && let Some(data_type) = projection + .cast_type + .as_deref() + .filter(|value| !value.is_empty() && *value != UNKNOWN_SQL_TYPE) + { + return Some(normalize_project_type(data_type)); + } + if let Some(data_type) = projection.transform_function.as_ref().and_then(|function| { + function_return_types + .get(&function.name.to_uppercase()) + .cloned() + }) { + return Some(data_type); + } + if let Some(data_type) = projection_cte_source_type(projection, cte_facts) { + return Some(data_type); + } + if projection.transform_kind == TransformKind::Direct { + let cte_fact = cte_column_fact(projection, cte_facts); + if cte_fact.is_some_and(|fact| fact.authoritative_type) + && let Some(data_type) = cte_fact.and_then(|fact| fact.data_type.clone()) + { + return Some(data_type); + } + if options.allow_direct_type_hint + && let Some(data_type) = project_type(projection, function_return_types, true) + { + return Some(data_type); + } + return cte_fact.and_then(|fact| fact.data_type.clone()); + } + if options.rich_type_inference { + return project_type(projection, function_return_types, true); + } + legacy_expression_type(projection) +} + +fn projection_cte_source_type( + projection: &polyglot_sql::ProjectionFact, + cte_facts: &CteColumnFacts, +) -> Option { + if projection.transform_kind != TransformKind::Direct { + return None; + } + let mut data_types: Vec = projection + .type_column_args + .iter() + .filter_map(|column| { + let source = column.source_name.as_ref().or(column.table.as_ref())?; + cte_facts + .get(&(source.to_lowercase(), column.column.to_lowercase())) + .and_then(|fact| fact.data_type.clone()) + }) + .collect(); + data_types.sort(); + data_types.dedup(); + let [data_type] = data_types.as_slice() else { + return None; + }; + Some(data_type.clone()) +} + +fn direct_projection_requires_legacy_fallback( + projection: &polyglot_sql::ProjectionFact, + inputs: &ProjectAnalysisInputs<'_>, + cte_facts: &CteColumnFacts, +) -> bool { + if projection.transform_kind == TransformKind::Aggregation + && !inputs.analysis.cte_facts.is_empty() + { + return true; + } + if inputs.rich_type_inference + || projection.transform_kind != TransformKind::Direct + || projection.passthrough_source.is_none() + { + return false; + } + let Some(fact) = cte_column_fact(projection, cte_facts) else { + return true; + }; + let Some(fact_type) = fact.data_type.as_deref() else { + return true; + }; + if !fact.authoritative_type { + return true; + } + project_type(projection, inputs.function_return_types, true).is_some_and(|root_type| { + canonical_project_type(&root_type) != canonical_project_type(fact_type) + }) +} + +fn canonical_project_type(value: &str) -> String { + let compact: String = normalize_project_type(value) + .chars() + .filter(|character| !character.is_ascii_whitespace() && *character != '_') + .flat_map(char::to_uppercase) + .collect(); + match compact.as_str() { + "INTEGER" => "INT".to_string(), + "VARCHAR" => "TEXT".to_string(), + _ => compact, + } +} + +fn projection_type_is_authoritative( + projection: &polyglot_sql::ProjectionFact, + function_return_types: &HashMap, + cte_facts: &CteColumnFacts, +) -> bool { + if projection + .cast_type + .as_deref() + .is_some_and(|value| !value.is_empty() && value != UNKNOWN_SQL_TYPE) + { + return true; + } + if projection + .transform_function + .as_ref() + .is_some_and(|function| function_return_types.contains_key(&function.name.to_uppercase())) + { + return true; + } + let source_facts: Vec<&CteColumnFact> = projection + .type_column_args + .iter() + .filter_map(|column| { + let source = column.source_name.as_ref().or(column.table.as_ref())?; + cte_facts.get(&(source.to_lowercase(), column.column.to_lowercase())) + }) + .collect(); + if projection.transform_kind == TransformKind::Direct + && !source_facts.is_empty() + && source_facts.iter().all(|fact| fact.authoritative_type) + { + return true; + } + if projection.transform_kind == TransformKind::Direct + && projection.passthrough_source.is_none() + && projection + .type_hint + .as_deref() + .is_some_and(|value| !value.is_empty() && value != UNKNOWN_SQL_TYPE) + { + return true; + } + projection.transform_kind == TransformKind::Direct + && cte_column_fact(projection, cte_facts).is_some_and(|fact| fact.authoritative_type) +} + +fn legacy_expression_type(projection: &polyglot_sql::ProjectionFact) -> Option { + if projection.transform_kind == TransformKind::Cast { + return projection.cast_type.as_deref().map(normalize_project_type); + } + if projection.transform_kind == TransformKind::Expression + && projection + .type_hint + .as_deref() + .is_some_and(|value| value.eq_ignore_ascii_case("BOOLEAN")) + { + return Some("BOOLEAN".to_string()); + } + None +} + +fn recovered_projection_nullability( + projection: &polyglot_sql::ProjectionFact, + cte_facts: &CteColumnFacts, +) -> ProjectNullability { + let direct = project_nullability(projection.nullability); + if !matches!(direct, ProjectNullability::Unknown) { + return direct; + } + if projection.transform_kind == TransformKind::Direct + && let Some(fact) = cte_column_fact(projection, cte_facts) + { + return fact.nullability; + } + ProjectNullability::Unknown +} + +fn cte_column_fact<'a>( + projection: &polyglot_sql::ProjectionFact, + cte_facts: &'a CteColumnFacts, +) -> Option<&'a CteColumnFact> { + let source = projection.passthrough_source.as_ref()?; + let column = projection.passthrough_column.as_ref()?; + cte_facts.get(&(source.to_lowercase(), column.to_lowercase())) +} + +fn recovered_projection_lineage( + projection: &polyglot_sql::ProjectionFact, + references: &HashMap, + cte_facts: &CteColumnFacts, +) -> (Vec, ProjectConfidence) { + if matches!( + projection.transform_kind, + TransformKind::Direct | TransformKind::Cast + ) && let Some(fact) = cte_column_fact(projection, cte_facts) + { + return (fact.upstream_columns.clone(), fact.confidence); + } + project_upstream_columns(projection, references) +} + +fn project_nullability(value: ProjectionNullability) -> ProjectNullability { + match value { + ProjectionNullability::NonNull => ProjectNullability::NonNull, + ProjectionNullability::Nullable => ProjectNullability::Nullable, + ProjectionNullability::Unknown => ProjectNullability::Unknown, + } +} + +fn project_type( + projection: &polyglot_sql::ProjectionFact, + function_return_types: &HashMap, + normalize_expression_type: bool, +) -> Option { + projection + .cast_type + .as_ref() + .filter(|value| !value.is_empty() && value.as_str() != UNKNOWN_SQL_TYPE) + .or_else(|| { + projection + .type_hint + .as_ref() + .filter(|value| !value.is_empty() && value.as_str() != UNKNOWN_SQL_TYPE) + }) + .map(|value| { + if normalize_expression_type { + normalize_project_type(value) + } else { + value.to_string() + } + }) + .or_else(|| { + projection.transform_function.as_ref().and_then(|function| { + function_return_types + .get(&function.name.to_uppercase()) + .cloned() + }) + }) +} + +fn normalize_project_type(value: &str) -> String { + let upper = value.to_uppercase(); + if upper == VARCHAR_SQL_TYPE { + return "TEXT".to_string(); + } + if let Some(length) = upper + .strip_prefix("TEXT(") + .and_then(|suffix| suffix.strip_suffix(')')) + { + return format!("VARCHAR({length})"); + } + value.to_string() +} + +fn project_upstream_columns( + projection: &polyglot_sql::ProjectionFact, + references: &HashMap, +) -> (Vec, ProjectConfidence) { + let mut columns: Vec = Vec::new(); + let mut seen: HashSet<(&str, &str, &str)> = HashSet::new(); + let mut confidence = ProjectConfidence::High; + for upstream in &projection.upstream { + if upstream.column.is_empty() || upstream.column == SQL_WILDCARD { + continue; + } + let Some(source_name) = upstream + .source_name + .as_deref() + .or(upstream.table.as_deref()) + else { + confidence = ProjectConfidence::Unknown; + continue; + }; + let Some(resource) = references.get(source_name) else { + continue; + }; + if upstream.confidence != ReferenceConfidence::Resolved && upstream.source_alias.is_none() { + confidence = ProjectConfidence::Medium; + } + let key = ( + resource.resource_type.as_str(), + resource.resource_name.as_str(), + upstream.column.as_str(), + ); + if !seen.insert(key) { + continue; + } + columns.push(ProjectLineageSource { + resource_type: resource.resource_type.clone(), + resource_name: resource.resource_name.clone(), + column_name: upstream.column.clone(), + }); + } + let columns = order_project_sources(columns); + (columns, confidence) +} + +fn order_project_sources(mut columns: Vec) -> Vec { + columns.sort_by(|left, right| { + ( + left.resource_type.as_str(), + left.resource_name.as_str(), + left.column_name.as_str(), + ) + .cmp(&( + right.resource_type.as_str(), + right.resource_name.as_str(), + right.column_name.as_str(), + )) + }); + columns.dedup_by(|left, right| { + left.resource_type == right.resource_type + && left.resource_name == right.resource_name + && left.column_name == right.column_name + }); + columns +} + +fn project_transform_kind(kind: TransformKind, has_upstream: bool) -> ProjectTransformKind { + match kind { + TransformKind::Star => ProjectTransformKind::Star, + TransformKind::Cast => ProjectTransformKind::Cast, + TransformKind::Aggregation => ProjectTransformKind::Aggregation, + TransformKind::Constant => ProjectTransformKind::Constant, + TransformKind::Direct => ProjectTransformKind::Direct, + TransformKind::Expression if !has_upstream => ProjectTransformKind::Constant, + TransformKind::Expression => ProjectTransformKind::Expression, + } +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main.rs b/crates/sqlbuild-rules-native/src/query_analysis/main.rs deleted file mode 100644 index 4061a47fb5..0000000000 --- a/crates/sqlbuild-rules-native/src/query_analysis/main.rs +++ /dev/null @@ -1,870 +0,0 @@ -//! Coarse, bounded native batch boundary for compact SQL query analysis. - -use polyglot_sql::{ - AnalyzeQueryOptions, DialectType, ProjectionNullability, QueryAnalysis, QueryShape, - ReferenceConfidence, TransformKind, ValidationSchema, analyze_query, - analyze_query_for_project_projections, -}; -use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator}; -use serde::{Deserialize, Serialize}; -use std::collections::{HashMap, HashSet}; - -use crate::constants::{SQL_WILDCARD, UNKNOWN_SQL_TYPE, VARCHAR_SQL_TYPE}; - -const DEFAULT_WORKERS: usize = 4; -const MAX_WORKERS: usize = 4; -const ANALYSIS_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; - -#[derive(Debug, Deserialize)] -struct AnalysisBatchRequest { - requests: Vec, - #[serde(default = "default_workers")] - workers: usize, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -struct AnalysisRequest { - sql: String, - #[serde(default = "default_dialect")] - dialect: String, - #[serde(default)] - schema: Option, -} - -#[derive(Debug, Serialize)] -#[serde(untagged)] -enum AnalysisResponse { - Success { analysis: QueryAnalysis }, - Failure { error: String }, -} - -#[derive(Debug, Deserialize)] -struct ProjectAnalysisBatchRequest { - requests: Vec, - #[serde(default = "default_workers")] - workers: usize, -} - -#[derive(Debug, Deserialize)] -struct CompactProjectAnalysisBatchRequest { - queries: Vec, - templates: Vec, - projections: Vec, - #[serde(default = "default_workers")] - workers: usize, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct CompactProjectAnalysisTemplateRequest { - query_index: usize, - #[serde(default)] - references: HashMap, - #[serde(default)] - function_return_types: HashMap, - #[serde(default)] - declared_column_order: Vec, - #[serde(default)] - recover_cte_facts: bool, - #[serde(default = "default_true")] - rich_type_inference: bool, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct CompactProjectAnalysisProjectionRequest { - template_index: usize, - #[serde(default)] - resource_names: HashMap, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct ProjectAnalysisRequest { - #[serde(flatten)] - query: AnalysisRequest, - #[serde(default)] - references: HashMap, - #[serde(default)] - function_return_types: HashMap, - #[serde(default)] - declared_column_order: Vec, - #[serde(default)] - recover_cte_facts: bool, - #[serde(default = "default_true")] - rich_type_inference: bool, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -struct LineageResource { - resource_type: String, - resource_name: String, -} - -#[derive(Debug, Serialize)] -#[serde(untagged)] -enum ProjectAnalysisResponse { - Success { analysis: ProjectAnalysis }, - Failure { error: String }, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct ProjectAnalysis { - columns: Vec, - lineage_columns: Vec, - has_star: bool, -} - -struct ProjectAnalysisInputs<'a> { - analysis: &'a QueryAnalysis, - references: &'a HashMap, - function_return_types: &'a HashMap, - declared_column_order: &'a [String], - recover_cte_facts: bool, - rich_type_inference: bool, -} - -struct ProjectAnalysisProjection { - analysis_index: usize, - references: HashMap, - function_return_types: HashMap, - declared_column_order: Vec, - recover_cte_facts: bool, - rich_type_inference: bool, -} - -struct ProjectAnalysisProjectionResult { - projection_index: usize, - resource_names: Vec<(String, String)>, -} - -#[derive(Debug, Serialize)] -struct ProjectColumn { - name: String, - #[serde(rename = "type")] - data_type: Option, - nullability: ProjectNullability, -} - -#[derive(Clone, Copy, Debug, Serialize)] -#[serde(rename_all = "snake_case")] -enum ProjectNullability { - NonNull, - Nullable, - Unknown, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct ProjectLineageColumn { - output_column: String, - upstream_columns: Vec, - transform_kind: ProjectTransformKind, - confidence: ProjectConfidence, -} - -#[derive(Clone, Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct ProjectLineageSource { - resource_type: String, - resource_name: String, - column_name: String, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] -#[serde(rename_all = "snake_case")] -enum ProjectTransformKind { - Direct, - Cast, - Expression, - Aggregation, - Star, - Constant, -} - -#[derive(Clone, Copy, Debug, Serialize)] -#[serde(rename_all = "snake_case")] -enum ProjectConfidence { - High, - Medium, - Unknown, -} - -fn default_workers() -> usize { - DEFAULT_WORKERS -} - -fn default_dialect() -> String { - "generic".to_string() -} - -fn default_true() -> bool { - true -} - -pub(crate) fn analyze_json(request_json: &str) -> Result { - let request: AnalysisBatchRequest = - serde_json::from_str(request_json).map_err(|error| error.to_string())?; - let workers = request.workers.clamp(1, MAX_WORKERS); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers.min(request.requests.len().max(1))) - .stack_size(ANALYSIS_WORKER_STACK_BYTES) - .build() - .map_err(|error| error.to_string())?; - let responses: Vec = pool.install(|| { - request - .requests - .into_par_iter() - .map(analyze_request) - .collect() - }); - serde_json::to_string(&responses).map_err(|error| error.to_string()) -} - -pub(crate) fn analyze_project_json(request_json: &str) -> Result { - let request: ProjectAnalysisBatchRequest = - serde_json::from_str(request_json).map_err(|error| error.to_string())?; - let workers = request.workers.clamp(1, MAX_WORKERS); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers.min(request.requests.len().max(1))) - .stack_size(ANALYSIS_WORKER_STACK_BYTES) - .build() - .map_err(|error| error.to_string())?; - let responses: Vec = pool.install(|| { - request - .requests - .into_par_iter() - .map(analyze_project_request) - .collect() - }); - serde_json::to_string(&responses).map_err(|error| error.to_string()) -} - -pub(crate) fn analyze_project_compact_json(request_json: &str) -> Result { - let request: CompactProjectAnalysisBatchRequest = - serde_json::from_str(request_json).map_err(|error| error.to_string())?; - let workers = request.workers.clamp(1, MAX_WORKERS); - if request - .templates - .iter() - .any(|template| template.query_index >= request.queries.len()) - { - return Err("compact project analysis received an invalid query index".to_string()); - } - if request - .projections - .iter() - .any(|projection| projection.template_index >= request.templates.len()) - { - return Err("compact project analysis received an invalid template index".to_string()); - } - let unique_query_count = request.queries.len(); - let unique_projection_count = request.templates.len(); - let mut project_projection_queries = vec![false; unique_query_count]; - for template in &request.templates { - if template.recover_cte_facts { - project_projection_queries[template.query_index] = true; - } - } - let unique_projections: Vec = request - .templates - .into_iter() - .map(|template| ProjectAnalysisProjection { - analysis_index: template.query_index, - references: template.references, - function_return_types: template.function_return_types, - declared_column_order: template.declared_column_order, - recover_cte_facts: template.recover_cte_facts, - rich_type_inference: template.rich_type_inference, - }) - .collect(); - let projection_results: Vec = request - .projections - .into_iter() - .map(|projection| { - let mut resource_names: Vec<(String, String)> = - projection.resource_names.into_iter().collect(); - resource_names.sort(); - ProjectAnalysisProjectionResult { - projection_index: projection.template_index, - resource_names, - } - }) - .collect(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers.min(unique_query_count.max(1))) - .stack_size(ANALYSIS_WORKER_STACK_BYTES) - .build() - .map_err(|error| error.to_string())?; - let query_analyses: Vec> = pool.install(|| { - request - .queries - .into_par_iter() - .zip(project_projection_queries.into_par_iter()) - .map(|(query, project_projections)| query_analysis(query, project_projections)) - .collect() - }); - let analyses: Vec> = pool.install(|| { - unique_projections - .into_par_iter() - .map(|projection| { - let analysis = query_analyses[projection.analysis_index] - .as_ref() - .map_err(Clone::clone)?; - Ok(project_analysis(ProjectAnalysisInputs { - analysis, - references: &projection.references, - function_return_types: &projection.function_return_types, - declared_column_order: &projection.declared_column_order, - recover_cte_facts: projection.recover_cte_facts, - rich_type_inference: projection.rich_type_inference, - })) - }) - .collect() - }); - serde_json::to_string(&compact_project_batch( - analyses, - projection_results, - unique_query_count, - unique_projection_count, - )) - .map_err(|error| error.to_string()) -} - -fn analyze_request(request: AnalysisRequest) -> AnalysisResponse { - match query_analysis(request, false) { - Ok(analysis) => AnalysisResponse::Success { analysis }, - Err(error) => AnalysisResponse::Failure { error }, - } -} - -fn query_analysis( - request: AnalysisRequest, - project_projections: bool, -) -> Result { - let dialect: DialectType = match request.dialect.parse() { - Ok(dialect) => dialect, - Err(error) => { - return Err(format!( - "unsupported Polyglot dialect '{}': {error}", - request.dialect - )); - } - }; - let options = AnalyzeQueryOptions { - complexity_guard: None, - dialect, - schema: request.schema, - }; - if project_projections { - analyze_query_for_project_projections(&request.sql, options) - } else { - analyze_query(&request.sql, options) - } - .map_err(|error| error.to_string()) -} - -fn analyze_project_request(request: ProjectAnalysisRequest) -> ProjectAnalysisResponse { - match analyze_project_result(request) { - Ok(analysis) => ProjectAnalysisResponse::Success { analysis }, - Err(error) => ProjectAnalysisResponse::Failure { error }, - } -} - -fn analyze_project_result(request: ProjectAnalysisRequest) -> Result { - let analysis = query_analysis(request.query, true)?; - Ok(project_analysis(ProjectAnalysisInputs { - analysis: &analysis, - references: &request.references, - function_return_types: &request.function_return_types, - declared_column_order: &request.declared_column_order, - recover_cte_facts: request.recover_cte_facts, - rich_type_inference: request.rich_type_inference, - })) -} - -type CompactProjectColumn = (usize, Option, u8, u8, u8, Vec<(usize, usize, usize)>); - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct CompactProjectBatch { - strings: Vec, - facts: Vec, - templates: Vec, - analyses: Vec, - unique_query_count: usize, - unique_projection_count: usize, -} - -#[derive(Debug, Serialize)] -#[serde(untagged)] -enum CompactProjectResponse { - Success((Vec, bool)), - Failure(String), -} - -type CompactProjectProjection = (usize, Vec<(usize, usize)>); - -#[derive(Default)] -struct StringInterner { - strings: Vec, - indexes: HashMap, -} - -impl StringInterner { - fn intern(&mut self, value: String) -> usize { - if let Some(index) = self.indexes.get(&value) { - return *index; - } - let index = self.strings.len(); - self.indexes.insert(value.clone(), index); - self.strings.push(value); - index - } -} - -fn compact_project_batch( - analyses: Vec>, - projections: Vec, - unique_query_count: usize, - unique_projection_count: usize, -) -> CompactProjectBatch { - let mut interner = StringInterner::default(); - let mut facts: Vec = Vec::new(); - let mut fact_indexes: HashMap = HashMap::new(); - let mut templates: Vec = Vec::with_capacity(analyses.len()); - for result in analyses { - let analysis = match result { - Ok(analysis) => analysis, - Err(error) => { - templates.push(CompactProjectResponse::Failure(error)); - continue; - } - }; - let mut columns = Vec::with_capacity(analysis.columns.len()); - for (column, lineage) in analysis.columns.into_iter().zip(analysis.lineage_columns) { - let upstream: Vec<(usize, usize, usize)> = lineage - .upstream_columns - .into_iter() - .map(|source| { - ( - interner.intern(source.resource_type), - interner.intern(source.resource_name), - interner.intern(source.column_name), - ) - }) - .collect(); - let fact = ( - interner.intern(column.name), - column.data_type.map(|value| interner.intern(value)), - nullability_code(column.nullability), - transform_code(lineage.transform_kind), - confidence_code(lineage.confidence), - upstream, - ); - let next_index = facts.len(); - let fact_index = *fact_indexes.entry(fact.clone()).or_insert_with(|| { - facts.push(fact); - next_index - }); - columns.push(fact_index); - } - templates.push(CompactProjectResponse::Success(( - columns, - analysis.has_star, - ))); - } - let mut compact_projections: Vec = - Vec::with_capacity(projections.len()); - for projection in projections { - let mut resource_names: Vec<(usize, usize)> = - Vec::with_capacity(projection.resource_names.len()); - for (canonical_name, resource_name) in projection.resource_names { - resource_names.push(( - interner.intern(canonical_name), - interner.intern(resource_name), - )); - } - compact_projections.push((projection.projection_index, resource_names)); - } - CompactProjectBatch { - strings: interner.strings, - facts, - templates, - analyses: compact_projections, - unique_query_count, - unique_projection_count, - } -} - -fn nullability_code(value: ProjectNullability) -> u8 { - match value { - ProjectNullability::Unknown => 0, - ProjectNullability::NonNull => 1, - ProjectNullability::Nullable => 2, - } -} - -fn transform_code(value: ProjectTransformKind) -> u8 { - match value { - ProjectTransformKind::Direct => 0, - ProjectTransformKind::Cast => 1, - ProjectTransformKind::Expression => 2, - ProjectTransformKind::Aggregation => 3, - ProjectTransformKind::Star => 4, - ProjectTransformKind::Constant => 5, - } -} - -fn confidence_code(value: ProjectConfidence) -> u8 { - match value { - ProjectConfidence::Unknown => 0, - ProjectConfidence::High => 1, - ProjectConfidence::Medium => 2, - } -} - -fn project_analysis(inputs: ProjectAnalysisInputs<'_>) -> ProjectAnalysis { - let cte_facts: CteColumnFacts = if inputs.recover_cte_facts { - recovered_cte_facts( - inputs.analysis, - inputs.function_return_types, - inputs.references, - inputs.rich_type_inference, - ) - } else { - HashMap::new() - }; - let infer_nullability = inputs.analysis.shape != QueryShape::SetOperation; - let mut columns: Vec = Vec::new(); - let mut lineage_columns: Vec = Vec::new(); - for projection in inputs - .analysis - .projections - .iter() - .filter(|projection| !projection.is_star) - { - let Some(output_column) = projection - .name - .as_ref() - .filter(|name| !name.is_empty() && name.as_str() != SQL_WILDCARD) - else { - continue; - }; - let (upstream_columns, confidence) = - recovered_projection_lineage(projection, inputs.references, &cte_facts); - let has_upstream = !upstream_columns.is_empty(); - let transform_kind = project_transform_kind(projection.transform_kind, has_upstream); - columns.push(ProjectColumn { - name: output_column.clone(), - data_type: recovered_projection_type( - projection, - inputs.function_return_types, - &cte_facts, - ProjectionTypeOptions { - allow_direct_type_hint: inputs.rich_type_inference, - rich_type_inference: inputs.rich_type_inference, - }, - ), - nullability: if infer_nullability { - recovered_projection_nullability(projection, &cte_facts) - } else { - ProjectNullability::Unknown - }, - }); - lineage_columns.push(ProjectLineageColumn { - output_column: output_column.clone(), - upstream_columns, - transform_kind, - confidence: if matches!(transform_kind, ProjectTransformKind::Constant) || has_upstream - { - confidence - } else { - ProjectConfidence::Unknown - }, - }); - } - if !inputs.analysis.star_projections.is_empty() && inputs.references.len() == 1 { - let order: HashMap<&str, usize> = inputs - .declared_column_order - .iter() - .enumerate() - .map(|(index, name)| (name.as_str(), index)) - .collect(); - let fallback = order.len(); - columns.sort_by_key(|column| order.get(column.name.as_str()).copied().unwrap_or(fallback)); - lineage_columns.sort_by_key(|column| { - order - .get(column.output_column.as_str()) - .copied() - .unwrap_or(fallback) - }); - } - ProjectAnalysis { - columns, - lineage_columns, - has_star: inputs.analysis.has_root_star, - } -} - -#[derive(Clone, Debug)] -struct CteColumnFact { - data_type: Option, - nullability: ProjectNullability, - upstream_columns: Vec, - confidence: ProjectConfidence, -} - -#[derive(Clone, Copy)] -struct ProjectionTypeOptions { - allow_direct_type_hint: bool, - rich_type_inference: bool, -} - -type CteColumnFacts = HashMap<(String, String), CteColumnFact>; - -fn recovered_cte_facts( - analysis: &QueryAnalysis, - function_return_types: &HashMap, - references: &HashMap, - rich_type_inference: bool, -) -> CteColumnFacts { - let mut facts: CteColumnFacts = HashMap::new(); - for cte in &analysis.cte_facts { - for projection in &cte.projections { - let Some(name) = projection.name.as_ref() else { - continue; - }; - let data_type = recovered_projection_type( - projection, - function_return_types, - &facts, - ProjectionTypeOptions { - allow_direct_type_hint: true, - rich_type_inference, - }, - ); - let nullability = if cte.shape == Some(QueryShape::SetOperation) { - project_nullability(projection.nullability) - } else { - recovered_projection_nullability(projection, &facts) - }; - let (upstream_columns, confidence) = - recovered_projection_lineage(projection, references, &facts); - facts.insert( - (cte.name.to_lowercase(), name.to_lowercase()), - CteColumnFact { - data_type, - nullability, - upstream_columns, - confidence, - }, - ); - } - } - facts -} - -fn recovered_projection_type( - projection: &polyglot_sql::ProjectionFact, - function_return_types: &HashMap, - cte_facts: &CteColumnFacts, - options: ProjectionTypeOptions, -) -> Option { - if let Some(data_type) = projection.transform_function.as_ref().and_then(|function| { - function_return_types - .get(&function.name.to_uppercase()) - .cloned() - }) { - return Some(data_type); - } - if projection.transform_kind == TransformKind::Direct { - if let Some(fact) = cte_column_fact(projection, cte_facts) - && fact.data_type.is_some() - { - return fact.data_type.clone(); - } - return options - .allow_direct_type_hint - .then(|| project_type(projection, function_return_types, false)) - .flatten(); - } - if options.rich_type_inference { - return project_type(projection, function_return_types, true); - } - legacy_expression_type(projection) -} - -fn legacy_expression_type(projection: &polyglot_sql::ProjectionFact) -> Option { - if projection.transform_kind == TransformKind::Cast { - return projection.cast_type.as_deref().map(normalize_project_type); - } - if projection.transform_kind == TransformKind::Expression - && projection - .type_hint - .as_deref() - .is_some_and(|value| value.eq_ignore_ascii_case("BOOLEAN")) - { - return Some("BOOLEAN".to_string()); - } - None -} - -fn recovered_projection_nullability( - projection: &polyglot_sql::ProjectionFact, - cte_facts: &CteColumnFacts, -) -> ProjectNullability { - let direct = project_nullability(projection.nullability); - if !matches!(direct, ProjectNullability::Unknown) { - return direct; - } - if projection.transform_kind == TransformKind::Direct - && let Some(fact) = cte_column_fact(projection, cte_facts) - { - return fact.nullability; - } - ProjectNullability::Unknown -} - -fn cte_column_fact<'a>( - projection: &polyglot_sql::ProjectionFact, - cte_facts: &'a CteColumnFacts, -) -> Option<&'a CteColumnFact> { - let source = projection.passthrough_source.as_ref()?; - let column = projection.passthrough_column.as_ref()?; - cte_facts.get(&(source.to_lowercase(), column.to_lowercase())) -} - -fn recovered_projection_lineage( - projection: &polyglot_sql::ProjectionFact, - references: &HashMap, - cte_facts: &CteColumnFacts, -) -> (Vec, ProjectConfidence) { - if projection.transform_kind == TransformKind::Direct - && let Some(fact) = cte_column_fact(projection, cte_facts) - { - return (fact.upstream_columns.clone(), fact.confidence); - } - project_upstream_columns(projection, references) -} - -fn project_nullability(value: ProjectionNullability) -> ProjectNullability { - match value { - ProjectionNullability::NonNull => ProjectNullability::NonNull, - ProjectionNullability::Nullable => ProjectNullability::Nullable, - ProjectionNullability::Unknown => ProjectNullability::Unknown, - } -} - -fn project_type( - projection: &polyglot_sql::ProjectionFact, - function_return_types: &HashMap, - normalize_expression_type: bool, -) -> Option { - projection - .cast_type - .as_ref() - .filter(|value| !value.is_empty() && value.as_str() != UNKNOWN_SQL_TYPE) - .or_else(|| { - projection - .type_hint - .as_ref() - .filter(|value| !value.is_empty() && value.as_str() != UNKNOWN_SQL_TYPE) - }) - .map(|value| { - if normalize_expression_type { - normalize_project_type(value) - } else { - value.to_string() - } - }) - .or_else(|| { - projection.transform_function.as_ref().and_then(|function| { - function_return_types - .get(&function.name.to_uppercase()) - .cloned() - }) - }) -} - -fn normalize_project_type(value: &str) -> String { - let upper = value.to_uppercase(); - if upper == VARCHAR_SQL_TYPE { - return "TEXT".to_string(); - } - if let Some(length) = upper - .strip_prefix("TEXT(") - .and_then(|suffix| suffix.strip_suffix(')')) - { - return format!("VARCHAR({length})"); - } - value.to_string() -} - -fn project_upstream_columns( - projection: &polyglot_sql::ProjectionFact, - references: &HashMap, -) -> (Vec, ProjectConfidence) { - let mut columns: Vec = Vec::new(); - let mut seen: HashSet<(&str, &str, &str)> = HashSet::new(); - let mut confidence = ProjectConfidence::High; - for upstream in &projection.upstream { - if upstream.column.is_empty() || upstream.column == SQL_WILDCARD { - continue; - } - let Some(source_name) = upstream - .source_name - .as_deref() - .or(upstream.table.as_deref()) - else { - confidence = ProjectConfidence::Unknown; - continue; - }; - let Some(resource) = references.get(source_name) else { - continue; - }; - if upstream.confidence != ReferenceConfidence::Resolved && upstream.source_alias.is_none() { - confidence = ProjectConfidence::Medium; - } - let key = ( - resource.resource_type.as_str(), - resource.resource_name.as_str(), - upstream.column.as_str(), - ); - if !seen.insert(key) { - continue; - } - columns.push(ProjectLineageSource { - resource_type: resource.resource_type.clone(), - resource_name: resource.resource_name.clone(), - column_name: upstream.column.clone(), - }); - } - columns.sort_by(|left, right| { - ( - left.resource_type.as_str(), - left.resource_name.as_str(), - left.column_name.as_str(), - ) - .cmp(&( - right.resource_type.as_str(), - right.resource_name.as_str(), - right.column_name.as_str(), - )) - }); - (columns, confidence) -} - -fn project_transform_kind(kind: TransformKind, has_upstream: bool) -> ProjectTransformKind { - match kind { - TransformKind::Star => ProjectTransformKind::Star, - TransformKind::Cast => ProjectTransformKind::Cast, - TransformKind::Aggregation => ProjectTransformKind::Aggregation, - TransformKind::Constant => ProjectTransformKind::Constant, - TransformKind::Direct => ProjectTransformKind::Direct, - TransformKind::Expression if !has_upstream => ProjectTransformKind::Constant, - TransformKind::Expression => ProjectTransformKind::Expression, - } -} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/analyze.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze.rs new file mode 100644 index 0000000000..51af080d5b --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze.rs @@ -0,0 +1,3 @@ +pub(crate) fn analyze_json(request_json: &str) -> Result { + crate::query_analysis::engine::analyze_json(request_json) +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project.rs new file mode 100644 index 0000000000..838dd720e6 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project.rs @@ -0,0 +1,3 @@ +pub(crate) fn analyze_project_json(request_json: &str) -> Result { + crate::query_analysis::engine::analyze_project_json(request_json) +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_compact.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_compact.rs new file mode 100644 index 0000000000..5b8f2a7da5 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_compact.rs @@ -0,0 +1,3 @@ +pub(crate) fn analyze_project_compact_json(request_json: &str) -> Result { + crate::query_analysis::engine::analyze_project_compact_json(request_json) +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs new file mode 100644 index 0000000000..dfdba3210b --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod analyze; +pub(crate) mod analyze_project; +pub(crate) mod analyze_project_compact; diff --git a/crates/sqlbuild-rules-native/src/query_analysis/mod.rs b/crates/sqlbuild-rules-native/src/query_analysis/mod.rs index c90d047f0b..6350d84af2 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/mod.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/mod.rs @@ -1,3 +1,9 @@ +#[path = "_helpers/borrowed_facts.rs"] +mod borrowed_facts; +#[path = "_helpers/compatibility_types.rs"] +mod compatibility_types; +#[path = "_helpers/engine.rs"] +mod engine; pub(crate) mod main; #[cfg(test)] diff --git a/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs b/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs index 4fb542fd46..4e849059da 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs @@ -1,6 +1,142 @@ -use crate::query_analysis::main::analyze_project_compact_json; +use crate::query_analysis::main::analyze_project_compact::analyze_project_compact_json; use serde_json::{Value, json}; +pub(crate) fn borrowed_facts_preserve_named_outputs_and_terminal_sources() -> bool { + let schema: polyglot_sql::ValidationSchema = serde_json::from_value(json!({ + "tables": [{"name": "orders", "columns": [ + {"name": "order_id", "type": "BIGINT", "nullable": false}, + {"name": "selected_id", "type": "BIGINT", "nullable": true} + ]}] + })) + .expect("valid schema"); + let expression = polyglot_sql::parse_one( + "WITH renamed AS (SELECT order_id AS selected_id FROM orders), nested AS (SELECT * FROM (SELECT selected_id FROM renamed) AS selected_orders(order_key)) SELECT order_key FROM nested", + 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.len(), 1); + assert_eq!(facts[0].name, "order_key"); + assert_eq!( + *facts[0].upstream, + std::collections::BTreeSet::from([("orders".to_string(), "order_id".to_string())]) + ); + assert_eq!( + facts[0].nullability, + polyglot_sql::ProjectionNullability::NonNull + ); + let expression = polyglot_sql::parse_one( + "WITH combined AS (SELECT order_id AS first_key FROM orders UNION ALL BY NAME SELECT order_id AS second_key FROM orders) SELECT first_key, second_key FROM combined", + 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.len(), 2); + assert!( + facts + .iter() + .all(|fact| fact.nullability == polyglot_sql::ProjectionNullability::Nullable) + ); + let types = crate::query_analysis::compatibility_types::infer( + &expression, + Some(&schema), + &std::collections::HashMap::new(), + polyglot_sql::DialectType::Snowflake, + ); + assert_eq!(types.get("first_key").map(String::as_str), Some("BIGINT")); + assert_eq!(types.get("second_key").map(String::as_str), Some("BIGINT")); + true +} + +pub(crate) fn combined_queries_preserve_standalone_binding() -> bool { + let schema = json!({"strict": true, "tables": [ + {"name": "orders", "columns": [{"name": "order_id", "type": "BIGINT"}]}, + {"name": "customers", "columns": [{"name": "customer_id", "type": "BIGINT"}]} + ]}); + for sql in [ + "SELECT order_id FROM orders", + "SELECT missing FROM orders", + "WITH selected AS (SELECT order_id FROM orders) SELECT order_id > 0 AS result FROM selected", + "WITH first_pass AS (SELECT CAST(order_id AS DOUBLE) AS amount, order_id AS unused FROM orders), second_pass AS (SELECT amount AS total FROM first_pass) SELECT total FROM second_pass", + "WITH first_pass (selected_id, unused_id) AS (SELECT order_id, order_id + 1 FROM orders) SELECT selected_id FROM first_pass", + "WITH first_pass AS (SELECT order_id FROM orders), combined AS (SELECT order_id FROM first_pass UNION ALL SELECT order_id FROM orders) SELECT order_id FROM combined", + "WITH selected AS (SELECT order_id FROM orders) SELECT order_id FROM selected WHERE order_id IS NOT NULL", + "WITH orders AS (SELECT * FROM orders), keys AS (SELECT order_id FROM orders WHERE order_id IS NOT NULL UNION ALL SELECT 0 AS order_id), selected AS (SELECT k.order_id FROM keys k), typed AS (SELECT CAST(s.order_id AS BIGINT) AS order_id FROM selected s) SELECT order_id FROM typed", + "WITH orders AS (SELECT * FROM orders), typed AS (SELECT CAST(CAST(o.order_id AS INT) AS NUMBER(38, 0)) AS order_id FROM orders o) SELECT order_id FROM typed", + "WITH combined AS (SELECT order_id FROM orders UNION ALL BY NAME SELECT customer_id AS order_id FROM customers), typed AS (SELECT CAST(order_id AS BIGINT) AS order_id FROM combined) SELECT order_id FROM typed", + "SELECT order_id FROM orders QUALIFY missing > 0", + "SELECT orders.order_id FROM orders JOIN customers USING (order_id)", + "SELECT FROM", + "SELECT order_id FROM orders; SELECT customer_id FROM customers", + ] { + let standalone: Value = serde_json::from_str( + &crate::semantic_validation::main::validation_json(&json!({ + "sql": sql, "dialect": "snowflake", "schema": schema, + "options": {"check_types": false, "check_references": true, "strict": true, "semantic": false, "strict_syntax": false} + }).to_string()).expect("standalone binding should return evidence"), + ).expect("standalone evidence should decode"); + let mut request = json!({ + "queries": [{"sql": sql, "dialect": "snowflake", "schema": schema}], + "templates": [{"queryIndex": 0, "recoverCteFacts": true, "references": { + "orders": {"resourceType": "model", "resourceName": "orders"}, + "customers": {"resourceType": "model", "resourceName": "customers"} + }}], + "projections": [{"templateIndex": 0}] + }); + let separate: Value = serde_json::from_str( + &analyze_project_compact_json(&request.to_string()) + .expect("separate analysis should return evidence"), + ) + .expect("separate evidence should decode"); + request["queries"][0]["binding_schema"] = schema.clone(); + let mut combined: Value = serde_json::from_str( + &analyze_project_compact_json(&request.to_string()) + .expect("combined analysis should return evidence"), + ) + .expect("combined evidence should decode"); + assert_eq!( + combined["validations"][0], standalone, + "binding mismatch for {sql}" + ); + combined + .as_object_mut() + .expect("response is an object") + .remove("validations"); + assert_eq!(combined, separate, "analysis mismatch for {sql}"); + } + true +} + +pub(crate) fn widening_aggregates_require_compatibility_recovery() -> bool { + for expression in [ + "AVG(order_id)", + "SUM(order_id)", + "MIN(order_id) > 0", + "MIN(order_id) / 2.0", + ] { + for rich in [false, true] { + let response: Value = serde_json::from_str( + &analyze_project_compact_json(&json!({ + "queries": [{"sql": format!("WITH selected AS (SELECT CAST(1 AS BIGINT) AS order_id) SELECT {expression} AS result FROM selected"), "dialect": "duckdb"}], + "templates": [{"queryIndex": 0, "recoverCteFacts": true, "richTypeInference": rich}], + "projections": [{"templateIndex": 0}] + }).to_string()).expect("query should analyze"), + ).expect("response should decode"); + assert_eq!( + response["templates"][0], + "native project type recovery requires legacy fallback" + ); + } + } + true +} + pub(crate) fn compact_project_query_interns_repeated_lineage_strings() -> bool { let response: Value = serde_json::from_str( &analyze_project_compact_json( @@ -67,6 +203,55 @@ pub(crate) fn repeated_project_facts_intern_complete_facts() -> bool { true } +pub(crate) fn interleaved_query_templates_preserve_template_order() -> bool { + let response: Value = serde_json::from_str( + &analyze_project_compact_json( + &json!({ + "queries": [ + {"sql": "SELECT 1 AS first_value", "dialect": "duckdb"}, + {"sql": "SELECT 2 AS second_value", "dialect": "duckdb"} + ], + "templates": [ + {"queryIndex": 1}, + {"queryIndex": 0}, + {"queryIndex": 1} + ], + "projections": [ + {"templateIndex": 0}, + {"templateIndex": 1}, + {"templateIndex": 2} + ] + }) + .to_string(), + ) + .expect("test assumption must hold"), + ) + .expect("test assumption must hold"); + + let strings = response["strings"] + .as_array() + .expect("compact strings should be an array"); + let facts = response["facts"] + .as_array() + .expect("compact facts should be an array"); + let templates = response["templates"] + .as_array() + .expect("compact templates should be an array"); + let names: Vec<&Value> = templates + .iter() + .map(|template| { + let fact_index = template[0][0] + .as_u64() + .expect("template should use a fact index") as usize; + let name_index = facts[fact_index][0] + .as_u64() + .expect("fact should use a name index") as usize; + &strings[name_index] + }) + .collect(); + names == vec!["second_value", "first_value", "second_value"] +} + pub(crate) fn canonical_queries_reuse_semantics_and_project_resources() -> bool { let response: Value = serde_json::from_str( &analyze_project_compact_json( @@ -135,3 +320,94 @@ pub(crate) fn canonical_queries_reuse_semantics_and_project_resources() -> bool assert_eq!(strings[second_resource_index], "archived_orders"); true } +pub(crate) fn native_compatibility_types_preserve_result_semantics() -> bool { + let expression = polyglot_sql::parse_one( + "WITH orders AS (SELECT CAST(1 AS INTEGER) AS order_id), values AS (SELECT CASE WHEN order_id > 0 THEN CAST(order_id AS BIGINT) ELSE NULL END AS result, order_id > 0 AS positive, SUM(order_id) AS total FROM orders GROUP BY order_id) SELECT result, positive, total FROM values", + polyglot_sql::DialectType::Snowflake, + ).expect("valid SQL"); + let functions = std::collections::HashMap::from([("SUM".to_string(), "FLOAT".to_string())]); + let types = crate::query_analysis::compatibility_types::infer( + &expression, + None, + &functions, + polyglot_sql::DialectType::Snowflake, + ); + assert_eq!(types.get("result").map(String::as_str), Some("BIGINT")); + assert_eq!(types.get("positive").map(String::as_str), Some("BOOLEAN")); + assert_eq!(types.get("total").map(String::as_str), Some("FLOAT")); + let expression = polyglot_sql::parse_one( + "WITH renamed(unknown_value, order_id) AS (SELECT unknown_function() AS original_value, CAST(1 AS BIGINT) AS original_id) SELECT * FROM (SELECT order_id FROM renamed) AS nested_orders", + polyglot_sql::DialectType::Snowflake, + ).expect("valid SQL"); + let types = crate::query_analysis::compatibility_types::infer( + &expression, + None, + &functions, + polyglot_sql::DialectType::Snowflake, + ); + assert_eq!(types.get("order_id").map(String::as_str), Some("BIGINT")); + true +} +pub(crate) fn borrowed_facts_preserve_union_dependencies_and_join_nullability() -> bool { + let schema: polyglot_sql::ValidationSchema = serde_json::from_value(serde_json::json!({ + "tables": [ + {"name": "orders", "columns": [{"name": "order_id", "type": "INTEGER", "nullable": false}, {"name": "payload", "type": "VARIANT"}]}, + {"name": "customers", "columns": [{"name": "customer_id", "type": "INTEGER", "nullable": false}]} + ] + })).expect("valid schema"); + let expression = polyglot_sql::parse_one( + "WITH combined AS (SELECT order_id FROM orders UNION ALL BY NAME SELECT customer_id AS order_id FROM customers), typed AS (SELECT CAST(order_id AS BIGINT) AS order_id FROM combined) SELECT order_id FROM typed", + 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[0].upstream, + std::collections::BTreeSet::from([ + ("orders".to_string(), "order_id".to_string()), + ("customers".to_string(), "customer_id".to_string()), + ]) + ); + assert_eq!( + facts[0].nullability, + polyglot_sql::ProjectionNullability::NonNull + ); + 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, + ).expect("valid SQL"); + let facts = crate::query_analysis::borrowed_facts::infer( + &expression, + Some(&schema), + polyglot_sql::DialectType::Snowflake, + ); + assert_eq!( + facts[0].nullability, + polyglot_sql::ProjectionNullability::NonNull + ); + assert_eq!( + facts[1].nullability, + polyglot_sql::ProjectionNullability::Nullable + ); + let expression = polyglot_sql::parse_one( + "WITH expanded AS (SELECT f.value AS item FROM orders o, LATERAL FLATTEN(INPUT => o.payload) f), renamed AS (SELECT item AS first_item, COALESCE(first_item, NULL) AS second_item FROM expanded) SELECT second_item FROM renamed", + 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[0].upstream, + std::collections::BTreeSet::from([("orders".to_string(), "payload".to_string())]) + ); + assert_eq!( + facts[0].nullability, + polyglot_sql::ProjectionNullability::Unknown + ); + true +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs b/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs index 238d99457c..cd78e51a87 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs @@ -1,8 +1,15 @@ -use crate::query_analysis::main::{analyze_json, analyze_project_json}; +use crate::query_analysis::main::analyze::analyze_json; +use crate::query_analysis::main::analyze_project::analyze_project_json; use crate::query_analysis::tests::helpers::{ + borrowed_facts_preserve_named_outputs_and_terminal_sources, + borrowed_facts_preserve_union_dependencies_and_join_nullability, canonical_queries_reuse_semantics_and_project_resources, + combined_queries_preserve_standalone_binding, compact_project_query_interns_repeated_lineage_strings, + interleaved_query_templates_preserve_template_order, + native_compatibility_types_preserve_result_semantics, repeated_project_facts_intern_complete_facts, + widening_aggregates_require_compatibility_recovery, }; use crate::query_analysis::tests::test_types::{ CompactQueryAnalysisTestCase, QueryAnalysisExpectedValue, QueryAnalysisTestCase, @@ -12,6 +19,31 @@ use serde_json::{Value, json}; #[test] fn given_compact_query_cases_when_analyzing_projects_then_expected_behavior_holds() { let test_cases = [ + CompactQueryAnalysisTestCase { + description: "borrowed facts preserve named outputs and actual terminal columns", + run: borrowed_facts_preserve_named_outputs_and_terminal_sources, + expected_success: true, + }, + CompactQueryAnalysisTestCase { + description: "borrowed output graph preserves union dependencies and outer-join nullability", + run: borrowed_facts_preserve_union_dependencies_and_join_nullability, + expected_success: true, + }, + CompactQueryAnalysisTestCase { + description: "borrowed native type evaluation preserves expression result semantics", + run: native_compatibility_types_preserve_result_semantics, + expected_success: true, + }, + CompactQueryAnalysisTestCase { + description: "combined compilation preserves standalone analysis and binding", + run: combined_queries_preserve_standalone_binding, + expected_success: true, + }, + CompactQueryAnalysisTestCase { + description: "widening aggregates preserve compatibility recovery", + run: widening_aggregates_require_compatibility_recovery, + expected_success: true, + }, CompactQueryAnalysisTestCase { description: "compact project queries intern repeated lineage strings", run: compact_project_query_interns_repeated_lineage_strings, @@ -27,6 +59,11 @@ fn given_compact_query_cases_when_analyzing_projects_then_expected_behavior_hold run: canonical_queries_reuse_semantics_and_project_resources, expected_success: true, }, + CompactQueryAnalysisTestCase { + description: "interleaved query templates preserve request order", + run: interleaved_query_templates_preserve_template_order, + expected_success: true, + }, ]; for test_case in test_cases { @@ -42,6 +79,44 @@ fn given_compact_query_cases_when_analyzing_projects_then_expected_behavior_hold #[test] fn given_query_analysis_cases_when_analyzing_batch_then_returns_expected_facts() { let test_cases = [ + QueryAnalysisTestCase { + description: "cast output retains both named union branches", + request: json!({"requests": [{ + "sql": "WITH combined AS (SELECT order_id FROM orders UNION ALL BY NAME SELECT customer_id AS order_id FROM customers), typed AS (SELECT CAST(order_id AS BIGINT) AS order_id FROM combined) SELECT order_id FROM typed", + "dialect": "snowflake", "recoverCteFacts": true, + "references": { + "orders": {"resourceType": "model", "resourceName": "orders"}, + "customers": {"resourceType": "model", "resourceName": "customers"} + } + }]}), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![ + QueryAnalysisExpectedValue { + pointer: "/0/analysis/lineageColumns/0/upstreamColumns/0/resourceName", + value: "customers", + }, + QueryAnalysisExpectedValue { + pointer: "/0/analysis/lineageColumns/0/upstreamColumns/1/resourceName", + value: "orders", + }, + ], + expected_nonempty_strings: vec![], + }, + QueryAnalysisTestCase { + description: "CTE comparisons preserve Boolean result types", + request: json!({"requests": [{ + "sql": "WITH selected AS (SELECT CAST(1 AS BIGINT) AS order_id) SELECT order_id > 0 AS result FROM selected", + "dialect": "duckdb", "recoverCteFacts": true, "richTypeInference": false + }]}), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![QueryAnalysisExpectedValue { + pointer: "/0/analysis/columns/0/type", + value: "BOOLEAN", + }], + expected_nonempty_strings: vec![], + }, QueryAnalysisTestCase { description: "raw batches preserve order and isolate parse errors", request: json!({ @@ -157,6 +232,110 @@ fn given_query_analysis_cases_when_analyzing_batch_then_returns_expected_facts() ], expected_nonempty_strings: vec![], }, + QueryAnalysisTestCase { + description: "root schema type takes precedence over normalized CTE recovery", + request: json!({ + "requests": [{ + "sql": "WITH selected AS (SELECT attributes FROM orders UNION ALL SELECT attributes FROM orders) SELECT attributes FROM selected", + "dialect": "snowflake", + "schema": {"tables": [{ + "name": "orders", + "columns": [{"name": "attributes", "type": "VARIANT"}] + }]}, + "richTypeInference": false, + "recoverCteFacts": true + }] + }), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![QueryAnalysisExpectedValue { + pointer: "/0/analysis/columns/0/type", + value: "VARIANT", + }], + expected_nonempty_strings: vec![], + }, + QueryAnalysisTestCase { + description: "specialized aggregate return types propagate through CTEs", + request: json!({ + "requests": [{ + "sql": "WITH aggregated AS (SELECT OBJECT_AGG(product_id, TO_VARIANT(quantity)) AS inventory FROM products), selected AS (SELECT inventory FROM aggregated) SELECT inventory FROM selected", + "dialect": "snowflake", + "schema": {"tables": [{ + "name": "products", + "columns": [ + {"name": "product_id", "type": "NUMBER"}, + {"name": "quantity", "type": "NUMBER"} + ] + }]}, + "functionReturnTypes": {"OBJECT_AGG": "OBJECT", "TO_VARIANT": "VARIANT"}, + "recoverCteFacts": true + }] + }), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![QueryAnalysisExpectedValue { + pointer: "/0/analysis/columns/0/type", + value: "OBJECT", + }], + expected_nonempty_strings: vec![], + }, + QueryAnalysisTestCase { + description: "outer casts take precedence over nested function return types", + request: json!({ + "requests": [{ + "sql": "WITH converted AS (SELECT CAST(TO_OBJECT(attributes) AS VARIANT) AS attributes FROM products) SELECT attributes FROM converted", + "dialect": "snowflake", + "schema": {"tables": [{ + "name": "products", + "columns": [{"name": "attributes", "type": "TEXT"}] + }]}, + "functionReturnTypes": {"TO_OBJECT": "OBJECT"}, + "recoverCteFacts": true + }] + }), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![QueryAnalysisExpectedValue { + pointer: "/0/analysis/columns/0/type", + value: "VARIANT", + }], + expected_nonempty_strings: vec![], + }, + QueryAnalysisTestCase { + description: "type-preserving aggregates inherit CTE argument types", + request: json!({ + "requests": [{ + "sql": "WITH typed AS (SELECT CAST(created_at AS TIMESTAMP_NTZ) AS created_at, CAST(is_active AS BOOLEAN) AS is_active FROM orders), aggregated AS (SELECT MAX(IFF(is_active, created_at, NULL)) AS latest_at FROM typed) SELECT latest_at FROM aggregated", + "dialect": "snowflake", + "recoverCteFacts": true + }] + }), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![QueryAnalysisExpectedValue { + pointer: "/0/analysis/columns/0/type", + value: "TIMESTAMPNTZ", + }], + expected_nonempty_strings: vec![], + }, + QueryAnalysisTestCase { + description: "conditional expressions inherit CTE value types", + request: json!({ + "requests": [{ + "sql": "WITH typed AS (SELECT CAST(fulfilled_at AS DATE) AS fulfilled_at FROM orders), selected AS (SELECT COALESCE(fulfilled_at, CAST(NULL AS DATE)) AS fulfilled_at FROM typed) SELECT fulfilled_at FROM selected", + "dialect": "snowflake", + "recoverCteFacts": true, + "richTypeInference": false + }] + }), + analyze: analyze_project_json, + expected_length: 1, + expected_values: vec![QueryAnalysisExpectedValue { + pointer: "/0/analysis/columns/0/type", + value: "DATE", + }], + expected_nonempty_strings: vec![], + }, QueryAnalysisTestCase { description: "three-stage set operation CTEs preserve expression types", request: json!({ diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/main.rs b/crates/sqlbuild-rules-native/src/semantic_validation/main.rs index 1405a850d6..fde57ff38d 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/main.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/main.rs @@ -22,6 +22,15 @@ struct ValidationRequest { options: SchemaValidationOptions, } +#[derive(Debug)] +pub(crate) struct ParsedValidationRequest<'a> { + pub(crate) sql: &'a str, + pub(crate) dialect: DialectType, + pub(crate) schema: &'a ValidationSchema, + pub(crate) result: ValidationResult, + pub(crate) expression: Option<&'a Expression>, +} + struct ClauseResolver<'a, 'b> { scope: &'a polyglot_sql::Scope, resolver: Resolver<'b>, @@ -67,12 +76,18 @@ fn validation_result(request: ValidationRequest) -> Result Result bool { + sql.as_bytes() + .windows(5) + .any(|word| word.eq_ignore_ascii_case(b"using")) + || sql + .as_bytes() + .windows(7) + .any(|word| word.eq_ignore_ascii_case(b"qualify")) +} + +pub(crate) fn complete_parsed_validation( + request: ParsedValidationRequest<'_>, +) -> Result { + let ParsedValidationRequest { + sql, + dialect, + schema, + mut result, + expression, + } = request; + if !result.valid || !may_have_extra_clause_checks(sql) { + return Ok(result); + } + let mapping = mapping_schema_from_validation_schema_with_dialect(schema, dialect); + if let Some(expression) = expression { + result + .errors + .extend(missing_clause_columns(&build_scope(expression), &mapping)); + } else { + for expression in Dialect::get(dialect) + .parse(sql) + .map_err(|error| error.to_string())? + { + result + .errors + .extend(missing_clause_columns(&build_scope(&expression), &mapping)); + } + } + result.valid = !result + .errors + .iter() + .any(|error| error.severity == polyglot_sql::ValidationSeverity::Error); + Ok(result) +} + fn missing_clause_columns( scope: &polyglot_sql::Scope, schema: &polyglot_sql::MappingSchema, diff --git a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/engine.rs b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/engine.rs index 4a692e1352..530e88f208 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/engine.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/engine.rs @@ -336,7 +336,7 @@ pub(crate) fn rule_metadata() -> Vec { ALL_RULE_METADATA.iter().map(|rule| **rule).collect() } -fn lint(request: LintRequest) -> Result { +pub(crate) fn lint(request: LintRequest) -> Result { if request.version != LINT_API_VERSION { return Err(format!( "unsupported native lint request version {}; expected {LINT_API_VERSION}", diff --git a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/mod.rs b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/mod.rs index 1a6749eaeb..3641ba2ff6 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/mod.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/mod.rs @@ -2,4 +2,5 @@ pub(crate) mod additional; pub(crate) mod engine; pub(crate) mod formatter; pub(crate) mod inline_relations; +pub(crate) mod preparation; pub(crate) mod terminal_shape; diff --git a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs new file mode 100644 index 0000000000..8a9ddc97c2 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs @@ -0,0 +1,183 @@ +//! One native lexical pass for interpolation provenance and CTE dependency evidence. + +use std::collections::HashSet; +use std::sync::LazyLock; + +use regex::Regex; + +use crate::sql_lint::types::{InterpolationSite, PreparedSql}; + +static SITES: LazyLock> = LazyLock::new(|| { + Regex::new( + r#"--[^\n]*(?:\n|\z)|/\*[\s\S]*?(?:\*/|\z)|'(?:\\.|''|[^'\\])*(?:'|\z)|"(?:\\.|""|[^"\\])*(?:"|\z)|(?P@@|\$\{|@|(?:__dbt_ref|__ref|__seed|__source|__table_fn|__udf)\s*\()"#, + ) +}); +static CTES: LazyLock> = LazyLock::new(|| { + Regex::new( + r#"(?is)(?:\bWITH(?:\s+RECURSIVE)?|,)\s*["`\[]?(?P[A-Za-z_][A-Za-z0-9_]*)["`\]]?(?:\s*\([^)]*\))?\s+AS\s*\("#, + ) +}); +static WORDS: LazyLock> = + LazyLock::new(|| Regex::new(r"\b[A-Za-z_][A-Za-z0-9_]*\b")); +static OPAQUE_CTE: LazyLock> = + LazyLock::new(|| Regex::new(r"(?i)\b[A-Za-z_][A-Za-z0-9_]*\s+AS\s*\(\s*(?:--[^\n]*\n\s*)?$")); + +pub(crate) fn prepare( + expanded: &str, + before_expansion: &str, + prior_sites: &[usize], +) -> Result, String> { + if !expanded.is_ascii() || !before_expansion.is_ascii() { + return Ok(None); + } + let site_pattern = SITES.as_ref().map_err(|error| error.to_string())?; + let cte_pattern = CTES.as_ref().map_err(|error| error.to_string())?; + let word_pattern = WORDS.as_ref().map_err(|error| error.to_string())?; + let opaque_pattern = OPAQUE_CTE.as_ref().map_err(|error| error.to_string())?; + let mut text = String::with_capacity(expanded.len()); + let mut sites: Vec = Vec::new(); + let mut copied_to = 0; + for captures in site_pattern.captures_iter(expanded) { + let Some(site) = captures.name("site") else { + continue; + }; + let start = site.start(); + if start < copied_to { + continue; + } + let Some(end) = site_end(expanded, start) else { + continue; + }; + text.push_str(&expanded[copied_to..start]); + let sentinel = format!("__sqb_lint_{}__", sites.len()); + let output_start = text.len(); + text.push_str(&sentinel); + sites.push(( + sentinel, + output_start, + text.len(), + start, + end, + expanded[start..end].to_string(), + )); + copied_to = end; + } + text.push_str(&expanded[copied_to..]); + let names: HashSet = cte_pattern + .captures_iter(expanded) + .map(|capture| capture["name"].to_ascii_lowercase()) + .collect(); + let mut referenced: HashSet = HashSet::new(); + for site in &sites { + referenced.extend( + word_pattern + .find_iter(&site.5) + .map(|word| word.as_str().to_ascii_lowercase()) + .filter(|word| names.contains(word)), + ); + } + for &start in prior_sites { + let Some(prefix) = before_expansion.get(..start) else { + return Ok(None); + }; + if opaque_pattern.is_match(prefix) { + referenced.extend( + cte_pattern + .captures_iter(prefix) + .map(|capture| capture["name"].to_ascii_lowercase()) + .filter(|name| names.contains(name)), + ); + } + } + let mut referenced: Vec = referenced.into_iter().collect(); + referenced.sort(); + Ok(Some((text, sites, referenced))) +} + +fn site_end(text: &str, start: usize) -> Option { + let bytes = text.as_bytes(); + match bytes[start] { + b'@' if bytes.get(start + 1) == Some(&b'@') => Some(scan_name(bytes, start + 2, true)), + b'@' if bytes.get(start + 1) == Some(&b'\'') => { + text[start + 2..].find('\'').map(|index| start + index + 3) + } + b'@' => { + let first = *bytes.get(start + 1)?; + if !first.is_ascii_alphabetic() && first != b'_' { + return None; + } + let end = scan_name(bytes, start + 1, false); + if bytes.get(end) == Some(&b'(') { + matching_paren(bytes, end) + } else { + Some(end) + } + } + b'$' => text[start + 2..].find('}').map(|index| start + index + 3), + b'_' => { + let end = scan_name(bytes, start, false); + (bytes.get(end) == Some(&b'(')) + .then(|| matching_paren(bytes, end)) + .flatten() + } + _ => None, + } +} + +fn scan_name(bytes: &[u8], mut index: usize, interpolation: bool) -> usize { + while bytes.get(index).is_some_and(|byte| { + byte.is_ascii_alphanumeric() + || *byte == b'_' + || interpolation && matches!(byte, b':' | b'.') + }) { + index += 1; + } + index +} + +fn matching_paren(bytes: &[u8], mut index: usize) -> Option { + let mut depth = 0; + let mut quote = None; + while index < bytes.len() { + let byte = bytes[index]; + if let Some(quoted) = quote { + if byte == b'\\' && index + 1 < bytes.len() { + index += 2; + continue; + } + if byte == quoted { + if bytes.get(index + 1) == Some("ed) { + index += 2; + continue; + } + quote = None; + } + } else if bytes[index..].starts_with(b"--") { + index = bytes[index + 2..] + .iter() + .position(|byte| *byte == b'\n') + .map_or(bytes.len(), |offset| index + offset + 3); + continue; + } else if bytes[index..].starts_with(b"/*") { + index = bytes[index + 2..] + .windows(2) + .position(|pair| pair == b"*/") + .map_or(bytes.len(), |offset| index + offset + 4); + continue; + } else { + match byte { + b'\'' | b'"' => quote = Some(byte), + b'(' => depth += 1, + b')' => { + depth -= 1; + if depth == 0 { + return Some(index + 1); + } + } + _ => {} + } + } + index += 1; + } + None +} diff --git a/crates/sqlbuild-rules-native/src/sql_lint/main/batch_engine.rs b/crates/sqlbuild-rules-native/src/sql_lint/main/batch_engine.rs new file mode 100644 index 0000000000..aea9fd8ae8 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_lint/main/batch_engine.rs @@ -0,0 +1,19 @@ +use crate::sql_lint::_helpers::engine::lint; +use crate::sql_lint::models::{LintRequest, LintResponse}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; + +pub(crate) fn lint_batch_json(request_json: &str) -> Result { + let requests: Vec = + serde_json::from_str(request_json).map_err(|error| error.to_string())?; + if requests.is_empty() { + return Ok("[]".to_owned()); + } + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(requests.len().min(4)) + .stack_size(16 * 1024 * 1024) + .build() + .map_err(|error| error.to_string())?; + let results: Vec> = + pool.install(|| requests.into_par_iter().map(lint).collect()); + serde_json::to_string(&results).map_err(|error| error.to_string()) +} diff --git a/crates/sqlbuild-rules-native/src/sql_lint/main/mod.rs b/crates/sqlbuild-rules-native/src/sql_lint/main/mod.rs index f4a89d1bea..506a01d172 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/main/mod.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/main/mod.rs @@ -1,4 +1,6 @@ +pub(crate) mod batch_engine; pub(crate) mod batch_formatter; pub(crate) mod catalogue; pub(crate) mod engine; pub(crate) mod formatter; +pub(crate) mod preparation; diff --git a/crates/sqlbuild-rules-native/src/sql_lint/main/preparation.rs b/crates/sqlbuild-rules-native/src/sql_lint/main/preparation.rs new file mode 100644 index 0000000000..6a3aa9e3ec --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_lint/main/preparation.rs @@ -0,0 +1,11 @@ +//! Native lexical preparation with Python-compatible ASCII offsets. + +use crate::sql_lint::types::PreparedSql; + +pub(crate) fn prepare( + expanded: &str, + before_expansion: &str, + prior_sites: &[usize], +) -> Result, String> { + crate::sql_lint::_helpers::preparation::prepare(expanded, before_expansion, prior_sites) +} diff --git a/crates/sqlbuild-rules-native/src/sql_lint/mod.rs b/crates/sqlbuild-rules-native/src/sql_lint/mod.rs index afb9a1ddcb..4d5986ec3e 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/mod.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/mod.rs @@ -4,6 +4,7 @@ pub(crate) mod _helpers; mod constants; pub(crate) mod main; mod models; +pub(crate) mod types; #[cfg(test)] mod tests; diff --git a/crates/sqlbuild-rules-native/src/sql_lint/tests.rs b/crates/sqlbuild-rules-native/src/sql_lint/tests.rs index 4aa3e9917a..1a632265ba 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/tests.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/tests.rs @@ -1,3 +1,5 @@ +#[path = "tests/test_batch_lint.rs"] +mod batch_lint; #[path = "tests/helpers.rs"] mod helpers; #[path = "tests/test_sql_lint.rs"] diff --git a/crates/sqlbuild-rules-native/src/sql_lint/tests/test_batch_lint.rs b/crates/sqlbuild-rules-native/src/sql_lint/tests/test_batch_lint.rs new file mode 100644 index 0000000000..5e4bf9ced7 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_lint/tests/test_batch_lint.rs @@ -0,0 +1,43 @@ +use crate::sql_lint::main::batch_engine::lint_batch_json; +use crate::sql_lint::main::engine::lint_json; +use crate::sql_lint::tests::test_types::BatchLintTestCase; +use serde_json::{Value, json}; + +#[test] +fn given_mixed_lint_inputs_when_batching_then_preserves_order_errors_and_findings() { + let test_cases = [BatchLintTestCase { + description: "mixed valid and invalid SQL retains standalone evidence", + queries: &[ + "SELECT 1", + "SELECT FROM", + "SELECT order_id FROM orders WHERE order_id = NULL", + ], + expected_responses: 3, + }]; + for test_case in test_cases { + let requests: Vec = test_case + .queries + .iter() + .map(|sql| json!({"version": 1, "sql": sql, "dialect": "duckdb"})) + .collect(); + let expected: Vec = requests + .iter() + .map(|request| { + serde_json::to_value(lint_json(&request.to_string()).map(|response| { + serde_json::from_str::(&response).expect("valid response") + })) + .expect("serializable result") + }) + .collect(); + let response = lint_batch_json(&serde_json::to_string(&requests).expect("valid requests")) + .expect("valid batch"); + let actual: Vec = serde_json::from_str(&response).expect("valid response"); + assert_eq!( + actual.len(), + test_case.expected_responses, + "{}", + test_case.description + ); + assert_eq!(actual, expected, "{}", test_case.description); + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_lint/tests/test_types.rs b/crates/sqlbuild-rules-native/src/sql_lint/tests/test_types.rs index db404b8d14..1ff58ad258 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/tests/test_types.rs @@ -5,6 +5,12 @@ pub(crate) struct LintTestCase { pub expected_anchors: &'static [(&'static str, &'static str)], } +pub(crate) struct BatchLintTestCase { + pub(crate) description: &'static str, + pub(crate) queries: &'static [&'static str], + pub(crate) expected_responses: usize, +} + pub(crate) struct LintMetadataTestCase { pub description: &'static str, pub sql: &'static str, diff --git a/crates/sqlbuild-rules-native/src/sql_lint/types.rs b/crates/sqlbuild-rules-native/src/sql_lint/types.rs new file mode 100644 index 0000000000..682c369422 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_lint/types.rs @@ -0,0 +1,2 @@ +pub(crate) type InterpolationSite = (String, usize, usize, usize, usize, String); +pub(crate) type PreparedSql = (String, Vec, Vec); diff --git a/docs/compiler-cache.md b/docs/compiler-cache.md index 8093e9e59c..7fc2ba4f2b 100644 --- a/docs/compiler-cache.md +++ b/docs/compiler-cache.md @@ -77,6 +77,18 @@ nullability, wildcard status, and any other fact consumed by downstream analysis invalidates downstream semantic analysis only when this exported interface changes. Its own compiled artifact still changes whenever its implementation changes. +For supported, schema-bound queries, the native compiler derives immutable CTE output facts from +one parsed query tree. It retains every projection position, including unnamed expressions used by +CTE column lists and positional unions. Known input columns take precedence over same-named SELECT +aliases. Queries with incomplete facts, unsupported source shapes, or unresolved types retain the +existing analysis path. Changes to this inference algorithm invalidate prior analysis identities. + +Compiler-expanded SQL and its source-location spans are also reused within one invocation. SQL lint +may start after attachment while semantic analysis continues. Compiler-proven dynamic-output paths +are checked again with their completed proof. Native and custom rules can execute concurrently; +their combined findings pass through the same final exception and suppression policy. These +invocation-local objects and workers are not persisted in the compiler cache. + ### 5. Shared derived indexes Project-wide indexes are built once per compile and shared by all consumers. They include model @@ -125,6 +137,9 @@ written, and phase timings. The current vertical slice reports model-analysis ba hits, misses, bypasses, and phase timings. Counts describe semantic records, not low-level database operations. +Concurrent phase timings can overlap. Use the total compile time and fresh-process wall time for +end-to-end comparisons rather than summing individual phase durations. + Fresh-process performance guards use representative semantic projects and cover: - cold compile with a clean target; diff --git a/docs/dense-compile-performance.md b/docs/dense-compile-performance.md new file mode 100644 index 0000000000..9570a7c6d6 --- /dev/null +++ b/docs/dense-compile-performance.md @@ -0,0 +1,55 @@ +# Dense all-rules compile guard + +The dense compile guard covers cold compilation of 1,000, 3,000, 5,000 and +10,000 models with every built-in compiler rule and one custom project rule +enabled. It complements the existing semantic and incremental-cache benchmarks. + +The deterministic fixture uses neutral order data and compiles offline with the +Snowflake dialect. Each 1,000-model group has 240 sources, 50 seeds, 25 SQL +functions, 1,700 audits and 1,000 SQL tests. Queries contain dependency import +CTEs, joins, named unions, explicit contracts, wide projections and Python macro +expansion. Independent declaration scopes preserve per-model macro visibility +as the project grows. SQL tests mock immediate inputs to keep test expansion +bounded. A smaller integration fixture executes its generated tests in DuckDB. + +The guard checks: + +- Complete built-in rule selection, one selected custom rule, and no ignored + rules, scoped exceptions or relaxed thresholds. +- Resource counts, minimum authored SQL volume, and zero diagnostics. +- Zero analysis/rule cache hits and analysis bypasses for every model. +- A fingerprint of semantic JSON and every compiled artifact. +- Fresh-process wall time and peak RSS, with phase timings in test logs. + +| Models | Wall-time ceiling | Process RSS ceiling | +| ---: | ---: | ---: | +| 1,000 | 14s | 1.5GiB | +| 3,000 | 34s | 2.75GiB | +| 5,000 | 55s | 3.25GiB | +| 10,000 | 115s | 4GiB | + +These are regression ceilings for the implemented compiler, with headroom for +host variation. They are not a claim of sub-10-second compilation at every size. + +CI additionally runs each size in a 4GiB cgroup with swap disabled, covering +the test harness and child processes together. Compiler, lint and rule-engine +changes select these jobs. Local process RSS measurements alone do not prove +the cgroup gate passes. + +Run an individual size with: + +```sh +make test-e2e-dense-compile-performance SQLBUILD_BENCHMARK_MODELS=1000 +``` + +For requests containing at least 32MiB of unique SQL, the compiler compacts native +analysis results in batches of 64. Smaller requests retain a single batch. Both +use up to four workers; local resolution views avoid copying CTE definitions +that are already available through the source catalogue. +These limits change scheduling, not analysis or fallback semantics. + +This fixture protects dense query analysis, declaration expansion, rules and +artifact generation. It is not an exact reproduction of every application: +dynamic plugins, incremental materialization mixes and deeply nested query +tails still require application-specific acceptance measurements. Keep those +measurements private; publish only neutral fixtures and their own results. diff --git a/fensu.toml b/fensu.toml index 072543bf37..528c17c8a2 100644 --- a/fensu.toml +++ b/fensu.toml @@ -241,11 +241,6 @@ tooling = [] rule_packs = ["rust"] select = ["FPRS"] -[[targets.rust.rule_exceptions]] -rule = "FPRSR307" -path = "crates/sqlbuild-rules-native/src/query_analysis/main.rs" -reason = "This module owns the single coarse native query-analysis boundary and its compact SQLBuild projection contract." - [[targets.rust.rule_exceptions]] rule = "FPRSR307" path = "crates/sqlbuild-rules-native/src/scope_metadata/main.rs" diff --git a/scripts/cold_compile_performance/_helpers/dense_project.py b/scripts/cold_compile_performance/_helpers/dense_project.py new file mode 100644 index 0000000000..a0ca920076 --- /dev/null +++ b/scripts/cold_compile_performance/_helpers/dense_project.py @@ -0,0 +1,330 @@ +"""Generate bounded CTE graphs with full compiler-rule evaluation.""" + +from bisect import bisect_right +from pathlib import Path + +from scripts.cold_compile_performance.constants import ( + DENSE_AMOUNT_AUDIT_SHARE, + DENSE_DEEP_CTE_PERCENTILE, + DENSE_ID_COLUMN, + DENSE_JOIN_COUNTS, + DENSE_JOIN_THRESHOLDS, + DENSE_MACRO_QUANTITY_COLUMNS, + DENSE_MACRO_SHARE, + DENSE_MEDIUM_CTE_PERCENTILE, + DENSE_UNION_MEDIUM_INDEX, + DENSE_UNION_TAIL_INDEX, +) +from tests.e2e.src.sqlbuild.cli.commands.main.compile.helpers import ( + _layered_write_functions, + _layered_write_macros, + _layered_write_seeds, +) + + +def dense_model_name(index: int) -> str: + return f"sales{index // 1000:03d}__int_clean__orders_{index:05d}" + + +def dense_column_count(index: int) -> int: + rank: int = (index * 17) % 1000 + for upper, width in ( + (100, 8), + (250, 12), + (450, 16), + (600, 20), + (730, 24), + (830, 32), + (900, 48), + (940, 64), + (970, 96), + (990, 160), + ): + if rank < upper: + return width + return 256 + (rank - 990) * 24 + + +def write_dense_compile_project(*, project_dir: Path, model_count: int) -> None: + """Write a deterministic query-heavy project without external inputs or services.""" + + project_dir.mkdir(parents=True) + domains: str = ", ".join( + f'"sales{cohort:03d}"' for cohort in range((model_count + 999) // 1000) + ) + (project_dir / "sqlbuild_project.toml").write_text( + 'name = "dense_orders"\nadapter = "snowflake"\n' + '[defaults]\ndatabase = "warehouse"\nschema = "analytics"\n' + f'[rules]\nselect = ["SQBR", "XSQBR"]\ndomains = [{domains}]\n', + encoding="utf-8", + ) + source_count: int = max(1, model_count * 24 // 100) + seed_count: int = max(1, model_count // 20) + function_count: int = max(1, model_count // 40) + _write_sources(project_dir=project_dir, source_count=source_count) + _layered_write_seeds(project_dir=project_dir, seed_count=seed_count) + _layered_write_functions(project_dir=project_dir, function_count=function_count) + for cohort in range((model_count + 999) // 1000): + models: Path = project_dir / "models" / f"sales{cohort:03d}" / "intermediate" / "clean" + models.mkdir(parents=True) + macro_count: int = max(1, min(1000, model_count - cohort * 1000) // 80) + _layered_write_macros(project_dir=project_dir, macro_count=macro_count) + macro_root: Path = project_dir / "macros" + for macro_index, macro_path in enumerate(sorted(macro_root.glob("*.py"))): + bucket: Path = models / "_sqlbuild" / "macros" / f"adjustment{macro_index // 8:03d}" + bucket.mkdir(parents=True, exist_ok=True) + content: str = macro_path.read_text(encoding="utf-8") + content = content.replace("macro_", f"macro_{cohort:03d}_").replace( + "base_offset", f"base_offset_{cohort:03d}" + ) + (bucket / f"offset{macro_index:05d}.py").write_text(content, encoding="utf-8") + macro_path.unlink() + macro_root.rmdir() + for index in range(model_count): + name: str = dense_model_name(index) + domain: str = f"sales{index // 1000:03d}" + model_owner: Path = ( + project_dir / "models" / domain / "intermediate" / "clean" / f"batch{index // 5:05d}" + ) + test_owner: Path = ( + project_dir + / "tests" + / "unit" + / domain + / "intermediate" + / "clean" + / f"batch{index // 5:05d}" + ) + macro_count = max(1, min(1000, model_count - index // 1000 * 1000) // 80) + model_owner.mkdir(parents=True, exist_ok=True) + test_owner.mkdir(parents=True, exist_ok=True) + (model_owner / f"{name}.sql").write_text( + _model_sql( + index=index, + source_count=source_count, + function_count=function_count, + macro_count=macro_count, + ), + encoding="utf-8", + ) + (test_owner / f"test_{name.replace('__', '_')}__preserves_order_id.sql").write_text( + _test_sql(index=index, source_count=source_count, macro_count=macro_count), + encoding="utf-8", + ) + rules: Path = project_dir / "rules" + rules.mkdir() + (rules / "orders.py").write_text( + "from sqlbuild.rules import Finding, Project, RuleContext, rule\n\n" + '@rule(code="XSQBRBENCH001", message="Order contracts must be present", ' + 'remediation="Declare the order contract.")\n' + "def order_contracts(*, project: Project, ctx: RuleContext) -> list[Finding]:\n" + " findings: list[Finding] = []\n" + " for model in ctx.project.models:\n" + " if not ctx.contracts.enforced(model):\n" + " findings.append(ctx.finding(subject=model))\n" + " return findings\n", + encoding="utf-8", + ) + (project_dir / "tests" / "test_rule_contracts.py").write_text( + "import pytest\n" + "from sqlbuild.rules.testing import RuleCase, evaluate_rule\n" + "from rules.orders import order_contracts\n\n" + "CASES = (\n" + ' RuleCase(description="enforced contract", source="MODEL (contract enforced, ' + 'columns (id (type INTEGER))); SELECT CAST(1 AS INTEGER) AS id", ' + "expected_finding_count=0),\n" + ' RuleCase(description="missing contract", source="MODEL (); SELECT 1 AS id", ' + "expected_finding_count=1),\n" + ")\n\n" + '@pytest.mark.parametrize("test_case", CASES, ids=lambda case: case.description)\n' + "def test_given_contract_when_checking_then_reports_expected_findings(test_case):\n" + " result = evaluate_rule(rule=order_contracts, test_case=test_case)\n" + " assert result.finding_count == test_case.expected_finding_count\n", + encoding="utf-8", + ) + + +def _write_sources(*, project_dir: Path, source_count: int) -> None: + sources: Path = project_dir / "sources" + sources.mkdir() + entries: list[str] = ["sources:"] + for index in range(source_count): + width: int = dense_column_count(index % 240) + extra: str = "".join( + f", CAST(1 AS DOUBLE) AS quantity_{column:03d}" for column in range(width - 2) + ) + entries.extend( + ( + f" - name: orders_{index:05d}", + " contract: enforced", + f' expression: "(SELECT 1 AS id, CAST(1 AS DOUBLE) AS amount{extra})"', + " columns:", + " - name: id", + " type: INTEGER", + " nullable: false", + " - name: amount", + " type: DOUBLE", + ) + ) + for column in range(width - 2): + entries.extend( + ( + f" - name: quantity_{column:03d}", + " type: DOUBLE", + " nullable: true", + ) + ) + (sources / "orders.yml").write_text("\n".join(entries) + "\n", encoding="utf-8") + + +def _join_count(index: int) -> int: + rank: int = (index * 37) % 100 + return DENSE_JOIN_COUNTS[bisect_right(DENSE_JOIN_THRESHOLDS, rank)] + + +def _branch_count(index: int) -> int: + return ( + 128 + if index % 1000 == DENSE_UNION_TAIL_INDEX + else 20 + if index % 20 == 1 + else 8 + if index % 10 == DENSE_UNION_MEDIUM_INDEX + else 1 + if index % 10 == 0 + else 2 + ) + + +def _test_sql(*, index: int, source_count: int, macro_count: int) -> str: + name: str = dense_model_name(index) + mocks: dict[str, str] = {} + base: str = ( + f"__source__orders_{_source_index(index=index, source_count=source_count):05d}" + if index % 1000 % 48 == 0 + else f"__ref__{dense_model_name(index - 1)}" + ) + mocks[base] = "SELECT CAST(1 AS INTEGER) AS id, CAST(1 AS DOUBLE) AS amount" + for lookup in range(_join_count(index)): + source_index: int = _source_index(index=index, source_count=source_count, offset=lookup + 1) + mocks[f"__source__orders_{source_index:05d}"] = mocks[base] + expected_id: int = 1 + if index % 10 < DENSE_MACRO_SHARE: + expected_id += index % 1000 % macro_count or 1 + adjustment: int = expected_id - 1 + amount: int = 1 + _join_count(index) + adjustment + int(index % 1000 % 43 == 0) + expected_columns: list[str] = [ + f"CAST({expected_id} AS INTEGER) AS id", + f"CAST({amount} AS DOUBLE) AS amount", + ] + for column in range(dense_column_count(index) - 2): + quantity: int = 2 + (adjustment if column < DENSE_MACRO_QUANTITY_COLUMNS else 0) + expected_columns.append(f"CAST({quantity} AS DOUBLE) AS quantity_{column:03d}") + expected: str = "SELECT " + ", ".join(expected_columns) + ctes: list[str] = [f"{mock} AS ({body})" for mock, body in mocks.items()] + ctes.append(f"__expected__{name} AS ({expected})") + return f'TEST (name "{name}__preserves_order_id");\nWITH ' + ",\n".join(ctes) + "\nSELECT 1\n" + + +def _source_index(*, index: int, source_count: int, offset: int = 0) -> int: + start: int = index // 1000 * 240 + return start + (index % 1000 + offset) % min(240, source_count - start) + + +def _model_sql(*, index: int, source_count: int, function_count: int, macro_count: int) -> str: + width: int = dense_column_count(index) + names: list[str] = ["id", "amount", *(f"quantity_{column:03d}" for column in range(width - 2))] + declarations: list[str] = [] + for column, name in enumerate(names): + data_type: str = "INTEGER" if column == 0 else "DOUBLE" + audit: str = ( + ", audits [not_null]" + if column == 0 or (column == 1 and index % 10 < DENSE_AMOUNT_AUDIT_SHARE) + else "" + ) + nullable: str = "false" if audit else "true" + declarations.append( + f" {name} (type {data_type}, nullable {nullable}, " + f'description "Order quantity measure"{audit})' + ) + header: str = ( + "MODEL (materialized table, contract enforced, columns (\n" + + ",\n".join(declarations) + + "\n));\n" + ) + source_index: int = _source_index(index=index, source_count=source_count) + relation: str = ( + f'__source("orders_{source_index:05d}")' + if index % 1000 % 48 == 0 + else f'__ref("{dense_model_name(index - 1)}")' + ) + ctes: list[str] = [f"imported AS (SELECT * FROM {relation})"] + rank: int = (index * 37) % 100 + joins: int = _join_count(index) + imported_sources: dict[int, str] = {source_index: "imported"} if index % 1000 % 48 == 0 else {} + lookup_names: list[str] = [] + for lookup in range(joins): + source_index = _source_index(index=index, source_count=source_count, offset=lookup + 1) + if source_index not in imported_sources: + imported_sources[source_index] = f"lookup_{lookup}" + ctes.append(f'lookup_{lookup} AS (SELECT * FROM __source("orders_{source_index:05d}"))') + lookup_names.append(imported_sources[source_index]) + amount: str = "b.amount" + "".join( + f" + COALESCE(l{lookup}.amount, 0)" for lookup in range(joins) + ) + if index % 10 < DENSE_MACRO_SHARE: + amount = f'@macro_{index // 1000:03d}_{index % 1000 % macro_count:05d}("{amount}")' + if index % 1000 % 43 == 0: + function_start: int = index // 1000 * 25 + function_index: int = function_start + index % 1000 % min( + 25, function_count - function_start + ) + amount = f'__udf("fn_{function_index:05d}")({amount})' + identity: str = ( + f'@macro_{index // 1000:03d}_{index % 1000 % macro_count:05d}("b.id")' + if index % 10 < DENSE_MACRO_SHARE + else "b.id" + ) + projections: list[str] = [ + f"CAST({identity} AS INTEGER) AS id", + f"CAST({amount} AS DOUBLE) AS amount", + ] + for column, name in enumerate(names[2:]): + expression: str = "COALESCE(b.amount, 0) + CAST(b.id AS DOUBLE)" + if index % 10 < DENSE_MACRO_SHARE and column < DENSE_MACRO_QUANTITY_COLUMNS: + expression = ( + f'@macro_{index // 1000:03d}_{index % 1000 % macro_count:05d}("{expression}")' + ) + projections.append(f"CAST({expression} AS DOUBLE) AS {name}") + join_sql: str = "".join( + f"\nLEFT JOIN {lookup_names[lookup]} AS l{lookup} ON l{lookup}.id = b.id" + for lookup in range(joins) + ) + ctes.append( + "projected AS (SELECT\n " + + ",\n ".join(projections) + + "\nFROM imported AS b" + + join_sql + + ")" + ) + columns: str = ", ".join(names) + branches: int = _branch_count(index) + union_branches: list[str] = [f"SELECT {columns} FROM projected"] + union_branches.extend( + f"SELECT {columns} FROM projected WHERE id = 0" for _ in range(branches - 1) + ) + ctes.append("combined AS (\n" + "\nUNION ALL BY NAME\n".join(union_branches) + "\n)") + previous: str = "combined" + depth: int = ( + 12 if rank >= DENSE_DEEP_CTE_PERCENTILE else 4 if rank >= DENSE_MEDIUM_CTE_PERCENTILE else 0 + ) + for step in range(depth): + current: str = f"calculated_{step}" + ctes.append(f"{current} AS (SELECT {columns} FROM {previous})") + previous = current + casts: str = ",\n ".join( + f"CAST({name} AS {'INTEGER' if name == DENSE_ID_COLUMN else 'DOUBLE'}) AS {name}" + for name in names + ) + ctes.append(f"final AS (SELECT\n {casts}\nFROM {previous})") + return header + "\nWITH " + ",\n".join(ctes) + f"\nSELECT {columns} FROM final\n" diff --git a/scripts/cold_compile_performance/constants.py b/scripts/cold_compile_performance/constants.py index d0c069de13..a8599c5db8 100644 --- a/scripts/cold_compile_performance/constants.py +++ b/scripts/cold_compile_performance/constants.py @@ -1,3 +1,14 @@ """Constants for fresh-process compile performance guards.""" UNLIMITED_CGROUP_MEMORY_VALUE: str = "max" + +DENSE_JOIN_THRESHOLDS: tuple[int, ...] = (55, 75, 90, 99) +DENSE_JOIN_COUNTS: tuple[int, ...] = (0, 1, 4, 8, 20) +DENSE_UNION_TAIL_INDEX: int = 54 +DENSE_UNION_MEDIUM_INDEX: int = 7 +DENSE_MACRO_SHARE: int = 3 +DENSE_MACRO_QUANTITY_COLUMNS: int = 2 +DENSE_AMOUNT_AUDIT_SHARE: int = 7 +DENSE_DEEP_CTE_PERCENTILE: int = 97 +DENSE_MEDIUM_CTE_PERCENTILE: int = 80 +DENSE_ID_COLUMN: str = "id" diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 09f29bf3c1..264de8d6fb 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -2,9 +2,14 @@ def catalogue_json() -> str: ... def evaluate_json(request_json: str) -> str: ... +def finalize_rule_findings_json(request_json: str) -> str: ... def format_sql_json(request_json: str) -> str: ... def format_sql_batch_json(request_json: str) -> str: ... def lint_sql_json(request_json: str) -> str: ... +def lint_sql_batch_json(request_json: str) -> str: ... +def prepare_lint_sql( + expanded: str, before_expansion: str, prior_sites: list[int] +) -> tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None: ... def load_config_json(project_dir: str) -> str: ... def selected_codes_json(request_json: str) -> str: ... def render_owned_skill(content: str, input_fingerprint: str) -> str: ... @@ -13,3 +18,9 @@ def parse_model_headers( headers: list[str], ) -> list[tuple[dict[str, object] | None, list[tuple[str, int, int]] | None, str | None]]: ... def tokenize_model_header(header: str) -> list[tuple[int, str, int]]: ... +def substitute_static_project_vars( + sqls: list[str], variables: list[tuple[str, str]] +) -> list[tuple[int, str | None]]: ... +def extract_static_sql_references( + sql: str, +) -> list[tuple[str, str, str | None, int | None]] | None: ... diff --git a/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py b/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py index 35daeabac4..8f82c27a20 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py @@ -3,6 +3,8 @@ from __future__ import annotations import time +from dataclasses import replace +from functools import partial from pathlib import Path from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -43,6 +45,7 @@ from sqlbuild.compiler.pipeline.models import ProjectGraph from sqlbuild.compiler.planner.main.selection.selection import resolve_project_selectors from sqlbuild.presentation.classes.transient_status_reporter import TransientStatusReporter +from sqlbuild.rule_engine.classes.early_sql_lint import EarlySqlLint from sqlbuild.rule_engine.main.load_config import load_rules_config from sqlbuild.rule_engine.main.run_rules import run_rules from sqlbuild.rule_engine.models import RulesRunResult @@ -67,6 +70,37 @@ def analyze_compile_project( ) -> CompileAnalysis: """Discover, compile, and validate the project into one analysis result.""" + with EarlySqlLint(enabled=not (select or exclude)) as early_lint: + return _analyze_compile_project( + project_dir=project_dir, + no_sql_validation=no_sql_validation, + selected_target=selected_target, + lineage_mode=lineage_mode, + cli_vars=cli_vars, + profile_flags=profile_flags, + analysis_selection=CompileAnalysisSelection( + select=select, exclude=exclude, no_cache=no_cache + ), + status=status, + early_lint=early_lint, + ) + + +def _analyze_compile_project( + *, + project_dir: Path, + no_sql_validation: bool, + selected_target: str | None, + lineage_mode: CompileLineageMode, + cli_vars: dict[str, object] | None, + profile_flags: CompileProfileFlags, + analysis_selection: CompileAnalysisSelection, + status: TransientStatusReporter | None, + early_lint: EarlySqlLint, +) -> CompileAnalysis: + select: tuple[str, ...] = analysis_selection.select + exclude: tuple[str, ...] = analysis_selection.exclude + discover_start: float = time.monotonic() _ = start_compile_phase(status=status, message="Discovering project...") discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( @@ -98,10 +132,11 @@ def analyze_compile_project( skip_column_inference=profile_flags.skip_column_inference, column_lineage_mode=compile_analysis_lineage_mode(lineage_mode), cli_vars=cli_vars, - analysis_selection=CompileAnalysisSelection( - select=select, - exclude=exclude, - no_cache=no_cache, + analysis_selection=replace( + analysis_selection, + on_inputs_ready=partial( + early_lint.start, dialect=adapter.sql_analysis_dialect() or "generic" + ), ), ) graph_ms: int = elapsed_ms(graph_start) @@ -159,6 +194,7 @@ def analyze_compile_project( project_dir=project_dir, dialect=adapter.sql_analysis_dialect() or "generic", selected_keys=selected_keys if select or exclude else None, + prepared_sql=early_lint.preparation, ) _ = complete_compile_phase( status=status, diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/cache.py b/src/sqlbuild/compiler/compile/_helpers/analysis/cache.py index bff28c575e..2a65186b3b 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/cache.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/cache.py @@ -44,7 +44,7 @@ from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic _ANALYSIS_CACHE_VERSION: int = 10 -_ANALYSIS_ALGORITHM_FINGERPRINT: str = "model-sql-analysis-v10-incremental-compact-batch" +_ANALYSIS_ALGORITHM_FINGERPRINT: str = "model-sql-analysis-v12-bound-borrowed-output-graph" _MAX_CACHE_ENTRY_BYTES: int = 10_000_000 _SHA256_HEX_LENGTH: int = 64 _CACHE_ENTRY_SEPARATOR: str = "\n" diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py index 0baafc2e41..1cb0db36fc 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py @@ -4,6 +4,7 @@ import logging import re +from dataclasses import replace from typing import Any, cast import orjson @@ -37,6 +38,7 @@ ) from sqlbuild.compiler.compile.exceptions import CompactAnalysisInputError from sqlbuild.compiler.compile.models import ( + CompactBatchExecutionOptions, CompactBatchPreparation, CompactLineageFacts, CompactProjectedFacts, @@ -50,12 +52,13 @@ PolyglotAnalysisResult, ProjectedAnalysisRequest, ) -from sqlbuild.compiler.compile.types import CompactBatchResponseCallback, CompiledResourceType +from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.compiler.lineage.types import ( ColumnLineageConfidence, ColumnTransformKind, InferredNullability, ) +from sqlbuild.compiler.profiling.main.record import record_compile_timing from sqlbuild.compiler.references.main._quoted_reference_call_pattern import ( quoted_reference_call_pattern, ) @@ -154,16 +157,20 @@ POLYGLOT_PAYLOAD_COLUMN as _POLYGLOT_PAYLOAD_COLUMN, ) from sqlbuild.compiler.sql_analysis.exceptions import SqlAnalysisBoundaryError +from sqlbuild.compiler.sql_analysis.main._decode_schema_validation import decode_schema_validation +from sqlbuild.compiler.sql_analysis.main._prepare_schema_validation import prepare_schema_validation 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 ( SqlBindingDiagnostic, + SqlBindingResult, SqlSchemaValidationRequest, ) from sqlbuild.compiler.sql_analysis.types import NativeQueryAnalysisModule from sqlbuild.diagnostics.main.log_debug_event import log_debug_event _DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile") +_NATIVE_LEGACY_FALLBACK: str = "native project type recovery requires legacy fallback" _REF_PATTERN: re.Pattern[str] = quoted_reference_call_pattern(SqlReferenceKind.REF) _SEED_PATTERN: re.Pattern[str] = quoted_reference_call_pattern(SqlReferenceKind.SEED) _SOURCE_PATTERN: re.Pattern[str] = quoted_reference_call_pattern(SqlReferenceKind.SOURCE) @@ -304,7 +311,14 @@ def analyze_columns_and_lineage_with_polyglot( column_nullability_by_table=column_nullability_by_table or {}, column_types_by_table=column_types_by_table or {}, inference_profile=profile, - allow_compact_analysis=allow_compact_analysis, + allow_compact_analysis=( + allow_compact_analysis + and not ( + precomputed is not None + and not precomputed.projected + and precomputed.analysis is None + ) + ), recover_cte_facts=recover_cte_facts, analysis=precomputed.analysis if precomputed is not None else None, ) @@ -374,43 +388,56 @@ def analyze_queries_with_compact_polyglot_batch( recover_cte_facts: tuple[bool, ...], rich_type_inference: bool = True, cached_batch: tuple[CompactBatchPreparation, object] | None = None, - on_response: CompactBatchResponseCallback | None = None, + execution: CompactBatchExecutionOptions | None = None, ) -> tuple[NativeCompactAnalysis, ...]: """Analyze rendered SQL in one bounded native call, preserving input order.""" + options: CompactBatchExecutionOptions = execution or CompactBatchExecutionOptions() + binding_schemas: tuple[dict[str, dict[str, str]] | None, ...] | None = options.binding_schemas + if not (len(query_sqls) == len(references) == len(placeholders) == len(recover_cte_facts)): raise CompactAnalysisInputError( "compact query-analysis batch inputs must have equal lengths" ) - preparation: CompactBatchPreparation = ( - cached_batch[0] - if cached_batch is not None - else _prepare_compact_analysis_batch( - query_sqls=query_sqls, - references=references, - placeholders=placeholders, - column_nullability_by_table=column_nullability_by_table, - column_types_by_table=column_types_by_table, - inference_profile=inference_profile, - recover_cte_facts=recover_cte_facts, - rich_type_inference=rich_type_inference, - ) - ) + if binding_schemas is not None and len(binding_schemas) != len(query_sqls): + raise CompactAnalysisInputError("binding schemas must match the query batch length") + if cached_batch is not None: + preparation: CompactBatchPreparation = cached_batch[0] + else: + with record_compile_timing("analysis_preparation_ms"): + preparation = _prepare_compact_analysis_batch( + query_sqls=query_sqls, + references=references, + placeholders=placeholders, + column_nullability_by_table=column_nullability_by_table, + column_types_by_table=column_types_by_table, + inference_profile=inference_profile, + recover_cte_facts=recover_cte_facts, + rich_type_inference=rich_type_inference, + binding_schemas=binding_schemas, + ) if len(preparation.cleaned_sql) != len(query_sqls): raise CompactAnalysisInputError( "cached compact query-analysis preparation has an invalid length" ) - response_payload: object = ( - cached_batch[1] - if cached_batch is not None - else _run_compact_analysis_batch(preparation=preparation) - ) - projected: tuple[NativeCompactAnalysis, ...] = _project_compact_analysis_batch( - preparation=preparation, - response_payload=response_payload, - ) - if cached_batch is None and on_response is not None: - on_response(preparation=preparation, response=response_payload) + if cached_batch is not None: + response_payload: object = cached_batch[1] + else: + with record_compile_timing("analysis_native_ms"): + response_payload = _run_compact_analysis_batch(preparation=preparation) + with record_compile_timing("analysis_projection_ms"): + projected: tuple[NativeCompactAnalysis, ...] = _project_compact_analysis_batch( + preparation=preparation, + response_payload=response_payload, + ) + if cached_batch is None: + projected = _attach_compiled_bindings( + preparation=preparation, + response=response_payload, + analyses=projected, + ) + if cached_batch is None and options.on_response is not None: + options.on_response(preparation=preparation, response=response_payload) return projected @@ -424,11 +451,12 @@ def _prepare_compact_analysis_batch( inference_profile: ExpressionInferenceProfile, recover_cte_facts: tuple[bool, ...], rich_type_inference: bool, + binding_schemas: tuple[dict[str, dict[str, str]] | None, ...] | None = None, ) -> CompactBatchPreparation: dialect: str | None = inference_profile.sql_analysis_dialect prepared: list[str] = [] queries: list[dict[str, object]] = [] - query_indexes: dict[tuple[str, str, bytes | None], int] = {} + query_indexes: dict[tuple[str, str, bytes | None, bytes | None], int] = {} templates: list[dict[str, object]] = [] template_indexes: dict[ tuple[ @@ -443,8 +471,19 @@ def _prepare_compact_analysis_batch( ] = {} projections: list[dict[str, object]] = [] function_return_types: dict[str, str] = dict(inference_profile.function_return_types) - for query_sql, query_references, query_placeholders, query_recover_cte_facts in zip( - query_sqls, references, placeholders, recover_cte_facts, strict=True + for ( + query_sql, + query_references, + query_placeholders, + query_recover_cte_facts, + binding_schema, + ) in zip( + query_sqls, + references, + placeholders, + recover_cte_facts, + binding_schemas if binding_schemas is not None else (None,) * len(query_sqls), + strict=True, ): cleaned_sql: str = _cleaned_analysis_sql( query_sql=query_sql, @@ -454,21 +493,31 @@ def _prepare_compact_analysis_batch( lineage_references: dict[str, tuple[CompiledResourceType, str]] = _lineage_reference_map( query_references ) - qualified_reference_names: frozenset[str] = _qualified_reference_names( - query_sql=cleaned_sql, - reference_names=lineage_references.keys(), + qualified_reference_names: frozenset[str] = ( + _qualified_reference_names( + query_sql=cleaned_sql, + reference_names=lineage_references.keys(), + ) + if binding_schema is None + else frozenset() ) canonical_stubs: dict[str, str] = { name: ( - name if name in qualified_reference_names else f"__sqlbuild_project_input_{index}" + name + if binding_schema is not None or name in qualified_reference_names + else f"__sqlbuild_project_input_{index}" ) for index, name in enumerate(lineage_references) } - analysis_sql: str = _cleaned_analysis_sql( - query_sql=query_sql, - placeholders=query_placeholders, - dialect=dialect, - relation_stubs=canonical_stubs, + analysis_sql: str = ( + 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, + ) ) query: dict[str, object] = { "sql": analysis_sql, @@ -484,10 +533,25 @@ def _prepare_compact_analysis_batch( ) if schema is not None: query["schema"] = schema - query_key: tuple[str, str, bytes | None] = ( + binding_payload: dict[str, object] | None = None + if binding_schema is not None: + binding_payload = prepare_schema_validation( + request=SqlSchemaValidationRequest( + sql=cleaned_sql, + dialect=dialect, + schema=binding_schema, + ) + ) + query["binding_schema"] = binding_payload["schema"] + query_key: tuple[str, str, bytes | None, bytes | None] = ( analysis_sql, dialect or "generic", (orjson.dumps(schema, option=orjson.OPT_SORT_KEYS) if schema is not None else None), + ( + orjson.dumps(binding_payload["schema"], option=orjson.OPT_SORT_KEYS) + if binding_payload is not None + else None + ), ) query_index: int | None = query_indexes.get(query_key) if query_index is None: @@ -574,6 +638,36 @@ def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> obje ) +def _attach_compiled_bindings( + *, + preparation: CompactBatchPreparation, + response: object, + analyses: tuple[NativeCompactAnalysis, ...], +) -> tuple[NativeCompactAnalysis, ...]: + if not isinstance(response, dict): + return analyses + validations: object = cast(dict[str, object], response).get("validations") + if validations is None: + return analyses + if not isinstance(validations, list) or len(validations) != len(preparation.queries): + raise SqlAnalysisBoundaryError("native compilation returned an invalid binding batch") + results: list[NativeCompactAnalysis] = [] + for analysis, projection in zip(analyses, preparation.projections, strict=True): + template_index: int = cast(int, projection["templateIndex"]) + query_index: int = cast(int, preparation.templates[template_index]["queryIndex"]) + validation: object = validations[query_index] + if validation is None: + results.append(analysis) + continue + binding: SqlBindingResult = decode_schema_validation( + sql=analysis.cleaned_sql, + dialect=cast(str, preparation.queries[query_index]["dialect"]), + response=validation, + ) + results.append(replace(analysis, binding_diagnostics=binding.diagnostics)) + return tuple(results) + + def _project_compact_analysis_batch( *, preparation: CompactBatchPreparation, response_payload: object ) -> tuple[NativeCompactAnalysis, ...]: @@ -631,7 +725,11 @@ def _project_compact_analysis_batch( sqlbuild_error=template, ) results.append( - NativeCompactAnalysis(cleaned_sql=cleaned_sql, analysis=None, projected=True) + NativeCompactAnalysis( + cleaned_sql=cleaned_sql, + analysis=None, + projected=template != _NATIVE_LEGACY_FALLBACK, + ) ) continue if not ( diff --git a/src/sqlbuild/compiler/compile/_helpers/assembly/project.py b/src/sqlbuild/compiler/compile/_helpers/assembly/project.py index 822adf065a..9998c664f2 100644 --- a/src/sqlbuild/compiler/compile/_helpers/assembly/project.py +++ b/src/sqlbuild/compiler/compile/_helpers/assembly/project.py @@ -82,6 +82,7 @@ CompactAnalysisCacheCandidate, CompactAnalysisCacheModel, CompactAnalysisCachePlan, + CompactBatchExecutionOptions, CompactBatchPreparation, CompileAuditInput, CompiledAudit, @@ -271,6 +272,11 @@ def assemble_compiled_project( run_id=inputs.run_id, ) return CompiledProject( + sql_expansions={ + model_input.model_file.file_path: model_input.sql_expansion + for model_input in inputs.model_inputs + if model_input.sql_expansion is not None + }, run_id=inputs.run_id, effective_target_name=inputs.effective_target_name, effective_connection=inputs.effective_connection, @@ -565,10 +571,36 @@ def _analyze_model_sql_in_parallel( else None ) if compact_candidate is not None: + completed_analyses: dict[str, PolyglotAnalysisResult] = {} + if analysis_cache is not None and len(compact_candidate.matching_indexes) == len(requests): + completed_analyses, _, _ = read_model_analyses( + context=analysis_cache, + cache_keys=request_cache_keys, + model_names=tuple(_model_name(request.model_input) for request in requests), + upstream_model_names_by_key={ + request.cache_key: _referenced_model_names( + model_input=request.model_input, + available_names=analyzed_model_names, + ) + for request in requests + if request.cache_key is not None + }, + ) + if all(key in completed_analyses for key in request_cache_keys): + record_analysis_cache_metrics( + batch_hits=len(requests), entry_hits=0, misses=0, bypasses=0 + ) + return { + _model_name(request.model_input): _ModelSqlAnalysis( + polyglot_analysis=completed_analyses[cache_key], + placeholders=request.placeholders, + ) + for request, cache_key in zip(requests, request_cache_keys, strict=True) + } try: compact_analyses: tuple[_ModelSqlAnalysis, ...] = _analyze_model_sql_requests( requests=requests, - cached_analyses={}, + cached_analyses=completed_analyses, column_nullability_by_table=column_nullability_by_table, column_types_by_table=column_types_by_table, inference_profile=inference_profile, @@ -832,38 +864,55 @@ def _analyze_model_sql_requests( prepared_by_index: dict[int, NativeCompactAnalysis] = {} diagnostics_by_index: dict[int, tuple[SqlBindingDiagnostic, ...]] = {} if uncached: + batch_requests: tuple[tuple[int, _ModelSqlAnalysisRequest], ...] = ( + tuple(enumerate(requests)) if cached_compact_batch is not None else uncached + ) prepared: tuple[NativeCompactAnalysis, ...] = ( analyze_queries_with_compact_polyglot_batch( - query_sqls=tuple(request.query_sql for _, request in uncached), - references=tuple(request.model_input.references for _, request in uncached), - placeholders=tuple(request.placeholders for _, request in uncached), + query_sqls=tuple(request.query_sql for _, request in batch_requests), + references=tuple( + request.model_input.references for _, request in batch_requests + ), + placeholders=tuple(request.placeholders for _, request in batch_requests), column_nullability_by_table=column_nullability_by_table, column_types_by_table=column_types_by_table, inference_profile=inference_profile, recover_cte_facts=tuple( - _should_recover_cte_facts(request.model_input) for _, request in uncached + _should_recover_cte_facts(request.model_input) + for _, request in batch_requests ), rich_type_inference=rich_type_inference, cached_batch=cached_compact_batch, - on_response=(on_compact_response if len(uncached) == len(requests) else None), + execution=CompactBatchExecutionOptions( + on_response=( + on_compact_response if len(uncached) == len(requests) else None + ), + binding_schemas=tuple( + request.binding_schema for _, request in batch_requests + ), + ), ) ) prepared_by_index = { - index: value for (index, _), value in zip(uncached, prepared, strict=True) + index: value for (index, _), value in zip(batch_requests, prepared, strict=True) } validation_indices: tuple[int, ...] = tuple( - index for index, request in uncached if request.binding_schema is not None + index + for index, request in uncached + if request.binding_schema is not None + and prepared_by_index[index].binding_diagnostics is None ) - validation_results: tuple[SqlBindingResult, ...] = get_schema_validations( - requests=tuple( - SqlSchemaValidationRequest( - sql=prepared_by_index[index].cleaned_sql, - dialect=inference_profile.sql_analysis_dialect, - schema=requests[index].binding_schema or {}, + with record_compile_timing("binding_validation_ms"): + validation_results: tuple[SqlBindingResult, ...] = get_schema_validations( + requests=tuple( + SqlSchemaValidationRequest( + sql=prepared_by_index[index].cleaned_sql, + dialect=inference_profile.sql_analysis_dialect, + schema=requests[index].binding_schema or {}, + ) + for index in validation_indices ) - for index in validation_indices ) - ) diagnostics_by_index = { index: result.diagnostics for index, result in zip(validation_indices, validation_results, strict=True) @@ -883,7 +932,9 @@ def _analyze_model_sql_requests( precomputed=( replace( prepared_by_index[index], - binding_diagnostics=diagnostics_by_index.get(index), + binding_diagnostics=diagnostics_by_index.get( + index, prepared_by_index[index].binding_diagnostics + ), ) if index in prepared_by_index else None @@ -976,9 +1027,10 @@ def _complete_inferred_bindings( for dependent_index in dependents_by_name.get(model_name, ()): if required_names_by_index[dependent_index] <= complete_schemas.keys(): ready.append(dependent_index) - validation_results: tuple[SqlBindingResult, ...] = get_schema_validations( - requests=tuple(deferred_validation_requests) - ) + with record_compile_timing("binding_validation_ms"): + validation_results: tuple[SqlBindingResult, ...] = get_schema_validations( + requests=tuple(deferred_validation_requests) + ) for index, validation_result in zip( deferred_validation_indices, validation_results, strict=True ): diff --git a/src/sqlbuild/compiler/compile/_helpers/attachment/core.py b/src/sqlbuild/compiler/compile/_helpers/attachment/core.py index a7cc7f5cda..10eba26326 100644 --- a/src/sqlbuild/compiler/compile/_helpers/attachment/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/attachment/core.py @@ -72,6 +72,7 @@ ) from sqlbuild.compiler.compile._helpers.render.sql_vars import ( expand_authored_sql_result, + prepare_static_project_vars_batch, substitute_sql_vars, ) from sqlbuild.compiler.compile._helpers.render.templating import ( @@ -86,6 +87,7 @@ from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.models import ( AuthoredSqlExpansionResult, + CompiledSqlExpansion, CompileModelConfig, CompileModelInput, CompileSeedInput, @@ -173,6 +175,7 @@ class _VisibleModelDeclarations: class _VisibleModelDeclarationCache: context: ModelInputBuildContext reusable_by_parent: bool + resource_specific: frozenset[ResourceIdentity] = frozenset() by_parent: dict[Path, _VisibleModelDeclarations] = field(default_factory=dict) @classmethod @@ -180,16 +183,17 @@ def build(cls, context: ModelInputBuildContext) -> _VisibleModelDeclarationCache resolver: DeclarationScopeResolver | None = context.declaration_resolver return cls( context=context, - reusable_by_parent=( - resolver is not None - and not any( - declaration.scope is ScopeKind.PRIVATE + reusable_by_parent=resolver is not None, + resource_specific=( + frozenset( + declaration.identity.owner for declaration in resolver.lookup.index.declarations + if declaration.scope is ScopeKind.PRIVATE + and declaration.identity.owner is not None ) - and not any( - resource.kind is ResourceKind.MODEL - for resource in resolver.lookup.grants_by_resource - ) + | frozenset(resolver.lookup.grants_by_resource) + if resolver is not None + else frozenset() ), ) @@ -197,6 +201,8 @@ def for_model( self, *, model_file: DiscoveredSqlModelFile, consumer: ResourceIdentity ) -> _VisibleModelDeclarations: parent: Path = model_file.file_path.parent + if consumer in self.resource_specific: + return _build_visible_declaration_indexes(model_file=model_file, context=self.context) cached: _VisibleModelDeclarations | None = self.by_parent.get(parent) if self.reusable_by_parent and cached is not None: return _rebind_visible_declarations(declarations=cached, consumer=consumer) @@ -430,8 +436,16 @@ def _build_model_inputs( target_config=target_config, ) declaration_cache: _VisibleModelDeclarationCache = _VisibleModelDeclarationCache.build(context) + model_files: tuple[DiscoveredSqlModelFile, ...] = discovered_inputs.model_files + prepared_var_substituted_sqls: tuple[str | None, ...] = prepare_static_project_vars_batch( + sqls=tuple(model_file.query_sql for model_file in model_files), + effective_vars=effective_vars, + ) model_file: DiscoveredSqlModelFile - for model_file in discovered_inputs.model_files: + prepared_var_substituted_sql: str | None + for model_file, prepared_var_substituted_sql in zip( + model_files, prepared_var_substituted_sqls, strict=True + ): model_identity: ResourceIdentity = ResourceIdentity( ResourceKind.MODEL, model_file.file_path.stem ) @@ -483,10 +497,14 @@ def _build_model_inputs( hook_functions=discovered_inputs.hook_functions, provider_names=frozenset(provider.name for provider in discovered_inputs.providers), ) - var_substituted_sql: str = substitute_sql_vars( - sql=model_file.query_sql, - file_path=model_file.file_path, - effective_vars=effective_vars, + var_substituted_sql: str = ( + prepared_var_substituted_sql + if prepared_var_substituted_sql is not None + else substitute_sql_vars( + sql=model_file.query_sql, + file_path=model_file.file_path, + effective_vars=effective_vars, + ) ) declaration_context: DeclarationResolutionContext = DeclarationResolutionContext( enums=declarations.enums, @@ -616,6 +634,15 @@ def _build_model_inputs( if header_schema_entry is None: model_inputs.append( CompileModelInput( + sql_expansion=( + CompiledSqlExpansion( + authored_sql=model_file.query_sql, + expanded_sql=macro_expansion.sql, + passes=(declaration_expansion.spans, macro_expansion.spans), + ) + if var_substituted_sql == model_file.query_sql + else None + ), model_file=model_file, config=model_config, query_sql=expanded_query_sql, @@ -634,6 +661,15 @@ def _build_model_inputs( model_inputs.append( CompileModelInput( + sql_expansion=( + CompiledSqlExpansion( + authored_sql=model_file.query_sql, + expanded_sql=macro_expansion.sql, + passes=(declaration_expansion.spans, macro_expansion.spans), + ) + if var_substituted_sql == model_file.query_sql + else None + ), model_file=model_file, config=model_config, query_sql=expanded_query_sql, diff --git a/src/sqlbuild/compiler/compile/_helpers/refs/references.py b/src/sqlbuild/compiler/compile/_helpers/refs/references.py index 5a56e70711..0e000b994b 100644 --- a/src/sqlbuild/compiler/compile/_helpers/refs/references.py +++ b/src/sqlbuild/compiler/compile/_helpers/refs/references.py @@ -4,6 +4,7 @@ import re +import sqlbuild._native as _native from sqlbuild.compiler.compile.constants import ( SQL_ARGUMENT_SEPARATOR_TOKEN, SQL_CLOSE_PAREN_TOKEN, @@ -39,6 +40,25 @@ def extract_sql_references(sql: str) -> tuple[CompileSqlReference, ...]: """Return logical SQL refs found outside comments and quoted text.""" + native_references: list[tuple[str, str, str | None, int | None]] | None = ( + _native.extract_static_sql_references(sql) + ) + if native_references is not None: + return tuple( + CompileSqlReference( + ref_kind=SqlReferenceKind(kind), + ref_name=name, + ref_package=package, + call_argument_count=call_argument_count, + ) + for kind, name, package, call_argument_count in native_references + ) + return _extract_sql_references_with_python(sql) + + +def _extract_sql_references_with_python(sql: str) -> tuple[CompileSqlReference, ...]: + """Return logical SQL refs through the authoritative general scanner.""" + references: list[CompileSqlReference] = [] index: int = 0 length: int = len(sql) diff --git a/src/sqlbuild/compiler/compile/_helpers/render/declarations.py b/src/sqlbuild/compiler/compile/_helpers/render/declarations.py index 6ac6091921..f150d64e1c 100644 --- a/src/sqlbuild/compiler/compile/_helpers/render/declarations.py +++ b/src/sqlbuild/compiler/compile/_helpers/render/declarations.py @@ -41,12 +41,13 @@ DeclarationIdentity, DeclarationRecord, ResourceIdentity, + ResourceRecord, ScopeIndex, UsageRecord, VisibilityRecord, VisibilityResolution, ) -from sqlbuild.compiler.scopes.types import DeclarationKind, ResourceKind, UsageKind +from sqlbuild.compiler.scopes.types import DeclarationKind, ResourceKind, ScopeKind, UsageKind from sqlbuild.compiler.sql_analysis.main._skip_block_comment import skip_block_comment from sqlbuild.compiler.sql_analysis.main._skip_line_comment import skip_line_comment from sqlbuild.compiler.sql_analysis.main._skip_quoted_text import skip_quoted_text @@ -64,6 +65,7 @@ r"(?P=quote)\s*\)" ) _DECLARATION_REFERENCE_START_PATTERN: re.Pattern[str] = re.compile(r"@(?Penum|const)\b") +_DECLARATION_SCAN_SPECIAL: re.Pattern[str] = re.compile(r"['\"`@/-]") _CONTEXT: str = "Enum and constant expansion" _ACCEPTED_VALUES_AUDIT: str = "accepted_values" _ENUM_REFERENCE_KIND: str = "enum" @@ -141,6 +143,12 @@ def build_declaration_scope_resolver( project_dir=discovered_inputs.project_dir, lookup=build_scope_lookup(index=scope_index), projection=DeclarationRuntimeProjection(declarations=MappingProxyType(declarations)), + resource_specific=frozenset( + declaration.identity.owner + for declaration in scope_index.declarations + if declaration.scope is ScopeKind.PRIVATE and declaration.identity.owner is not None + ) + | frozenset(grant.resource for grant in scope_index.grants), ) @@ -158,6 +166,21 @@ def resolve_declaration_context( target_path = file_path.relative_to(resolver.project_dir) except ValueError: target_path = file_path + resources: tuple[ResourceRecord, ...] = ( + resolver.lookup.resources.get(resource, ()) + if resource is not None + else resolver.lookup.resources_by_path.get(target_path.as_posix(), ()) + ) + cache_key: tuple[str, str] | None = None + if ( + len(resources) == 1 + and resolver.resource_specific is not None + and resources[0].identity not in resolver.resource_specific + ): + cache_key = (Path(resources[0].path).parent.as_posix(), resources[0].ownership_root.path) + cached: DeclarationResolutionContext | None = resolver.contexts_by_directory.get(cache_key) + if cached is not None: + return _rebind_declaration_context(context=cached, consumer=resources[0].identity) resolution: VisibilityResolution = resolve_scope_visibility( lookup=resolver.lookup, target=resource or target_path ) @@ -227,7 +250,7 @@ def resolve_declaration_context( constant_visibility[record.identity.name] = records elif record.identity.kind is DeclarationKind.MACRO: macro_visibility[record.identity.name] = records - return DeclarationResolutionContext( + context: DeclarationResolutionContext = DeclarationResolutionContext( enums=enums, constants=constants, inaccessible_enums=inaccessible_enums, @@ -243,6 +266,43 @@ def resolve_declaration_context( or (resolution.target.matches[0].identity if resolution.target.matches else None) ), ) + if cache_key is not None: + resolver.cache_context(key=cache_key, context=context) + return context + + +def _rebind_declaration_context( + *, context: DeclarationResolutionContext, consumer: ResourceIdentity +) -> DeclarationResolutionContext: + return replace( + context, + consumer=consumer, + enums=dict(context.enums), + constants=dict(context.constants), + macros=dict(context.macros), + macro_records=dict(context.macro_records), + inaccessible_enums=dict(context.inaccessible_enums), + inaccessible_constants=dict(context.inaccessible_constants), + inaccessible_macros=dict(context.inaccessible_macros), + enum_visibility=_rebind_visibility( + records_by_name=context.enum_visibility, consumer=consumer + ), + constant_visibility=_rebind_visibility( + records_by_name=context.constant_visibility, consumer=consumer + ), + macro_visibility=_rebind_visibility( + records_by_name=context.macro_visibility, consumer=consumer + ), + ) + + +def _rebind_visibility( + *, records_by_name: dict[str, tuple[VisibilityRecord, ...]], consumer: ResourceIdentity +) -> dict[str, tuple[VisibilityRecord, ...]]: + result: dict[str, tuple[VisibilityRecord, ...]] = {} + for name, records in records_by_name.items(): + result[name] = tuple(replace(record, resource=consumer) for record in records) + return result def declaration_usage_records( @@ -776,6 +836,10 @@ def _find_next_reference_start(*, sql: str, start: int) -> int | None: return None index: int = start while index < len(sql): + special: re.Match[str] | None = _DECLARATION_SCAN_SPECIAL.search(sql, index) + if special is None: + return None + index = special.start() character: str = sql[index] if character in SQL_QUOTE_TOKENS: index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) diff --git a/src/sqlbuild/compiler/compile/_helpers/render/sql_vars.py b/src/sqlbuild/compiler/compile/_helpers/render/sql_vars.py index 0466ff909b..561c891295 100644 --- a/src/sqlbuild/compiler/compile/_helpers/render/sql_vars.py +++ b/src/sqlbuild/compiler/compile/_helpers/render/sql_vars.py @@ -7,6 +7,7 @@ from collections.abc import Mapping from pathlib import Path +import sqlbuild._native as _native from sqlbuild.compiler.compile._helpers.render.declarations import ( expand_declaration_references_result, expand_declaration_references_with_spans, @@ -48,6 +49,10 @@ _SQL_INTERPOLATION_SPECIAL: re.Pattern[str] = re.compile(r"['\"`\-/@]") _CONTEXT: str = "SQL interpolation" +_NATIVE_UNCHANGED: int = 0 +_NATIVE_SUBSTITUTED: int = 1 +_NATIVE_FALLBACK: int = 2 +_NATIVE_RESULT_LENGTH: int = 2 def validate_var_macro_collision( @@ -208,6 +213,42 @@ def substitute_sql_vars( return rendered_sql +def prepare_static_project_vars_batch( + *, + sqls: tuple[str, ...], + effective_vars: dict[str, object], +) -> tuple[str | None, ...]: + """Prepare safe static substitutions, leaving dynamic fallbacks in authored order.""" + + scalar_variables: list[tuple[str, str]] = [ + (name, render_project_var_text(value=value, label=f"SQL variable '@@{name}'")) + for name, value in effective_vars.items() + if value is None or isinstance(value, str | int | float | bool) + ] + raw_results: object = _native.substitute_static_project_vars(list(sqls), scalar_variables) + if not isinstance(raw_results, list) or len(raw_results) != len(sqls): + raise CompileInputError("native static SQL interpolation returned an invalid batch") + results: list[str | None] = [] + for sql, raw_result in zip(sqls, raw_results, strict=True): + if not ( + isinstance(raw_result, tuple) + and len(raw_result) == _NATIVE_RESULT_LENGTH + and type(raw_result[0]) is int + and (raw_result[1] is None or isinstance(raw_result[1], str)) + ): + raise CompileInputError("native static SQL interpolation returned an invalid result") + status, rendered = raw_result + if status == _NATIVE_UNCHANGED and rendered is None: + results.append(sql) + elif status == _NATIVE_SUBSTITUTED and rendered is not None: + results.append(rendered) + elif status == _NATIVE_FALLBACK and rendered is None: + results.append(None) + else: + raise CompileInputError("native static SQL interpolation returned an invalid status") + return tuple(results) + + def substitute_sql_vars_with_spans( *, sql: str, @@ -219,7 +260,6 @@ def substitute_sql_vars_with_spans( if SQL_INTERPOLATION_TOKEN not in sql: return sql, () - parts: list[str] = [] spans: list[ExpansionSpan] = [] output_length: int = 0 diff --git a/src/sqlbuild/compiler/compile/models.py b/src/sqlbuild/compiler/compile/models.py index 7a627778d9..aba90f651b 100644 --- a/src/sqlbuild/compiler/compile/models.py +++ b/src/sqlbuild/compiler/compile/models.py @@ -13,7 +13,9 @@ from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.types import ( AttachedAuditTargetKind, + CompactBatchResponseCallback, CompiledResourceType, + CompileInputsReadyCallback, DiagnosticPhase, DiagnosticSeverity, FunctionLanguage, @@ -291,6 +293,14 @@ class DeclarationScopeResolver: project_dir: Path | None lookup: ScopeLookup projection: DeclarationRuntimeProjection + resource_specific: frozenset[ResourceIdentity] | None = None + contexts_by_directory: dict[tuple[str, str], DeclarationResolutionContext] = field( + default_factory=dict, compare=False, repr=False + ) + + def cache_context(self, *, key: tuple[str, str], context: DeclarationResolutionContext) -> None: + """Retain a process-local context for equivalent lexical consumers.""" + self.contexts_by_directory[key] = context @dataclass(frozen=True) @@ -613,12 +623,15 @@ class AnalysisCacheContext: @dataclass(frozen=True) class CompileAnalysisSelection: - """Selection inputs used to limit deep model SQL analysis.""" + """Deep-analysis selection and invocation-local preparation observer.""" select: tuple[str, ...] = () exclude: tuple[str, ...] = () auto_load_sources: bool = False no_cache: bool = False + on_inputs_ready: CompileInputsReadyCallback | None = field( + default=None, compare=False, repr=False + ) @dataclass(frozen=True) @@ -717,6 +730,7 @@ class CompileModelInput: macro_deps: tuple[str, ...] = field(default_factory=tuple) macro_usages: tuple[UsageRecord, ...] = field(default_factory=tuple) declaration_usages: tuple[UsageRecord, ...] = field(default_factory=tuple) + sql_expansion: CompiledSqlExpansion | None = field(default=None, compare=False, repr=False) @dataclass(frozen=True) @@ -1006,6 +1020,9 @@ class CompiledProject: diagnostics: tuple[CompilerDiagnostic, ...] = field(default_factory=tuple) external_sql_reference_resolver: ExternalSqlReferenceResolver | None = None scope_index: ScopeIndex = field(default_factory=ScopeIndex) + sql_expansions: dict[Path, CompiledSqlExpansion] = field( + default_factory=dict, compare=False, repr=False + ) @dataclass(frozen=True) @@ -1141,6 +1158,15 @@ class ExpansionSpan: output_end: int +@dataclass(frozen=True) +class CompiledSqlExpansion: + """Process-local authored-to-expanded SQL evidence shared with compiler checks.""" + + authored_sql: str + expanded_sql: str + passes: tuple[tuple[ExpansionSpan, ...], ...] + + @dataclass(frozen=True) class MappedOffset: """An offset in rendered SQL resolved back onto the text that produced it.""" @@ -1226,6 +1252,14 @@ class CompactProjectedFacts: ] +@dataclass(frozen=True) +class CompactBatchExecutionOptions: + """Optional binding work and cache publication for one native batch.""" + + binding_schemas: tuple[dict[str, dict[str, str]] | None, ...] | None = None + on_response: CompactBatchResponseCallback | None = None + + @dataclass(frozen=True) class CompactBatchPreparation: """Deduplicated native request data and ordered cleaned SQL inputs.""" diff --git a/src/sqlbuild/compiler/compile/types.py b/src/sqlbuild/compiler/compile/types.py index 9c7a2d2a58..1ff7649d5b 100644 --- a/src/sqlbuild/compiler/compile/types.py +++ b/src/sqlbuild/compiler/compile/types.py @@ -8,7 +8,7 @@ from sqlbuild.sql_values.models import SqlValue if TYPE_CHECKING: - from sqlbuild.compiler.compile.models import CompactBatchPreparation + from sqlbuild.compiler.compile.models import CompactBatchPreparation, CompileProjectInputs class TypedSqlValueRenderer(Protocol): @@ -31,6 +31,12 @@ class CompactBatchResponseCallback(Protocol): def __call__(self, *, preparation: CompactBatchPreparation, response: object) -> None: ... +class CompileInputsReadyCallback(Protocol): + """Invocation-local observer for the completed pre-semantic input snapshot.""" + + def __call__(self, *, inputs: CompileProjectInputs) -> None: ... + + class AttachedAuditTargetKind(StrEnum): MODEL = "model" SOURCE = "source" diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py b/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py index d265ebe064..3cad384065 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py @@ -112,7 +112,7 @@ class _ModelHeaderToken: @dataclass(frozen=True) class _ModelHeaderTokenization: values: dict[str, object] | None - column_offsets: tuple[tuple[str, int, int], ...] | None + column_relative_locations: tuple[tuple[str, int, int, int], ...] | None error: str | None @@ -264,21 +264,25 @@ def header_column_locations( """Return authored column locations from a parsed SQLBuild header.""" tokenization: _ModelHeaderTokenization = _model_header_parse(header) - if tokenization.column_offsets is None: - raise ModelHeaderSyntaxError("Native MODEL header parser returned no column offsets") - line_starts: tuple[int, ...] | None = None - locations: dict[str, SourceLocation] = {} - for name, position, length in tokenization.column_offsets: - if line_starts is None: - line_starts = _line_starts(contents) - locations[name] = _location_for_absolute_span( - contents=contents, - start=header_start + position, - end=header_start + position + length, - relative_path=relative_path, - line_starts=line_starts, + if tokenization.column_relative_locations is None: + raise ModelHeaderSyntaxError("Native MODEL header parser returned no column locations") + header_line: int = contents.count("\n", 0, header_start) + 1 + preceding_newline: int = contents.rfind("\n", 0, header_start) + header_column: int = ( + header_start + 1 if preceding_newline < 0 else header_start - preceding_newline + ) + return { + name: SourceLocation( + path=relative_path, + line=header_line + relative_line - 1, + column=relative_column + (header_column - 1 if relative_line == 1 else 0), + end_line=header_line + relative_line - 1, + end_column=( + relative_column + (header_column - 1 if relative_line == 1 else 0) + length + ), ) - return locations + for name, relative_line, relative_column, length in tokenization.column_relative_locations + } def model_output_column_locations( @@ -643,12 +647,39 @@ def prepare_model_header_tokens(headers: list[str]) -> None: values=( _project_native_header_map(native_values) if native_values is not None else None ), - column_offsets=(tuple(column_offsets) if column_offsets is not None else None), + column_relative_locations=( + _header_column_relative_locations(header=header, offsets=column_offsets) + if column_offsets is not None + else None + ), error=error, ), ) +def _header_column_relative_locations( + *, header: str, offsets: list[tuple[str, int, int]] +) -> tuple[tuple[str, int, int, int], ...]: + """Project ordered header offsets to one-based relative source locations in one pass.""" + + locations: list[tuple[str, int, int, int]] = [] + line: int = 1 + column: int = 1 + cursor: int = 0 + for name, position, length in offsets: + while cursor < position: + newline: int = header.find("\n", cursor, position) + if newline < 0: + column += position - cursor + cursor = position + else: + line += 1 + column = 1 + cursor = newline + 1 + locations.append((name, line, column, length)) + return tuple(locations) + + def prepare_model_file_headers(contents: list[str]) -> None: """Batch tokenization for syntactically recognizable MODEL file headers.""" diff --git a/src/sqlbuild/compiler/pipeline/main/compiled_project.py b/src/sqlbuild/compiler/pipeline/main/compiled_project.py index c088c7bca0..84d343e16e 100644 --- a/src/sqlbuild/compiler/pipeline/main/compiled_project.py +++ b/src/sqlbuild/compiler/pipeline/main/compiled_project.py @@ -86,6 +86,8 @@ def build_compiled_project( external_sql_reference_resolver=external_sql_reference_resolver, no_cache=no_cache, ) + if analysis_selection is not None and analysis_selection.on_inputs_ready is not None: + analysis_selection.on_inputs_ready(inputs=compile_inputs) inference_profile: ExpressionInferenceProfile = adapter.expression_inference_profile() analysis_model_names: frozenset[str] | None = _resolve_analysis_model_names( compile_inputs=compile_inputs, diff --git a/src/sqlbuild/compiler/profiling/constants.py b/src/sqlbuild/compiler/profiling/constants.py index ba6a20cddc..2794341f76 100644 --- a/src/sqlbuild/compiler/profiling/constants.py +++ b/src/sqlbuild/compiler/profiling/constants.py @@ -5,6 +5,10 @@ COMPILE_TIMING_PHASES: tuple[CompileTimingPhase, ...] = ( "attachment_ms", "model_analysis_ms", + "analysis_preparation_ms", + "analysis_native_ms", + "analysis_projection_ms", + "binding_validation_ms", "test_input_compile_ms", "test_planning_ms", "comparison_render_ms", diff --git a/src/sqlbuild/compiler/profiling/types.py b/src/sqlbuild/compiler/profiling/types.py index 35337a97a3..64d8b0dd73 100644 --- a/src/sqlbuild/compiler/profiling/types.py +++ b/src/sqlbuild/compiler/profiling/types.py @@ -5,6 +5,10 @@ type CompileTimingPhase = Literal[ "attachment_ms", "model_analysis_ms", + "analysis_preparation_ms", + "analysis_native_ms", + "analysis_projection_ms", + "binding_validation_ms", "test_input_compile_ms", "test_planning_ms", "comparison_render_ms", diff --git a/src/sqlbuild/compiler/sql_analysis/_helpers/polyglot_normalization.py b/src/sqlbuild/compiler/sql_analysis/_helpers/polyglot_normalization.py index ee054f8514..12dd8a3202 100644 --- a/src/sqlbuild/compiler/sql_analysis/_helpers/polyglot_normalization.py +++ b/src/sqlbuild/compiler/sql_analysis/_helpers/polyglot_normalization.py @@ -56,15 +56,21 @@ def normalize_sql_for_polyglot_impl(*, sql: str, dialect: str | None) -> str: if sql[index] in SQL_QUOTE_CHARACTERS: index = _skip_snowflake_quoted_text(sql=sql, start=index) continue + if sql[index] != SQL_OPEN_BRACKET: + index += 1 + continue close_bracket: int | None = _matching_bracket(sql=sql, open_bracket=index) if close_bracket is None: index += 1 continue key_sql: str = sql[index + 1 : close_bracket] + if _SIMPLE_BRACKET_KEY.fullmatch(key_sql) is not None: + index = close_bracket + 1 + continue normalized_key: str = normalize_sql_for_polyglot_impl(sql=key_sql, dialect=dialect) base_end: int = _trailing_trivia_start(sql=sql, end=index) base_match: re.Match[str] | None = _VARIANT_PATH_SUFFIX.search(sql[:base_end]) - if base_match is not None and _SIMPLE_BRACKET_KEY.fullmatch(key_sql) is None: + if base_match is not None: base_sql: str = base_match.group("base") trivia: str = sql[base_end:index] replacements.append( diff --git a/src/sqlbuild/compiler/sql_analysis/main/_decode_schema_validation.py b/src/sqlbuild/compiler/sql_analysis/main/_decode_schema_validation.py new file mode 100644 index 0000000000..d72928817c --- /dev/null +++ b/src/sqlbuild/compiler/sql_analysis/main/_decode_schema_validation.py @@ -0,0 +1,11 @@ +"""Decode binding evidence from native compilation.""" + +from sqlbuild.compiler.sql_analysis._helpers.schema_validation import _binding_result +from sqlbuild.compiler.sql_analysis.models import SqlBindingResult + + +def decode_schema_validation( + *, sql: str, dialect: str | None, response: object +) -> SqlBindingResult: + """Decode binding evidence produced by either native compilation entry point.""" + return _binding_result(sql=sql, dialect=dialect, response=response) diff --git a/src/sqlbuild/compiler/sql_analysis/main/_prepare_schema_validation.py b/src/sqlbuild/compiler/sql_analysis/main/_prepare_schema_validation.py new file mode 100644 index 0000000000..b7cae89a6c --- /dev/null +++ b/src/sqlbuild/compiler/sql_analysis/main/_prepare_schema_validation.py @@ -0,0 +1,9 @@ +"""Prepare shared native schema-validation inputs.""" + +from sqlbuild.compiler.sql_analysis._helpers.schema_validation import _request_payload +from sqlbuild.compiler.sql_analysis.models import SqlSchemaValidationRequest + + +def prepare_schema_validation(*, request: SqlSchemaValidationRequest) -> dict[str, object]: + """Prepare the schema-validation protocol for native compilation.""" + return _request_payload(request=request) diff --git a/src/sqlbuild/lint/_helpers/expansion.py b/src/sqlbuild/lint/_helpers/expansion.py index 68e76d1cf5..c34329f481 100644 --- a/src/sqlbuild/lint/_helpers/expansion.py +++ b/src/sqlbuild/lint/_helpers/expansion.py @@ -5,6 +5,7 @@ import re from pathlib import Path +import sqlbuild._native as _native from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.discovery.main.resolve_adapter import resolve_adapter from sqlbuild.compiler.compile.constants import ( @@ -16,7 +17,11 @@ from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.main.expand_sql_with_spans import expand_sql_with_spans from sqlbuild.compiler.compile.main.sql_expansion_context import build_sql_expansion_context -from sqlbuild.compiler.compile.models import ExpansionSpan, SqlExpansionContext +from sqlbuild.compiler.compile.models import ( + CompiledSqlExpansion, + ExpansionSpan, + SqlExpansionContext, +) from sqlbuild.compiler.compile.types import TypedSqlValueRenderer from sqlbuild.compiler.discovery.exceptions import DiscoveryError from sqlbuild.compiler.discovery.main.discover import discover_project_inputs @@ -133,6 +138,7 @@ def prepare_lint_body( external_identifiers: tuple[str, ...] = (), allows_ceremonial_select: bool = False, allows_dynamic_output_star: bool = False, + compiled_expansion: CompiledSqlExpansion | None = None, ) -> LintBody: """Expand one authored body and neutralize whatever interpolation remains.""" @@ -149,6 +155,15 @@ def prepare_lint_body( expansion_passes: tuple[tuple[ExpansionSpan, ...], ...] try: if ( + compiled_expansion is not None + and authored_body.strip() == compiled_expansion.authored_sql + ): + body_start += len(authored_body) - len(authored_body.lstrip()) + body_end = body_start + len(compiled_expansion.authored_sql) + expansion_input = compiled_expansion.authored_sql + expanded = compiled_expansion.expanded_sql + expansion_passes = compiled_expansion.passes + elif ( MACRO_TOKEN not in expansion_input and TEMPLATE_INTERPOLATION_START not in expansion_input ): @@ -164,18 +179,28 @@ def prepare_lint_body( ) from error neutralized: str sites: tuple[InterpolationSite, ...] - neutralized, sites = neutralize_interpolation(body=expanded) + native_prepared: tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None = ( + _native.prepare_lint_sql( + expanded, expansion_input, [site.neutralized_start for site in pre_expansion_sites] + ) + ) + if native_prepared is None: + neutralized, sites = neutralize_interpolation(body=expanded) + externally_referenced_ctes: tuple[str, ...] = _externally_referenced_ctes( + expanded=expanded, + interpolation_sites=sites, + pre_expansion_body=expansion_input, + pre_expansion_sites=pre_expansion_sites, + ) + else: + neutralized = native_prepared[0] + sites = tuple(InterpolationSite(*site) for site in native_prepared[1]) + externally_referenced_ctes = tuple(native_prepared[2]) dependency_identifiers: tuple[str, ...] = tuple( site.sentinel for site in sites if _DEPENDENCY_INTRINSIC_PATTERN.match(site.original_text) is not None ) - externally_referenced_ctes: tuple[str, ...] = _externally_referenced_ctes( - expanded=expanded, - interpolation_sites=sites, - pre_expansion_body=expansion_input, - pre_expansion_sites=pre_expansion_sites, - ) return LintBody( file_path=file_path, body_start=body_start, diff --git a/src/sqlbuild/lint/_helpers/native_sql.py b/src/sqlbuild/lint/_helpers/native_sql.py index da91aa2672..a47b599724 100644 --- a/src/sqlbuild/lint/_helpers/native_sql.py +++ b/src/sqlbuild/lint/_helpers/native_sql.py @@ -2,15 +2,13 @@ from __future__ import annotations -import json -import os import re from bisect import bisect_right -from concurrent.futures import ThreadPoolExecutor -from itertools import chain from pathlib import Path from typing import Any, cast +import orjson + import sqlbuild._native as _native from sqlbuild.compiler.compile.constants import ( ASSERT_TEST_CTE_PREFIX, @@ -41,7 +39,6 @@ from sqlbuild.lint.models import LintBody, LintConfig, LintEdit, LintViolation _NATIVE_LINT_API_VERSION: int = 1 -_MAX_NATIVE_LINT_WORKERS: int = 8 _NEWLINE_CHARACTER: str = "\n" _UNUSED_CTE_CODE: str = "SQBRSQL005" _PARSE_ERROR_POSITION_PATTERN: re.Pattern[str] = re.compile( @@ -155,40 +152,37 @@ def _native_cache_key(*, body: LintBody, config: LintConfig) -> _NativeCacheKey: def _native_responses( *, requests: dict[_NativeCacheKey, dict[str, object]] ) -> dict[_NativeCacheKey, _NativeResult]: - """Analyze unique SQL bodies concurrently while preserving request order.""" + """Analyze unique SQL bodies in one bounded native batch.""" keys: tuple[_NativeCacheKey, ...] = tuple(requests) payloads: tuple[dict[str, object], ...] = tuple(requests.values()) - worker_count: int = min( - len(payloads), - _MAX_NATIVE_LINT_WORKERS, - os.cpu_count() or 1, - ) - if worker_count <= 1: - results: tuple[_NativeResult, ...] = tuple( - _native_response_or_error(payload=payload) for payload in payloads - ) - else: - chunk_size: int = (len(payloads) + worker_count - 1) // worker_count - chunks: tuple[tuple[dict[str, object], ...], ...] = tuple( - payloads[start : start + chunk_size] for start in range(0, len(payloads), chunk_size) + try: + decoded: object = orjson.loads( + _native.lint_sql_batch_json( + orjson.dumps(payloads, option=orjson.OPT_SORT_KEYS).decode() + ) ) - with ThreadPoolExecutor(max_workers=worker_count) as executor: - results = tuple(chain.from_iterable(executor.map(_native_response_chunk, chunks))) + except (TypeError, ValueError) as error: + raise NativeLintError(str(error)) from error + if not isinstance(decoded, list) or len(decoded) != len(keys): + raise NativeLintError("native lint engine returned an invalid batch response") + results: tuple[_NativeResult, ...] = tuple(_native_result(response=value) for value in decoded) return dict(zip(keys, results, strict=True)) -def _native_response_or_error(payload: dict[str, object]) -> _NativeResult: - try: - return _native_response(payload=payload) - except NativeLintError as error: - return error - - -def _native_response_chunk( - payloads: tuple[dict[str, object], ...], -) -> tuple[_NativeResult, ...]: - return tuple(_native_response_or_error(payload) for payload in payloads) +def _native_result(*, response: object) -> _NativeResult: + if not isinstance(response, dict): + return NativeLintError("native lint engine returned a non-object response") + payload: dict[str, object] = cast(dict[str, object], response) + if isinstance(payload.get("Err"), str): + return NativeLintError(cast(str, payload["Err"])) + decoded: object = payload.get("Ok") + if not isinstance(decoded, dict): + return NativeLintError("native lint engine returned a non-object response") + result: dict[str, Any] = cast(dict[str, Any], decoded) + if result.get("version") != _NATIVE_LINT_API_VERSION: + return NativeLintError("native lint engine returned an unsupported response version") + return result def _parse_failure_violation( @@ -221,20 +215,6 @@ def _parse_failure_violation( ) -def _native_response(*, payload: dict[str, object]) -> dict[str, Any]: - try: - decoded: object = json.loads( - _native.lint_sql_json(json.dumps(payload, sort_keys=True, separators=(",", ":"))) - ) - except (TypeError, ValueError) as error: - raise NativeLintError(str(error)) from error - if not isinstance(decoded, dict): - raise NativeLintError("native lint engine returned a non-object response") - if decoded.get("version") != _NATIVE_LINT_API_VERSION: - raise NativeLintError("native lint engine returned an unsupported response version") - return {str(key): value for key, value in decoded.items()} - - def _authored_violation( *, raw_diagnostic: object, body: LintBody, contents: str ) -> LintViolation | None: diff --git a/src/sqlbuild/lint/main/build_expansion_context.py b/src/sqlbuild/lint/main/build_expansion_context.py new file mode 100644 index 0000000000..d01bff6005 --- /dev/null +++ b/src/sqlbuild/lint/main/build_expansion_context.py @@ -0,0 +1,15 @@ +"""Prepare invocation-local SQL expansion evidence for compiler-integrated lint.""" + +from pathlib import Path + +from sqlbuild.compiler.compile.models import SqlExpansionContext +from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs +from sqlbuild.lint._helpers.expansion import build_lint_expansion_context + + +def build_expansion_context( + *, project_dir: Path, discovered_inputs: DiscoveredProjectInputs +) -> SqlExpansionContext: + return build_lint_expansion_context( + project_dir=project_dir, discovered_inputs=discovered_inputs + ) diff --git a/src/sqlbuild/lint/main/run_lint.py b/src/sqlbuild/lint/main/run_lint.py index afd632134c..8c36a82172 100644 --- a/src/sqlbuild/lint/main/run_lint.py +++ b/src/sqlbuild/lint/main/run_lint.py @@ -5,7 +5,7 @@ from pathlib import Path from sqlbuild.compiler.compile.constants import MODEL_DIRECTORY_NAME -from sqlbuild.compiler.compile.models import SqlExpansionContext +from sqlbuild.compiler.compile.models import CompiledSqlExpansion, SqlExpansionContext from sqlbuild.compiler.compile.types import TypedSqlValueRenderer from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.compiler.scopes.constants import ( @@ -30,6 +30,8 @@ def run_lint( selected_paths: frozenset[Path] | None = None, discovered_inputs: DiscoveredProjectInputs | None = None, dynamic_output_paths: frozenset[Path] = frozenset(), + compiled_expansions: dict[Path, CompiledSqlExpansion] | None = None, + expansion_context: SqlExpansionContext | None = None, ) -> LintRunResult: """Lint all DSL files in the project without modifying anything.""" @@ -38,7 +40,7 @@ def run_lint( ) violations: list[LintViolation] = [] bodies: list[LintBody] = [] - context: SqlExpansionContext | None = _expansion_context( + context: SqlExpansionContext | None = expansion_context or _expansion_context( project_dir=project_dir, native_enabled=config.native_enabled, value_renderer=value_renderer, @@ -58,14 +60,15 @@ def run_lint( and not declaration_directories.intersection(relative_parts[1:-1]) ), ) - violations.extend( - lint_native_headers( - contents=contents, - file_path=file_path, - headers=headers, - config=config, + if config.header_rules_enabled: + violations.extend( + lint_native_headers( + contents=contents, + file_path=file_path, + headers=headers, + config=config, + ) ) - ) if context is not None: bodies.extend( _prepared_bodies( @@ -75,6 +78,7 @@ def run_lint( context=context, project_dir=project_dir, allows_dynamic_output_star=file_path.resolve() in dynamic_output_paths, + compiled_expansions=compiled_expansions, ) ) @@ -120,8 +124,10 @@ def _prepared_bodies( context: SqlExpansionContext, project_dir: Path, allows_dynamic_output_star: bool, + compiled_expansions: dict[Path, CompiledSqlExpansion] | None, ) -> tuple[LintBody, ...]: bodies: list[LintBody] = [] + compiled_expansion: CompiledSqlExpansion | None = (compiled_expansions or {}).get(file_path) external_identifiers: tuple[str, ...] = external_identifiers_for_headers( contents=contents, headers=headers ) @@ -147,6 +153,7 @@ def _prepared_bodies( external_identifiers=external_identifiers, allows_ceremonial_select=allows_ceremonial_select, allows_dynamic_output_star=allows_dynamic_output_star, + compiled_expansion=compiled_expansion, ) ) return tuple(bodies) diff --git a/src/sqlbuild/lint/models.py b/src/sqlbuild/lint/models.py index 680b371403..ccebbe0b77 100644 --- a/src/sqlbuild/lint/models.py +++ b/src/sqlbuild/lint/models.py @@ -116,6 +116,7 @@ class LintConfig: dialect: str = "generic" enabled_native_rules: tuple[str, ...] | None = None ignored_native_rules: tuple[str, ...] = () + header_rules_enabled: bool = True @dataclass(frozen=True) diff --git a/src/sqlbuild/rule_engine/_helpers/engine/native.py b/src/sqlbuild/rule_engine/_helpers/engine/native.py index 605662494a..b850a6e9f1 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/native.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/native.py @@ -9,7 +9,8 @@ import re import sys import tempfile -from dataclasses import asdict, replace +from collections.abc import Sequence +from dataclasses import asdict, fields, replace from decimal import Decimal from pathlib import Path from typing import Any, cast @@ -21,7 +22,6 @@ from sqlbuild.compiler.compile.models import ( CompactLineageFacts, CompiledLineageColumnFact, - CompiledLineageSourceFact, CompiledModel, CompiledProject, CompiledSqlScenario, @@ -70,6 +70,7 @@ def evaluate_native( catalogue: tuple[Rule, ...], dialect: str = "generic", initial_findings: tuple[Finding, ...] = (), + defer_suppressions: bool = False, ) -> RulesResult: """Evaluate one compiled model batch through the native engine.""" @@ -98,6 +99,7 @@ def evaluate_native( ], "scope_index": scope_metadata_projection(index=project.scope_index), "initial_findings": [_finding_payload(finding) for finding in initial_findings], + "defer_suppressions": defer_suppressions, "custom_rules": _custom_rule_payloads( catalogue=catalogue, project=project, project_dir=project_dir ), @@ -106,6 +108,15 @@ def evaluate_native( retry_native_misses: int = 0 request["custom_host"] = None try: + if not config.cache.enabled or not (project_dir / "target" / "rules-cache").exists(): + custom_host, custom_host_input = _custom_host_payload( + project=project, + config=config, + project_dir=project_dir, + catalogue=tuple(selected_catalogue), + dialect=dialect, + ) + request["custom_host"] = custom_host try: response_json: str = _evaluate_request(request) except ValueError as error: @@ -152,6 +163,33 @@ def _evaluate_request(request: dict[str, object]) -> str: ) +def finalize_native_findings( + *, + config: RulesConfig, + project_dir: Path, + evaluated_codes: tuple[str, ...], + findings: tuple[Finding, ...], +) -> tuple[Finding, ...]: + """Apply exception policy once to completed SQL, native, and custom findings.""" + + request: dict[str, object] = { + "version": RULES_NATIVE_API_VERSION, + "project_dir": str(project_dir.resolve()), + "config": _config_payload(config), + "evaluated_codes": evaluated_codes, + "findings": [_finding_payload(finding) for finding in findings], + } + try: + payload: object = orjson.loads( + _native.finalize_rule_findings_json(orjson.dumps(request).decode()) + ) + except (ValueError, TypeError) as error: + raise RulesError(str(error)) from error + if not isinstance(payload, list): + raise RulesError("native rules engine returned invalid finalized findings") + return tuple(_decode_finding(value) for value in payload) + + def load_native_config(project_dir: Path) -> dict[str, object]: """Load strict rules TOML through the native configuration owner.""" @@ -582,42 +620,56 @@ def _custom_fact_fingerprint(*, project: CompiledProject, attributes: frozenset[ return hashlib.sha256(encoded).hexdigest() -def _project_fact_models(*, project: CompiledProject) -> tuple[CompiledModel, ...]: +def _project_fact_models(*, project: CompiledProject) -> tuple[dict[str, object], ...]: """Project models with lazy lineage represented by stable semantic data only.""" - projected: list[CompiledModel] = [] + projected: list[dict[str, object]] = [] for model in project.models: - projected_model: CompiledModel = model - lineage: object = model.fast_lineage_columns - if isinstance(lineage, CompactLineageFacts): - lineage = _compact_lineage_fact_payloads(lineage=lineage) - projected_model = replace(model, fast_lineage_columns=cast(Any, lineage)) - projected.append(projected_model) + payload: dict[str, object] = { + field.name: getattr(model, field.name) for field in fields(model) + } + payload["fast_lineage_columns"] = _lineage_fingerprint_payload( + lineage=model.fast_lineage_columns + ) + projected.append(payload) return tuple(projected) -def _compact_lineage_fact_payloads( - *, lineage: CompactLineageFacts -) -> tuple[CompiledLineageColumnFact, ...]: - """Project compact lineage into the same semantic shape as materialized facts.""" - - payloads: list[CompiledLineageColumnFact] = [] - for name_index, transform_code, confidence_code, sources in lineage.rows: - payloads.append( - CompiledLineageColumnFact( - output_column=lineage.string_pool[name_index], - upstream_columns=tuple( - CompiledLineageSourceFact( - resource_type=lineage.string_pool[source[0]], - resource_name=lineage.resource_name(source[1]), - column_name=lineage.string_pool[source[2]], - ) - for source in sources - ), - transform_kind=lineage.transform_kind(transform_code), - confidence=lineage.confidence(confidence_code), +def _lineage_fingerprint_payload( + *, lineage: Sequence[CompiledLineageColumnFact] | None +) -> tuple[tuple[object, ...], ...] | None: + """Hash semantic lineage without rehydrating compact facts into Python objects.""" + + if lineage is None: + return None + payloads: list[tuple[object, ...]] = [] + if isinstance(lineage, CompactLineageFacts): + for name_index, transform_code, confidence_code, sources in lineage.rows: + upstream: tuple[tuple[str, str, str], ...] = tuple( + ( + lineage.string_pool[source[0]], + lineage.resource_name(source[1]), + lineage.string_pool[source[2]], + ) + for source in sources + ) + payloads.append( + ( + lineage.string_pool[name_index], + upstream, + lineage.transform_kind(transform_code), + lineage.confidence(confidence_code), + ) + ) + else: + for column in lineage: + upstream = tuple( + (source.resource_type, source.resource_name, source.column_name) + for source in column.upstream_columns + ) + payloads.append( + (column.output_column, upstream, column.transform_kind, column.confidence) ) - ) return tuple(payloads) @@ -653,6 +705,7 @@ def _custom_host_payload( return None, None serializable_project: CompiledProject = replace( project, + sql_expansions={}, loaded_macros={}, loader_functions=(), hook_functions=(), diff --git a/src/sqlbuild/rule_engine/_helpers/run/rules.py b/src/sqlbuild/rule_engine/_helpers/run/rules.py index b03fb82f8b..3149e33871 100644 --- a/src/sqlbuild/rule_engine/_helpers/run/rules.py +++ b/src/sqlbuild/rule_engine/_helpers/run/rules.py @@ -6,24 +6,46 @@ import json import os import time +from concurrent.futures import Executor, Future, ThreadPoolExecutor from dataclasses import asdict, dataclass, replace from importlib.metadata import version from pathlib import Path from typing import Any -from sqlbuild.compiler.compile.models import CompiledModel, CompiledObjectKey, CompiledProject +from sqlbuild.compiler.compile.models import ( + CompiledModel, + CompiledObjectKey, + CompiledProject, + CompiledSqlExpansion, + CompileProjectInputs, + SqlExpansionContext, +) from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.compiler.pipeline.models import ProjectGraph +from sqlbuild.lint.main.build_expansion_context import build_expansion_context from sqlbuild.lint.main.collect_project_files import collect_project_files from sqlbuild.lint.main.run_lint import run_lint from sqlbuild.lint.models import LintConfig, LintRunResult, LintViolation from sqlbuild.rule_engine._helpers.engine.catalogue import build_catalogue, select_rules from sqlbuild.rule_engine._helpers.engine.config import resolve_rule_ignore_selectors from sqlbuild.rule_engine._helpers.engine.hermeticity import verify_custom_rules -from sqlbuild.rule_engine._helpers.engine.native import evaluate_native +from sqlbuild.rule_engine._helpers.engine.native import ( + evaluate_native, + finalize_native_findings, + native_catalogue, +) from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.models import Finding, Rule, RulesConfig, RulesResult, RulesRunResult +from sqlbuild.rule_engine.main.load_config import load_rules_config +from sqlbuild.rule_engine.models import ( + Finding, + PreparedSqlLint, + PreparedSqlLintResult, + Rule, + RulesConfig, + RulesResult, + RulesRunResult, +) _SQL_RULE_CACHE_VERSION: str = "sql-rules-v2" _SQLBUILD_VERSION: str = version("sqlbuild") @@ -45,8 +67,9 @@ def evaluate_rules( project_dir: Path, dialect: str, selected_keys: frozenset[CompiledObjectKey] | None = None, + prepared_sql: PreparedSqlLint | None = None, ) -> RulesRunResult: - """Run selected native built-ins before selected custom Python rules.""" + """Evaluate independent rule phases concurrently, then finalize their combined findings.""" effective_config: RulesConfig = resolve_rule_ignore_selectors( config=config, project=graph.project ) @@ -69,32 +92,43 @@ def evaluate_rules( graph=graph, selected_keys=selected_keys ) selected_project: CompiledProject = _selected_project(graph=graph, model_paths=model_paths) - sql_started: float = time.monotonic() - sql_result: _SqlRulesEvaluation = _run_sql_rules( - rules=native_rules, - project_dir=resolved_project_dir, - discovered_inputs=discovered_inputs, - project=selected_project, - dialect=dialect, - selected_model_paths=model_paths, - cache_enabled=effective_config.cache.enabled, + selected_config: RulesConfig = replace( + _selection_rules_config(config=effective_config, model_paths=model_paths), + select=tuple(rule.code for rule in selected), + ignore=(), ) - sql_ms: int = round((time.monotonic() - sql_started) * 1000) - result: RulesResult = evaluate_native( - project=selected_project, - config=replace( - _selection_rules_config(config=effective_config, model_paths=model_paths), - select=tuple(rule.code for rule in selected), - ignore=(), - ), + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-rules") as executor: + native_result: Future[RulesResult] = executor.submit( + evaluate_native, + project=selected_project, + config=selected_config, + project_dir=resolved_project_dir, + catalogue=catalogue, + dialect=dialect, + defer_suppressions=True, + ) + sql_started: float = time.monotonic() + sql_result: _SqlRulesEvaluation = _run_sql_rules( + rules=native_rules, + project_dir=resolved_project_dir, + discovered_inputs=discovered_inputs, + project=selected_project, + dialect=dialect, + selected_model_paths=model_paths, + cache_enabled=effective_config.cache.enabled, + prepared_sql=prepared_sql, + ) + sql_ms: int = round((time.monotonic() - sql_started) * 1000) + result: RulesResult = native_result.result() + finalized: tuple[Finding, ...] = finalize_native_findings( + config=selected_config, project_dir=resolved_project_dir, - catalogue=catalogue, - dialect=dialect, - initial_findings=sql_result.findings, + evaluated_codes=tuple(rule.code for rule in selected), + findings=(*sql_result.findings, *result.findings), ) findings: tuple[Finding, ...] = tuple( sorted( - result.findings, + finalized, key=lambda item: (item.path.as_posix(), item.line, item.column, item.code), ) ) @@ -122,6 +156,59 @@ def _config_references_custom_rules(config: RulesConfig) -> bool: return any(selector.startswith("XSQBR") for selector in selectors) +def prepare_sql_rules( + *, inputs: CompileProjectInputs, executor: Executor, dialect: str +) -> PreparedSqlLint | None: + project_dir: Path | None = inputs.discovered_inputs.project_dir + if project_dir is None or any(diagnostic.is_error for diagnostic in inputs.diagnostics): + return None + try: + config: RulesConfig = load_rules_config(project_dir=project_dir) + except RulesError: + return None + if not config.select or (config.cache.enabled and _read_sql_rule_cache(project_dir)): + return None + codes: tuple[str, ...] = _selected_sql_codes(config) + if not codes: + return None + return PreparedSqlLint( + codes=codes, + future=executor.submit( + _prepare_sql_lint, + project_dir=project_dir, + config=LintConfig( + dialect=dialect, enabled_native_rules=codes, header_rules_enabled=False + ), + discovered_inputs=inputs.discovered_inputs, + compiled_expansions={ + model.model_file.file_path: model.sql_expansion + for model in inputs.model_inputs + if model.sql_expansion is not None + }, + ), + ) + + +def _prepare_sql_lint( + *, + project_dir: Path, + config: LintConfig, + discovered_inputs: DiscoveredProjectInputs, + compiled_expansions: dict[Path, CompiledSqlExpansion], +) -> PreparedSqlLintResult: + context: SqlExpansionContext = build_expansion_context( + project_dir=project_dir, discovered_inputs=discovered_inputs + ) + result: LintRunResult = run_lint( + project_dir=project_dir, + config=config, + discovered_inputs=discovered_inputs, + compiled_expansions=compiled_expansions, + expansion_context=context, + ) + return PreparedSqlLintResult(result=result, context=context) + + def _selection_rules_config( *, config: RulesConfig, model_paths: frozenset[str] | None ) -> RulesConfig: @@ -135,6 +222,19 @@ def _selection_rules_config( ) +def _selected_sql_codes(config: RulesConfig) -> tuple[str, ...]: + codes: list[str] = [] + for rule in native_catalogue(): + code: str = str(rule["code"]) + if not code.startswith("SQBRSQL"): + continue + selected: bool = any(code.startswith(selector) for selector in config.select) + ignored: bool = any(code.startswith(selector) for selector in config.ignore) + if selected and not ignored: + codes.append(code) + return tuple(codes) + + def _selected_project( *, graph: ProjectGraph, model_paths: frozenset[str] | None ) -> CompiledProject: @@ -148,6 +248,47 @@ def _selected_project( ) +def _run_prepared_lint( + *, + project_dir: Path, + config: LintConfig, + selected_paths: frozenset[Path], + discovered_inputs: DiscoveredProjectInputs, + compiled_expansions: dict[Path, CompiledSqlExpansion], + dynamic_output_paths: frozenset[Path], + prepared_sql: PreparedSqlLint | None, +) -> LintRunResult: + if prepared_sql is None or set(prepared_sql.codes) != set(config.enabled_native_rules or ()): + return run_lint( + project_dir=project_dir, + config=config, + selected_paths=selected_paths, + discovered_inputs=discovered_inputs, + compiled_expansions=compiled_expansions, + dynamic_output_paths=dynamic_output_paths, + ) + completed: PreparedSqlLintResult = prepared_sql.future.result() + prepared: LintRunResult = completed.result + proof_paths: frozenset[Path] = selected_paths & dynamic_output_paths + violations: tuple[LintViolation, ...] = tuple( + violation + for violation in prepared.violations + if violation.file_path in selected_paths and violation.file_path not in proof_paths + ) + if proof_paths: + proven: LintRunResult = run_lint( + project_dir=project_dir, + config=config, + selected_paths=proof_paths, + discovered_inputs=discovered_inputs, + compiled_expansions=compiled_expansions, + dynamic_output_paths=proof_paths, + expansion_context=completed.context, + ) + violations = (*violations, *proven.violations) + return replace(prepared, violations=violations, files_checked=len(selected_paths)) + + def _run_sql_rules( *, rules: tuple[Rule, ...], @@ -157,6 +298,7 @@ def _run_sql_rules( dialect: str, selected_model_paths: frozenset[str] | None, cache_enabled: bool, + prepared_sql: PreparedSqlLint | None = None, ) -> _SqlRulesEvaluation: codes: tuple[str, ...] = tuple(rule.code for rule in rules if rule.code.startswith("SQBRSQL")) if not codes: @@ -190,17 +332,21 @@ def _run_sql_rules( if path not in model_paths ) result: LintRunResult | None = ( - run_lint( + _run_prepared_lint( project_dir=project_dir, - config=LintConfig(dialect=dialect, enabled_native_rules=codes), + config=LintConfig( + dialect=dialect, enabled_native_rules=codes, header_rules_enabled=False + ), selected_paths=frozenset(selected_paths), discovered_inputs=discovered_inputs, + compiled_expansions=project.sql_expansions, dynamic_output_paths=frozenset( (project_dir / model.relative_path).resolve() for model in models_by_path.values() if model.dynamic_column_contract is not None and model.dynamic_column_contract.output_proven ), + prepared_sql=prepared_sql, ) if selected_paths else None diff --git a/src/sqlbuild/rule_engine/classes/early_sql_lint.py b/src/sqlbuild/rule_engine/classes/early_sql_lint.py new file mode 100644 index 0000000000..064f91e9a7 --- /dev/null +++ b/src/sqlbuild/rule_engine/classes/early_sql_lint.py @@ -0,0 +1,36 @@ +"""Own speculative SQL checks for one compilation invocation.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from types import TracebackType + +from sqlbuild.compiler.compile.models import CompileProjectInputs +from sqlbuild.rule_engine.main._prepare_sql_rules import prepare_sql_rules +from sqlbuild.rule_engine.models import PreparedSqlLint + + +class EarlySqlLint: + def __init__(self, *, enabled: bool) -> None: + self.enabled: bool = enabled + self.preparation: PreparedSqlLint | None = None + self.executor: ThreadPoolExecutor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="sqlbuild-lint" + ) + + def __enter__(self) -> EarlySqlLint: + return self + + def __exit__( + self, + exception_type: type[BaseException] | None, + exception: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.executor.shutdown(wait=True, cancel_futures=True) + + def start(self, *, inputs: CompileProjectInputs, dialect: str) -> None: + if self.enabled: + self.preparation = prepare_sql_rules( + inputs=inputs, executor=self.executor, dialect=dialect + ) diff --git a/src/sqlbuild/rule_engine/main/_prepare_sql_rules.py b/src/sqlbuild/rule_engine/main/_prepare_sql_rules.py new file mode 100644 index 0000000000..ebdebfc73a --- /dev/null +++ b/src/sqlbuild/rule_engine/main/_prepare_sql_rules.py @@ -0,0 +1,17 @@ +"""Start reusable SQL lint work after all compiler inputs have been expanded.""" + +from __future__ import annotations + +from concurrent.futures import Executor + +from sqlbuild.compiler.compile.models import CompileProjectInputs +from sqlbuild.rule_engine._helpers.run.rules import prepare_sql_rules as prepare +from sqlbuild.rule_engine.models import PreparedSqlLint + + +def prepare_sql_rules( + *, inputs: CompileProjectInputs, executor: Executor, dialect: str +) -> PreparedSqlLint | None: + """Prepare configured SQL checks without publishing findings before compilation finishes.""" + + return prepare(inputs=inputs, executor=executor, dialect=dialect) diff --git a/src/sqlbuild/rule_engine/main/run_rules.py b/src/sqlbuild/rule_engine/main/run_rules.py index 26dced6bbe..df98639d44 100644 --- a/src/sqlbuild/rule_engine/main/run_rules.py +++ b/src/sqlbuild/rule_engine/main/run_rules.py @@ -6,7 +6,7 @@ from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.compiler.pipeline.models import ProjectGraph from sqlbuild.rule_engine._helpers.run.rules import evaluate_rules -from sqlbuild.rule_engine.models import RulesConfig, RulesRunResult +from sqlbuild.rule_engine.models import PreparedSqlLint, RulesConfig, RulesRunResult def run_rules( @@ -17,6 +17,7 @@ def run_rules( project_dir: Path, dialect: str, selected_keys: frozenset[CompiledObjectKey] | None = None, + prepared_sql: PreparedSqlLint | None = None, ) -> RulesRunResult: """Run selected native built-ins before selected custom Python rules.""" @@ -27,4 +28,5 @@ def run_rules( project_dir=project_dir, dialect=dialect, selected_keys=selected_keys, + prepared_sql=prepared_sql, ) diff --git a/src/sqlbuild/rule_engine/models.py b/src/sqlbuild/rule_engine/models.py index 936652b0b2..79cfa8c786 100644 --- a/src/sqlbuild/rule_engine/models.py +++ b/src/sqlbuild/rule_engine/models.py @@ -2,10 +2,13 @@ from __future__ import annotations +from concurrent.futures import Future from dataclasses import dataclass, field from pathlib import Path from typing import Any +from sqlbuild.compiler.compile.models import SqlExpansionContext +from sqlbuild.lint.models import LintRunResult from sqlbuild.rule_engine.types import RuleCheck, RuleOptionValue, RuleSubject @@ -293,6 +296,22 @@ class RulesResult: custom_ms: int = 0 +@dataclass(frozen=True) +class PreparedSqlLintResult: + """Reusable expansion context and completed preliminary SQL checks.""" + + result: LintRunResult + context: SqlExpansionContext + + +@dataclass(frozen=True) +class PreparedSqlLint: + """Process-local SQL checks started after compiler expansion has completed.""" + + codes: tuple[str, ...] + future: Future[PreparedSqlLintResult] + + @dataclass(frozen=True) class RulesRunResult: """Unified findings and phase measurements for configured rules.""" diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_compile_json_behavior.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_compile_json_behavior.py index b17f1277b8..1e2d4724ac 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_compile_json_behavior.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_compile_json_behavior.py @@ -70,6 +70,10 @@ def test_given_waffle_shop_when_running_compile_json_then_it_reports_offline_que "built_in_rules_ms", "custom_rules_ms", "model_analysis_ms", + "analysis_preparation_ms", + "analysis_native_ms", + "analysis_projection_ms", + "binding_validation_ms", "rule_cache_hits", "rule_cache_misses", "test_input_compile_ms", 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 6b90899436..62c8f3a354 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 @@ -163,3 +163,12 @@ class DeepSqlAnalysisCompileTestCase: description: str function_depth: int expected_stdout_fragment: str + + +@dataclass(frozen=True) +class DenseCompileGuardTestCase: + description: str + model_count: int + expected_max_wall_seconds: float + expected_max_rss_bytes: int + expected_fingerprint: str diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py index c2900457eb..f726b3d421 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py @@ -566,14 +566,19 @@ def _run_fresh_process_compile_benchmark( output_path.open("wb") as output_file, stderr_path.open("wb") as stderr_file, ): - completed: subprocess.CompletedProcess[bytes] = subprocess.run( + with subprocess.Popen( command, - check=False, stdout=output_file, stderr=stderr_file, - timeout=expected_max_wall_seconds + 10.0, - ) - assert completed.returncode == 0, stderr_path.read_text(encoding="utf-8") + start_new_session=True, + ) as process: + try: + returncode: int = process.wait(timeout=expected_max_wall_seconds + 10.0) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + raise + assert returncode == 0, stderr_path.read_text(encoding="utf-8") elapsed_text, peak_rss_kib_text = measurement_path.read_text(encoding="utf-8").split() elapsed_seconds: float = float(elapsed_text) payload_object: object = json.loads(output_path.read_bytes()) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py new file mode 100644 index 0000000000..17cf08cd1f --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py @@ -0,0 +1,135 @@ +"""Fresh-process guards for dense query graphs and the complete built-in ruleset.""" + +import logging +from itertools import filterfalse +from operator import attrgetter +from pathlib import Path +from typing import cast + +import pytest + +from scripts.cold_compile_performance._helpers.dense_project import ( + write_dense_compile_project, +) +from scripts.cold_compile_performance.main.assert_required_cgroup_memory_limit import ( + assert_required_cgroup_memory_limit, +) +from sqlbuild.rule_engine._helpers.engine.catalogue import build_catalogue, select_rules +from sqlbuild.rule_engine.main.load_config import load_rules_config +from sqlbuild.rule_engine.models import Rule, RulesConfig +from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import ( + DenseCompileGuardTestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.main.compile.helpers import ( + FreshProcessCompileBenchmarkResult, + _run_fresh_process_compile_benchmark, + measure_model_sql_bytes, +) + +_LOGGER: logging.Logger = logging.getLogger(__name__) +_GIB: int = 1024 * 1024 * 1024 + + +@pytest.mark.performance +@pytest.mark.cold_compile_performance +@pytest.mark.parametrize( + "test_case", + ( + DenseCompileGuardTestCase( + "dense_models_1000_all_rules", + 1000, + 14.0, + 3 * _GIB // 2, + "24d355f005d1e4c027ad06bc80678e6dede7c3d6a55e5e9bcae3ba4cfc1de922", + ), + DenseCompileGuardTestCase( + "dense_models_3000_all_rules", + 3000, + 34.0, + 11 * _GIB // 4, + "554dbb7e4aa130abc1e8801482aff123a746d382ed1f5b70d5c0d80fe2559edb", + ), + DenseCompileGuardTestCase( + "dense_models_5000_all_rules", + 5000, + 55.0, + 13 * _GIB // 4, + "79013b434bc8efdc0a6ebbcc6e23af24ba4b459a6f2ca99fcdd75fd9f5042e85", + ), + DenseCompileGuardTestCase( + "dense_models_10000_all_rules", + 10000, + 115.0, + 4 * _GIB, + "8125ddde6d3ae172bb15878f95048d66398453c516bbd837b1fd83c04c5982b5", + ), + ), + ids=lambda case: case.description, +) +def test_given_dense_project_when_compiling_cold_then_preserves_rules_semantics_and_budgets( + test_case: DenseCompileGuardTestCase, + tmp_path: Path, +) -> None: + assert_required_cgroup_memory_limit() + project_dir: Path = tmp_path / "dense_orders" + write_dense_compile_project(project_dir=project_dir, model_count=test_case.model_count) + config: RulesConfig = load_rules_config(project_dir=project_dir) + catalogue: tuple[Rule, ...] = build_catalogue(config=config, project_dir=project_dir) + selected: tuple[Rule, ...] = select_rules( + catalogue=catalogue, config=config, project_dir=project_dir + ) + builtin_codes: set[str] = {rule.code for rule in filterfalse(attrgetter("custom"), catalogue)} + assert len(builtin_codes) >= 75 + assert {rule.code for rule in filterfalse(attrgetter("custom"), selected)} == builtin_codes + assert sum(rule.custom for rule in selected) == 1 + assert not config.rule_exceptions + assert not config.rule_ignores + assert not config.ignore + assert not config.thresholds + assert measure_model_sql_bytes(project_dir) >= 7000 * test_case.model_count + result: FreshProcessCompileBenchmarkResult = _run_fresh_process_compile_benchmark( + project_dir=project_dir, + label=f"dense-{test_case.model_count}", + expected_max_wall_seconds=test_case.expected_max_wall_seconds, + compile_args=("--no-cache",), + ) + _LOGGER.info( + "dense compile models=%d builtin_rules=%d wall=%.3fs peak_rss_bytes=%d fingerprint=%s timings=%s", + test_case.model_count, + len(builtin_codes), + result.elapsed_seconds, + result.peak_rss_bytes, + result.semantic_fingerprint, + result.payload["compile_timings"], + ) + assert result.payload["diagnostics"] == [] + assert result.payload["has_errors"] is False + assert result.payload["summary"] == { + "models": test_case.model_count, + "selected_models": test_case.model_count, + "sources": test_case.model_count * 24 // 100, + "seeds": test_case.model_count // 20, + "selected_seeds": test_case.model_count // 20, + "functions": test_case.model_count // 40, + "selected_functions": test_case.model_count // 40, + "audits": test_case.model_count * 17 // 10, + "tests": test_case.model_count, + "hooks": 0, + "execution_layers": 48, + "errors": 0, + "warnings": 0, + } + timings: dict[str, int] = cast(dict[str, int], result.payload["compile_timings"]) + assert timings["analysis_cache_bypasses"] == test_case.model_count + assert timings["analysis_entry_cache_hits"] == timings["analysis_batch_cache_hits"] == 0 + assert timings["rule_cache_hits"] == 0 + assert timings["rule_cache_misses"] == 2 * test_case.model_count + 1 + assert timings["built_in_rules_ms"] > 0 + assert timings["custom_rules_ms"] > 0 + assert result.semantic_fingerprint == test_case.expected_fingerprint + assert result.elapsed_seconds < test_case.expected_max_wall_seconds + assert result.peak_rss_bytes < test_case.expected_max_rss_bytes + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py index 79db1575f6..e803bb6d5c 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py @@ -47,16 +47,16 @@ expected_max_rss_bytes=2 * _GIB, expected_max_cache_bytes=96 * _MIB, expected_cold_fingerprint=( - "02af81fb7723d30ec5082ed06fb039e45acba3d60b2d8a30fcbccb5a990bb9fc" + "6841a94fc2d3a740cb0db14fe6cee46bfe32171b9580949134cd40379a31bf94" ), expected_leaf_edit_fingerprint=( - "eb9c3fce74888baad3818aa3f732e8786e8306b2a10c0ff82c15bb74a5a1620a" + "a58095a0cd5c2866cd99df92c2503174b400a136379b68dc5889272cef30f17e" ), expected_macro_edit_fingerprint=( - "1de325ddf1ed816c797218def79b23a45e4c020c96e564b71a8ec3cc86112d50" + "fdef695e0e1d08cb3809bb0589b0c535b1bbb604bf81a428627dce09317911f3" ), expected_project_config_fingerprint=( - "a320a971740013ccfbd09870864892635d5f5fca6a34ede117f3100537df8cc3" + "990efe5eee76ba8e0690e5203264aa948e59c39e12463120a3661ea976f9ec8d" ), ), FreshProcessCompileCachePerformanceGuardTestCase( @@ -74,16 +74,16 @@ expected_max_rss_bytes=2 * _GIB, expected_max_cache_bytes=160 * _MIB, expected_cold_fingerprint=( - "0c9258635fe4f9d31ec2dd1c99dbc63f5c8551ca7fea494003c5ab1c3a2ad57c" + "32dac15ee9c15a2aea2e4783e888231e402f15a8fe83270ef55af1aa75924b6c" ), expected_leaf_edit_fingerprint=( - "d9f1e6840338b4ec75a105666a37e255a6c0eedb1edbcb924cc9c07ec06a7666" + "4b1d9da0a0141ded2db0cfb7719fb0ccc12353f129638822e451c11e98f55724" ), expected_macro_edit_fingerprint=( - "68ef78100da39c45b9a28751abff295c2f01a246f4515935cb288a648629ed6a" + "3a8d7e4ebe24a354f4d9f470e4b37b0286cf12a2f36f850053651c17d9f8e3b3" ), expected_project_config_fingerprint=( - "539f538bdd4387d1aa01cc9e36f378e2a0e61854229d8a482a8cfae4d7c123b6" + "c1efca198bc601d5135898ad059cfb6e016872e15f216cc357662bb185ba29f3" ), ), FreshProcessCompileCachePerformanceGuardTestCase( @@ -101,16 +101,16 @@ expected_max_rss_bytes=2 * _GIB, expected_max_cache_bytes=320 * _MIB, expected_cold_fingerprint=( - "6887ddd2547137efd8fe67d13fb07142f6287c0911d152d69eca0d495310e02b" + "7e1410c35f6026c7520711224aea2a95d8fe46515a2d2d370a3ee21180a36fdc" ), expected_leaf_edit_fingerprint=( - "6447fccd4a45b5cfe959780cd9c3577d19e97bf7ad53b31cd30c31e1bdba0f34" + "059dc35c437adacfe0098e0cc19ccfeaae9fc5ca9aa682721cd3f37cbff92f77" ), expected_macro_edit_fingerprint=( - "999eddf213ff4f95b4b10b11890fc852c610410b730c38c81721d466a6534a68" + "e2cfd1b086c42ce4d181380a3be1d7a3d1360ee7ccea6f5dfbb5a8d2af5a76ce" ), expected_project_config_fingerprint=( - "17ae2532010bef51a653859929c3dc26ad88fdd77725cc791a2102184c86cbbc" + "38323964889d18747f10614c8442d9a2cdd9d38efed88369cb10c3bf1099ed85" ), ), ], diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_performance.py index 9bdaf8a679..065d48883d 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_performance.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_performance.py @@ -39,7 +39,7 @@ expected_max_wall_seconds=11.0, expected_max_rss_bytes=int(1.5 * _GIB), expected_semantic_fingerprint=( - "02af81fb7723d30ec5082ed06fb039e45acba3d60b2d8a30fcbccb5a990bb9fc" + "6841a94fc2d3a740cb0db14fe6cee46bfe32171b9580949134cd40379a31bf94" ), ), FreshProcessCompilePerformanceGuardTestCase( @@ -54,7 +54,7 @@ expected_max_wall_seconds=15.5, expected_max_rss_bytes=int(1.75 * _GIB), expected_semantic_fingerprint=( - "0c9258635fe4f9d31ec2dd1c99dbc63f5c8551ca7fea494003c5ab1c3a2ad57c" + "32dac15ee9c15a2aea2e4783e888231e402f15a8fe83270ef55af1aa75924b6c" ), ), FreshProcessCompilePerformanceGuardTestCase( @@ -69,7 +69,7 @@ expected_max_wall_seconds=29.5, expected_max_rss_bytes=2 * _GIB, expected_semantic_fingerprint=( - "6887ddd2547137efd8fe67d13fb07142f6287c0911d152d69eca0d495310e02b" + "7e1410c35f6026c7520711224aea2a95d8fe46515a2d2d370a3ee21180a36fdc" ), ), ], @@ -92,11 +92,13 @@ def test_given_scaled_semantic_project_when_compiling_fresh_then_preserves_resou expected_max_wall_seconds=test_case.expected_max_wall_seconds, ) _LOGGER.info( - "fresh-process semantic compile models=%d wall=%.3fs peak_rss_bytes=%d fingerprint=%s", + "fresh-process semantic compile models=%d wall=%.3fs peak_rss_bytes=%d " + "fingerprint=%s timings=%s", test_case.model_count, result.elapsed_seconds, result.peak_rss_bytes, result.semantic_fingerprint, + result.payload["compile_timings"], ) assert result.payload["has_errors"] is False diff --git a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py index c2ee589180..9f27e2bdfa 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -5,6 +5,46 @@ from sqlbuild.compiler.planner.models import CursorOverrides +@dataclass(frozen=True) +class BoundProjectionCompileTestCase: + description: str + query_sql: str + expected_rows: tuple[tuple[int | None, ...], ...] + expected_edges: int + output_contract: str = "order_id (type BIGINT, nullable false)" + + +@dataclass(frozen=True) +class AliasSourceCompileTestCase: + """Lexical source shape exercised by an input-column precedence regression.""" + + description: str + query_prefix: str + source_relation: str + expected_edge_count: int + + +@dataclass(frozen=True) +class DerivedNativeCompileTestCase: + """Expected compiled and executed output from the native query graph.""" + + description: str + query_sql: str + expected_exit_code: int + expected_diagnostics: tuple[str, ...] + expected_rows: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True) +class CombinedCompilationTestCase: + """Expected CLI behavior for schema-bound CTE compilation.""" + + description: str + projection: str + expected_exit_code: int + expected_diagnostics: tuple[str, ...] + + @dataclass(frozen=True) class RulesIntegrationTestCase: """One compiler-integrated Rules command expectation.""" @@ -423,3 +463,11 @@ class ExpectedCountTestCase: class ExpectedMessageTestCase: description: str expected_message: str + + +@dataclass(frozen=True) +class DenseCompileFixtureTestCase: + description: str + model_count: int + expected_rule_misses: int + expected_tail_lineage_edges: int diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_compile.py b/tests/integration/src/sqlbuild/cli/commands/main/test_compile.py index f6fe1ae38c..63a25020f4 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_compile.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_compile.py @@ -11,6 +11,7 @@ from sqlbuild.cli.commands.main.entrypoint.entry import main from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + CombinedCompilationTestCase, ContractNullabilityCompileIntegrationTestCase, ExpectedCountTestCase, ExpectedMessageTestCase, @@ -19,6 +20,53 @@ ) +@pytest.mark.parametrize( + "test_case", + ( + CombinedCompilationTestCase( + description="typed CTE comparison compiles", + projection="order_id > 0", + expected_exit_code=0, + expected_diagnostics=(), + ), + CombinedCompilationTestCase( + description="missing CTE input retains binding diagnostic", + projection="missing > 0", + expected_exit_code=1, + expected_diagnostics=("B002",), + ), + ), + ids=lambda case: case.description, +) +def test_given_typed_dependency_when_compiling_cte_then_preserves_binding_and_output_contract( + test_case: CombinedCompilationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text('name = "orders"\nadapter = "duckdb"\n') + models: Path = tmp_path / "models" + models.mkdir() + sources: Path = tmp_path / "sources" + sources.mkdir() + (sources / "orders.yml").write_text( + "sources:\n - name: orders\n contract: enforced\n" + ' expression: "(SELECT CAST(1 AS BIGINT) AS order_id)"\n' + " columns:\n - name: order_id\n type: BIGINT\n" + ) + (models / "selected_orders.sql").write_text( + "MODEL (columns (result (type BOOLEAN)));\n" + 'WITH selected AS (SELECT order_id FROM __source("orders"))\n' + f"SELECT {test_case.projection} AS result FROM selected\n" + ) + exit_code: int = main( + ["--project-dir", str(tmp_path), "--no-color", "compile", "--json", "--no-cache"] + ) + result: dict[str, object] = json.loads(capsys.readouterr().out) + diagnostics: list[dict[str, object]] = cast(list[dict[str, object]], result["diagnostics"]) + assert exit_code == test_case.expected_exit_code + assert tuple(item["code"] for item in diagnostics) == test_case.expected_diagnostics + + @pytest.mark.parametrize( "test_case", [ExpectedCountTestCase(description="typed headers compile", expected_count=2)], @@ -65,6 +113,45 @@ def test_given_multiple_typed_model_headers_when_compiling_then_cli_preserves_he assert (tmp_path / "target" / "compiled" / "models" / "orders.sql").is_file() +@pytest.mark.parametrize( + "test_case", + [ + ExpectedCountTestCase( + description="compiler rules reuse expanded macro SQL", expected_count=1 + ) + ], + ids=lambda case: case.description, +) +def test_given_model_macro_when_compiling_with_sql_rules_then_expands_once_and_preserves_location( + test_case: ExpectedCountTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n[rules]\nselect = ["SQBRSQL001"]\n' + ) + macro_dir: Path = tmp_path / "models" / "orders" / "_sqlbuild" / "_macros" + macro_dir.mkdir(parents=True) + (macro_dir / "missing_value.py").write_text( + "from pathlib import Path\n" + "def missing_value() -> str:\n" + " counter = Path(__file__).with_suffix('.calls')\n" + " count = int(counter.read_text()) if counter.exists() else 0\n" + " counter.write_text(str(count + 1))\n" + " return 'NULL'\n" + ) + (tmp_path / "models" / "orders" / "orders.sql").write_text( + "MODEL (materialized table);\n\nSELECT 1 AS order_id WHERE 1 = @missing_value()\n" + ) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + result: dict[str, object] = json.loads(capsys.readouterr().out) + diagnostics: list[dict[str, object]] = cast(list[dict[str, object]], result["diagnostics"]) + assert exit_code == 1 + assert int((macro_dir / "missing_value.calls").read_text()) == test_case.expected_count + assert tuple(item["code"] for item in diagnostics) == ("SQBRSQL001",) + assert diagnostics[0]["line"] == 3 + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_compile_native_facts.py b/tests/integration/src/sqlbuild/cli/commands/main/test_compile_native_facts.py new file mode 100644 index 0000000000..29ee9e7177 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_compile_native_facts.py @@ -0,0 +1,289 @@ +"""Real CLI and DuckDB coverage for the borrowed query graph.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import cast + +import duckdb +import pytest + +from sqlbuild.cli.commands.main.entrypoint.entry import main +from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + AliasSourceCompileTestCase, + BoundProjectionCompileTestCase, + DerivedNativeCompileTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + ( + DerivedNativeCompileTestCase( + description="derived-table column aliases preserve the declared contract", + query_sql=( + "WITH typed AS (SELECT CAST(1 AS BIGINT) AS original_id) " + "SELECT order_id FROM (SELECT original_id FROM typed) AS nested_orders(order_id)" + ), + expected_exit_code=0, + expected_diagnostics=(), + expected_rows=((1,),), + ), + DerivedNativeCompileTestCase( + description="CTE column aliases survive a derived-table star", + query_sql=( + "WITH typed(order_id) AS (SELECT CAST(1 AS BIGINT) AS original_id) " + "SELECT * FROM (SELECT order_id FROM typed) AS nested_orders" + ), + expected_exit_code=0, + expected_diagnostics=(), + expected_rows=((1,),), + ), + DerivedNativeCompileTestCase( + description="both named-union branches survive a derived table", + query_sql=( + "WITH combined AS (SELECT CAST(1 AS BIGINT) AS order_id " + "UNION ALL BY NAME SELECT CAST(2 AS BIGINT) AS order_id) " + "SELECT * FROM (SELECT order_id FROM combined) AS nested_orders ORDER BY order_id" + ), + expected_exit_code=0, + expected_diagnostics=(), + expected_rows=((1,), (2,)), + ), + ), + ids=lambda case: case.description, +) +def test_given_derived_native_facts_when_compiling_then_contract_and_execution_agree( + test_case: DerivedNativeCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text('name = "orders"\nadapter = "duckdb"\n') + models: Path = tmp_path / "models" + models.mkdir() + (models / "orders.sql").write_text( + "MODEL (contract enforced, columns (order_id (type BIGINT, nullable false)));\n" + + test_case.query_sql + ) + + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + result: dict[str, object] = json.loads(capsys.readouterr().out) + diagnostics: list[dict[str, object]] = cast(list[dict[str, object]], result["diagnostics"]) + + assert exit_code == test_case.expected_exit_code + assert tuple(item["code"] for item in diagnostics) == test_case.expected_diagnostics + compiled_sql: str = (tmp_path / "target" / "compiled" / "models" / "orders.sql").read_text() + with duckdb.connect() as connection: + rows: tuple[tuple[int, ...], ...] = tuple(connection.execute(compiled_sql).fetchall()) + assert rows == test_case.expected_rows + + +@pytest.mark.parametrize( + "test_case", + ( + AliasSourceCompileTestCase( + description="direct source", + query_prefix="", + source_relation='__source("orders")', + expected_edge_count=5, + ), + AliasSourceCompileTestCase( + description="CTE source", + query_prefix='WITH source_rows AS (SELECT * FROM __source("orders")) ', + source_relation="source_rows", + expected_edge_count=5, + ), + ), + ids=lambda case: case.description, +) +def test_given_input_and_output_share_name_when_compiling_then_lineage_keeps_input_dependency( + test_case: AliasSourceCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text('name = "orders"\nadapter = "duckdb"\n') + models: Path = tmp_path / "models" + models.mkdir() + sources: Path = tmp_path / "sources" + sources.mkdir() + (sources / "orders.yml").write_text( + "sources:\n - name: orders\n contract: enforced\n" + " expression: orders_input\n columns:\n" + " - name: order_id\n type: INTEGER\n" + " - name: amount\n type: DOUBLE\n" + " - name: status\n type: VARCHAR\n" + ) + (models / "order_totals.sql").write_text( + "MODEL (contract enforced, columns (order_id (type INTEGER), " + "status (type VARCHAR), total (type DOUBLE)));\n" + f"{test_case.query_prefix}" + "SELECT order_id, " + "CAST(CASE WHEN order_id % 2 = 0 THEN 'even' ELSE 'odd' END AS VARCHAR) AS status, " + "CAST(CASE WHEN order_id % 2 = 0 THEN amount + 2 " + "WHEN status = 'priority' THEN amount * 2 ELSE amount - 1 END AS DOUBLE) AS total " + f"FROM {test_case.source_relation} AS input" + ) + + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + result: dict[str, object] = json.loads(capsys.readouterr().out) + assert exit_code == 0 + assert result["diagnostics"] == [] + resources: dict[str, object] = cast(dict[str, object], result["resources"]) + compiled_models: list[dict[str, object]] = cast(list[dict[str, object]], resources["models"]) + lineage: dict[str, object] = cast(dict[str, object], compiled_models[0]["lineage"]) + # order_id and the status alias each depend on order_id; total depends on + # order_id, amount, and the physical status column, for five edges in total. + assert lineage["edge_count"] == test_case.expected_edge_count + compiled_sql: str = ( + tmp_path / "target" / "compiled" / "models" / "order_totals.sql" + ).read_text() + with duckdb.connect() as connection: + connection.execute( + "CREATE TABLE orders_input AS SELECT 1 AS order_id, " + "CAST(10 AS DOUBLE) AS amount, 'priority' AS status" + ) + # Compiled artifacts retain logical source intrinsics. Bind this test's + # single source to its physical table before executing the artifact. + connection.execute("CREATE MACRO __source(name) AS TABLE SELECT * FROM orders_input") + assert connection.execute(compiled_sql).fetchall() == [(1, "odd", 20.0)] + connection.execute("UPDATE orders_input SET status = 'standard'") + assert connection.execute(compiled_sql).fetchall() == [(1, "odd", 9.0)] + + +@pytest.mark.parametrize( + "test_case", + ( + BoundProjectionCompileTestCase( + description="qualified star exclusions preserve the output shape", + query_sql='WITH selected AS (SELECT o.* EXCLUDE (category, amount) FROM __source("orders") o) SELECT * FROM selected ORDER BY order_id', + expected_rows=((1,), (2,)), + expected_edges=1, + ), + BoundProjectionCompileTestCase( + description="positional CTE aliases retain unnamed nullable expressions", + query_sql='WITH selected(a, b) AS (SELECT CAST(NULL AS VARCHAR), order_id FROM __source("orders")) SELECT a AS order_id FROM selected', + expected_rows=((None,), (None,)), + expected_edges=0, + output_contract="order_id (type VARCHAR, nullable true)", + ), + BoundProjectionCompileTestCase( + description="unnamed union branches retain every input dependency", + query_sql='WITH selected AS (SELECT MAX(order_id) AS order_id FROM __source("orders") UNION ALL SELECT MAX(CAST(amount AS BIGINT)) FROM __source("orders")) SELECT order_id FROM selected ORDER BY order_id', + expected_rows=((2,), (20,)), + expected_edges=2, + output_contract="order_id (type BIGINT, nullable true)", + ), + BoundProjectionCompileTestCase( + description="named unions align inputs with different projection orders", + query_sql='WITH selected AS (SELECT CAST(NULL AS BIGINT) AS order_id, amount FROM __source("orders") UNION ALL BY NAME SELECT amount, order_id FROM __source("orders")) SELECT order_id FROM selected ORDER BY order_id', + expected_rows=((1,), (2,), (None,), (None,)), + expected_edges=1, + output_contract="order_id (type BIGINT, nullable true)", + ), + BoundProjectionCompileTestCase( + description="value rows preserve positional aliases and nullable constants", + query_sql='WITH selected AS (SELECT v.id AS order_id FROM (VALUES (CAST(1 AS BIGINT)), (CAST(NULL AS BIGINT))) AS v(id) CROSS JOIN __source("orders") AS o WHERE o.order_id = 1) SELECT order_id FROM selected ORDER BY order_id', + expected_rows=((1,), (None,)), + expected_edges=0, + output_contract="order_id (type BIGINT, nullable true)", + ), + BoundProjectionCompileTestCase( + description="star replacements preserve constant provenance", + query_sql='WITH selected AS (SELECT * EXCLUDE (category, amount) REPLACE (CAST(0 AS BIGINT) AS order_id) FROM __source("orders")) SELECT * FROM selected', + expected_rows=((0,), (0,)), + expected_edges=0, + ), + BoundProjectionCompileTestCase( + description="full using join coalesces its non-null input keys", + query_sql='WITH joined AS (SELECT order_id FROM (SELECT order_id FROM __source("orders") WHERE order_id = 1) a FULL JOIN (SELECT order_id FROM __source("orders") WHERE order_id = 2) b USING (order_id)) SELECT order_id FROM joined ORDER BY order_id', + expected_rows=((1,), (2,)), + expected_edges=1, + ), + ), + ids=lambda case: case.description, +) +def test_given_bound_wildcard_or_using_join_when_compiling_then_facts_match_execution( + test_case: BoundProjectionCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text('name = "orders"\nadapter = "duckdb"\n') + sources: Path = tmp_path / "sources" + sources.mkdir() + (sources / "orders.yml").write_text( + "sources:\n - name: orders\n contract: enforced\n" + " expression: orders_input\n columns:\n" + " - name: order_id\n type: BIGINT\n nullable: false\n" + " - name: category\n type: VARCHAR\n" + " - name: amount\n type: DOUBLE\n" + ) + models: Path = tmp_path / "models" + models.mkdir() + (models / "selected_orders.sql").write_text( + f"MODEL (contract enforced, columns ({test_case.output_contract}));\n" + test_case.query_sql + ) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + result: dict[str, object] = json.loads(capsys.readouterr().out) + assert exit_code == 0 + assert result["diagnostics"] == [] + resources: dict[str, object] = cast(dict[str, object], result["resources"]) + compiled_models: list[dict[str, object]] = cast(list[dict[str, object]], resources["models"]) + lineage: dict[str, object] = cast(dict[str, object], compiled_models[0]["lineage"]) + assert lineage["edge_count"] == test_case.expected_edges + sql: str = (tmp_path / "target" / "compiled" / "models" / "selected_orders.sql").read_text() + with duckdb.connect() as connection: + connection.execute( + "CREATE TABLE orders_input (order_id BIGINT, category VARCHAR, amount DOUBLE)" + ) + connection.execute("INSERT INTO orders_input VALUES (1, 'ready', 10), (2, 'queued', 20)") + connection.execute("CREATE MACRO __source(name) AS TABLE SELECT * FROM orders_input") + assert tuple(connection.execute(sql).fetchall()) == test_case.expected_rows + + +@pytest.mark.parametrize( + "test_case", + ( + BoundProjectionCompileTestCase( + description="object wildcard retains all possible payload inputs", + query_sql='WITH packed AS (SELECT OBJECT_CONSTRUCT_KEEP_NULL(*) AS payload FROM __source("orders")) SELECT CAST(payload:order_id AS BIGINT) AS order_id FROM packed', + expected_rows=(), + expected_edges=2, + ), + ), + ids=lambda case: case.description, +) +def test_given_object_wildcard_when_compiling_then_lineage_preserves_payload_inputs( + test_case: BoundProjectionCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "snowflake"\n[defaults]\ndatabase = "warehouse"\nschema = "analytics"\n' + ) + sources: Path = tmp_path / "sources" + sources.mkdir() + (sources / "orders.yml").write_text( + "sources:\n - name: orders\n contract: enforced\n" + " expression: orders_input\n columns:\n" + " - name: order_id\n type: BIGINT\n" + " - name: amount\n type: DOUBLE\n" + ) + models: Path = tmp_path / "models" + models.mkdir() + (models / "selected_orders.sql").write_text( + "MODEL (contract enforced, columns (order_id (type BIGINT, nullable true)));\n" + + test_case.query_sql + ) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + result: dict[str, object] = json.loads(capsys.readouterr().out) + assert exit_code == 0 + assert result["diagnostics"] == [] + resources: dict[str, object] = cast(dict[str, object], result["resources"]) + compiled_models: list[dict[str, object]] = cast(list[dict[str, object]], resources["models"]) + lineage: dict[str, object] = cast(dict[str, object], compiled_models[0]["lineage"]) + assert lineage["edge_count"] == test_case.expected_edges + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_dense_compile_fixture.py b/tests/integration/src/sqlbuild/cli/commands/main/test_dense_compile_fixture.py new file mode 100644 index 0000000000..e07879e43e --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_dense_compile_fixture.py @@ -0,0 +1,64 @@ +"""Exercise the dense guard's generated SQL tests with real DuckDB execution.""" + +import json +from pathlib import Path + +import pytest + +from scripts.cold_compile_performance._helpers.dense_project import ( + write_dense_compile_project, +) +from sqlbuild.cli.commands.main.entrypoint.entry import main +from tests.integration.src.sqlbuild.cli.commands.main._test_types import DenseCompileFixtureTestCase + + +@pytest.mark.parametrize( + "test_case", + (DenseCompileFixtureTestCase("joins unions macros and functions", 64, 129, 134),), + ids=lambda case: case.description, +) +def test_given_dense_fixture_when_compiling_and_testing_then_all_rules_and_rows_pass( + test_case: DenseCompileFixtureTestCase, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + project: Path = tmp_path / "orders" + write_dense_compile_project(project_dir=project, model_count=test_case.model_count) + config: Path = project / "sqlbuild_project.toml" + config.write_text( + config.read_text() + .replace('adapter = "snowflake"', 'adapter = "duckdb"') + .replace('database = "warehouse"\n', "") + .replace('schema = "analytics"', 'schema = "main"') + + f'\n[connection]\ndatabase = "{project / "orders.duckdb"}"\n', + encoding="utf-8", + ) + compiled: int = main(["--project-dir", str(project), "compile", "--no-cache", "--json"]) + output: str = capsys.readouterr().out + assert compiled == 0, output + payload: dict = json.loads(output) + assert payload["diagnostics"] == [] + assert payload["compile_timings"]["rule_cache_hits"] == 0 + assert payload["compile_timings"]["rule_cache_misses"] == test_case.expected_rule_misses + assert ( + payload["resources"]["models"][54]["lineage"]["edge_count"] + == test_case.expected_tail_lineage_edges + ) + built: int = main( + [ + "--project-dir", + str(project), + "build", + "--select", + "fn_00000", + "--no-tests", + "--no-audits", + ] + ) + captured: str = "".join(capsys.readouterr()) + assert built == 0, captured + tested: int = main(["--project-dir", str(project), "test"]) + captured = "".join(capsys.readouterr()) + assert tested == 0, captured + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules_concurrent.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules_concurrent.py new file mode 100644 index 0000000000..0344d39ee5 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules_concurrent.py @@ -0,0 +1,65 @@ +"""Combined rule-policy coverage through the real compiler and custom-rule host.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import cast + +import pytest + +from sqlbuild.cli.commands.main.entrypoint.entry import main +from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + RulePassIntegrationTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + (RulePassIntegrationTestCase("SQL and custom exceptions finalize together", 0, ()),), + ids=lambda case: case.description, +) +def test_given_sql_and_custom_exceptions_when_compiling_cold_and_warm_then_both_apply( + test_case: RulePassIntegrationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n' + '[rules]\nselect = ["SQBRSQL004", "XSQBRARCH001"]\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n" + '[[rules.rule_exceptions]]\nrule = "SQBRSQL004"\n' + 'path = "models/orders.sql"\nreason = "A single synthetic row is intentional."\n' + '[[rules.rule_exceptions]]\nrule = "XSQBRARCH001"\n' + 'path = "models/orders.sql"\nreason = "The synthetic model name is intentional."\n' + ) + models: Path = tmp_path / "models" + models.mkdir() + (models / "orders.sql").write_text("MODEL ();\nSELECT 1 AS order_id LIMIT 1\n") + rules: Path = tmp_path / "rules" + rules.mkdir() + (rules / "architecture.py").write_text( + "from sqlbuild.rules import Finding, Model, RuleContext, rule\n\n" + '@rule(code="XSQBRARCH001", message="Review model name", remediation="Rename the model.")\n' + "def review_name(*, model: Model, ctx: RuleContext) -> list[Finding]:\n" + " return [ctx.finding(subject=model)]\n" + ) + + cold_exit: int = main(["--project-dir", str(tmp_path), "compile", "--json"]) + cold: dict[str, object] = json.loads(capsys.readouterr().out) + warm_exit: int = main(["--project-dir", str(tmp_path), "compile", "--json"]) + warm: dict[str, object] = json.loads(capsys.readouterr().out) + + assert cold_exit == warm_exit == test_case.expected_exit_code + cold_diagnostics: list[dict[str, object]] = cast(list[dict[str, object]], cold["diagnostics"]) + warm_diagnostics: list[dict[str, object]] = cast(list[dict[str, object]], warm["diagnostics"]) + assert tuple(item["code"] for item in cold_diagnostics) == test_case.expected_diagnostics + assert warm_diagnostics == cold_diagnostics + cold_timings: dict[str, int] = cast(dict[str, int], cold["compile_timings"]) + warm_timings: dict[str, int] = cast(dict[str, int], warm["compile_timings"]) + assert cold_timings["rule_cache_misses"] > 0 + assert warm_timings["rule_cache_hits"] > 0 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) 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 25eba08297..4e05eff095 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py @@ -12,6 +12,7 @@ ) from sqlbuild.compiler.compile.types import AttachedAuditTargetKind, SqlTestMode from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.compiler.references.types import SqlReferenceKind from sqlbuild.compiler.scopes.models import UsageRecord from sqlbuild.compiler.scopes.types import ScopeKind from sqlbuild.spec.contracts.models import ResolvedTimeTravelRetention @@ -24,6 +25,28 @@ class UnexpectedAnalysisFailureTestCase: expected_error: str +@dataclass(frozen=True) +class SqlReferenceExtractionTestCase: + description: str + sql: str + expected_references: tuple[tuple[SqlReferenceKind, str, str | None], ...] + + +@dataclass(frozen=True) +class SqlReferenceExtractionErrorTestCase: + description: str + sql: str + expected_error: str + + +@dataclass(frozen=True) +class StaticProjectVarBatchTestCase: + description: str + sqls: tuple[str, ...] + effective_vars: dict[str, object] + expected_sqls: tuple[str | None, ...] + + @dataclass(frozen=True) class AuditFactoryAttachmentTestCase: description: str diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_analysis_cache.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_analysis_cache.py index b37ea2b430..d1fe68db7f 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_analysis_cache.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_analysis_cache.py @@ -9,6 +9,7 @@ import pytest from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.compile._helpers.analysis import cache as analysis_cache from sqlbuild.compiler.compile._helpers.analysis.cache import ( build_analysis_cache_context, model_analysis_cache_key, @@ -42,6 +43,7 @@ "sources/raw.yml": """ sources: - name: raw_orders + contract: enforced expression: "(SELECT 1 AS order_id)" columns: - name: order_id @@ -65,6 +67,40 @@ } +@pytest.mark.parametrize( + "test_case", + ( + AnalysisCacheTestCase( + description="changed compiler semantics invalidate prior analysis", + expected_count=1, + ), + ), + ids=lambda case: case.description, +) +def test_given_changed_analysis_algorithm_when_compiling_then_refreshes_cached_facts( + test_case: AnalysisCacheTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + write_repo_files: Callable[[Path, dict[str, str]], None], +) -> None: + write_repo_files(tmp_path, _CACHE_REPO_FILES) + cold_project: CompiledProject = compile_project_with_cache(project_dir=tmp_path) + monkeypatch.setattr( + analysis_cache, "_ANALYSIS_ALGORITHM_FINGERPRINT", "synthetic-next-algorithm" + ) + with collect_compile_timings() as changed_metrics: + changed_project: CompiledProject = compile_project_with_cache(project_dir=tmp_path) + with collect_compile_timings() as normal_metrics: + normal_project: CompiledProject = compile_project_with_cache(project_dir=tmp_path) + + assert changed_metrics.as_milliseconds()["analysis_cache_misses"] == test_case.expected_count + assert changed_metrics.as_milliseconds()["analysis_entry_cache_hits"] == 0 + assert changed_metrics.as_milliseconds()["analysis_batch_cache_hits"] == 0 + assert normal_metrics.as_milliseconds()["analysis_entry_cache_hits"] == test_case.expected_count + assert normal_project.models == changed_project.models == cold_project.models + assert normal_project.diagnostics == cold_project.diagnostics + + @pytest.mark.parametrize( "test_case", (AnalysisCacheTestCase(description="successful cache hit", expected_count=1),), @@ -103,6 +139,35 @@ def test_given_successful_analysis_when_compiling_again_then_reuses_identical_ca ) +@pytest.mark.parametrize( + "test_case", + (AnalysisCacheTestCase(description="completed compact cache reuse", expected_count=1),), + ids=lambda case: case.description, +) +def test_given_completed_compact_cache_when_compiling_again_then_skips_analysis_and_binding( + test_case: AnalysisCacheTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + write_repo_files: Callable[[Path, dict[str, str]], None], +) -> None: + write_repo_files(tmp_path, _CACHE_REPO_FILES) + monkeypatch.setattr( + assembly_project, "_COMPACT_BATCH_CACHE_MIN_MODEL_COUNT", test_case.expected_count + ) + cold_project: CompiledProject = compile_project_with_cache(project_dir=tmp_path) + analyzer: Mock = Mock(side_effect=AssertionError("completed facts must not be reanalyzed")) + validator: Mock = Mock(side_effect=AssertionError("completed bindings must not be revalidated")) + monkeypatch.setattr(assembly_project, "_analyze_model_sql_requests", analyzer) + monkeypatch.setattr(assembly_project, "get_schema_validations", validator) + + warm_project: CompiledProject = compile_project_with_cache(project_dir=tmp_path) + + assert warm_project.models == cold_project.models + assert warm_project.diagnostics == cold_project.diagnostics + analyzer.assert_not_called() + validator.assert_not_called() + + @pytest.mark.parametrize( "test_case", (AnalysisCacheTestCase(description="exact compact batch reuse", expected_count=1),), 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 4fafde5e17..22e7f48a90 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 @@ -175,6 +175,61 @@ def analyze_project_queries_compact_json(request_json: str) -> str: assert failed_analysis.analysis_succeeded is False +@pytest.mark.parametrize( + "test_case", + ( + InferColumnsTestCase( + description="cast expression overrides input type through CTE", + query_sql='WITH typed AS (SELECT COALESCE(CAST(amount AS FLOAT), CAST(0 AS FLOAT)) AS amount FROM __ref("orders")) SELECT amount FROM typed', + expected_columns=( + InferredColumn( + name="amount", type="FLOAT", nullability=InferredNullability.NON_NULL + ), + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_cte_cast_when_batch_analyzing_then_cast_type_overrides_input_type( + test_case: InferColumnsTestCase, +) -> None: + query_sql: str = test_case.query_sql + references: tuple[CompileSqlReference, ...] = ( + CompileSqlReference(ref_kind=SqlReferenceKind.REF, ref_name="orders"), + ) + nullability: dict[str, dict[str, InferredNullability]] = { + "orders": {"amount": InferredNullability.NULLABLE} + } + types: dict[str, dict[str, str]] = {"orders": {"amount": "VARCHAR"}} + profile: ExpressionInferenceProfile = ExpressionInferenceProfile( + sql_analysis_dialect="snowflake" + ) + + prepared: NativeCompactAnalysis = analyze_queries_with_compact_polyglot_batch( + query_sqls=(query_sql,), + references=(references,), + placeholders=(None,), + column_nullability_by_table=nullability, + column_types_by_table=types, + inference_profile=profile, + recover_cte_facts=(True,), + rich_type_inference=False, + )[0] + result: PolyglotAnalysisResult = analyze_columns_and_lineage_with_polyglot( + query_sql=query_sql, + references=references, + column_nullability_by_table=nullability, + column_types_by_table=types, + inference_profile=profile, + allow_compact_analysis=True, + recover_cte_facts=True, + precomputed=prepared, + ) + + assert result.analysis_succeeded + assert result.columns == test_case.expected_columns + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_references.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_references.py new file mode 100644 index 0000000000..9084472122 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_references.py @@ -0,0 +1,78 @@ +"""Tests for native-assisted logical SQL reference extraction.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile._helpers.refs.references import extract_sql_references +from sqlbuild.compiler.compile.exceptions import CompileInputError +from sqlbuild.compiler.compile.models import CompileSqlReference +from sqlbuild.compiler.references.types import SqlReferenceKind +from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + SqlReferenceExtractionErrorTestCase, + SqlReferenceExtractionTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + [ + SqlReferenceExtractionTestCase( + description="simple references preserve authored order", + sql=( + "SELECT * FROM __source('orders') " + 'UNION ALL SELECT * FROM __dbt_ref("shop", "customers")' + ), + expected_references=( + (SqlReferenceKind.SOURCE, "orders", None), + (SqlReferenceKind.DBT_REF, "customers", "shop"), + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_simple_references_when_extracting_then_returns_authored_order( + test_case: SqlReferenceExtractionTestCase, +) -> None: + references: tuple[CompileSqlReference, ...] = extract_sql_references(test_case.sql) + + assert ( + tuple( + (reference.ref_kind, reference.ref_name, reference.ref_package) + for reference in references + ) + == test_case.expected_references + ) + + +@pytest.mark.parametrize( + "test_case", + [ + SqlReferenceExtractionErrorTestCase( + description="unclosed reference parenthesis preserves diagnostic", + sql='SELECT * FROM __ref("orders"', + expected_error="unclosed parenthesis", + ), + SqlReferenceExtractionErrorTestCase( + description="expression reference name preserves diagnostic", + sql="SELECT * FROM __ref(concat('ord', 'ers'))", + expected_error="name argument must be a quoted string or identifier", + ), + SqlReferenceExtractionErrorTestCase( + description="unclosed block comment preserves diagnostic", + sql='SELECT * FROM __ref("orders") /* unterminated', + expected_error="unclosed block comment", + ), + SqlReferenceExtractionErrorTestCase( + description="unclosed quoted string preserves diagnostic", + sql='SELECT * FROM __ref("orders") WHERE note = \'unterminated', + expected_error="unclosed quoted string", + ), + ], + ids=lambda case: case.description, +) +def test_given_unsupported_reference_sql_when_extracting_then_preserves_python_diagnostic( + test_case: SqlReferenceExtractionErrorTestCase, +) -> None: + with pytest.raises(CompileInputError, match=test_case.expected_error): + extract_sql_references(test_case.sql) diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_vars.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_vars.py index 5ad41f0952..06783b1633 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_vars.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_sql_vars.py @@ -7,10 +7,12 @@ import pytest from sqlbuild.compiler.compile._helpers.render.sql_vars import ( + prepare_static_project_vars_batch, substitute_sql_vars, ) from sqlbuild.compiler.compile.exceptions import CompileInputError from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + StaticProjectVarBatchTestCase, SubstituteSqlVarsErrorTestCase, SubstituteSqlVarsTestCase, ) @@ -192,6 +194,18 @@ def test_given_sql_and_vars_when_substituting_then_returns_expected( r'\["analyst","reporter"\]. Use a macro to consume structured vars\.' ), ), + SubstituteSqlVarsErrorTestCase( + description="deferred placeholder preserves unclosed quote diagnostics", + sql="SELECT @@@partition_start, 'unterminated", + effective_vars={}, + expected_error_fragment="SQL interpolation contains an unclosed quoted string", + ), + SubstituteSqlVarsErrorTestCase( + description="deferred placeholder preserves unclosed comment diagnostics", + sql="SELECT @@@partition_start /* unterminated", + effective_vars={}, + expected_error_fragment="SQL interpolation contains an unclosed block comment", + ), ], ids=lambda case: case.description, ) @@ -205,3 +219,40 @@ def test_given_missing_var_when_substituting_then_raises( effective_vars=test_case.effective_vars, context_values=test_case.context_values, ) + + +@pytest.mark.parametrize( + "test_case", + [ + StaticProjectVarBatchTestCase( + description="static scalars use native results and dynamic values fall back", + sqls=( + "SELECT @@revision, '@@status', @@@window_start", + "-- @@revision\nSELECT 1", + "SELECT '@@ENV:SQLBUILD_TEST_BATCH_USER'", + ), + effective_vars={"revision": 7, "status": "ready"}, + expected_sqls=( + "SELECT 7, 'ready', @@@window_start", + "-- @@revision\nSELECT 1", + None, + ), + ), + StaticProjectVarBatchTestCase( + description="structured project variables fall back", + sqls=("SELECT @@grants",), + effective_vars={"grants": {"role": "analyst"}}, + expected_sqls=(None,), + ), + ], + ids=lambda case: case.description, +) +def test_given_static_model_sql_when_batch_preparing_then_returns_safe_results_and_fallbacks( + test_case: StaticProjectVarBatchTestCase, +) -> None: + result: tuple[str | None, ...] = prepare_static_project_vars_batch( + sqls=test_case.sqls, + effective_vars=test_case.effective_vars, + ) + + assert result == test_case.expected_sqls diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_for_polyglot.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_for_polyglot.py index b43a94e2e0..d92dd9aa3e 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_for_polyglot.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_for_polyglot.py @@ -13,6 +13,18 @@ @pytest.mark.parametrize( "test_case", ( + PolyglotSqlNormalizationTestCase( + description="normalizes dynamic access after arithmetic and variable tokens", + sql=( + "SELECT amount / quantity - $discount, " + "payload:labels[TO_VARCHAR(product_id)] FROM orders" + ), + dialect="snowflake", + expected_sql=( + "SELECT amount / quantity - $discount, " + "GET(payload:labels, TO_VARCHAR(product_id)) FROM orders" + ), + ), PolyglotSqlNormalizationTestCase( description="rewrites a dynamic key after a Snowflake variant path", sql="SELECT payload:labels[TO_VARCHAR(payload:item_id)] FROM events", diff --git a/tests/unit/src/sqlbuild/lint/_helpers/test_native_sql.py b/tests/unit/src/sqlbuild/lint/_helpers/test_native_sql.py index 3d202a6826..e2a0a18ff3 100644 --- a/tests/unit/src/sqlbuild/lint/_helpers/test_native_sql.py +++ b/tests/unit/src/sqlbuild/lint/_helpers/test_native_sql.py @@ -76,8 +76,8 @@ def test_given_invalid_native_response_when_linting_then_boundary_fails_closed( ) -> None: monkeypatch.setattr( native_sql._native, - "lint_sql_json", - lambda _request: test_case.response, + "lint_sql_batch_json", + lambda _request: json.dumps([{"Ok": json.loads(test_case.response)}]), ) target: Path = tmp_path / "model.sql" body: LintBody = LintBody( @@ -113,10 +113,11 @@ def test_given_one_unparseable_body_when_linting_then_other_files_still_complete tmp_path: Path, ) -> None: lint_response: Mock = Mock( - side_effect=[ValueError(test_case.error_message), '{"version":1,"diagnostics":[]}'] + return_value=json.dumps( + [{"Err": test_case.error_message}, {"Ok": {"version": 1, "diagnostics": []}}] + ) ) - monkeypatch.setattr(native_sql._native, "lint_sql_json", lint_response) - monkeypatch.setattr(native_sql, "_MAX_NATIVE_LINT_WORKERS", 1) + monkeypatch.setattr(native_sql._native, "lint_sql_batch_json", lint_response) failed_path: Path = tmp_path / "failed.sql" healthy_path: Path = tmp_path / "healthy.sql" failed_text: str = "SELECT category\nGROUP BY GROUPING SETS ((category))" @@ -166,10 +167,10 @@ def test_given_identical_expanded_bodies_when_linting_then_native_analysis_is_re calls: list[str] = [] def lint_sql_json(request: str) -> str: - calls.append(request) - return '{"version":1,"diagnostics":[]}' + calls.extend(json.dumps(payload) for payload in json.loads(request)) + return json.dumps([{"Ok": {"version": 1, "diagnostics": []}} for _ in json.loads(request)]) - monkeypatch.setattr(native_sql._native, "lint_sql_json", lint_sql_json) + monkeypatch.setattr(native_sql._native, "lint_sql_batch_json", lint_sql_json) paths: tuple[Path, Path] = (tmp_path / "first.sql", tmp_path / "second.sql") bodies: tuple[LintBody, ...] = tuple( LintBody( @@ -210,10 +211,10 @@ def test_given_different_dependency_identities_when_linting_then_native_cache_is calls: list[str] = [] def lint_sql_json(request: str) -> str: - calls.append(request) - return '{"version":1,"diagnostics":[]}' + calls.extend(json.dumps(payload) for payload in json.loads(request)) + return json.dumps([{"Ok": {"version": 1, "diagnostics": []}} for _ in json.loads(request)]) - monkeypatch.setattr(native_sql._native, "lint_sql_json", lint_sql_json) + monkeypatch.setattr(native_sql._native, "lint_sql_batch_json", lint_sql_json) sql: str = "WITH imported_orders AS (SELECT * FROM __sqb_lint_0__) SELECT 1" paths: tuple[Path, Path] = (tmp_path / "dependency.sql", tmp_path / "dynamic.sql") bodies: tuple[LintBody, LintBody] = ( @@ -412,7 +413,11 @@ def test_given_diagnostic_crossing_generated_sql_when_mapping_then_range_falls_b ], } ) - monkeypatch.setattr(native_sql._native, "lint_sql_json", lambda _request: response) + monkeypatch.setattr( + native_sql._native, + "lint_sql_batch_json", + lambda _request: json.dumps([{"Ok": json.loads(response)}]), + ) target: Path = tmp_path / "model.sql" body: LintBody = LintBody( file_path=target, diff --git a/tests/unit/src/sqlbuild/lint/_helpers/test_sqlbuild_tokens.py b/tests/unit/src/sqlbuild/lint/_helpers/test_sqlbuild_tokens.py index c9605b2852..f3a00e4407 100644 --- a/tests/unit/src/sqlbuild/lint/_helpers/test_sqlbuild_tokens.py +++ b/tests/unit/src/sqlbuild/lint/_helpers/test_sqlbuild_tokens.py @@ -4,6 +4,7 @@ import pytest +import sqlbuild._native as _native from sqlbuild.lint._helpers.sqlbuild_tokens import ( map_neutralized_offset, neutralize_interpolation, @@ -149,6 +150,12 @@ def test_given_body_when_neutralizing_then_sentinels_replace_interpolation( neutralized, sites = neutralize_interpolation(body=test_case.body) assert neutralized == test_case.expected_neutralized assert tuple(site.original_text for site in sites) == test_case.expected_original_texts + native: tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None = ( + _native.prepare_lint_sql(test_case.body, test_case.body, []) + ) + assert native is not None + assert native[0] == test_case.expected_neutralized + assert tuple(InterpolationSite(*site) for site in native[1]) == sites @pytest.mark.parametrize( diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py index 125feec286..b2502f0467 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py @@ -111,6 +111,12 @@ class ScopePayloadTestCase: expected_result: bool +@dataclass(frozen=True) +class ProjectFingerprintRepresentationTestCase: + description: str + expected_equivalent: bool + + @dataclass(frozen=True) class NativeFactPayloadTestCase: description: str diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_facts.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_facts.py index 29ce88b1cb..0154526824 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_facts.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_facts.py @@ -5,6 +5,7 @@ import pytest from sqlbuild.compiler.compile.models import ( + CompactLineageFacts, CompiledObjectKey, CompiledProject, CompiledSqlScenario, @@ -14,8 +15,11 @@ from sqlbuild.compiler.discovery.models import ( DiscoveredSqlScenarioFile, ) +from sqlbuild.rule_engine._helpers.engine.native import _custom_fact_fingerprint +from sqlbuild.rule_engine.constants import RULE_CONTEXT_PROJECT_FACT from tests.unit.src.sqlbuild.rule_engine._helpers.engine._test_types import ( NativeFactPayloadTestCase, + ProjectFingerprintRepresentationTestCase, ) from tests.unit.src.sqlbuild.rule_engine._helpers.engine.helpers import ( captured_native_request, @@ -132,3 +136,56 @@ def test_given_compiled_sql_facts_when_evaluating_native_then_exact_safe_rows_ar assert request["models"][0]["targeting_test_count"] == 1 assert "/private/project" not in str(request["sql_tests"] + request["sql_scenarios"]) assert "secret fixture value" not in str(request["sql_tests"] + request["sql_scenarios"]) + + +@pytest.mark.parametrize( + "test_case", + ( + ProjectFingerprintRepresentationTestCase( + "lineage representation preserves cache identity", True + ), + ), + ids=lambda case: case.description, +) +def test_given_compact_or_materialized_lineage_when_fingerprinting_then_semantics_define_identity( + test_case: ProjectFingerprintRepresentationTestCase, +) -> None: + project: CompiledProject = build_project( + name="selected_orders", + relative_path="models/selected_orders.sql", + sql="SELECT order_id FROM orders", + config_values={}, + ) + lineage: CompactLineageFacts = CompactLineageFacts( + string_pool=("order_id", "source", "canonical_input", "orders", "other_orders"), + rows=((0, 0, 1, ((1, 2, 0),)),), + resource_name_indexes={2: 3}, + ) + compact: CompiledProject = replace( + project, models=(replace(project.models[0], fast_lineage_columns=lineage),) + ) + materialized: CompiledProject = replace( + project, models=(replace(project.models[0], fast_lineage_columns=tuple(lineage)),) + ) + changed: CompiledProject = replace( + project, + models=( + replace( + project.models[0], + fast_lineage_columns=replace(lineage, resource_name_indexes={2: 4}), + ), + ), + ) + attributes: frozenset[str] = frozenset({RULE_CONTEXT_PROJECT_FACT}) + compact_fingerprint: str = _custom_fact_fingerprint(project=compact, attributes=attributes) + materialized_fingerprint: str = _custom_fact_fingerprint( + project=materialized, attributes=attributes + ) + changed_fingerprint: str = _custom_fact_fingerprint(project=changed, attributes=attributes) + + assert (compact_fingerprint == materialized_fingerprint) is test_case.expected_equivalent + assert changed_fingerprint != compact_fingerprint + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"]))