From 08f086ad060cc5e5c02f87200427c90f8bab83bd Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:27:47 +0100 Subject: [PATCH 01/11] feat: bound unit test, scenario and model execution cost Limit the real models one SQL unit test may execute ([settings] unit_test_max_real_models, default 10) with a compile error that names every real model, its dependency path, the inputs to mock, when a scenario may fit better and the exact key to raise. Evaluate each compared actual output once into a temporary relation (DuckDB, MotherDuck, PostgreSQL; Snowflake reads RESULT_SCAN) so counts, both difference directions, column probes and samples read it. Add per-item time limits: test_timeout (default 30s) and scenario_timeout in [settings] and target execution_limits, TEST()/SCENARIO() timeout, --test-timeout and --scenario-timeout, and per-model max_duration in MODEL(), [path_defaults] and [defaults]. At expiry SQLBuild cancels the running statement, reports TIMEOUT with the limit, its source, the cancellation outcome and query IDs, and continues with independent work. Every built-in adapter now enforces deadlines, including the target's whole-build max_duration: Snowflake, PostgreSQL, BigQuery and Databricks bound statements server-side, and all adapters cancel in-flight work. --- .../src/compiler/_helpers/sql_tests/mod.rs | 1 + .../compiler/_helpers/sql_tests/planning.rs | 95 +++-- .../_helpers/sql_tests/real_model_limit.rs | 136 +++++++ .../compiler/_helpers/sql_tests/rendering.rs | 112 +++++- .../sqlbuild-analysis/src/compiler/models.rs | 2 + .../sqlbuild-analysis/src/compiler/tests.rs | 2 + .../src/compiler/tests/single_read/helpers.rs | 80 ++++ .../compiler/tests/test_sql_test_rendering.rs | 6 + .../sqlbuild-cli/src/argv/command_table.json | 30 ++ crates/sqlbuild-cli/src/dispatch/constants.rs | 9 +- .../src/effective/constants.rs | 28 ++ .../src/project/_helpers/defaults.rs | 3 + .../src/project/_helpers/sections.rs | 70 ++++ .../src/project/_helpers/targets.rs | 26 +- .../sqlbuild-config/src/project/constants.rs | 24 +- .../src/build_progress/_helpers/model_rows.rs | 4 +- .../src/build_progress/_helpers/text.rs | 12 +- .../src/build_progress/constants.rs | 2 + .../sqlbuild-core/src/cli_text/constants.rs | 2 +- .../sqlbuild-core/src/progress/constants.rs | 1 + .../main/format_expectation_detail.rs | 6 +- .../src/test_command/constants.rs | 1 + .../main/render_test_completion.rs | 24 +- .../test_command/main/test_outcome_status.rs | 3 +- .../src/execution_limits/constants.rs | 8 + .../src/execution_limits/main/mod.rs | 1 + .../main/parse_execution_limit.rs | 18 + .../src/execution_limits/mod.rs | 6 + .../src/execution_limits/models.rs | 8 + .../src/execution_limits/tests.rs | 4 + .../tests/test_parse_execution_limit.rs | 48 +++ .../src/execution_limits/tests/test_types.rs | 8 + crates/sqlbuild-model-config/src/lib.rs | 1 + .../_helpers/materialization.rs | 16 + .../main/fixed_duration_seconds.rs | 12 + .../src/model_validation/main/mod.rs | 1 + .../main/validate_model_config.rs | 5 +- .../src/sql_tests/_helpers/header.rs | 20 + .../src/sql_tests/main/sql_test_block.rs | 4 +- .../src/sql_tests/models.rs | 1 + .../_helpers/analysis/sql_test_comparisons.rs | 24 +- .../_helpers/build_command/execution.rs | 13 + .../_helpers/contract_records/constants.rs | 10 + .../bindings/_helpers/discovery/constants.rs | 8 +- .../_helpers/discovery/sql_test_records.rs | 14 + .../_helpers/model_config/project_defaults.rs | 3 +- .../bindings/_helpers/native_host/executor.rs | 12 + .../_helpers/planner_records/constants.rs | 6 + .../planner_strategy/model_entries.rs | 36 ++ .../bindings/_helpers/records/constants.rs | 3 + .../bindings/_helpers/records/construct.rs | 17 +- .../_helpers/scenario_command/test_command.rs | 25 +- .../bindings/_helpers/sql_test_command/mod.rs | 1 + .../_helpers/sql_test_command/registration.rs | 7 +- .../_helpers/sql_test_command/test_command.rs | 36 +- .../_helpers/sql_test_command/timeouts.rs | 163 ++++++++ .../_helpers/sql_test_glue/plan_driver.rs | 5 + .../_helpers/sql_test_glue/plan_entries.rs | 23 +- .../_helpers/sql_test_glue/planning.rs | 2 + src/sqlbuild/_native.pyi | 42 +++ src/sqlbuild/_native_host.py | 2 + src/sqlbuild/_native_host.pyi | 2 + .../contract/_helpers/raw_statement_audit.py | 6 + .../adapter/contract/classes/base_adapter.py | 12 + .../adapter/contract/classes/connection.py | 13 +- .../contract/classes/duckdb_backed_adapter.py | 14 + .../contract/classes/strict_adapter.py | 10 + src/sqlbuild/adapter/contract/constants.py | 1 + .../bigquery/classes/bigquery_adapter.py | 20 +- .../bigquery/classes/bigquery_connection.py | 44 ++- .../databricks/classes/databricks_adapter.py | 20 +- .../classes/databricks_connection.py | 52 ++- .../postgres/classes/postgres_adapter.py | 13 + .../postgres/classes/postgres_connection.py | 22 ++ .../classes/snowflake_active_statements.py | 45 +++ .../snowflake/classes/snowflake_adapter.py | 21 +- .../snowflake/classes/snowflake_connection.py | 8 +- .../snowflake/classes/snowflake_cursor.py | 26 +- src/sqlbuild/adapters/snowflake/exceptions.py | 10 + .../sqlserver/classes/sqlserver_adapter.py | 19 +- .../sqlserver/classes/sqlserver_connection.py | 27 +- .../executor/build/_helpers/model_timeouts.py | 45 +++ .../executor/build/classes/build_scheduler.py | 26 +- src/sqlbuild/executor/build/models.py | 2 + .../executor/pipeline/_helpers/scenario.py | 41 +- .../pipeline/_helpers/scenario_timeouts.py | 54 +++ .../executor/pipeline/_helpers/testing.py | 16 +- .../scenario/_helpers/execution/run.py | 27 +- src/sqlbuild/executor/scenario/models.py | 10 + .../executor/testing/_helpers/comparison.py | 72 ++++ .../testing/_helpers/difference_samples.py | 21 +- .../testing/_helpers/expected_column_probe.py | 21 +- .../testing/_helpers/timeout_policy.py | 19 + .../testing/_helpers/timeout_result.py | 103 +++++ .../executor/testing/classes/__init__.py | 0 .../classes/sql_test_actual_outputs.py | 104 ++++++ src/sqlbuild/executor/testing/constants.py | 6 + .../executor/testing/main/_execute.py | 114 +++--- src/sqlbuild/executor/testing/models.py | 3 + src/sqlbuild/executor/testing/types.py | 1 + .../execution_limits/_helpers/cancellation.py | 30 ++ .../execution_limits/_helpers/context.py | 10 + .../execution_limits/classes/__init__.py | 0 .../classes/active_connections.py | 34 ++ .../execution_limits/classes/item_deadline.py | 114 ++++++ .../runtime/execution_limits/constants.py | 21 ++ .../runtime/execution_limits/exceptions.py | 11 + .../main/current_execution_deadline.py | 18 +- .../main/current_item_deadline.py | 10 + .../main/enforce_execution_deadline.py | 11 +- .../main/execution_deadline_scope.py | 24 +- .../main/item_deadline_scope.py | 42 +++ .../main/item_deadline_suspended.py | 21 ++ .../main/timeout_report_lines.py | 29 ++ .../main/tracked_statement_connection.py | 34 ++ .../runtime/execution_limits/models.py | 53 +++ .../runtime/execution_limits/types.py | 23 ++ .../classes/statement_lifecycle.py | 14 + .../cli/commands/main/test/_test_types.py | 51 +++ .../cli/commands/main/test/helpers.py | 58 +++ .../commands/main/test/test_cost_guards.py | 352 ++++++++++++++++++ .../unit/src/sqlbuild/adapters/_test_types.py | 7 + tests/unit/src/sqlbuild/adapters/helpers.py | 122 ++++++ .../adapters/test_statement_cancellation.py | 189 ++++++++++ .../execution/test_build_process_resources.py | 10 +- .../pipeline/_helpers/test_scenario.py | 5 +- .../pipeline/_helpers/test_testing.py | 12 +- .../runtime/execution_limits/__init__.py | 0 .../runtime/execution_limits/_test_types.py | 54 +++ .../runtime/execution_limits/helpers.py | 51 +++ .../test_item_deadline_scope.py | 222 +++++++++++ .../test_run_item_timeouts.py | 88 +++++ 132 files changed, 3879 insertions(+), 203 deletions(-) create mode 100644 crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/real_model_limit.rs create mode 100644 crates/sqlbuild-analysis/src/compiler/tests/single_read/helpers.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/constants.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/main/mod.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/main/parse_execution_limit.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/mod.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/models.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/tests.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/tests/test_parse_execution_limit.rs create mode 100644 crates/sqlbuild-model-config/src/execution_limits/tests/test_types.rs create mode 100644 crates/sqlbuild-model-config/src/model_validation/main/fixed_duration_seconds.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/timeouts.rs create mode 100644 src/sqlbuild/adapters/snowflake/classes/snowflake_active_statements.py create mode 100644 src/sqlbuild/adapters/snowflake/exceptions.py create mode 100644 src/sqlbuild/executor/build/_helpers/model_timeouts.py create mode 100644 src/sqlbuild/executor/pipeline/_helpers/scenario_timeouts.py create mode 100644 src/sqlbuild/executor/testing/_helpers/comparison.py create mode 100644 src/sqlbuild/executor/testing/_helpers/timeout_policy.py create mode 100644 src/sqlbuild/executor/testing/_helpers/timeout_result.py create mode 100644 src/sqlbuild/executor/testing/classes/__init__.py create mode 100644 src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py create mode 100644 src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py create mode 100644 src/sqlbuild/runtime/execution_limits/classes/__init__.py create mode 100644 src/sqlbuild/runtime/execution_limits/classes/active_connections.py create mode 100644 src/sqlbuild/runtime/execution_limits/classes/item_deadline.py create mode 100644 src/sqlbuild/runtime/execution_limits/constants.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/current_item_deadline.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/timeout_report_lines.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py create mode 100644 src/sqlbuild/runtime/execution_limits/types.py create mode 100644 tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py create mode 100644 tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py create mode 100644 tests/unit/src/sqlbuild/runtime/execution_limits/__init__.py create mode 100644 tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py create mode 100644 tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py create mode 100644 tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py create mode 100644 tests/unit/src/sqlbuild/runtime/execution_limits/test_run_item_timeouts.py diff --git a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/mod.rs b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/mod.rs index eefb8cb50c..0513121193 100644 --- a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/mod.rs +++ b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/mod.rs @@ -16,6 +16,7 @@ pub mod markers; pub mod mock_reads; pub mod planning; pub mod reader_references; +pub(crate) mod real_model_limit; pub mod reference_calls; pub mod relation_markers; pub mod rendering; diff --git a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/planning.rs b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/planning.rs index e0d564e389..0f2f4dfe4f 100644 --- a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/planning.rs +++ b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/planning.rs @@ -31,6 +31,9 @@ use crate::compiler::_helpers::sql_tests::reader_references::{ CompilerReads, assertion_ref_targets, calls_reference, unresolvable_call, unresolved_reference_error, }; +use crate::compiler::_helpers::sql_tests::real_model_limit::{ + RealModelChain, real_model_limit_message, +}; use crate::compiler::_helpers::sql_tests::relation_markers::relation_marker_calls; use crate::compiler::_helpers::sql_tests::rendering::{ RenderRequest, render_comparison_sql, render_dialect, rendered_chain_steps, @@ -94,6 +97,7 @@ struct ProjectContext { patterns: SqlTestPatterns, render_dialect: Arc, rejects_nested_with: bool, + max_real_models: Option, /// Each model query's reserved CTE name, since every test reading the model checks it. reserved_ctes: Arc>>>, } @@ -469,7 +473,7 @@ pub(crate) fn resolve_chains(request: SqlTestChainBatch) -> Result Result = ordered_names + .iter() + .filter(|name| context.models.contains_key(*name)) + .cloned() + .collect(); + let chain = RealModelChain { + test_name: &plan.test_name, + file_label: &plan.file_label, + roots: &expected_names, + real_models: &real_models, + limit, + }; + if let Some(message) = + real_model_limit_message(&chain, |name| real_model_inputs(&topo_request, name)) + { + return Ok(PlannedResponse { + request: RenderRequest { + chain: Vec::new(), + assertions: Vec::new(), + sql_analysis_enabled: context.sql_analysis_enabled, + set_difference_operator: context.set_difference_operator.clone(), + sql_analysis_dialect: Some(context.dialect.clone()), + probe_step_index: None, + actual_relations: HashMap::new(), + output_step_index: None, + }, + model_names: real_models, + warnings: vec![SqlTestPlanWarning { + model_name: None, + severity: "error", + message, + }], + }); + } + } let mut warnings: Vec = Vec::new(); let mut reported_missing_mocks: HashSet<(&'static str, String)> = HashSet::new(); let mut reachable_mocks: HashSet = HashSet::new(); @@ -928,6 +972,8 @@ fn plan_model_test( set_difference_operator: context.set_difference_operator.clone(), sql_analysis_dialect: Some(context.dialect.clone()), probe_step_index: None, + actual_relations: HashMap::new(), + output_step_index: None, }; Ok(PlannedResponse { request, @@ -1302,28 +1348,31 @@ fn classify_fixtures( fixtures } -fn topo_sort_model_chain(request: TopoSortRequest<'_>) -> Vec { +fn topo_sort_model_chain(request: &TopoSortRequest<'_>) -> Vec { let mut roots: Vec = request.expected_names.to_vec(); roots.sort(); - depth_first_postorder(roots, |name| { - let Some(model) = request.models.get(name) else { - return Vec::new(); - }; - let mut dependencies: Vec = match request.overrides.get(name) { - Some(override_sql) => marker_names( - &request.patterns.reference, - &request.patterns.lexical, - override_sql, - ), - None => model.model_dependencies.clone(), - }; - dependencies.retain(|dependency| { - !request.mock_refs.contains_key(dependency) && request.models.contains_key(dependency) - }); - dependencies.sort(); - dependencies.dedup(); - dependencies - }) + depth_first_postorder(roots, |name| real_model_inputs(request, name)) +} + +/// The unmocked project models one chain model reads, sorted and deduplicated. +fn real_model_inputs(request: &TopoSortRequest<'_>, name: &str) -> Vec { + let Some(model) = request.models.get(name) else { + return Vec::new(); + }; + let mut dependencies: Vec = match request.overrides.get(name) { + Some(override_sql) => marker_names( + &request.patterns.reference, + &request.patterns.lexical, + override_sql, + ), + None => model.model_dependencies.clone(), + }; + dependencies.retain(|dependency| { + !request.mock_refs.contains_key(dependency) && request.models.contains_key(dependency) + }); + dependencies.sort(); + dependencies.dedup(); + dependencies } fn analyze_and_resolve_sql( diff --git a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/real_model_limit.rs b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/real_model_limit.rs new file mode 100644 index 0000000000..78adbbb5fa --- /dev/null +++ b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/real_model_limit.rs @@ -0,0 +1,136 @@ +//! The compile error for a unit test that executes more real models than the project allows. + +use std::collections::{BTreeSet, HashMap, VecDeque}; + +/// The project setting that bounds how many real models one unit test may execute. +pub(crate) const REAL_MODEL_LIMIT_KEY: &str = "unit_test_max_real_models"; + +/// One unit test's chain, measured against the configured real-model limit. +pub(crate) struct RealModelChain<'a> { + pub(crate) test_name: &'a str, + pub(crate) file_label: &'a str, + pub(crate) roots: &'a [String], + pub(crate) real_models: &'a [String], + pub(crate) limit: usize, +} + +/// The explanation of an oversized chain, or `None` when the chain is within the limit. +pub(crate) fn real_model_limit_message(chain: &RealModelChain<'_>, inputs: F) -> Option +where + F: Fn(&str) -> Vec, +{ + let count = chain.real_models.len(); + if count <= chain.limit { + return None; + } + let roots: Vec<&String> = chain + .roots + .iter() + .filter(|root| chain.real_models.contains(root)) + .collect(); + let paths = dependency_paths(&roots, &inputs); + let width = chain.real_models.iter().map(String::len).max().unwrap_or(0); + let mut lines: Vec = vec![ + format!( + "SQL unit test '{}' executes {count} real models (limit {}).", + chain.test_name, chain.limit + ), + format!(" Test file: {}", chain.file_label), + format!( + " Model under test: {}", + roots + .iter() + .map(|root| root.as_str()) + .collect::>() + .join(", ") + ), + String::new(), + "Real models executed, and the dependency path that pulls each one in:".to_owned(), + ]; + for model in chain.real_models.iter().rev() { + let reason = match paths.get(model.as_str()) { + Some(path) if path.len() == 1 => "model under test".to_owned(), + Some(path) => path.join(" -> "), + None => "read by the test".to_owned(), + }; + lines.push(format!(" {model: [String; 4] { + [ + " 2. If this test checks behaviour across several models, consider a scenario".to_owned(), + " (tests/scenarios/). A unit test inlines every real model into one statement," + .to_owned(), + " so the warehouse compiles and optimises the whole chain at once; a scenario" + .to_owned(), + " builds each model step by step, compiling each statement separately.".to_owned(), + ] +} + +/// The direct real inputs of each model under test, which a test mocks to cut the chain. +fn mock_advice(roots: &[&String], real_models: &[String], inputs: &F) -> Vec +where + F: Fn(&str) -> Vec, +{ + let mut seen: BTreeSet = BTreeSet::new(); + let mut mocks: Vec = Vec::new(); + for root in roots { + for input in inputs(root) { + let is_root = roots.iter().any(|candidate| **candidate == input); + if !is_root && real_models.contains(&input) && seen.insert(input.clone()) { + mocks.push(input); + } + } + } + if mocks.is_empty() { + return vec![ + " 1. Mock models further up the paths above with __ref__ CTEs.".to_owned(), + ]; + } + let mut lines: Vec = + vec![" 1. Mock the direct inputs of the model under test to cut the chain:".to_owned()]; + for mock in mocks { + lines.push(format!(" __ref__{mock} AS (SELECT ...)")); + } + lines +} + +/// The shortest read path from a model under test to each real model, breadth first. +fn dependency_paths(roots: &[&String], inputs: &F) -> HashMap> +where + F: Fn(&str) -> Vec, +{ + let mut paths: HashMap> = HashMap::new(); + let mut queue: VecDeque = VecDeque::new(); + for root in roots { + if !paths.contains_key(root.as_str()) { + paths.insert((*root).clone(), vec![(*root).clone()]); + queue.push_back((*root).clone()); + } + } + while let Some(model) = queue.pop_front() { + let path = paths.get(&model).cloned().unwrap_or_default(); + for input in inputs(&model) { + if paths.contains_key(&input) { + continue; + } + let mut extended = path.clone(); + extended.push(input.clone()); + paths.insert(input.clone(), extended); + queue.push_back(input); + } + } + paths +} diff --git a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/rendering.rs b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/rendering.rs index ff2599b2b3..606e29a7c0 100644 --- a/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/rendering.rs +++ b/crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/rendering.rs @@ -44,6 +44,8 @@ struct DifferenceSampleRequest { sample_limit: usize, #[serde(default)] use_top_clause: bool, + #[serde(default)] + actual_relation: Option, } #[derive(Debug, Deserialize)] @@ -68,6 +70,10 @@ pub(crate) struct RenderRequest { pub(crate) sql_analysis_dialect: Option, #[serde(default)] pub(crate) probe_step_index: Option, + #[serde(default)] + pub(crate) actual_relations: HashMap, + #[serde(default)] + pub(crate) output_step_index: Option, } #[derive(Debug, Serialize)] @@ -460,6 +466,15 @@ impl RenderCteState { Ok(()) } + /// Read already materialised model outputs instead of recomputing their CTEs. + fn substitute(&mut self, substitutions: &[(String, String)]) { + for (key, body) in substitutions { + if let Some(cte) = self.lifted.iter_mut().find(|cte| cte.key == *key) { + cte.body.clone_from(body); + } + } + } + fn definitions(&mut self) -> Vec { std::mem::take(&mut self.lifted) .iter() @@ -613,7 +628,10 @@ fn render_difference_sample_sql(request: &DifferenceSampleRequest) -> Result relation_select(relation), + None => cte_state.actual_step_sql(step, &suffix, enabled)?, + }; let expected_sql = cte_state.expected_step_sql(step, expected_input, enabled)?; let mut cte_parts = cte_state.definitions(); cte_parts.push(cte_definition_sql("__actual", &actual_sql)); @@ -634,11 +652,46 @@ fn render_difference_sample_sql(request: &DifferenceSampleRequest) -> Result String { + format!("SELECT * FROM {relation}") +} + +/// The statement without top-level CTEs nothing reads, or unchanged when it cannot be sliced. +fn pruned_statement(statement: &str, dialect: SliceDialect) -> String { + let Ok(Some(split)) = split_top_level_with(statement, dialect) else { + return statement.to_string(); + }; + let used = used_ctes(&split, dialect); + if used.iter().all(|used| *used) { + return statement.to_string(); + } + let definitions: Vec = split + .ctes + .iter() + .zip(used) + .filter(|(_, used)| *used) + .map(|(cte, _)| cte_definition_sql(cte.header, cte.body)) + .collect(); + if definitions.is_empty() { + return split.body.trim_start().to_string(); + } + format!( + "WITH {}\n{}", + definitions.join(",\n"), + split.body.trim_start() + ) } /// Columns compared for an expected-output step: the listed expected columns, else every column. @@ -667,6 +720,8 @@ pub(crate) fn render_comparison_sql(request: &RenderRequest) -> Result = Vec::new(); let rendered_steps = rendered_chain_steps(&request.chain, &request.assertions); let mut probe_actual_cte: Option = None; + let mut output_actual_cte: Option = None; + let mut substitutions: Vec<(String, String)> = Vec::new(); for (step_index, step) in request.chain.iter().enumerate() { let suffix = cte_state.unique_suffix(&step.model_name); @@ -675,11 +730,24 @@ pub(crate) fn render_comparison_sql(request: &RenderRequest) -> Result { + substitutions.push(( + cte_key(&model_cte_name(&step.model_name, &suffix)), + relation_select(relation), + )); + relation_select(relation) + } + None => cte_state.actual_step_sql(step, &suffix, enabled)?, + }; comparison_ctes.push((actual_cte.clone(), actual_sql)); if request.probe_step_index == Some(step_index) { probe_actual_cte = Some(actual_cte.clone()); } + if request.output_step_index == Some(step_index) { + output_actual_cte = Some(actual_cte.clone()); + break; + } let Some(expected_input) = step.expected_cte_sql.as_deref() else { continue; }; @@ -698,6 +766,22 @@ pub(crate) fn render_comparison_sql(request: &RenderRequest) -> Result Result String { diff --git a/crates/sqlbuild-analysis/src/compiler/models.rs b/crates/sqlbuild-analysis/src/compiler/models.rs index 4902677c6c..04ab21109c 100644 --- a/crates/sqlbuild-analysis/src/compiler/models.rs +++ b/crates/sqlbuild-analysis/src/compiler/models.rs @@ -32,6 +32,8 @@ pub struct SqlTestPlanBatch { #[serde(default = "crate::compiler::_helpers::sql_tests::glue::default_true")] pub include_plan: bool, pub lexical_syntax: LexicalSyntax, + #[serde(default)] + pub max_unit_test_real_models: Option, } /// The tests whose model chains to resolve; chain ordering reads only declared dependencies. diff --git a/crates/sqlbuild-analysis/src/compiler/tests.rs b/crates/sqlbuild-analysis/src/compiler/tests.rs index 5700adc60d..2037327ce9 100644 --- a/crates/sqlbuild-analysis/src/compiler/tests.rs +++ b/crates/sqlbuild-analysis/src/compiler/tests.rs @@ -2,6 +2,8 @@ mod helpers; #[path = "tests/simple_template/helpers.rs"] mod simple_template_helpers; +#[path = "tests/single_read/helpers.rs"] +mod single_read_helpers; #[path = "tests/test_sql_test_cte_rename.rs"] mod sql_test_cte_rename; #[path = "tests/test_sql_test_cte_slices.rs"] diff --git a/crates/sqlbuild-analysis/src/compiler/tests/single_read/helpers.rs b/crates/sqlbuild-analysis/src/compiler/tests/single_read/helpers.rs new file mode 100644 index 0000000000..c41381ef0b --- /dev/null +++ b/crates/sqlbuild-analysis/src/compiler/tests/single_read/helpers.rs @@ -0,0 +1,80 @@ +use serde_json::{Value, json}; + +fn render_single(request: Value) -> String { + let response: Value = serde_json::from_str( + &crate::compiler::main::sql_test_rendering::render_json( + &json!({"requests": [request]}).to_string(), + ) + .expect("test assumption must hold"), + ) + .expect("test assumption must hold"); + response[0]["sql"] + .as_str() + .expect("test assumption must hold") + .to_string() +} + +fn lifted_step() -> Value { + json!({ + "modelName": "orders", + "resolvedSql": "WITH __ref__stg_orders AS (SELECT 1 AS order_id) SELECT * FROM __ref__stg_orders", + "expectedCteSql": "SELECT 2 AS order_id", + "liftedCtes": [["__ref__stg_orders", "SELECT 1 AS order_id"]], + "comparisonBodySql": "SELECT * FROM __ref__stg_orders" + }) +} + +pub(crate) fn stored_actual_output_is_read_instead_of_the_chain() -> bool { + let output = render_single(json!({ + "chain": [lifted_step()], + "sqlAnalysisEnabled": false, + "sqlAnalysisDialect": "duckdb", + "outputStepIndex": 0 + })); + assert!( + output.ends_with("\nSELECT * FROM __actual__orders"), + "{output}" + ); + assert!( + output.contains("__ref__stg_orders AS (SELECT 1 AS order_id)"), + "{output}" + ); + assert!(!output.contains("__expected__orders"), "{output}"); + + let comparison = render_single(json!({ + "chain": [lifted_step()], + "sqlAnalysisEnabled": false, + "sqlAnalysisDialect": "duckdb", + "actualRelations": {"0": "stored_orders"} + })); + assert!( + comparison.contains("__actual__orders AS (SELECT * FROM stored_orders)"), + "{comparison}" + ); + assert!(!comparison.contains("__ref__stg_orders"), "{comparison}"); + assert!(comparison.contains("AS unexpected_count"), "{comparison}"); + + let sample: Value = serde_json::from_str( + &crate::compiler::main::sql_test_difference_sampling::render_difference_sample_json( + &json!({ + "step": lifted_step(), + "sqlAnalysisEnabled": false, + "setDifferenceOperator": "EXCEPT", + "sqlAnalysisDialect": "duckdb", + "direction": "missing", + "sampleLimit": 3, + "actualRelation": "stored_orders" + }) + .to_string(), + ) + .expect("test assumption must hold"), + ) + .expect("test assumption must hold"); + let sample_sql = sample["sql"].as_str().expect("test assumption must hold"); + assert!( + sample_sql.contains("__actual AS (SELECT * FROM stored_orders)"), + "{sample_sql}" + ); + assert!(!sample_sql.contains("__ref__stg_orders"), "{sample_sql}"); + true +} diff --git a/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_rendering.rs b/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_rendering.rs index d7baa3a076..de727b2c3c 100644 --- a/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_rendering.rs +++ b/crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_rendering.rs @@ -14,6 +14,7 @@ use crate::compiler::tests::helpers::{ tsql_model_cte_collisions_are_renamed_by_token_span, tsql_unprovable_cte_renames_are_refused_with_named_ctes, }; +use crate::compiler::tests::single_read_helpers::stored_actual_output_is_read_instead_of_the_chain; use crate::compiler::tests::test_types::SqlTestRenderingTestCase; #[test] @@ -104,6 +105,11 @@ fn given_sql_rendering_cases_when_rendering_native_batches_then_expected_behavio run: difference_sample_lifts_expected_helper_ctes, expected_success: true, }, + SqlTestRenderingTestCase { + description: "stored actual outputs are read instead of recomputing the chain", + run: stored_actual_output_is_read_instead_of_the_chain, + expected_success: true, + }, ]; for test_case in test_cases { diff --git a/crates/sqlbuild-cli/src/argv/command_table.json b/crates/sqlbuild-cli/src/argv/command_table.json index 86e9ecdfb6..f0af9a7345 100644 --- a/crates/sqlbuild-cli/src/argv/command_table.json +++ b/crates/sqlbuild-cli/src/argv/command_table.json @@ -1237,6 +1237,16 @@ "type": "int", "group": 1 }, + { + "kind": "store", + "options": [ + "--test-timeout" + ], + "dest": "test_timeout", + "metavar": "DURATION", + "help": "Per-test timeout for unit tests this build runs, such as '30s', or 'disabled'", + "group": 1 + }, { "kind": "store_true", "options": [ @@ -1690,6 +1700,16 @@ "type": "int", "group": 1 }, + { + "kind": "store", + "options": [ + "--test-timeout" + ], + "dest": "test_timeout", + "metavar": "DURATION", + "help": "Per-test timeout for this run, such as '30s' or '2m', or 'disabled'; overrides TEST() and project settings", + "group": 1 + }, { "kind": "store", "options": [ @@ -5532,6 +5552,16 @@ "type": "int", "group": 1 }, + { + "kind": "store", + "options": [ + "--scenario-timeout" + ], + "dest": "scenario_timeout", + "metavar": "DURATION", + "help": "Per-scenario timeout for this run, such as '10m', or 'disabled'", + "group": 1 + }, { "kind": "store_true", "options": [ diff --git a/crates/sqlbuild-cli/src/dispatch/constants.rs b/crates/sqlbuild-cli/src/dispatch/constants.rs index 0e28af1fe5..16ecc2f8af 100644 --- a/crates/sqlbuild-cli/src/dispatch/constants.rs +++ b/crates/sqlbuild-cli/src/dispatch/constants.rs @@ -135,7 +135,7 @@ pub(crate) const PLAN_ROUTE: Route = routed( REQUEST_KEYWORD_SHAPE, ); -const BUILD_FIELDS: [Field; 34] = [ +const BUILD_FIELDS: [Field; 35] = [ PROJECT_DIR, NO_SQL_VALIDATION, same("no_cache"), @@ -170,6 +170,7 @@ const BUILD_FIELDS: [Field; 34] = [ attr("event_output_path", "event_output"), same("max_microbatches"), same("selection_diagnostics"), + same("test_timeout"), ]; const BUILD_ARGUMENTS: [Field; 1] = [request(class("BuildCommandRequest"), &BUILD_FIELDS)]; pub(crate) const BUILD_ROUTE: Route = routed( @@ -202,7 +203,7 @@ pub(crate) const FRESHNESS_ROUTE: Route = routed( REQUEST_KEYWORD_SHAPE, ); -const TEST_FIELDS: [Field; 12] = [ +const TEST_FIELDS: [Field; 13] = [ PROJECT_DIR, NO_SQL_VALIDATION, NO_COLOR, @@ -215,6 +216,7 @@ const TEST_FIELDS: [Field; 12] = [ same("concurrency"), attr("case_name", "case"), same("inspect"), + same("test_timeout"), ]; const TEST_ARGUMENTS: [Field; 1] = [request(class("TestCommandRequest"), &TEST_FIELDS)]; pub(crate) const TEST_ROUTE: Route = routed( @@ -682,7 +684,7 @@ const SCENARIO_COMMON: [Field; 7] = [ EXCLUDE, attr("retain", "scenario_retain"), ]; -const SCENARIO_TEST_FIELDS: [Field; 15] = [ +const SCENARIO_TEST_FIELDS: [Field; 16] = [ SCENARIO_COMMON[0], SCENARIO_COMMON[1], SCENARIO_COMMON[2], @@ -698,6 +700,7 @@ const SCENARIO_TEST_FIELDS: [Field; 15] = [ SNAPSHOT_LIMITS, JSON_OUTPUT, JSON_OUTPUT_PATH, + same("scenario_timeout"), ]; const SCENARIO_TEST_ARGUMENTS: [Field; 1] = [request( class("ScenarioTestCommandRequest"), diff --git a/crates/sqlbuild-config/src/effective/constants.rs b/crates/sqlbuild-config/src/effective/constants.rs index 6767150114..840e99a048 100644 --- a/crates/sqlbuild-config/src/effective/constants.rs +++ b/crates/sqlbuild-config/src/effective/constants.rs @@ -161,6 +161,34 @@ pub(crate) const TARGET_FIELDS: &[TargetField] = &[ decided_by: &["remediation"], }, }, + TargetField { + group: LIMITS, + name: "test_timeout", + rule: MergeRule::Overlay { + decided_by: &["test_timeout"], + }, + }, + TargetField { + group: LIMITS, + name: "test_timeout_seconds", + rule: MergeRule::Overlay { + decided_by: &["test_timeout"], + }, + }, + TargetField { + group: LIMITS, + name: "scenario_timeout", + rule: MergeRule::Overlay { + decided_by: &["scenario_timeout"], + }, + }, + TargetField { + group: LIMITS, + name: "scenario_timeout_seconds", + rule: MergeRule::Overlay { + decided_by: &["scenario_timeout"], + }, + }, TargetField { group: Some("diff"), name: "max_full_rows", diff --git a/crates/sqlbuild-config/src/project/_helpers/defaults.rs b/crates/sqlbuild-config/src/project/_helpers/defaults.rs index e6d118dc9c..07474bafe8 100644 --- a/crates/sqlbuild-config/src/project/_helpers/defaults.rs +++ b/crates/sqlbuild-config/src/project/_helpers/defaults.rs @@ -10,6 +10,7 @@ use crate::project::_helpers::fields::{ }; use crate::project::_helpers::help::{change_policy_problem, change_policy_toml_help}; use crate::project::_helpers::keys::{Place, reject_removed_virtual_keys, validate_allowed_keys}; +use crate::project::_helpers::sections::item_timeout; use crate::project::constants::{ BATCH_CONCURRENCY_KEY, CHANGE_POLICY_KEYS, CONTRACT_POLICIES, DEFAULTS_KEYS, DISABLED_VALUE, GLOB_SEGMENTS, INHERIT_VALUE, MATERIALIZATION_DEFAULT_KEYS, MAX_MICROBATCHES_KEY, @@ -176,6 +177,8 @@ pub(crate) fn load_defaults(payload: Option<&ConfigValue>, file: ConfigFile) -> ), ("tags", tags), ]); + let (max_duration, _) = item_timeout(mapping, "max_duration", ("defaults", file), None)?; + fields.push(("max_duration", max_duration)); Ok(object("DefaultsConfig", fields)) } diff --git a/crates/sqlbuild-config/src/project/_helpers/sections.rs b/crates/sqlbuild-config/src/project/_helpers/sections.rs index e6d2fb5489..ecb135feaf 100644 --- a/crates/sqlbuild-config/src/project/_helpers/sections.rs +++ b/crates/sqlbuild-config/src/project/_helpers/sections.rs @@ -8,6 +8,7 @@ use crate::project::_helpers::fields::{ optional_text_field, string_mapping, string_sequence, text_or_none, value_or_none, }; use crate::project::_helpers::keys::{Place, reject_removed_virtual_keys, validate_allowed_keys}; +use crate::project::_helpers::targets::seconds_value; use crate::project::_helpers::values::{int_below, python_equal}; use crate::project::constants::{ AUDIT_SEVERITIES, COLLECTION_RENDERINGS, COLUMN_CONTRACT_MODES, CONCURRENCY_KEY, @@ -19,6 +20,10 @@ use crate::project::constants::{ SNAPSHOT_SCHEMA_CHANGE_POLICIES, SQL_ANALYSIS_KEY, UNACCOUNTED_PARTITION_POLICIES, USD_PER_CREDIT_KEY, ZERO_FIXED_DURATION, ZERO_FLOAT, }; +use crate::project::constants::{ + DEFAULT_TEST_TIMEOUT, DEFAULT_UNIT_TEST_MAX_REAL_MODELS, DISABLED_VALUE, SCENARIO_TIMEOUT_KEY, + TEST_TIMEOUT_KEY, UNIT_TEST_MAX_REAL_MODELS_KEY, +}; use crate::project::errors::ConfigFailure; use crate::project::models::{Built, ConfigFile}; use crate::project::types::{Entries, Fields, Loaded}; @@ -38,6 +43,54 @@ pub(crate) fn fixed_duration_seconds(value: &str, allow_zero: bool) -> Loaded, +) -> Loaded<(Built, Built)> { + let Some(text) = optional_str(mapping, key)?.or_else(|| default.map(str::to_owned)) else { + return Ok(( + Built::Value(ConfigValue::Null), + Built::Value(ConfigValue::Null), + )); + }; + if text == DISABLED_VALUE { + return Ok(( + Built::Value(ConfigValue::String(text)), + Built::Value(ConfigValue::Null), + )); + } + match fixed_duration_seconds(&text, false)? { + Some(seconds) => Ok(( + Built::Value(ConfigValue::String(text)), + Built::Value(seconds_value(&seconds)), + )), + None => Err(ConfigFailure::at( + file, + format!( + "{label}.{key} must be a positive fixed duration such as '30s' or '10m', or \ + 'disabled'" + ), + )), + } +} + +fn real_model_limit(mapping: &Entries) -> Loaded { + let value = optional_int( + mapping, + UNIT_TEST_MAX_REAL_MODELS_KEY, + DEFAULT_UNIT_TEST_MAX_REAL_MODELS, + )?; + if int_below(&value, 1) { + return Err(ConfigFailure::plain(format!( + "settings.{UNIT_TEST_MAX_REAL_MODELS_KEY} must be >= 1" + ))); + } + Ok(value) +} + fn settings_mapping(payload: Option<&ConfigValue>, file: ConfigFile) -> Loaded<&Entries> { let mapping = coerce_mapping(payload, place(file, "settings"))?; reject_removed_virtual_keys( @@ -185,7 +238,24 @@ pub(crate) fn load_settings(payload: Option<&ConfigValue>, file: ConfigFile) -> "default_audit_run_scope", text_or_none(optional_str(mapping, "default_audit_run_scope")?), ), + ( + UNIT_TEST_MAX_REAL_MODELS_KEY, + Built::Value(real_model_limit(mapping)?), + ), ]; + let mut fields = fields; + for (key, seconds_key, default) in [ + ( + TEST_TIMEOUT_KEY, + "test_timeout_seconds", + Some(DEFAULT_TEST_TIMEOUT), + ), + (SCENARIO_TIMEOUT_KEY, "scenario_timeout_seconds", None), + ] { + let (text, seconds) = item_timeout(mapping, key, ("settings", file), default)?; + fields.push((key, text)); + fields.push((seconds_key, seconds)); + } Ok(object("SettingsConfig", fields)) } diff --git a/crates/sqlbuild-config/src/project/_helpers/targets.rs b/crates/sqlbuild-config/src/project/_helpers/targets.rs index 6ee083be09..63c3c214b4 100644 --- a/crates/sqlbuild-config/src/project/_helpers/targets.rs +++ b/crates/sqlbuild-config/src/project/_helpers/targets.rs @@ -12,7 +12,7 @@ use crate::project::_helpers::fields::{ use crate::project::_helpers::keys::{ Place, reject_removed_virtual_keys, reject_unknown_mapping_keys, validate_allowed_keys, }; -use crate::project::_helpers::sections::fixed_duration_seconds; +use crate::project::_helpers::sections::{fixed_duration_seconds, item_timeout}; use crate::project::_helpers::values::{int_below, is_int}; use crate::project::constants::{ CLONE_KEYS, CONFIRMATION_POLICIES, EXECUTION_LIMIT_KEYS, MAX_FULL_ROWS_KEY, @@ -100,15 +100,21 @@ fn execution_limits(payload: Option<&ConfigValue>, label: &str, file: ConfigFile } }, }; - Ok(object( - "ExecutionLimitsConfig", - vec![ - ("max_models", value_or_none(max_models)), - ("max_duration", text_or_none(max_duration)), - ("max_duration_seconds", seconds), - ("remediation", optional_text_field(mapping, "remediation")?), - ], - )) + let mut fields: Fields = vec![ + ("max_models", value_or_none(max_models)), + ("max_duration", text_or_none(max_duration)), + ("max_duration_seconds", seconds), + ("remediation", optional_text_field(mapping, "remediation")?), + ]; + for (key, seconds_key) in [ + ("test_timeout", "test_timeout_seconds"), + ("scenario_timeout", "scenario_timeout_seconds"), + ] { + let (text, seconds) = item_timeout(mapping, key, (label, file), None)?; + fields.push((key, text)); + fields.push((seconds_key, seconds)); + } + Ok(object("ExecutionLimitsConfig", fields)) } /// Whole seconds as a Python integer. diff --git a/crates/sqlbuild-config/src/project/constants.rs b/crates/sqlbuild-config/src/project/constants.rs index 373d44b4c7..a8e53fead1 100644 --- a/crates/sqlbuild-config/src/project/constants.rs +++ b/crates/sqlbuild-config/src/project/constants.rs @@ -37,7 +37,12 @@ pub(crate) const SETTING_SNIPPET_INDENT: &str = " "; pub(crate) const ADDITIONAL_HELP_SEPARATOR: &str = "\n = help: "; /// The `SettingsConfig` fields, which `[settings]` may set. -pub(crate) const SETTINGS_KEYS: [&str; 11] = [ +pub(crate) const UNIT_TEST_MAX_REAL_MODELS_KEY: &str = "unit_test_max_real_models"; +pub(crate) const DEFAULT_UNIT_TEST_MAX_REAL_MODELS: i64 = 10; +pub(crate) const TEST_TIMEOUT_KEY: &str = "test_timeout"; +pub(crate) const SCENARIO_TIMEOUT_KEY: &str = "scenario_timeout"; +pub(crate) const DEFAULT_TEST_TIMEOUT: &str = "30s"; +pub(crate) const SETTINGS_KEYS: [&str; 14] = [ "sql_analysis", "require_sql_analysis", "query_change_tracking", @@ -49,8 +54,11 @@ pub(crate) const SETTINGS_KEYS: [&str; 11] = [ "table_promotion_mode", "default_audit_severity", "default_audit_run_scope", + "unit_test_max_real_models", + "test_timeout", + "scenario_timeout", ]; -pub(crate) const DEFAULTS_KEYS: [&str; 32] = [ +pub(crate) const DEFAULTS_KEYS: [&str; 33] = [ "append_cursor_inclusive", "batch_concurrency", "batch_size", @@ -70,6 +78,7 @@ pub(crate) const DEFAULTS_KEYS: [&str; 32] = [ "incremental_strategy", "lookback", "materialized", + "max_duration", "max_microbatches", "merge_exclude_columns", "microbatch_strategy", @@ -85,7 +94,7 @@ pub(crate) const DEFAULTS_KEYS: [&str; 32] = [ "unaccounted_partition_policy", ]; /// The MODEL header keys, which `[path_defaults]` entries may set. -pub(crate) const MODEL_HEADER_KEYS: [&str; 67] = [ +pub(crate) const MODEL_HEADER_KEYS: [&str; 68] = [ "alias", "append_cursor_inclusive", "audit_factories", @@ -121,6 +130,7 @@ pub(crate) const MODEL_HEADER_KEYS: [&str; 67] = [ "invalidate_hard_deletes", "lookback", "materialized", + "max_duration", "max_microbatches", "merge_exclude_columns", "microbatch_limit", @@ -175,7 +185,13 @@ pub(crate) const TARGET_KEYS: [&str; 18] = [ "warehouses", ]; pub(crate) const CLONE_KEYS: [&str; 2] = ["allow_as_clone_origin", "allow_as_clone_destination"]; -pub(crate) const EXECUTION_LIMIT_KEYS: [&str; 3] = ["max_models", "max_duration", "remediation"]; +pub(crate) const EXECUTION_LIMIT_KEYS: [&str; 5] = [ + "max_models", + "max_duration", + "remediation", + "test_timeout", + "scenario_timeout", +]; pub(crate) const SNAPSHOT_KEYS: [&str; 4] = [ "current_state_full_refresh", "historical_full_refresh", diff --git a/crates/sqlbuild-core/src/build_progress/_helpers/model_rows.rs b/crates/sqlbuild-core/src/build_progress/_helpers/model_rows.rs index 152aa2def0..7d78004d30 100644 --- a/crates/sqlbuild-core/src/build_progress/_helpers/model_rows.rs +++ b/crates/sqlbuild-core/src/build_progress/_helpers/model_rows.rs @@ -4,7 +4,7 @@ use crate::build_progress::_helpers::audits::{AuditDisplayEntry, aggregate_audit use crate::build_progress::_helpers::rows::{TopLevelRow, log_block, sql_block, top_level_line}; use crate::build_progress::_helpers::text::{ abbreviated_rows, audit_outcome_display, counter_text, execution_status_display, - format_duration, pad, skip_detail, test_outcome_display, truncate_name, + format_duration, model_status_display, pad, skip_detail, test_outcome_display, truncate_name, }; use crate::build_progress::constants::{ BATCH_INDENT, FAILED, HOOK_MAX_LABEL_WIDTH, HOOK_MIN_LABEL_WIDTH, HOOK_PHASE_WIDTH, @@ -178,7 +178,7 @@ pub(crate) fn model_line( counter: &counter_text(counter, layout.total), resource_type: &facts.display_type, name: &truncate_name(&name, layout.name_width), - status: &execution_status_display(&model.status), + status: &model_status_display(&model.status, model.error.code.as_deref()), duration: &duration, detail: &detail, }, diff --git a/crates/sqlbuild-core/src/build_progress/_helpers/text.rs b/crates/sqlbuild-core/src/build_progress/_helpers/text.rs index ce726f2a01..67ec55d183 100644 --- a/crates/sqlbuild-core/src/build_progress/_helpers/text.rs +++ b/crates/sqlbuild-core/src/build_progress/_helpers/text.rs @@ -2,7 +2,8 @@ use crate::build_progress::constants::{ ERROR, FAILED, INSUFFICIENT, PASS, ROW_COUNT_ABBREVIATION_THRESHOLD, ROW_COUNT_BILLION, - ROW_COUNT_MILLION, ROW_COUNT_THOUSAND, SECONDS_PER_MINUTE, SKIPPED, SUCCESS, WARN, + ROW_COUNT_MILLION, ROW_COUNT_THOUSAND, SECONDS_PER_MINUTE, SKIPPED, SUCCESS, TIMEOUT_CODE, + TIMEOUT_OUTCOME, WARN, }; use crate::build_progress::models::SkipText; use crate::text::main::python_isspace::python_isspace; @@ -98,10 +99,19 @@ pub(crate) fn test_outcome_display(outcome: &str) -> &'static str { match outcome { PASS => "PASS", ERROR => "ERROR", + TIMEOUT_OUTCOME => "TIMEOUT", _ => "FAIL", } } +/// `TIMEOUT` for a model attempt stopped by its time limit, else its execution status word. +pub(crate) fn model_status_display(status: &str, error_code: Option<&str>) -> String { + if status == FAILED && error_code == Some(TIMEOUT_CODE) { + return "TIMEOUT".to_owned(); + } + execution_status_display(status) +} + /// ` skip: ` for a skipped row. pub(crate) fn skip_detail(skip: &SkipText) -> String { let label: String = skip diff --git a/crates/sqlbuild-core/src/build_progress/constants.rs b/crates/sqlbuild-core/src/build_progress/constants.rs index 648027a09d..73f5985fd3 100644 --- a/crates/sqlbuild-core/src/build_progress/constants.rs +++ b/crates/sqlbuild-core/src/build_progress/constants.rs @@ -35,3 +35,5 @@ pub(crate) const WARN: &str = "warn"; pub(crate) const ERROR: &str = "error"; pub(crate) const INSUFFICIENT: &str = "insufficient"; pub(crate) const DELTA_AND_FINAL: &str = "delta_and_final"; +pub(crate) const TIMEOUT_OUTCOME: &str = "timeout"; +pub(crate) const TIMEOUT_CODE: &str = "C417"; diff --git a/crates/sqlbuild-core/src/cli_text/constants.rs b/crates/sqlbuild-core/src/cli_text/constants.rs index e84eb2526f..c356a7ab7b 100644 --- a/crates/sqlbuild-core/src/cli_text/constants.rs +++ b/crates/sqlbuild-core/src/cli_text/constants.rs @@ -10,7 +10,7 @@ pub(crate) const ORANGE: &str = "\x1b[38;5;208m"; pub(crate) const RESET: &str = "\x1b[0m"; pub(crate) const SUCCESS_STATUS_LABELS: [&str; 3] = ["OK", "PASS", "SUCCESS"]; pub(crate) const WARNING_STATUS_LABELS: [&str; 2] = ["WARN", "WARNING"]; -pub(crate) const ERROR_STATUS_LABELS: [&str; 3] = ["ERROR", "FAIL", "FAILED"]; +pub(crate) const ERROR_STATUS_LABELS: [&str; 4] = ["ERROR", "FAIL", "FAILED", "TIMEOUT"]; pub(crate) const SKIPPED_STATUS_LABELS: [&str; 2] = ["SKIP", "SKIPPED"]; pub(crate) const PHASE_OK_GLYPH: &str = "\u{2713}"; pub(crate) const PHASE_FAIL_GLYPH: &str = "\u{2717}"; diff --git a/crates/sqlbuild-core/src/progress/constants.rs b/crates/sqlbuild-core/src/progress/constants.rs index aa0ba9ec5e..31492f197a 100644 --- a/crates/sqlbuild-core/src/progress/constants.rs +++ b/crates/sqlbuild-core/src/progress/constants.rs @@ -99,6 +99,7 @@ pub(crate) const PLANNING_COMPLETION_PREFIXES: [&str; 14] = [ pub(crate) const ASSERTION_PREFIX: &str = "assertion "; pub(crate) const PASS_OUTCOME: &str = "pass"; pub(crate) const ERROR_OUTCOME: &str = "error"; +pub(crate) const TIMEOUT_OUTCOME: &str = "timeout"; /// Inspection SQL longer than this many characters is cut. pub(crate) const SQL_PREVIEW_LENGTH: usize = 160; pub(crate) const EXECUTION_TITLE: &str = "Execution"; diff --git a/crates/sqlbuild-core/src/progress/main/format_expectation_detail.rs b/crates/sqlbuild-core/src/progress/main/format_expectation_detail.rs index fb9ed9533e..85f4d54266 100644 --- a/crates/sqlbuild-core/src/progress/main/format_expectation_detail.rs +++ b/crates/sqlbuild-core/src/progress/main/format_expectation_detail.rs @@ -1,6 +1,6 @@ //! The failing row detail of one SQL expectation row. -use crate::progress::constants::{ASSERTION_PREFIX, ERROR_OUTCOME, PASS_OUTCOME}; +use crate::progress::constants::{ASSERTION_PREFIX, ERROR_OUTCOME, PASS_OUTCOME, TIMEOUT_OUTCOME}; use crate::progress::models::ExpectationStep; use crate::progress::types::DifferenceSample; @@ -21,10 +21,10 @@ fn sample_lines(label: &str, samples: &[DifferenceSample]) -> Vec { lines } -/// Row counts, samples and differing columns of a failing step; empty when it passed or errored. +/// Row counts, samples and differing columns of a failing step; empty when it passed, errored or timed out. #[must_use] pub fn format_expectation_detail(step: &ExpectationStep) -> String { - if step.outcome == PASS_OUTCOME || step.outcome == ERROR_OUTCOME { + if [PASS_OUTCOME, ERROR_OUTCOME, TIMEOUT_OUTCOME].contains(&step.outcome.as_str()) { return String::new(); } if step.model_name.starts_with(ASSERTION_PREFIX) { diff --git a/crates/sqlbuild-core/src/test_command/constants.rs b/crates/sqlbuild-core/src/test_command/constants.rs index 883fcc2cf2..5fb390f8eb 100644 --- a/crates/sqlbuild-core/src/test_command/constants.rs +++ b/crates/sqlbuild-core/src/test_command/constants.rs @@ -7,3 +7,4 @@ pub(crate) const UNKNOWN_GROUP: &str = "(unknown)"; pub(crate) const NAME_MIN_WIDTH: usize = 50; pub(crate) const CHILD_NAME_INDENT: usize = 4; pub(crate) const ERROR_SEVERITY: &str = "error"; +pub(crate) const TIMEOUT_CODE: &str = "C417"; diff --git a/crates/sqlbuild-core/src/test_command/main/render_test_completion.rs b/crates/sqlbuild-core/src/test_command/main/render_test_completion.rs index fcf1d9decd..aa33b13fec 100644 --- a/crates/sqlbuild-core/src/test_command/main/render_test_completion.rs +++ b/crates/sqlbuild-core/src/test_command/main/render_test_completion.rs @@ -3,7 +3,7 @@ use crate::cli_text::main::format_coded_error::format_coded_error; use crate::cli_text::main::format_summary_footer::format_summary_footer; use crate::cli_text::models::{CliStyle, SummaryCount}; -use crate::test_command::constants::{DEFAULT_ERROR, ERROR_LABEL, ERROR_LABEL_WIDTH}; +use crate::test_command::constants::{DEFAULT_ERROR, ERROR_LABEL, ERROR_LABEL_WIDTH, TIMEOUT_CODE}; use crate::test_command::models::TestFailure; use crate::text::main::python_splitlines::python_splitlines; @@ -51,15 +51,19 @@ pub fn render_test_completion( failures: &[TestFailure], style: CliStyle, ) -> String { - let summary: String = format_summary_footer( - style, - &[ - count("PASS", passed), - count("FAIL", total - passed), - count("TOTAL", total), - ], - None, - ); + let timed_out: usize = failures + .iter() + .filter(|failure| failure.error_code.as_deref() == Some(TIMEOUT_CODE)) + .count(); + let mut counts: Vec = vec![ + count("PASS", passed), + count("FAIL", total - passed - timed_out.min(total - passed)), + ]; + if timed_out > 0 { + counts.push(count("TIMEOUT", timed_out)); + } + counts.push(count("TOTAL", total)); + let summary: String = format_summary_footer(style, &counts, None); let mut rendered: String = format!("\n{summary}"); if !failures.is_empty() { let mut lines: Vec = vec![style.error_strong("Failures:"), String::new()]; diff --git a/crates/sqlbuild-core/src/test_command/main/test_outcome_status.rs b/crates/sqlbuild-core/src/test_command/main/test_outcome_status.rs index 42941e1198..04a9f2869f 100644 --- a/crates/sqlbuild-core/src/test_command/main/test_outcome_status.rs +++ b/crates/sqlbuild-core/src/test_command/main/test_outcome_status.rs @@ -1,11 +1,12 @@ //! The status word of a test outcome. -/// `PASS`, `ERROR` for execution errors, else `FAIL`. +/// `PASS`, `ERROR` for execution errors, `TIMEOUT` past the time limit, else `FAIL`. #[must_use] pub fn test_outcome_status(outcome: &str) -> &'static str { match outcome { "pass" => "PASS", "error" => "ERROR", + "timeout" => "TIMEOUT", _ => "FAIL", } } diff --git a/crates/sqlbuild-model-config/src/execution_limits/constants.rs b/crates/sqlbuild-model-config/src/execution_limits/constants.rs new file mode 100644 index 0000000000..ccddab1d33 --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/constants.rs @@ -0,0 +1,8 @@ +//! Execution-limit header keys and values. + +/// The TEST() and SCENARIO() header key for a per-item timeout. +pub const TIMEOUT_KEY: &str = "timeout"; +/// The MODEL() header and defaults key for a per-model execution limit. +pub const MODEL_MAX_DURATION_KEY: &str = "max_duration"; +/// The value that turns an inherited limit off. +pub const DISABLED_LIMIT: &str = "disabled"; diff --git a/crates/sqlbuild-model-config/src/execution_limits/main/mod.rs b/crates/sqlbuild-model-config/src/execution_limits/main/mod.rs new file mode 100644 index 0000000000..10b84c46d5 --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/main/mod.rs @@ -0,0 +1 @@ +pub mod parse_execution_limit; diff --git a/crates/sqlbuild-model-config/src/execution_limits/main/parse_execution_limit.rs b/crates/sqlbuild-model-config/src/execution_limits/main/parse_execution_limit.rs new file mode 100644 index 0000000000..e58b85e3db --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/main/parse_execution_limit.rs @@ -0,0 +1,18 @@ +//! Parse a per-item execution limit such as `30s`, `10m` or `disabled`. + +use crate::execution_limits::constants::DISABLED_LIMIT; +use crate::execution_limits::models::ExecutionLimit; +use crate::model_validation::main::fixed_duration_seconds::fixed_duration_seconds; + +/// The limit `text` spells, or `None` when it is not a positive fixed duration or `disabled`. +#[must_use] +pub fn parse_execution_limit(text: &str) -> Option { + let text = text.trim(); + if text == DISABLED_LIMIT { + return Some(ExecutionLimit::Disabled); + } + match u64::try_from(fixed_duration_seconds(text)?) { + Ok(seconds) if seconds > 0 => Some(ExecutionLimit::Seconds(seconds)), + Ok(_) | Err(_) => None, + } +} diff --git a/crates/sqlbuild-model-config/src/execution_limits/mod.rs b/crates/sqlbuild-model-config/src/execution_limits/mod.rs new file mode 100644 index 0000000000..042268f15e --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/mod.rs @@ -0,0 +1,6 @@ +pub mod constants; +pub mod main; +pub mod models; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-model-config/src/execution_limits/models.rs b/crates/sqlbuild-model-config/src/execution_limits/models.rs new file mode 100644 index 0000000000..f6e04e0eff --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/models.rs @@ -0,0 +1,8 @@ +//! A parsed per-item execution limit. + +/// An execution limit: whole seconds, or explicitly disabled. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExecutionLimit { + Seconds(u64), + Disabled, +} diff --git a/crates/sqlbuild-model-config/src/execution_limits/tests.rs b/crates/sqlbuild-model-config/src/execution_limits/tests.rs new file mode 100644 index 0000000000..f5e6fab257 --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/tests.rs @@ -0,0 +1,4 @@ +#[path = "tests/test_parse_execution_limit.rs"] +mod parse_execution_limit; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-model-config/src/execution_limits/tests/test_parse_execution_limit.rs b/crates/sqlbuild-model-config/src/execution_limits/tests/test_parse_execution_limit.rs new file mode 100644 index 0000000000..834a2a0c5f --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/tests/test_parse_execution_limit.rs @@ -0,0 +1,48 @@ +use crate::execution_limits::main::parse_execution_limit::parse_execution_limit; +use crate::execution_limits::models::ExecutionLimit; +use crate::execution_limits::tests::test_types::ParseExecutionLimitTestCase; + +#[test] +fn given_limit_text_when_parsing_then_expected_limit_is_returned() { + let test_cases = [ + ParseExecutionLimitTestCase { + description: "seconds", + text: "30s", + expected_limit: Some(ExecutionLimit::Seconds(30)), + }, + ParseExecutionLimitTestCase { + description: "compound fixed duration", + text: "1h30m", + expected_limit: Some(ExecutionLimit::Seconds(5_400)), + }, + ParseExecutionLimitTestCase { + description: "disabled", + text: "disabled", + expected_limit: Some(ExecutionLimit::Disabled), + }, + ParseExecutionLimitTestCase { + description: "calendar duration", + text: "1mo", + expected_limit: None, + }, + ParseExecutionLimitTestCase { + description: "zero", + text: "0s", + expected_limit: None, + }, + ParseExecutionLimitTestCase { + description: "bare number", + text: "30", + expected_limit: None, + }, + ]; + + for test_case in test_cases { + assert_eq!( + parse_execution_limit(test_case.text), + test_case.expected_limit, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-model-config/src/execution_limits/tests/test_types.rs b/crates/sqlbuild-model-config/src/execution_limits/tests/test_types.rs new file mode 100644 index 0000000000..e262756763 --- /dev/null +++ b/crates/sqlbuild-model-config/src/execution_limits/tests/test_types.rs @@ -0,0 +1,8 @@ +use crate::execution_limits::models::ExecutionLimit; + +/// One limit text and the limit it parses to, if any. +pub(super) struct ParseExecutionLimitTestCase { + pub(super) description: &'static str, + pub(super) text: &'static str, + pub(super) expected_limit: Option, +} diff --git a/crates/sqlbuild-model-config/src/lib.rs b/crates/sqlbuild-model-config/src/lib.rs index a34311abe3..130e9d67fa 100644 --- a/crates/sqlbuild-model-config/src/lib.rs +++ b/crates/sqlbuild-model-config/src/lib.rs @@ -5,6 +5,7 @@ pub mod config_presence; pub mod declarations; pub mod errors; +pub mod execution_limits; pub mod header_metadata; pub mod model_validation; pub mod name_patterns; diff --git a/crates/sqlbuild-model-config/src/model_validation/_helpers/materialization.rs b/crates/sqlbuild-model-config/src/model_validation/_helpers/materialization.rs index 9c815c2d80..01c6543312 100644 --- a/crates/sqlbuild-model-config/src/model_validation/_helpers/materialization.rs +++ b/crates/sqlbuild-model-config/src/model_validation/_helpers/materialization.rs @@ -1,5 +1,7 @@ //! Contract, materialization, storage, migration and placeholder rules with Python's errors. +use crate::execution_limits::constants::MODEL_MAX_DURATION_KEY; +use crate::execution_limits::main::parse_execution_limit::parse_execution_limit; use std::collections::{BTreeSet, HashMap}; use sqlbuild_core::text::main::is_word_char::is_word_char; @@ -75,6 +77,20 @@ pub(crate) fn check_contract(config: &ConfigView<'_, N>) -> Che ))) } +/// Check that `max_duration` is a positive fixed duration or `disabled`. +pub(crate) fn check_execution_limit(config: &ConfigView<'_, N>) -> Check { + let Some(text) = config.string(MODEL_MAX_DURATION_KEY)? else { + return Ok(()); + }; + if parse_execution_limit(&text).is_some() { + return Ok(()); + } + Err(config.error(format!( + "{MODEL_MAX_DURATION_KEY} must be a positive fixed duration such as '15m' or '1h', or \ + 'disabled', got '{text}'" + ))) +} + /// Check `validate_non_incremental_config`. pub(crate) fn check_non_incremental(config: &ConfigView<'_, N>) -> Check { if config.string("materialized")?.as_deref() == Some(INCREMENTAL_MATERIALIZATION) { diff --git a/crates/sqlbuild-model-config/src/model_validation/main/fixed_duration_seconds.rs b/crates/sqlbuild-model-config/src/model_validation/main/fixed_duration_seconds.rs new file mode 100644 index 0000000000..a35e2110d8 --- /dev/null +++ b/crates/sqlbuild-model-config/src/model_validation/main/fixed_duration_seconds.rs @@ -0,0 +1,12 @@ +//! The whole seconds of a fixed duration such as `90s` or `1h30m`. + +use crate::model_validation::_helpers::durations::parse_duration; + +/// The duration's seconds, or `None` when it does not parse or includes months or years. +#[must_use] +pub fn fixed_duration_seconds(text: &str) -> Option { + match parse_duration(text) { + Ok(Some(duration)) if duration.total_months == 0 => Some(duration.fixed_seconds), + Ok(_) | Err(_) => None, + } +} diff --git a/crates/sqlbuild-model-config/src/model_validation/main/mod.rs b/crates/sqlbuild-model-config/src/model_validation/main/mod.rs index 8c47253af8..2724bc0337 100644 --- a/crates/sqlbuild-model-config/src/model_validation/main/mod.rs +++ b/crates/sqlbuild-model-config/src/model_validation/main/mod.rs @@ -1,3 +1,4 @@ +pub mod fixed_duration_seconds; pub mod retention_override; pub mod table_type_override; pub mod validate_column_migrations; diff --git a/crates/sqlbuild-model-config/src/model_validation/main/validate_model_config.rs b/crates/sqlbuild-model-config/src/model_validation/main/validate_model_config.rs index 74b962e9a0..5b261cbea5 100644 --- a/crates/sqlbuild-model-config/src/model_validation/main/validate_model_config.rs +++ b/crates/sqlbuild-model-config/src/model_validation/main/validate_model_config.rs @@ -3,8 +3,8 @@ use crate::model_validation::_helpers::config::ConfigView; use crate::model_validation::_helpers::incremental::check_incremental; use crate::model_validation::_helpers::materialization::{ - check_contract, check_custom_materialization, check_migration, check_non_incremental, - check_placeholders, check_project_capability, check_storage_policies, + check_contract, check_custom_materialization, check_execution_limit, check_migration, + check_non_incremental, check_placeholders, check_project_capability, check_storage_policies, }; use crate::model_validation::_helpers::references::check_references; use crate::model_validation::_helpers::snapshot::check_snapshot; @@ -24,6 +24,7 @@ pub fn validate_model_config( check_incremental(&config, facts)?; check_project_capability(&config, project)?; check_contract(&config)?; + check_execution_limit(&config)?; check_non_incremental(&config)?; check_snapshot(&config, facts)?; check_custom_materialization(&config, project)?; diff --git a/crates/sqlbuild-model-config/src/sql_tests/_helpers/header.rs b/crates/sqlbuild-model-config/src/sql_tests/_helpers/header.rs index 752bb37b15..8f0160a287 100644 --- a/crates/sqlbuild-model-config/src/sql_tests/_helpers/header.rs +++ b/crates/sqlbuild-model-config/src/sql_tests/_helpers/header.rs @@ -1,5 +1,7 @@ //! TEST header fields: name, mode, analysis flag and cursor bounds. +use crate::execution_limits::constants::TIMEOUT_KEY; +use crate::execution_limits::main::parse_execution_limit::parse_execution_limit; use crate::sql_tests::_helpers::mode::defined_modes; use crate::sql_tests::_helpers::parameters::entry; use crate::sql_tests::constants::{ @@ -25,6 +27,24 @@ pub(crate) fn cursor_bound( } } +/// A header's validated `timeout`: a positive fixed duration or `disabled`. +pub(crate) fn header_timeout( + header: &[(RawKey, RawSqlValue)], + header_name: &str, + file_path: &str, +) -> Result, String> { + match entry(header, TIMEOUT_KEY) { + None | Some(RawSqlValue::Null) => Ok(None), + Some(RawSqlValue::Str(text)) if parse_execution_limit(text).is_some() => { + Ok(Some(text.trim().to_owned())) + } + Some(_) => Err(format!( + "{header_name} timeout in '{file_path}' must be a positive fixed duration such as \ + \"45s\" or \"2m\", or \"disabled\"" + )), + } +} + pub(crate) fn header_checks( header: &[(RawKey, RawSqlValue)], file_path: &str, diff --git a/crates/sqlbuild-model-config/src/sql_tests/main/sql_test_block.rs b/crates/sqlbuild-model-config/src/sql_tests/main/sql_test_block.rs index 99ada140f3..4395512b16 100644 --- a/crates/sqlbuild-model-config/src/sql_tests/main/sql_test_block.rs +++ b/crates/sqlbuild-model-config/src/sql_tests/main/sql_test_block.rs @@ -1,6 +1,6 @@ //! Python's `build_sql_test_block`: validate one TEST block's header against its SQL. -use crate::sql_tests::_helpers::header::{cursor_bound, header_checks, test_mode}; +use crate::sql_tests::_helpers::header::{cursor_bound, header_checks, header_timeout, test_mode}; use crate::sql_tests::_helpers::parameters::{entry, parameters_and_cases}; use crate::sql_tests::constants::{CURSOR_END_KEY, CURSOR_START_KEY, MODEL_MODE, NAME_KEY}; use crate::sql_tests::models::SqlTestBlock; @@ -32,7 +32,9 @@ pub fn sql_test_block( supported for model tests, not mode '{mode}'" )); } + let timeout = header_timeout(header, "TEST()", file_path)?; Ok(SqlTestBlock { + timeout, name, mode, parameters, diff --git a/crates/sqlbuild-model-config/src/sql_tests/models.rs b/crates/sqlbuild-model-config/src/sql_tests/models.rs index f8b064e738..32b6a29e54 100644 --- a/crates/sqlbuild-model-config/src/sql_tests/models.rs +++ b/crates/sqlbuild-model-config/src/sql_tests/models.rs @@ -26,4 +26,5 @@ pub struct SqlTestBlock { pub cases: Vec, pub cursor_start: Option, pub cursor_end: Option, + pub timeout: Option, } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis/sql_test_comparisons.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis/sql_test_comparisons.rs index 8d0f62592e..3698831d1d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis/sql_test_comparisons.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis/sql_test_comparisons.rs @@ -1,5 +1,7 @@ //! Render executable SQL test and scenario comparison queries from planned entries. +use std::collections::HashMap; + use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; use pyo3::types::{PyDict, PyDictMethods}; use pyo3::{FromPyObject, pyfunction, wrap_pyfunction}; @@ -23,6 +25,8 @@ pub(crate) struct RenderSettings { set_difference_operator: String, sql_analysis_dialect: Option, probe_step_index: Option, + output_step_index: Option, + actual_relations: HashMap, } impl RenderSettings { @@ -32,6 +36,8 @@ impl RenderSettings { .unwrap_or_else(|| DEFAULT_SET_DIFFERENCE.to_owned()), sql_analysis_dialect: optional_value(options, "sql_analysis_dialect")?, probe_step_index: optional_value(options, "probe_step_index")?, + output_step_index: optional_value(options, "output_step_index")?, + actual_relations: optional_value(options, "actual_relations")?.unwrap_or_default(), }) } } @@ -76,8 +82,19 @@ fn rendered_comparisons( &settings.set_difference_operator, settings.sql_analysis_dialect.as_deref(), )?; - if let (Some(index), Some(fields)) = (settings.probe_step_index, request.as_object_mut()) { - fields.insert("probeStepIndex".to_owned(), json!(index)); + if let Some(fields) = request.as_object_mut() { + if let Some(index) = settings.probe_step_index { + fields.insert("probeStepIndex".to_owned(), json!(index)); + } + if let Some(index) = settings.output_step_index { + fields.insert("outputStepIndex".to_owned(), json!(index)); + } + if !settings.actual_relations.is_empty() { + fields.insert( + "actualRelations".to_owned(), + json!(settings.actual_relations), + ); + } } requests.push(request); } @@ -110,6 +127,8 @@ pub(crate) fn single_comparison_sql( .extract()?, sql_analysis_dialect: adapter.call_method0("sql_analysis_dialect")?.extract()?, probe_step_index: None, + output_step_index: None, + actual_relations: HashMap::new(), }; rendered_comparisons(entry.py(), std::slice::from_ref(entry), &settings)? .into_iter() @@ -183,6 +202,7 @@ fn render_sql_test_difference_sample_sql( "direction": request.get("direction")?.extract::()?, "sampleLimit": request.get("sample_limit")?.extract::()?, "useTopClause": settings.sql_analysis_dialect.as_deref() == Some(TOP_CLAUSE_DIALECT), + "actualRelation": optional_value::(options, "actual_relation")?, })) .map_err(value_error)?; let response: String = py diff --git a/crates/sqlbuild-python/src/bindings/_helpers/build_command/execution.rs b/crates/sqlbuild-python/src/bindings/_helpers/build_command/execution.rs index 2e7343de8f..ec7ba2f74d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/build_command/execution.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/build_command/execution.rs @@ -36,6 +36,8 @@ use crate::bindings::_helpers::records::construct::new_record; use crate::bindings::_helpers::records::fields::field as native_field; use crate::bindings::_helpers::records::fields::mapping_get; use crate::bindings::_helpers::records::ids::RecordId; +use crate::bindings::_helpers::sql_test_command::test_command::timeout_record; +use crate::bindings::_helpers::sql_test_command::timeouts::{TimeoutKind, resolve_run_timeout}; const EXECUTION_COMMAND: &str = "sqb build"; const SUCCESS_STATUS: &str = "success"; @@ -331,6 +333,17 @@ pub(crate) fn execute_build_plan<'py>( ("end_cursor_int", end_int.clone()), ("use_color", boolean(py, context.use_color)), ("providers", providers.clone()), + ( + "test_timeouts", + timeout_record( + py, + &resolve_run_timeout( + (&project, &context.get("discovered_inputs")?), + request("test_timeout")?.extract()?, + TimeoutKind::Test, + )?, + )?, + ), ( "project_relations", build_project_relation_targets( diff --git a/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs index 59205d3677..2654ff7f31 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs @@ -85,6 +85,10 @@ const R_EXECUTION_LIMITS_CONFIG: RecordSpec = RecordSpec { FieldSpec::new("max_duration", FieldDefault::None), FieldSpec::new("max_duration_seconds", FieldDefault::None), FieldSpec::new("remediation", FieldDefault::None), + FieldSpec::new("test_timeout", FieldDefault::None), + FieldSpec::new("test_timeout_seconds", FieldDefault::None), + FieldSpec::new("scenario_timeout", FieldDefault::None), + FieldSpec::new("scenario_timeout_seconds", FieldDefault::None), ], }; const R_TARGET_DIFF_CONFIG: RecordSpec = RecordSpec { @@ -203,6 +207,11 @@ const R_SETTINGS_CONFIG: RecordSpec = RecordSpec { FieldSpec::new("table_promotion_mode", FieldDefault::None), FieldSpec::new("default_audit_severity", FieldDefault::None), FieldSpec::new("default_audit_run_scope", FieldDefault::None), + FieldSpec::new("unit_test_max_real_models", FieldDefault::Int(10)), + FieldSpec::new("test_timeout", FieldDefault::Text("30s")), + FieldSpec::new("test_timeout_seconds", FieldDefault::Int(30)), + FieldSpec::new("scenario_timeout", FieldDefault::None), + FieldSpec::new("scenario_timeout_seconds", FieldDefault::None), ], }; const R_SCOPES_CONFIG: RecordSpec = RecordSpec { @@ -317,6 +326,7 @@ const R_DEFAULTS_CONFIG: RecordSpec = RecordSpec { FieldSpec::new("batch_size", FieldDefault::None), FieldSpec::new("batch_concurrency", FieldDefault::None), FieldSpec::new("max_microbatches", FieldDefault::None), + FieldSpec::new("max_duration", FieldDefault::None), FieldSpec::new("unaccounted_partition_policy", FieldDefault::None), FieldSpec::new("replay_on_change", FieldDefault::None), FieldSpec::new("row_diff_exclude_columns", FieldDefault::EmptyTuple), diff --git a/crates/sqlbuild-python/src/bindings/_helpers/discovery/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/discovery/constants.rs index ecba49d408..5503d6a77f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/discovery/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/discovery/constants.rs @@ -1,6 +1,6 @@ //! Header keys native discovery accepts, sorted by code point as Python sorts them. -pub(crate) const SQL_MODEL_HEADER_KEYS: [&str; 67] = [ +pub(crate) const SQL_MODEL_HEADER_KEYS: [&str; 68] = [ "alias", "append_cursor_inclusive", "audit_factories", @@ -36,6 +36,7 @@ pub(crate) const SQL_MODEL_HEADER_KEYS: [&str; 67] = [ "invalidate_hard_deletes", "lookback", "materialized", + "max_duration", "max_microbatches", "merge_exclude_columns", "microbatch_limit", @@ -92,7 +93,7 @@ pub(crate) const SQL_AUDIT_HEADER_KEYS: [&str; 11] = [ "value", ]; pub(crate) const SQL_HOOK_HEADER_KEYS: [&str; 1] = ["description"]; -pub(crate) const SQL_TEST_HEADER_KEYS: [&str; 7] = [ +pub(crate) const SQL_TEST_HEADER_KEYS: [&str; 8] = [ "cases", "cursor_end", "cursor_start", @@ -100,8 +101,9 @@ pub(crate) const SQL_TEST_HEADER_KEYS: [&str; 7] = [ "name", "parameters", "sql_analysis", + "timeout", ]; -pub(crate) const SQL_SCENARIO_HEADER_KEYS: [&str; 2] = ["description", "tags"]; +pub(crate) const SQL_SCENARIO_HEADER_KEYS: [&str; 3] = ["description", "tags", "timeout"]; /// The fields a SQL hook entry reports, in output order. pub(crate) const SQL_HOOK_OUTPUT_FIELDS: [&str; 6] = [ "statement", diff --git a/crates/sqlbuild-python/src/bindings/_helpers/discovery/sql_test_records.rs b/crates/sqlbuild-python/src/bindings/_helpers/discovery/sql_test_records.rs index 4ae7d8acca..e464542475 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/discovery/sql_test_records.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/discovery/sql_test_records.rs @@ -14,6 +14,7 @@ use sqlbuild_discovery::sql_tests::models::{ DiscoveredScenarioFile, DiscoveredSqlTestFile, SqlTestBlock, SqlTestFileOptions, }; use sqlbuild_discovery::tree::models::ProjectTree; +use sqlbuild_model_config::execution_limits::main::parse_execution_limit::parse_execution_limit; use crate::bindings::_helpers::compile_driver::python_calls::imported; use crate::bindings::_helpers::discovery::constants; @@ -39,6 +40,7 @@ const PATHLIB_MODULE: &str = "pathlib"; const SCENARIO_KIND: &str = "sql_scenario"; const DESCRIPTION_KEY: &str = "description"; const TAGS_KEY: &str = "tags"; +const TIMEOUT_KEY: &str = "timeout"; /// A native discoverer of one kind of SQL test file. type Discoverer = fn( @@ -177,6 +179,18 @@ fn scenario_file<'py>( )); } } + if let Some(timeout) = header_values.get_item(TIMEOUT_KEY)? { + let valid: bool = match timeout.extract::() { + Ok(text) => parse_execution_limit(&text).is_some(), + Err(_) => false, + }; + if !valid { + return failure(format!( + "SCENARIO() timeout in '{shown}' must be a positive fixed duration such as \ + \"10m\", or \"disabled\"" + )); + } + } if file.sql_body.is_empty() { return failure(format!( "SQL scenario '{shown}' must define SQL after SCENARIO(...)" diff --git a/crates/sqlbuild-python/src/bindings/_helpers/model_config/project_defaults.rs b/crates/sqlbuild-python/src/bindings/_helpers/model_config/project_defaults.rs index 1746e11126..882cd71ca8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/model_config/project_defaults.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/model_config/project_defaults.rs @@ -5,7 +5,7 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult}; use pyo3::types::{PyDict, PyDictMethods, PyList}; /// Each default in `project_defaults_to_mapping`'s order, and whether a falsy value is unset. -const DEFAULT_KEYS: [(&str, bool); 30] = [ +const DEFAULT_KEYS: [(&str, bool); 31] = [ ("materialized", false), ("database", false), ("schema", false), @@ -27,6 +27,7 @@ const DEFAULT_KEYS: [(&str, bool); 30] = [ ("batch_size", false), ("batch_concurrency", false), ("max_microbatches", false), + ("max_duration", false), ("unaccounted_partition_policy", false), ("replay_on_change", false), ("row_diff_exclude_columns", true), diff --git a/crates/sqlbuild-python/src/bindings/_helpers/native_host/executor.rs b/crates/sqlbuild-python/src/bindings/_helpers/native_host/executor.rs index a2d1968c07..fe6bf9c1e5 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/native_host/executor.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/native_host/executor.rs @@ -78,6 +78,8 @@ pub(crate) enum ExecutorHost { RunTestPipeline, SelectScenarioSnapshotCaptureCandidates, TestPipelineCallbacks, + RunItemTimeouts, + ScenarioTestSettings, ExecutePythonChecks, ExecuteIngressPythonLoaderNodes, RunIngressPythonLoaderNodes, @@ -534,6 +536,16 @@ pub(crate) const ENTRIES: &[HostEntry] = &[ kind: HostKind::Class, edge: HostEdge::Executor, }, + HostEntry { + name: "RunItemTimeouts", + kind: HostKind::Class, + edge: HostEdge::Executor, + }, + HostEntry { + name: "ScenarioTestSettings", + kind: HostKind::Class, + edge: HostEdge::Executor, + }, HostEntry { name: "execute_python_checks", kind: HostKind::Function, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs index 9cbc14ad68..90d4e7c6cf 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs @@ -651,6 +651,9 @@ const R_MODEL_PLAN_ENTRY: RecordSpec = RecordSpec { FieldSpec::new("custom_config", FieldDefault::EmptyDict), FieldSpec::new("custom_placeholders", FieldDefault::EmptyDict), FieldSpec::new("column_rename_hints", FieldDefault::EmptyTuple), + FieldSpec::new("max_duration", FieldDefault::None), + FieldSpec::new("max_duration_seconds", FieldDefault::None), + FieldSpec::new("max_duration_source", FieldDefault::None), ], }; const R_CLONE_SOURCE_PLAN_ENTRY: RecordSpec = RecordSpec { @@ -808,6 +811,9 @@ const R_SQL_TEST_PLAN_ENTRY: RecordSpec = RecordSpec { FieldSpec::new("scope_deps", FieldDefault::EmptyTuple), FieldSpec::new("function_deps", FieldDefault::EmptyTuple), FieldSpec::new("sql_analysis_enabled", FieldDefault::True), + FieldSpec::new("real_model_count", FieldDefault::None), + FieldSpec::new("timeout", FieldDefault::None), + FieldSpec::new("timeout_seconds", FieldDefault::None), ], }; const R_SCENARIO_RELATION_MAP: RecordSpec = RecordSpec { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/model_entries.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/model_entries.rs index 1b89faf1d9..85509850c1 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/model_entries.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/model_entries.rs @@ -8,6 +8,9 @@ use pyo3::{IntoPyObject, pyfunction}; use sqlbuild_analysis::planner_strategy::main::microbatch_lookback::microbatch_lookback; use sqlbuild_analysis::planner_strategy::main::table_type_ddl::table_type_ddl; use sqlbuild_analysis::planner_strategy::models::{CursorTypeCheck, LookbackConfig}; +use sqlbuild_model_config::execution_limits::constants::MODEL_MAX_DURATION_KEY; +use sqlbuild_model_config::execution_limits::main::parse_execution_limit::parse_execution_limit; +use sqlbuild_model_config::execution_limits::models::ExecutionLimit; use crate::bindings::_helpers::analysis::model_reads; use crate::bindings::_helpers::compile_driver::python_calls::call_with; @@ -366,6 +369,38 @@ fn ddl_cursor_bounds<'py>(request: &ModelEntryDecisions<'py>) -> PyResult( + model: &Bound<'py, PyAny>, + values: &Bound<'py, PyAny>, + entry: &Bound<'py, PyDict>, +) -> PyResult<()> { + let Some(text) = config_str(values, MODEL_MAX_DURATION_KEY)? else { + return Ok(()); + }; + let text: String = text.extract()?; + let seconds: Option = match parse_execution_limit(&text) { + Some(ExecutionLimit::Seconds(seconds)) => Some(seconds), + Some(ExecutionLimit::Disabled) | None => None, + }; + let config: Bound<'py, PyAny> = native_field(model, "config")?; + let from_header: bool = + native_field(&config, "model_header_keys")?.contains(MODEL_MAX_DURATION_KEY)?; + let matched: Option = native_field(&config, "matched_path_default")?.extract()?; + let source: String = match (from_header, matched) { + (true, _) => "MODEL() header".to_owned(), + (false, Some(glob)) => format!("[path_defaults.\"{glob}\"] or the [defaults] it inherits"), + (false, None) => "[defaults]".to_owned(), + }; + entry.set_item("max_duration", text)?; + entry.set_item("max_duration_seconds", seconds)?; + entry.set_item( + "max_duration_source", + format!("{MODEL_MAX_DURATION_KEY} in {source}"), + )?; + Ok(()) +} + /// Copy the model's authored strategy, cursor, batch and snapshot settings onto the entry. fn config_fields<'py>( request: &ModelEntryDecisions<'py>, @@ -538,6 +573,7 @@ pub(crate) fn assemble_model_plan_entry<'py>( entry.set_item(field, native_field(change, field)?)?; } config_fields(request, &values, &entry)?; + execution_limit_fields(model, &values, &entry)?; entry.set_item("unique_key", unique_key)?; entry.set_item("merge_exclude_columns", merge_exclude)?; entry.set_item("table_type", table_type)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/records/constants.rs index 65b19fa0c6..9fe49eb2dd 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/records/constants.rs @@ -267,6 +267,7 @@ const BUILD_COMMAND_REQUEST: RecordSpec = RecordSpec { FieldSpec::new("max_microbatches", FieldDefault::None), FieldSpec::new("selection_diagnostics", FieldDefault::False), FieldSpec::new("selector_files", FieldDefault::EmptyTuple), + FieldSpec::new("test_timeout", FieldDefault::None), ], }; const BUILD_INVOCATION: RecordSpec = RecordSpec { @@ -590,6 +591,7 @@ const SCENARIO_TEST_COMMAND_REQUEST: RecordSpec = RecordSpec { ), FieldSpec::new("json_output", FieldDefault::False), FieldSpec::new("json_output_path", FieldDefault::None), + FieldSpec::new("scenario_timeout", FieldDefault::None), ], }; const SCENARIO_CAPTURE_COMMAND_REQUEST: RecordSpec = RecordSpec { @@ -653,6 +655,7 @@ const TEST_COMMAND_REQUEST: RecordSpec = RecordSpec { FieldSpec::new("concurrency", FieldDefault::None), FieldSpec::new("case_name", FieldDefault::None), FieldSpec::new("inspect", FieldDefault::False), + FieldSpec::new("test_timeout", FieldDefault::None), ], }; const REFACTOR_COMMAND_REQUEST: RecordSpec = RecordSpec { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/records/construct.rs b/crates/sqlbuild-python/src/bindings/_helpers/records/construct.rs index 7b2e3ddb86..d70b479ce1 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/records/construct.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/records/construct.rs @@ -3,9 +3,9 @@ use pyo3::IntoPyObject; use pyo3::call::PyCallArgs; use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyModule, PyResult, Python}; +use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyErr, PyModule, PyResult, Python}; use pyo3::sync::PyOnceLock; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyDictMethods}; use crate::bindings::_helpers::compile_driver::python_calls::keywords; use crate::bindings::_helpers::records::ids::{ErrorId, RecordId}; @@ -41,6 +41,19 @@ pub(crate) fn error_class(py: Python<'_>, id: ErrorId) -> PyResult, id: ErrorId, message: String, code: &str) -> PyErr { + let raised = (|| { + let options: Bound<'_, PyDict> = PyDict::new(py); + options.set_item("code", code)?; + error_class(py, id)?.call((message,), Some(&options)) + })(); + match raised { + Ok(error) => PyErr::from_value(error), + Err(error) => error, + } +} + /// The native record class or string enum `id` as registered, ignoring later attribute replacement. pub(crate) fn record_class(py: Python<'_>, id: RecordId) -> PyResult> { let classes: &Vec> = CLASSES diff --git a/crates/sqlbuild-python/src/bindings/_helpers/scenario_command/test_command.rs b/crates/sqlbuild-python/src/bindings/_helpers/scenario_command/test_command.rs index ab9aa1935a..14018b8ddf 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/scenario_command/test_command.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/scenario_command/test_command.rs @@ -32,6 +32,8 @@ use crate::bindings::_helpers::scenario_command::setup::{ require_sql_analysis, resolve_adapter, selected_scenarios, style, text, write, write_headers, }; use crate::bindings::_helpers::scenario_command::snapshot_sync::{LocalSync, sync_local_snapshots}; +use crate::bindings::_helpers::sql_test_command::test_command::timeout_record; +use crate::bindings::_helpers::sql_test_command::timeouts::{TimeoutKind, resolve_run_timeout}; const LOCAL_ADAPTER: &str = "duckdb"; const LOCAL_STATUSES: [&str; 4] = ["PASS", "FAIL", "ERROR", "SKIP"]; @@ -254,8 +256,27 @@ fn run_scenarios( } else { arguments.extend([ ("connection_config", run.connection_config.clone()), - ("retain", request.getattr("retain")?), - ("concurrency", int(py, run.concurrency)), + ( + "settings", + call_with( + &host::entry(py, ExecutorHost::ScenarioTestSettings)?, + &[ + ("retain", request.getattr("retain")?), + ("concurrency", int(py, run.concurrency)), + ( + "timeouts", + timeout_record( + py, + &resolve_run_timeout( + (&native_field(pipeline, "project")?, &run.project.discovered), + request.getattr("scenario_timeout")?.extract()?, + TimeoutKind::Scenario, + )?, + )?, + ), + ], + )?, + ), ( "connection_hooks", connection_hooks(&run.stream, (&run.adapter_name, run.style))?, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/mod.rs index 1e38419e09..5db811030d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/mod.rs @@ -1,2 +1,3 @@ pub(crate) mod registration; pub(crate) mod test_command; +pub(crate) mod timeouts; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/registration.rs index daf61f4830..0d79653596 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/registration.rs @@ -1,9 +1,10 @@ -//! Register the native `sqb test` command and its progress labels. +//! Register the native `sqb test` command, its progress labels and execution-limit parsing. use pyo3::prelude::{Bound, PyModule, PyResult}; -use crate::bindings::_helpers::sql_test_command::test_command; +use crate::bindings::_helpers::sql_test_command::{test_command, timeouts}; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - test_command::register(module) + test_command::register(module)?; + timeouts::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/test_command.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/test_command.rs index 6b3dc8a185..67b23cf31d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/test_command.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/test_command.rs @@ -49,6 +49,9 @@ use crate::bindings::_helpers::result_output::documents::write_execution_json_ou use crate::bindings::_helpers::result_output::event_writer::ExecutionEventWriter; use crate::bindings::_helpers::result_output::integration_values::json_value; use crate::bindings::_helpers::result_output::resource_ids::sql_test_identity; +use crate::bindings::_helpers::sql_test_command::timeouts::{ + RunTimeout, TimeoutKind, resolve_run_timeout, +}; const DECIMAL_KIND: &str = "decimal"; const PASS_OUTCOME: &str = "pass"; @@ -185,6 +188,16 @@ fn expectation_rows<'py>(result: &Bound<'py, PyAny>) -> PyResult) -> PyResult { + let count: Option = native_field(result, "real_model_count")?.extract()?; + Ok(match count { + Some(1) => " 1 real model".to_owned(), + Some(count) => format!(" {count} real models"), + None => String::new(), + }) +} + /// The model a planned test's progress rows are grouped under. fn entry_group(entry: &Bound<'_, PyAny>) -> PyResult { let mut chain: Vec<(String, bool)> = Vec::new(); @@ -240,6 +253,7 @@ impl TestCallbackRelay { text(py, outcome_status(&result.getattr("outcome")?)?), ), ("child_rows", expectation_rows(result)?.into_any()), + ("detail", text(py, &real_model_detail(result)?)), ("error_code", result.getattr("error_code")?), ("error_help", result.getattr("error_help")?), ("error_message", result.getattr("error_message")?), @@ -465,7 +479,7 @@ fn inspect(stream: &Bound<'_, PyAny>, plan_output: &Bound<'_, PyAny>) -> PyResul fn run_tests<'py>( invocation: &Bound<'py, PyAny>, pipeline: &Bound<'py, PyAny>, - concurrency: i64, + (concurrency, timeout): (i64, &RunTimeout), ) -> PyResult>> { let py: Python<'py> = invocation.py(); let stream: Bound<'py, PyAny> = invocation.getattr("progress_stream")?; @@ -578,6 +592,7 @@ fn run_tests<'py>( ("max_concurrency", PyInt::new(py, concurrency).into_any()), ("callbacks", callbacks), ("run_id", pipeline.getattr("project")?.getattr("run_id")?), + ("timeouts", timeout_record(py, timeout)?), ], )? .try_iter()? @@ -589,6 +604,14 @@ fn run_tests<'py>( Ok(results) } +/// The executor's `RunItemTimeouts` for a resolved run limit. +pub(crate) fn timeout_record<'py>( + py: Python<'py>, + timeout: &RunTimeout, +) -> PyResult> { + host::entry(py, ExecutorHost::RunItemTimeouts)?.call((), Some(&timeout.keywords(py)?)) +} + /// Write the summary, failures and JSON output; returns the exit code. fn complete( request: &Bound<'_, PyAny>, @@ -669,7 +692,16 @@ pub(crate) fn run_test_command(py: Python<'_>, request: &Bound<'_, PyAny>) -> Py header.set_item("concurrency", concurrency)?; header.set_item("use_color", invocation.getattr("use_color")?)?; write_execution_header(py, Some(header))?; - let results: Vec> = run_tests(&invocation, &pipeline, concurrency)?; + let timeout: RunTimeout = resolve_run_timeout( + ( + &native_field(&pipeline, "project")?, + &invocation.getattr("discovered_inputs")?, + ), + request.getattr("test_timeout")?.extract()?, + TimeoutKind::Test, + )?; + let results: Vec> = + run_tests(&invocation, &pipeline, (concurrency, &timeout))?; write_test_runtime_target( ( &native_field(&invocation, "effective_project_dir")? diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/timeouts.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/timeouts.rs new file mode 100644 index 0000000000..1b90e9350a --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_command/timeouts.rs @@ -0,0 +1,163 @@ +//! Resolve the per-test and per-scenario time limit a run applies from CLI, target and settings. + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::{ + Bound, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; +use pyo3::types::{PyDict, PyDictMethods}; +use pyo3::{pyfunction, wrap_pyfunction}; +use sqlbuild_model_config::execution_limits::main::parse_execution_limit::parse_execution_limit; +use sqlbuild_model_config::execution_limits::models::ExecutionLimit; + +use crate::bindings::_helpers::contracts::spec_resolution::resolve_target_config; +use crate::bindings::_helpers::records::construct::new_error; +use crate::bindings::_helpers::records::fields::field as native_field; +use crate::bindings::_helpers::records::ids::ErrorId; + +const INVALID_TIMEOUT_CODE: &str = "C418"; +const LOCAL_FILE: &str = "sqlbuild_local.toml"; +const DEFAULT_TEST_TIMEOUT: &str = "30s"; + +/// Which per-item limit a command resolves. +#[derive(Clone, Copy, Debug)] +pub(crate) enum TimeoutKind { + Test, + Scenario, +} + +impl TimeoutKind { + fn key(self) -> &'static str { + match self { + Self::Test => "test_timeout", + Self::Scenario => "scenario_timeout", + } + } + + fn flag(self) -> &'static str { + match self { + Self::Test => "--test-timeout", + Self::Scenario => "--scenario-timeout", + } + } +} + +/// The run's limit for each item: text, whole seconds (`None` when disabled) and its source. +#[derive(Clone, Debug)] +pub(crate) struct RunTimeout { + pub(crate) limit: Option, + pub(crate) seconds: Option, + pub(crate) source: Option, + pub(crate) overrides_items: bool, +} + +impl RunTimeout { + /// The executor keyword arguments that carry this limit. + pub(crate) fn keywords<'py>(&self, py: Python<'py>) -> PyResult> { + let keywords: Bound<'py, PyDict> = PyDict::new(py); + keywords.set_item("timeout_limit", self.limit.as_deref())?; + keywords.set_item("timeout_seconds", self.seconds)?; + keywords.set_item("timeout_source", self.source.as_deref())?; + keywords.set_item("timeout_overrides_items", self.overrides_items)?; + Ok(keywords) + } +} + +fn invalid(py: Python<'_>, message: String) -> PyErr { + new_error(py, ErrorId::CompileInput, message, INVALID_TIMEOUT_CODE) +} + +fn limit_seconds(text: &str) -> Option> { + match parse_execution_limit(text)? { + ExecutionLimit::Seconds(seconds) => Some(Some(seconds)), + ExecutionLimit::Disabled => Some(None), + } +} + +/// The limit for `kind`: the CLI value, else the target's, else the project or local setting. +pub(crate) fn resolve_run_timeout( + (project, discovered): (&Bound<'_, PyAny>, &Bound<'_, PyAny>), + cli: Option, + kind: TimeoutKind, +) -> PyResult { + let py: Python<'_> = project.py(); + if let Some(text) = cli { + let Some(seconds) = limit_seconds(&text) else { + return Err(invalid( + py, + format!( + "{} must be a positive fixed duration such as '30s' or '2m', or 'disabled'; \ + got '{text}'", + kind.flag() + ), + )); + }; + return Ok(RunTimeout { + limit: Some(text), + seconds, + source: Some(kind.flag().to_owned()), + overrides_items: true, + }); + } + let target_name: Bound<'_, PyAny> = native_field(project, "effective_target_name")?; + if !target_name.is_none() { + let target: Bound<'_, PyAny> = resolve_target_config( + py, + &native_field(discovered, "project_config")?, + &native_field(discovered, "local_config")?, + &target_name, + )?; + let limits: Bound<'_, PyAny> = native_field(&target, "execution_limits")?; + let text: Option = native_field(&limits, kind.key())?.extract()?; + if let Some(text) = text { + let name: String = target_name.extract()?; + return Ok(RunTimeout { + seconds: limit_seconds(&text).flatten(), + limit: Some(text), + source: Some(format!("[targets.{name}.execution_limits] {}", kind.key())), + overrides_items: false, + }); + } + } + let settings: Bound<'_, PyAny> = native_field(project, "settings")?; + let text: Option = native_field(&settings, kind.key())?.extract()?; + let Some(text) = text else { + return Ok(RunTimeout { + limit: None, + seconds: None, + source: None, + overrides_items: false, + }); + }; + let local: bool = native_field( + &native_field(discovered, "local_config")?, + "setting_overrides", + )? + .contains(kind.key())?; + let source: String = if local { + format!("[settings] {} in {LOCAL_FILE}", kind.key()) + } else if matches!(kind, TimeoutKind::Test) && text == DEFAULT_TEST_TIMEOUT { + format!("default {DEFAULT_TEST_TIMEOUT} ([settings] {})", kind.key()) + } else { + format!("[settings] {}", kind.key()) + }; + Ok(RunTimeout { + seconds: limit_seconds(&text).flatten(), + limit: Some(text), + source: Some(source), + overrides_items: false, + }) +} + +/// Whole seconds of an execution limit such as `30s`, or `None` for `disabled`. +#[pyfunction] +fn execution_limit_seconds(text: &str) -> PyResult> { + limit_seconds(text).ok_or_else(|| { + PyValueError::new_err(format!( + "'{text}' is not a positive fixed duration such as '30s' or 'disabled'" + )) + }) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(execution_limit_seconds, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs index 428f284d3a..f21ff48196 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs @@ -96,6 +96,11 @@ pub(crate) fn planned<'py>( .extract()?, render_sql: artifact || options.render_sql, include_plan: !artifact && options.include_plan, + max_unit_test_real_models: native_field( + &native_field(project, "settings")?, + "unit_test_max_real_models", + )? + .extract()?, }; let (rows, planning_ns, rendering_ns) = match plan_compiled_sql_tests(py, request) { Ok(value) => value, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_entries.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_entries.rs index 7bdb0110f8..7bf7ead6dc 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_entries.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_entries.rs @@ -2,6 +2,9 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult}; use pyo3::types::{PyDict, PyDictMethods, PyListMethods, PyTuple}; +use sqlbuild_model_config::execution_limits::constants::TIMEOUT_KEY as TIMEOUT_HEADER_KEY; +use sqlbuild_model_config::execution_limits::main::parse_execution_limit::parse_execution_limit; +use sqlbuild_model_config::execution_limits::models::ExecutionLimit; use crate::bindings::_helpers::planner_strategy::constants::MODEL_REFERENCE_KIND; use crate::bindings::_helpers::planner_strategy::models::RelationFixturePlanningContext; @@ -126,7 +129,25 @@ pub(super) fn entry<'py>( "function_deps", function_deps(test, plan, inputs.0, (&mocks, inputs.1))?, )?; - fields.set_item("chain", native_field(plan, "chain")?)?; + let chain = native_field(plan, "chain")?; + let direct: bool = native_field(test, "payload")?.hasattr("actual_cte")?; + fields.set_item( + "real_model_count", + if direct { None } else { Some(chain.len()?) }, + )?; + let header: Bound<'py, PyDict> = + native_field(&native_field(test, "test_block")?, "header_values")?.cast_into::()?; + let timeout: Option = match header.get_item(TIMEOUT_HEADER_KEY)? { + Some(value) => value.extract()?, + None => None, + }; + let timeout_seconds: Option = match timeout.as_deref().and_then(parse_execution_limit) { + Some(ExecutionLimit::Seconds(seconds)) => Some(seconds), + Some(ExecutionLimit::Disabled) | None => None, + }; + fields.set_item("timeout", timeout)?; + fields.set_item("timeout_seconds", timeout_seconds)?; + fields.set_item("chain", chain)?; fields.set_item("assertions", native_field(plan, "assertions")?)?; fields.set_item("sql_analysis_enabled", analysis)?; new_record(py, RecordId::SqlTestPlanEntry, (), Some(&fields)) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/planning.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/planning.rs index 82d8ff5e76..df61cf369a 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/planning.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/planning.rs @@ -62,6 +62,7 @@ pub(crate) struct PlanningRequest<'py> { pub(crate) lexical_syntax: LexicalSyntaxInput, pub(crate) render_sql: bool, pub(crate) include_plan: bool, + pub(crate) max_unit_test_real_models: Option, } /// The chain request: compiled models and the tests whose chains to resolve. @@ -139,6 +140,7 @@ fn planning_batch(mut request: PlanningRequest<'_>) -> PyResult tuple[MicrobatchInterval, ...]: ... +def execution_limit_seconds(text: str) -> int | None: ... def render_sql_test_comparison_sql( test_entries: Sequence[SqlTestPlanEntry], *, set_difference_operator: str = "EXCEPT", sql_analysis_dialect: str | None = None, probe_step_index: int | None = None, + output_step_index: int | None = None, + actual_relations: Mapping[int, str] | None = None, ) -> list[str]: ... def render_sql_test_difference_sample_sql( *, @@ -5088,6 +5091,7 @@ def render_sql_test_difference_sample_sql( sample_limit: int, set_difference_operator: str = "EXCEPT", sql_analysis_dialect: str | None = None, + actual_relation: str | None = None, ) -> str: ... def describe_missing_expected_columns( *, model_name: str, expected_columns: Sequence[str], available_columns: Iterable[str] @@ -7655,6 +7659,7 @@ class BuildCommandRequest(NativeRecord): max_microbatches: int | None selection_diagnostics: bool selector_files: tuple[SelectorFileSummary, ...] + test_timeout: str | None def __init__( self, project_dir: Path | None = None, @@ -7691,6 +7696,7 @@ class BuildCommandRequest(NativeRecord): max_microbatches: int | None = None, selection_diagnostics: bool = False, selector_files: tuple[SelectorFileSummary, ...] = (), + test_timeout: str | None = None, ) -> None: ... class BuildInvocation(NativeRecord): @@ -8213,6 +8219,7 @@ class ScenarioTestCommandRequest(NativeRecord): limit_inputs: ScenarioSnapshotLimitInputs json_output: bool json_output_path: Path | None + scenario_timeout: str | None def __init__( self, project_dir: Path | None = None, @@ -8230,6 +8237,7 @@ class ScenarioTestCommandRequest(NativeRecord): limit_inputs: ScenarioSnapshotLimitInputs = ..., json_output: bool = False, json_output_path: Path | None = None, + scenario_timeout: str | None = None, ) -> None: ... class ScenarioCaptureCommandRequest(NativeRecord): @@ -8316,6 +8324,7 @@ class TestCommandRequest(NativeRecord): concurrency: int | None case_name: str | None inspect: bool + test_timeout: str | None def __init__( self, project_dir: Path | None = None, @@ -8330,6 +8339,7 @@ class TestCommandRequest(NativeRecord): concurrency: int | None = None, case_name: str | None = None, inspect: bool = False, + test_timeout: str | None = None, ) -> None: ... class RefactorCommandRequest(NativeRecord): @@ -13575,12 +13585,20 @@ class ExecutionLimitsConfig(NativeRecord): max_duration: str | None max_duration_seconds: int | None remediation: str | None + test_timeout: str | None + test_timeout_seconds: int | None + scenario_timeout: str | None + scenario_timeout_seconds: int | None def __init__( self, max_models: int | None = ..., max_duration: str | None = ..., max_duration_seconds: int | None = ..., remediation: str | None = ..., + test_timeout: str | None = ..., + test_timeout_seconds: int | None = ..., + scenario_timeout: str | None = ..., + scenario_timeout_seconds: int | None = ..., ) -> None: ... class TargetDiffConfig(NativeRecord): @@ -13712,6 +13730,11 @@ class SettingsConfig(NativeRecord): table_promotion_mode: str | None default_audit_severity: AuditSeverity | None default_audit_run_scope: str | None + unit_test_max_real_models: int + test_timeout: str + test_timeout_seconds: int | None + scenario_timeout: str | None + scenario_timeout_seconds: int | None def __init__( self, sql_analysis: bool = ..., @@ -13725,6 +13748,11 @@ class SettingsConfig(NativeRecord): table_promotion_mode: str | None = ..., default_audit_severity: AuditSeverity | None = ..., default_audit_run_scope: str | None = ..., + unit_test_max_real_models: int = ..., + test_timeout: str = ..., + test_timeout_seconds: int | None = ..., + scenario_timeout: str | None = ..., + scenario_timeout_seconds: int | None = ..., ) -> None: ... @property def sql_validation(self) -> bool: ... @@ -13847,6 +13875,7 @@ class DefaultsConfig(NativeRecord): batch_size: str | int | None batch_concurrency: int | None max_microbatches: int | None + max_duration: str | None unaccounted_partition_policy: str | None replay_on_change: str | None row_diff_exclude_columns: tuple[str, ...] @@ -13883,6 +13912,7 @@ class DefaultsConfig(NativeRecord): batch_size: str | int | None = ..., batch_concurrency: int | None = ..., max_microbatches: int | None = ..., + max_duration: str | None = ..., unaccounted_partition_policy: str | None = ..., replay_on_change: str | None = ..., row_diff_exclude_columns: tuple[str, ...] = ..., @@ -15679,6 +15709,9 @@ class ModelPlanEntry(NativeRecord): custom_config: dict[str, object] custom_placeholders: dict[str, str] column_rename_hints: tuple[ColumnRenameHint, ...] + max_duration: str | None + max_duration_seconds: int | None + max_duration_source: str | None def __init__( self, key: CompiledObjectKey, @@ -15759,6 +15792,9 @@ class ModelPlanEntry(NativeRecord): custom_config: dict[str, object] = ..., custom_placeholders: dict[str, str] = ..., column_rename_hints: tuple[ColumnRenameHint, ...] = ..., + max_duration: str | None = ..., + max_duration_seconds: int | None = ..., + max_duration_source: str | None = ..., ) -> None: ... class CloneSourcePlanEntry(NativeRecord): @@ -15972,6 +16008,9 @@ class SqlTestPlanEntry(NativeRecord): scope_deps: tuple[CompiledObjectKey, ...] function_deps: tuple[CompiledObjectKey, ...] sql_analysis_enabled: bool + real_model_count: int | None + timeout: str | None + timeout_seconds: int | None def __init__( self, key: CompiledObjectKey, @@ -15994,6 +16033,9 @@ class SqlTestPlanEntry(NativeRecord): scope_deps: tuple[CompiledObjectKey, ...] = ..., function_deps: tuple[CompiledObjectKey, ...] = ..., sql_analysis_enabled: bool = ..., + real_model_count: int | None = ..., + timeout: str | None = ..., + timeout_seconds: int | None = ..., ) -> None: ... class ScenarioRelationMap(NativeRecord): diff --git a/src/sqlbuild/_native_host.py b/src/sqlbuild/_native_host.py index b7574a6e37..3ee235e68e 100644 --- a/src/sqlbuild/_native_host.py +++ b/src/sqlbuild/_native_host.py @@ -406,6 +406,8 @@ "select_scenario_snapshot_capture_candidates", ), "TestPipelineCallbacks": ("sqlbuild.executor.pipeline.models", "TestPipelineCallbacks"), + "RunItemTimeouts": ("sqlbuild.runtime.execution_limits.models", "RunItemTimeouts"), + "ScenarioTestSettings": ("sqlbuild.executor.scenario.models", "ScenarioTestSettings"), "execute_python_checks": ( "sqlbuild.executor.python_nodes.main._checks", "execute_python_checks", diff --git a/src/sqlbuild/_native_host.pyi b/src/sqlbuild/_native_host.pyi index 1e53f628e7..1dde9af711 100644 --- a/src/sqlbuild/_native_host.pyi +++ b/src/sqlbuild/_native_host.pyi @@ -318,6 +318,7 @@ from sqlbuild.executor.scenario.models import ScenarioLocalReplaySource as Scena from sqlbuild.executor.scenario.models import ( ScenarioSnapshotCaptureLimits as ScenarioSnapshotCaptureLimits, ) +from sqlbuild.executor.scenario.models import ScenarioTestSettings as ScenarioTestSettings from sqlbuild.executor.scheduling.types import ExecutionStatus as ExecutionStatus from sqlbuild.executor.testing.models import SqlTestExecutionResult as SqlTestExecutionResult from sqlbuild.integrations.dagster.constants import ( @@ -518,6 +519,7 @@ from sqlbuild.runtime.execution_limits.exceptions import ( from sqlbuild.runtime.execution_limits.main.execution_deadline_scope import ( execution_deadline_scope as execution_deadline_scope, ) +from sqlbuild.runtime.execution_limits.models import RunItemTimeouts as RunItemTimeouts from sqlbuild.runtime.observability.classes.event_dispatcher import ( EventDispatcher as EventDispatcher, ) diff --git a/src/sqlbuild/adapter/contract/_helpers/raw_statement_audit.py b/src/sqlbuild/adapter/contract/_helpers/raw_statement_audit.py index 0eb81c97e8..a8c18616ac 100644 --- a/src/sqlbuild/adapter/contract/_helpers/raw_statement_audit.py +++ b/src/sqlbuild/adapter/contract/_helpers/raw_statement_audit.py @@ -24,6 +24,12 @@ "adapters/bigquery/classes/bigquery_connection.py": frozenset( {("_BigQueryConnection.query_job", "query")} ), + "adapters/databricks/classes/databricks_connection.py": frozenset( + {("_DatabricksConnection._run_setting", "execute")} + ), + "adapters/postgres/classes/postgres_connection.py": frozenset( + {("_PostgresConnection._apply_statement_timeout", "execute")} + ), } _RAW_FACTORY_METHODS: frozenset[str] = frozenset({"cursor"}) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index f915886a37..2ea80590f8 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -111,6 +111,7 @@ class BaseAdapter(RetentionAdapterMixin, StrictAdapter): views_read_with_reader_access: ClassVar[bool] = False allows_implicit_managed_write_schema: ClassVar[bool] = False execution_duration_limit_seconds: ClassVar[int | None] = None + statement_cancellation_is_asynchronous: ClassVar[bool] = False metadata_inspection_concurrency: ClassVar[int] = 1 supports_session_warehouse: ClassVar[bool] = False _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( @@ -155,6 +156,17 @@ def interrupt_connection(self, connection: Any) -> bool: del connection return False + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """A relation holding a unit test's evaluated output, or `None` to compare inline.""" + + del connection, sql, name + return None + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Remove a relation `create_sql_test_output` returned.""" + + del connection, relation + def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapter/contract/classes/connection.py b/src/sqlbuild/adapter/contract/classes/connection.py index 59a8821cd2..656a050917 100644 --- a/src/sqlbuild/adapter/contract/classes/connection.py +++ b/src/sqlbuild/adapter/contract/classes/connection.py @@ -9,10 +9,16 @@ from sqlbuild.adapter.contract.main._attribute_failed_sql import attribute_failed_sql from sqlbuild.adapter.contract.models import QueryResult -from sqlbuild.runtime.execution_limits.exceptions import ExecutionDurationLimitError +from sqlbuild.runtime.execution_limits.exceptions import ( + ExecutionDurationLimitError, + ItemTimeoutError, +) from sqlbuild.runtime.execution_limits.main.enforce_execution_deadline import ( enforce_execution_deadline, ) +from sqlbuild.runtime.execution_limits.main.tracked_statement_connection import ( + tracked_statement_connection, +) from sqlbuild.runtime.observability.classes.statement_lifecycle import StatementLifecycle @@ -32,11 +38,12 @@ def execute(self, *, connection: Any, sql: str) -> Any: with StatementLifecycle(adapter=self.adapter_name, sql=sql, intent="execute") as lifecycle: try: enforce_execution_deadline() - result: Any = self._execute(connection=connection, sql=sql) + with tracked_statement_connection(adapter=self, connection=connection): + result: Any = self._execute(connection=connection, sql=sql) except Exception as error: try: enforce_execution_deadline() - except ExecutionDurationLimitError as deadline_error: + except (ExecutionDurationLimitError, ItemTimeoutError) as deadline_error: raise deadline_error from error attributed_error: Exception = attribute_failed_sql(error=error, sql=sql) if attributed_error is error: diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index a0d3a8c3f9..e57bbd8024 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -37,6 +37,7 @@ from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( + CANCELLABLE_DURATION_LIMIT_SECONDS, DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, DUCKDB_MISSING_OBJECT_MARKER, @@ -113,6 +114,19 @@ def get_columns_for_relations( ) relation_grants_supported: ClassVar[bool] = False + execution_duration_limit_seconds: ClassVar[int | None] = CANCELLABLE_DURATION_LIMIT_SECONDS + + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Store a unit test's actual output once in a session temporary table.""" + + _ = self.execute(connection=connection, sql=f"CREATE TEMPORARY TABLE {name} AS {sql}") + return name + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Drop the temporary table `create_sql_test_output` created.""" + + _ = self.execute(connection=connection, sql=f"DROP TABLE IF EXISTS {relation}") + def interrupt_connection(self, connection: Any) -> bool: """Cancel a connection's running statement from another thread; False if unsupported.""" diff --git a/src/sqlbuild/adapter/contract/classes/strict_adapter.py b/src/sqlbuild/adapter/contract/classes/strict_adapter.py index 7b676a97bb..ebddef75ae 100644 --- a/src/sqlbuild/adapter/contract/classes/strict_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/strict_adapter.py @@ -105,6 +105,16 @@ def interrupt_connection(self, connection: Any) -> bool: """Cancel a connection's running statement from another thread; False if unsupported.""" ... + @abstractmethod + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Evaluate a unit test's actual output once and return a relation reading it.""" + ... + + @abstractmethod + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Remove a relation `create_sql_test_output` returned.""" + ... + @abstractmethod def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapter/contract/constants.py b/src/sqlbuild/adapter/contract/constants.py index db92dd0498..2f323717dd 100644 --- a/src/sqlbuild/adapter/contract/constants.py +++ b/src/sqlbuild/adapter/contract/constants.py @@ -33,3 +33,4 @@ {"CatalogException", "BinderException"} ) DUCKDB_MISSING_OBJECT_MARKER: str = "does not exist" +CANCELLABLE_DURATION_LIMIT_SECONDS: int = 2_147_483 diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index e0356b7fd3..e8e55487c9 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -41,6 +41,7 @@ from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( + CANCELLABLE_DURATION_LIMIT_SECONDS, DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, ) @@ -158,6 +159,8 @@ def get_columns_for_relations( adapter=self, connection=connection, relations=relations ) + execution_duration_limit_seconds: ClassVar[int | None] = CANCELLABLE_DURATION_LIMIT_SECONDS + statement_cancellation_is_asynchronous: ClassVar[bool] = True adapter_name: ClassVar[str] = BuiltinAdapter.BIGQUERY.value sql_analysis_dialect_name: ClassVar[str | None] = "bigquery" sql_lexical_syntax: ClassVar[SqlLexicalSyntax] = SqlLexicalSyntax( @@ -1505,11 +1508,22 @@ def close(self, connection: _BigQueryConnection) -> None: def supports_transactions(self) -> bool: return False + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Keep unit-test comparisons inline; this adapter stores no test outputs.""" + + del connection, sql, name + return None + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Nothing to drop: no test outputs are stored.""" + + del connection, relation + def interrupt_connection(self, connection: Any) -> bool: - """Cancel a connection's running statement from another thread; False if unsupported.""" + """Request cancellation of the connection's jobs; BigQuery stops them asynchronously.""" - del connection - return False + connection.cancel_active_jobs() + return True def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py index 789a0cb8c4..4b352a0d7b 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py @@ -1,10 +1,17 @@ """BigQuery connection wrapper.""" +import math +import threading from typing import Any from sqlbuild.adapters.bigquery._helpers.statement_telemetry import affected_rows from sqlbuild.adapters.bigquery.classes.bigquery_cursor import _BigQueryCursor from sqlbuild.adapters.bigquery.constants import SELECT_STATEMENT_TYPE +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( + current_execution_deadline, +) +from sqlbuild.runtime.execution_limits.models import ExecutionDeadline from sqlbuild.runtime.observability.classes.statement_lifecycle import StatementLifecycle @@ -14,6 +21,8 @@ class _BigQueryConnection: def __init__(self, *, client: Any, location: str | None) -> None: self.client: Any = client self.location: str | None = location + self._active_jobs: dict[int, Any] = {} + self._jobs_lock: threading.Lock = threading.Lock() def execute(self, sql: str) -> _BigQueryCursor: with StatementLifecycle(adapter="bigquery", sql=sql, intent="execute") as lifecycle: @@ -21,7 +30,7 @@ def execute(self, sql: str) -> _BigQueryCursor: job_id: str | None = _job_id(job=job) lifecycle.submitted(job_id=job_id) try: - cursor: _BigQueryCursor = _BigQueryCursor(job.result()) + cursor: _BigQueryCursor = _BigQueryCursor(self.job_result(job)) except Exception as error: lifecycle.failed(error=error, job_id=job_id) raise @@ -34,7 +43,38 @@ def execute(self, sql: str) -> _BigQueryCursor: return cursor def query_job(self, sql: str) -> Any: - return self.client.query(sql, location=self.location) + deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() + if deadline is None: + return self.client.query(sql, location=self.location) + deadline.enforce() + from google.cloud import bigquery + + return self.client.query( + sql, + location=self.location, + job_config=bigquery.QueryJobConfig( + job_timeout_ms=max(1, math.ceil(deadline.remaining_seconds * 1000)) + ), + ) + + def job_result(self, job: Any) -> Any: + """Wait for a job while it is registered for cancellation at a deadline.""" + + with self._jobs_lock: + self._active_jobs[id(job)] = job + try: + return job.result() + finally: + with self._jobs_lock: + self._active_jobs.pop(id(job), None) + + def cancel_active_jobs(self) -> None: + """Request cancellation of every job this connection is waiting on.""" + + with self._jobs_lock: + jobs: list[Any] = list(self._active_jobs.values()) + for job in jobs: + job.cancel() def close(self) -> None: close: object | None = getattr(self.client, "close", None) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 00caaef22e..164f32b272 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -40,6 +40,7 @@ from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( + CANCELLABLE_DURATION_LIMIT_SECONDS, DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, ) @@ -147,6 +148,8 @@ def get_columns_for_relations( adapter=self, connection=connection, relations=relations ) + execution_duration_limit_seconds: ClassVar[int | None] = CANCELLABLE_DURATION_LIMIT_SECONDS + statement_cancellation_is_asynchronous: ClassVar[bool] = True adapter_name: ClassVar[str] = BuiltinAdapter.DATABRICKS.value sql_analysis_dialect_name: ClassVar[str | None] = "databricks" sql_lexical_syntax: ClassVar[SqlLexicalSyntax] = SqlLexicalSyntax( @@ -1163,11 +1166,22 @@ def render_apply_historical_check_snapshot_changes( invalidate_hard_deletes=invalidate_hard_deletes, ) + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Keep unit-test comparisons inline; this adapter stores no test outputs.""" + + del connection, sql, name + return None + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Nothing to drop: no test outputs are stored.""" + + del connection, relation + def interrupt_connection(self, connection: Any) -> bool: - """Cancel a connection's running statement from another thread; False if unsupported.""" + """Request cancellation of the connection's running statements.""" - del connection - return False + connection.cancel_active_statements() + return True def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_connection.py b/src/sqlbuild/adapters/databricks/classes/databricks_connection.py index bf0efc8e7d..3938413448 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_connection.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_connection.py @@ -1,8 +1,17 @@ """Databricks connection wrapper.""" +import math +import threading from typing import Any from sqlbuild.adapter.contract.classes.observed_cursor import ObservedCursor +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( + current_execution_deadline, +) +from sqlbuild.runtime.execution_limits.models import ExecutionDeadline + +_MAX_STATEMENT_TIMEOUT_SECONDS: int = 172_800 class _DatabricksConnection: @@ -10,13 +19,54 @@ class _DatabricksConnection: def __init__(self, raw_connection: Any) -> None: self.raw_connection: Any = raw_connection + self._statement_timeout_set: bool = False + self._active_cursors: dict[int, Any] = {} + self._cursors_lock: threading.Lock = threading.Lock() def execute(self, sql: str) -> Any: + self._apply_statement_timeout() cursor: ObservedCursor = self.cursor() - return cursor.execute(sql) + with self._cursors_lock: + self._active_cursors[id(cursor.raw_cursor)] = cursor.raw_cursor + try: + return cursor.execute(sql) + finally: + with self._cursors_lock: + self._active_cursors.pop(id(cursor.raw_cursor), None) def close(self) -> None: self.raw_connection.close() def cursor(self) -> ObservedCursor: return ObservedCursor(raw_cursor=self.raw_connection.cursor(), adapter="databricks") + + def cancel_active_statements(self) -> None: + """Ask Databricks to cancel each statement this connection is running.""" + + with self._cursors_lock: + cursors: list[Any] = list(self._active_cursors.values()) + for raw_cursor in cursors: + raw_cursor.cancel() + + def _apply_statement_timeout(self) -> None: + """Bound the next statement by the remaining deadline, or restore the session default.""" + + deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() + if deadline is None: + if self._statement_timeout_set: + self._run_setting("RESET STATEMENT_TIMEOUT") + self._statement_timeout_set = False + return + deadline.enforce() + seconds: int = min( + _MAX_STATEMENT_TIMEOUT_SECONDS, max(1, math.ceil(deadline.remaining_seconds)) + ) + self._run_setting(f"SET STATEMENT_TIMEOUT = {seconds}") + self._statement_timeout_set = True + + def _run_setting(self, sql: str) -> None: + raw_cursor: Any = self.raw_connection.cursor() + try: + raw_cursor.execute(sql) + finally: + raw_cursor.close() diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 793d5fef84..26d4744c89 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -41,6 +41,7 @@ from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( + CANCELLABLE_DURATION_LIMIT_SECONDS, DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, QUALIFIED_NAME_SEPARATOR, @@ -126,6 +127,7 @@ def get_columns_for_relations( ) adapter_name: ClassVar[str] = BuiltinAdapter.POSTGRES.value + execution_duration_limit_seconds: ClassVar[int | None] = CANCELLABLE_DURATION_LIMIT_SECONDS sql_analysis_dialect_name: ClassVar[str | None] = "postgres" sql_lexical_syntax: ClassVar[SqlLexicalSyntax] = SqlLexicalSyntax( escape_string_prefix=True, nested_block_comments=True @@ -140,6 +142,17 @@ def get_columns_for_relations( historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, ) + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Store a unit test's actual output once in a session temporary table.""" + + _ = self.execute(connection=connection, sql=f"CREATE TEMPORARY TABLE {name} AS {sql}") + return name + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Drop the temporary table `create_sql_test_output` created.""" + + _ = self.execute(connection=connection, sql=f"DROP TABLE IF EXISTS {relation}") + def interrupt_connection(self, connection: Any) -> bool: """Cancel a connection's running statement from another thread; False if unsupported.""" diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_connection.py b/src/sqlbuild/adapters/postgres/classes/postgres_connection.py index cee1a1f71f..eb61f542f7 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_connection.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_connection.py @@ -1,8 +1,14 @@ """PostgreSQL connection wrapper.""" +import math from typing import Any from sqlbuild.adapter.contract.classes.observed_cursor import ObservedCursor +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( + current_execution_deadline, +) +from sqlbuild.runtime.execution_limits.models import ExecutionDeadline class _PostgresConnection: @@ -10,8 +16,10 @@ class _PostgresConnection: def __init__(self, raw_connection: Any) -> None: self.raw_connection: Any = raw_connection + self._statement_timeout_set: bool = False def execute(self, sql: str) -> Any: + self._apply_statement_timeout() cursor: ObservedCursor = self.cursor() cursor.execute(sql) return cursor @@ -21,3 +29,17 @@ def cursor(self) -> ObservedCursor: def close(self) -> None: self.raw_connection.close() + + def _apply_statement_timeout(self) -> None: + """Bound the next statement by the remaining deadline, or restore the session default.""" + + deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() + if deadline is None: + if self._statement_timeout_set: + self.raw_connection.execute("RESET statement_timeout") + self._statement_timeout_set = False + return + deadline.enforce() + milliseconds: int = max(1, math.ceil(deadline.remaining_seconds * 1000)) + self.raw_connection.execute(f"SET statement_timeout = {milliseconds}") + self._statement_timeout_set = True diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_active_statements.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_active_statements.py new file mode 100644 index 0000000000..dd44d0e05d --- /dev/null +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_active_statements.py @@ -0,0 +1,45 @@ +"""The Snowflake statements one connection is running, so a deadline can cancel them.""" + +from __future__ import annotations + +import threading +from collections.abc import Generator +from contextlib import contextmanager +from typing import Any + +from sqlbuild.adapters.snowflake.exceptions import SnowflakeCancellationError + + +class SnowflakeActiveStatements: + """Running cursors of one connection, with the SQL each submitted.""" + + def __init__(self) -> None: + self._running: dict[int, tuple[Any, str]] = {} + self._lock: threading.Lock = threading.Lock() + + @contextmanager + def running(self, *, raw_cursor: Any, sql: str) -> Generator[None]: + """Record `sql` running on `raw_cursor` for the duration of the block.""" + + with self._lock: + self._running[id(raw_cursor)] = (raw_cursor, sql) + try: + yield + finally: + with self._lock: + _ = self._running.pop(id(raw_cursor), None) + + def cancel(self, *, raw_connection: Any) -> None: + """Ask Snowflake to abort every running statement; raise if it refuses.""" + + with self._lock: + running: list[tuple[Any, str]] = list(self._running.values()) + for raw_cursor, sql in running: + request_id: object | None = getattr(raw_cursor, "_request_id", None) + if request_id is None: + continue + response: object = raw_connection._cancel_query(sql, request_id) + if isinstance(response, dict) and response.get("success") is False: + raise SnowflakeCancellationError( + request_id=request_id, message=response.get("message") + ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index c6affbc356..98145c003f 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -2837,10 +2837,25 @@ def render_cursor_bound_literal(self, *, value: str, cursor_type: str | None) -> return f"'{value}'" def interrupt_connection(self, connection: Any) -> bool: - """Cancel a connection's running statement from another thread; False if unsupported.""" + """Abort the connection's running statements through Snowflake's cancel request.""" - del connection - return False + connection.active_statements.cancel(raw_connection=connection.raw_connection) + return True + + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Run the output query once and read its persisted result through RESULT_SCAN.""" + + del name + cursor: Any = self.execute(connection=connection, sql=sql) + query_id: object | None = getattr(cursor, "sfqid", None) + if not isinstance(query_id, str) or not query_id: + return None + return f"TABLE(RESULT_SCAN('{query_id}'))" + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Persisted query results expire on their own; there is nothing to drop.""" + + del connection, relation def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py index 1ee72c1948..9c1b2ba851 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py @@ -3,6 +3,9 @@ from typing import Any from sqlbuild.adapters.snowflake._helpers.session_context import may_switch_session_database +from sqlbuild.adapters.snowflake.classes.snowflake_active_statements import ( + SnowflakeActiveStatements, +) from sqlbuild.adapters.snowflake.classes.snowflake_cursor import _SnowflakeCursor from sqlbuild.adapters.snowflake.constants import CURRENT_DATABASE_ATTRIBUTE @@ -12,6 +15,7 @@ class _SnowflakeConnection: def __init__(self, raw_connection: Any) -> None: self.raw_connection: Any = raw_connection + self.active_statements: SnowflakeActiveStatements = SnowflakeActiveStatements() def execute(self, sql: str, *, statement_params: dict[str, str] | None = None) -> Any: kwargs: dict[str, object] = {} @@ -24,7 +28,9 @@ def close(self) -> None: def cursor(self) -> _SnowflakeCursor: return _SnowflakeCursor( - self.raw_connection.cursor(), on_executed=self._forget_changed_session_context + self.raw_connection.cursor(), + on_executed=self._forget_changed_session_context, + statements=self.active_statements, ) def _forget_changed_session_context(self, *, sql: str) -> None: diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py index 94d6b96685..593e4c19c9 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py @@ -13,7 +13,11 @@ from uuid import uuid4 from sqlbuild._native import SQLBUILD_QUERY_TAG_APP, CostContext, record_statement +from sqlbuild.adapters.snowflake.classes.snowflake_active_statements import ( + SnowflakeActiveStatements, +) from sqlbuild.cost.models import CostResourceContext, StatementExecutionTelemetry +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( current_execution_deadline, ) @@ -31,7 +35,10 @@ _STATEMENT_PARAMETER_ERROR_FRAGMENT: str = "STATEMENT PARAMETER" _RAW_CURSOR_ATTRIBUTE: str = "raw_cursor" _ON_EXECUTED_ATTRIBUTE: str = "on_executed" -_OWN_ATTRIBUTES: frozenset[str] = frozenset({_RAW_CURSOR_ATTRIBUTE, _ON_EXECUTED_ATTRIBUTE}) +_STATEMENTS_ATTRIBUTE: str = "statements" +_OWN_ATTRIBUTES: frozenset[str] = frozenset( + {_RAW_CURSOR_ATTRIBUTE, _ON_EXECUTED_ATTRIBUTE, _STATEMENTS_ATTRIBUTE} +) _EXECUTEMANY_INTENT: str = "executemany" @@ -45,9 +52,16 @@ def _statement_id() -> str: class _SnowflakeCursor: """Delegate cursor operations while tagging and recording executed statements.""" - def __init__(self, raw_cursor: Any, *, on_executed: StatementListener | None = None) -> None: + def __init__( + self, + raw_cursor: Any, + *, + on_executed: StatementListener | None = None, + statements: SnowflakeActiveStatements | None = None, + ) -> None: object.__setattr__(self, _RAW_CURSOR_ATTRIBUTE, raw_cursor) object.__setattr__(self, _ON_EXECUTED_ATTRIBUTE, on_executed) + object.__setattr__(self, _STATEMENTS_ATTRIBUTE, statements or SnowflakeActiveStatements()) def __setattr__(self, name: str, value: Any) -> None: if name in _OWN_ATTRIBUTES: @@ -129,7 +143,11 @@ def _execute_with_telemetry( ), ) as lifecycle: try: - result: Any = operation(sql, *args, **kwargs) + statements: SnowflakeActiveStatements = object.__getattribute__( + self, _STATEMENTS_ATTRIBUTE + ) + with statements.running(raw_cursor=self.raw_cursor, sql=sql): + result: Any = operation(sql, *args, **kwargs) except Exception as error: if query_tag_injected and _can_retry_without_query_tag( cursor=self.raw_cursor, @@ -308,7 +326,7 @@ def _with_query_tag( def _with_execution_deadline(*, kwargs: dict[str, Any]) -> dict[str, Any]: - deadline: ExecutionDeadline | None = current_execution_deadline() + deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() if deadline is None: return dict(kwargs) deadline.enforce() diff --git a/src/sqlbuild/adapters/snowflake/exceptions.py b/src/sqlbuild/adapters/snowflake/exceptions.py new file mode 100644 index 0000000000..d9d61f99c6 --- /dev/null +++ b/src/sqlbuild/adapters/snowflake/exceptions.py @@ -0,0 +1,10 @@ +"""Snowflake adapter exceptions.""" + + +class SnowflakeCancellationError(Exception): + """Raised when Snowflake refuses a request to cancel a running statement.""" + + code: str = "C417" + + def __init__(self, *, request_id: object, message: object) -> None: + super().__init__(f"Snowflake refused to cancel request {request_id}: {message}") diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index af95fd3413..c7628b588b 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -42,6 +42,7 @@ from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( + CANCELLABLE_DURATION_LIMIT_SECONDS, DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE, ) @@ -139,6 +140,7 @@ def get_columns_for_relations( adapter=self, connection=connection, relations=relations ) + execution_duration_limit_seconds: ClassVar[int | None] = CANCELLABLE_DURATION_LIMIT_SECONDS adapter_name: ClassVar[str] = BuiltinAdapter.SQLSERVER.value sql_analysis_dialect_name: ClassVar[str | None] = "tsql" sql_lexical_syntax: ClassVar[SqlLexicalSyntax] = SqlLexicalSyntax(nested_block_comments=True) @@ -2797,11 +2799,22 @@ def supports_table_functions(self) -> bool: def supports_unqualified_function_fingerprints(self) -> bool: return False + def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: + """Keep unit-test comparisons inline; this adapter stores no test outputs.""" + + del connection, sql, name + return None + + def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: + """Nothing to drop: no test outputs are stored.""" + + del connection, relation + def interrupt_connection(self, connection: Any) -> bool: - """Cancel a connection's running statement from another thread; False if unsupported.""" + """Cancel the connection's running statement through the driver.""" - del connection - return False + connection.cancel_active_statements() + return True def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_connection.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_connection.py index 47f9019295..f8c0bbc59e 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_connection.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_connection.py @@ -1,8 +1,14 @@ """SQL Server connection wrapper.""" +import threading from typing import Any from sqlbuild.adapter.contract.classes.observed_cursor import ObservedCursor +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( + current_execution_deadline, +) +from sqlbuild.runtime.execution_limits.models import ExecutionDeadline class _SqlServerConnection: @@ -10,14 +16,33 @@ class _SqlServerConnection: def __init__(self, raw_connection: Any) -> None: self.raw_connection: Any = raw_connection + self._running: int = 0 + self._running_lock: threading.Lock = threading.Lock() def execute(self, sql: str) -> Any: + deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() + if deadline is not None: + deadline.enforce() cursor: ObservedCursor = self.cursor() - cursor.execute(sql) + with self._running_lock: + self._running += 1 + try: + cursor.execute(sql) + finally: + with self._running_lock: + self._running -= 1 return cursor def cursor(self) -> ObservedCursor: return ObservedCursor(raw_cursor=self.raw_connection.cursor(), adapter="sqlserver") + def cancel_active_statements(self) -> None: + """Cancel the running statement through the driver's DB-Library cancel.""" + + with self._running_lock: + running: bool = self._running > 0 + if running: + self.raw_connection._conn.cancel() + def close(self) -> None: self.raw_connection.close() diff --git a/src/sqlbuild/executor/build/_helpers/model_timeouts.py b/src/sqlbuild/executor/build/_helpers/model_timeouts.py new file mode 100644 index 0000000000..22fed6f5a7 --- /dev/null +++ b/src/sqlbuild/executor/build/_helpers/model_timeouts.py @@ -0,0 +1,45 @@ +"""Per-model execution limits: the policy a model attempt runs under and its timeout result.""" + +from __future__ import annotations + +import dataclasses + +from sqlbuild._native import ModelPlanEntry +from sqlbuild.executor.run.models import ModelExecutionResult +from sqlbuild.runtime.execution_limits.constants import ITEM_TIMEOUT_CODE +from sqlbuild.runtime.execution_limits.main.timeout_report_lines import timeout_report_lines +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, ItemTimeoutReport + + +def model_timeout_policy(entry: ModelPlanEntry) -> ItemTimeoutPolicy | None: + """The model's `max_duration`, or `None` when it has none or it is disabled.""" + + if entry.max_duration is None or entry.max_duration_seconds is None: + return None + return ItemTimeoutPolicy( + limit=entry.max_duration, + seconds=entry.max_duration_seconds, + source=entry.max_duration_source or "max_duration", + ) + + +def timed_out_model_result( + *, result: ModelExecutionResult, report: ItemTimeoutReport +) -> ModelExecutionResult: + """The failed attempt reported as a timeout, keeping its phase and failed SQL.""" + + lines: list[str] = [ + *timeout_report_lines(report=report, label="Model"), + "", + "SQLBuild stopped this attempt and ran no further statements for it; models that", + "depend on it are skipped as for any failure.", + "", + "To allow more time, raise max_duration in the model's MODEL() header, its", + '[path_defaults] entry or [defaults] in sqlbuild_project.toml, or set it to "disabled".', + ] + return dataclasses.replace( + result, + error_code=ITEM_TIMEOUT_CODE, + error_message="\n".join(lines), + error_help=None, + ) diff --git a/src/sqlbuild/executor/build/classes/build_scheduler.py b/src/sqlbuild/executor/build/classes/build_scheduler.py index 8eb61acfeb..b3d85ed1f4 100644 --- a/src/sqlbuild/executor/build/classes/build_scheduler.py +++ b/src/sqlbuild/executor/build/classes/build_scheduler.py @@ -50,6 +50,10 @@ from sqlbuild.executor.build._helpers.blocking import downstream_blocked_keys from sqlbuild.executor.build._helpers.end_audits import run_end_audits from sqlbuild.executor.build._helpers.indexes import build_execution_indexes +from sqlbuild.executor.build._helpers.model_timeouts import ( + model_timeout_policy, + timed_out_model_result, +) from sqlbuild.executor.build._helpers.retention import ( apply_table_type_conversion, materialization_recreates_relation, @@ -122,6 +126,8 @@ from sqlbuild.runtime.execution_limits.main.enforce_execution_deadline import ( enforce_execution_deadline, ) +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope +from sqlbuild.runtime.execution_limits.types import ItemKind from sqlbuild.runtime.observability.classes.resource_attempt_lifecycle import ( ResourceAttemptLifecycle, ) @@ -921,6 +927,7 @@ def _execute_test_node( adapter=self._adapter, connection=connection, quality_scope="model", + timeouts=self._runtime.test_timeouts, ) duration: int = int((time.monotonic() - start) * 1000) log_debug_event( @@ -1014,10 +1021,19 @@ def _execute_model_node( sqlbuild_event="start", sqlbuild_kind=model_entry.materialization_type, ) - with diagnostics_context( - sqlbuild_subject="model", - sqlbuild_name=model_entry.name, - sqlbuild_kind=model_entry.materialization_type, + with ( + diagnostics_context( + sqlbuild_subject="model", + sqlbuild_name=model_entry.name, + sqlbuild_kind=model_entry.materialization_type, + ), + item_deadline_scope( + kind=ItemKind.MODEL, + name=model_entry.name, + policy=model_timeout_policy(model_entry), + adapter=self._adapter, + connection=connection, + ) as deadline, ): try: table_type_entry: TableTypePlanEntry | None = self._table_type_entries.get( @@ -1111,6 +1127,8 @@ def _execute_model_node( error_code=BUILD_WORKER_FAILED_CODE, error_message=str(error), ) + if deadline is not None and deadline.expired and result.status != ExecutionStatus.SUCCESS: + result = timed_out_model_result(result=result, report=deadline.report()) duration: int = int((time.monotonic() - start) * 1000) completed_result: ModelExecutionResult = dataclasses.replace(result, duration_ms=duration) log_debug_event( diff --git a/src/sqlbuild/executor/build/models.py b/src/sqlbuild/executor/build/models.py index 6ef91accb1..c56c637503 100644 --- a/src/sqlbuild/executor/build/models.py +++ b/src/sqlbuild/executor/build/models.py @@ -35,6 +35,7 @@ from sqlbuild.provider.main.runtime import ProviderContainer from sqlbuild.python_nodes.models import SqlResourceRef from sqlbuild.runtime.contracts.types import ConnectionElapsedCallback, NodeStartCallback +from sqlbuild.runtime.execution_limits.models import RunItemTimeouts @dataclass(frozen=True) @@ -94,6 +95,7 @@ class BuildRuntimeParams: ) = None microbatch_lease_check: Callable[[], None] | None = None project_relations: Mapping[SqlResourceRef, str] | None = None + test_timeouts: RunItemTimeouts | None = None @dataclass(frozen=True) diff --git a/src/sqlbuild/executor/pipeline/_helpers/scenario.py b/src/sqlbuild/executor/pipeline/_helpers/scenario.py index e7e8ec9329..e4d884f4d6 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/scenario.py +++ b/src/sqlbuild/executor/pipeline/_helpers/scenario.py @@ -37,6 +37,10 @@ close_connections, open_worker_connections, ) +from sqlbuild.executor.pipeline._helpers.scenario_timeouts import ( + scenario_timeout, + timed_out_scenario_result, +) from sqlbuild.executor.pipeline._helpers.settings import resolve_promotion_mode from sqlbuild.executor.scenario.classes.prepared_scenario_schemas import PreparedScenarioSchemas from sqlbuild.executor.scenario.classes.running_scenarios import RunningScenarios @@ -59,10 +63,14 @@ ScenarioRunResult, ScenarioSnapshotCaptureRunResult, ScenarioSnapshotStateResult, + ScenarioTestSettings, ) from sqlbuild.executor.scenario.types import ScenarioLocalRunStatus, ScenarioSnapshotState from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.runtime.contracts.models import ConnectionHooks +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy +from sqlbuild.runtime.execution_limits.types import ItemKind from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.classes.resource_attempt_lifecycle import ( ResourceAttemptLifecycle, @@ -123,8 +131,7 @@ def run_scenario_test_pipeline( connection_config: dict[str, object], adapter: BaseAdapter, project_name: str, - retain: bool, - concurrency: int = 1, + settings: ScenarioTestSettings, connection_hooks: ConnectionHooks | None = None, on_scenario_start: Callable[[CompiledSqlScenario], None] | None = None, on_scenario_complete: Callable[ @@ -134,7 +141,13 @@ def run_scenario_test_pipeline( ) -> tuple[ScenarioRunResult, ...]: """Execute selected scenarios from a compiled project.""" - worker_count: int = max(1, min(concurrency, len(scenarios))) + retain: bool = settings.retain + policies: dict[str, ItemTimeoutPolicy | None] = { + scenario.name: scenario_timeout(scenario=scenario, timeouts=settings.timeouts) + for scenario in scenarios + } + + worker_count: int = max(1, min(settings.concurrency, len(scenarios))) running: RunningScenarios = RunningScenarios(adapter=adapter) scheduling: _ScenarioScheduling = _ScenarioScheduling( worker_count=worker_count, @@ -162,14 +175,24 @@ def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: connection: Any = connection_pool.get() running.start(scenario_plan=scenario_plan, connection=connection) try: - return execute_scenario_run( - scenario_plan=scenario_plan, + with item_deadline_scope( + kind=ItemKind.SCENARIO, + name=scenario_plan.name, + policy=policies.get(scenario_plan.name), adapter=adapter, connection=connection, - run_id=pipeline_result.project.run_id, - retain=retain, - options=options, - ) + ) as deadline: + result: ScenarioRunResult = execute_scenario_run( + scenario_plan=scenario_plan, + adapter=adapter, + connection=connection, + run_id=pipeline_result.project.run_id, + retain=retain, + options=options, + ) + if deadline is not None and deadline.expired: + return timed_out_scenario_result(result=result, report=deadline.report()) + return result finally: running.finish(connection=connection) connection_pool.put(connection) diff --git a/src/sqlbuild/executor/pipeline/_helpers/scenario_timeouts.py b/src/sqlbuild/executor/pipeline/_helpers/scenario_timeouts.py new file mode 100644 index 0000000000..95e50be276 --- /dev/null +++ b/src/sqlbuild/executor/pipeline/_helpers/scenario_timeouts.py @@ -0,0 +1,54 @@ +"""Per-scenario time limits and the result of a scenario that ran out of time.""" + +from __future__ import annotations + +import dataclasses + +from sqlbuild._native import CompiledSqlScenario, execution_limit_seconds +from sqlbuild.executor.scenario.models import ScenarioRunResult +from sqlbuild.executor.scheduling.types import ExecutionStatus +from sqlbuild.runtime.execution_limits.constants import ITEM_TIMEOUT_CODE +from sqlbuild.runtime.execution_limits.main.timeout_report_lines import timeout_report_lines +from sqlbuild.runtime.execution_limits.models import ( + ItemTimeoutPolicy, + ItemTimeoutReport, + RunItemTimeouts, +) + +_TIMEOUT_KEY: str = "timeout" + + +def scenario_timeout( + *, scenario: CompiledSqlScenario, timeouts: RunItemTimeouts | None +) -> ItemTimeoutPolicy | None: + """The CLI limit, else the SCENARIO() timeout, else the target's or project's.""" + + declared: object = scenario.scenario_file.header_values.get(_TIMEOUT_KEY) + text: str | None = declared if isinstance(declared, str) else None + location: str = scenario.source_path.as_posix() if scenario.source_path else scenario.name + return (timeouts or RunItemTimeouts()).resolve( + declared=text, + declared_seconds=None if text is None else execution_limit_seconds(text), + declared_source=f"SCENARIO() timeout in {location}", + ) + + +def timed_out_scenario_result( + *, result: ScenarioRunResult, report: ItemTimeoutReport +) -> ScenarioRunResult: + """The scenario reported as timed out, keeping the step results it reached.""" + + lines: list[str] = [ + *timeout_report_lines(report=report, label="Scenario"), + "", + 'To allow more time, set SCENARIO(timeout "...") in the scenario file,', + "[settings] scenario_timeout in sqlbuild_project.toml (or sqlbuild_local.toml), or", + "--scenario-timeout for one run.", + ] + return dataclasses.replace( + result, + status=ExecutionStatus.FAILED, + error_code=ITEM_TIMEOUT_CODE, + error_help=None, + error_message="\n".join(lines), + ) diff --git a/src/sqlbuild/executor/pipeline/_helpers/testing.py b/src/sqlbuild/executor/pipeline/_helpers/testing.py index 0a8fd02066..be3e33156e 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/testing.py +++ b/src/sqlbuild/executor/pipeline/_helpers/testing.py @@ -36,6 +36,7 @@ from sqlbuild.executor.testing.main._execute import execute_sql_test from sqlbuild.executor.testing.models import SqlTestExecutionResult, StepResult from sqlbuild.executor.testing.types import SqlTestOutcome +from sqlbuild.runtime.execution_limits.models import RunItemTimeouts from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.classes.resource_attempt_lifecycle import ( ResourceAttemptLifecycle, @@ -63,6 +64,7 @@ def run_test_pipeline( callbacks: TestPipelineCallbacks | None = None, run_id: str | None = None, max_concurrency: int = 1, + timeouts: RunItemTimeouts | None = None, ) -> tuple[SqlTestExecutionResult, ...]: """Execute all SQL unit tests from a compiled plan.""" @@ -122,6 +124,7 @@ def run_test_pipeline( on_test_start=resolved_callbacks.on_test_start, on_test_complete=resolved_callbacks.on_test_complete, run_id=canonical_run_id, + timeouts=timeouts or RunItemTimeouts(), ) return results except BaseException as error: @@ -141,6 +144,7 @@ def _run_tests( on_test_start: Callable[[SqlTestPlanEntry], None] | None, on_test_complete: Callable[[SqlTestExecutionResult], None] | None, run_id: str, + timeouts: RunItemTimeouts, ) -> tuple[SqlTestExecutionResult, ...]: if worker_count == 1: serial_results: list[SqlTestExecutionResult] = [] @@ -153,6 +157,7 @@ def _run_tests( on_test_start=on_test_start, before_test_start=None, run_id=run_id, + timeouts=timeouts, ) serial_results.append(result) if on_test_complete is not None: @@ -180,6 +185,7 @@ def _run_tests( on_test_start=on_test_start, before_test_start=lambda index=index: start_coordinator.wait(index=index), run_id=run_id, + timeouts=timeouts, ), ) ) @@ -201,6 +207,7 @@ def _execute_test_entry_with_pooled_connection( on_test_start: Callable[[SqlTestPlanEntry], None] | None, before_test_start: Callable[[], None] | None, run_id: str, + timeouts: RunItemTimeouts, ) -> SqlTestExecutionResult: connection: Any = connection_pool.get() try: @@ -212,6 +219,7 @@ def _execute_test_entry_with_pooled_connection( on_test_start=on_test_start, before_test_start=before_test_start, run_id=run_id, + timeouts=timeouts, ) finally: connection_pool.put(connection) @@ -226,6 +234,7 @@ def _execute_test_entry( on_test_start: Callable[[SqlTestPlanEntry], None] | None, before_test_start: Callable[[], None] | None, run_id: str, + timeouts: RunItemTimeouts, ) -> SqlTestExecutionResult: with ResourceAttemptLifecycle( resource_id=sql_test_resource_id( @@ -248,7 +257,12 @@ def _execute_test_entry( missing_functions=missing_functions, ) if missing_functions - else execute_sql_test(test_entry=entry, adapter=adapter, connection=connection) + else execute_sql_test( + test_entry=entry, + adapter=adapter, + connection=connection, + timeouts=timeouts, + ) ) if result.outcome != SqlTestOutcome.PASS: lifecycle.failed(error_code=result.error_code) diff --git a/src/sqlbuild/executor/scenario/_helpers/execution/run.py b/src/sqlbuild/executor/scenario/_helpers/execution/run.py index 4eb7161627..d1ab571421 100644 --- a/src/sqlbuild/executor/scenario/_helpers/execution/run.py +++ b/src/sqlbuild/executor/scenario/_helpers/execution/run.py @@ -36,6 +36,9 @@ ScenarioStepResults, ) from sqlbuild.executor.scheduling.types import ExecutionStatus +from sqlbuild.runtime.execution_limits.main.item_deadline_suspended import ( + item_deadline_suspended, +) from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle _STEP_FAILED_MESSAGE: str = "scenario step failed" @@ -66,12 +69,13 @@ def execute_scenario_run_steps( if isinstance(error, (KeyboardInterrupt, ScenarioStopRequested)): options.interrupts.announce_stop() if not retain: - cleanup_interrupted_scenarios( - scenario_plans=(scenario_plan,), - adapter=adapter, - connection=connection, - interrupts=options.interrupts, - ) + with item_deadline_suspended(): + cleanup_interrupted_scenarios( + scenario_plans=(scenario_plan,), + adapter=adapter, + connection=connection, + interrupts=options.interrupts, + ) raise @@ -243,11 +247,12 @@ def _finish_scenario( if interrupts.stop_requested.is_set(): interrupts.announce_stop() if not retain and not interrupts.abandoned.is_set(): - cleanup_result = execute_scenario_cleanup( - scenario_plan=scenario_plan, - adapter=adapter, - connection=connection, - ) + with item_deadline_suspended(): + cleanup_result = execute_scenario_cleanup( + scenario_plan=scenario_plan, + adapter=adapter, + connection=connection, + ) if cleanup_result.status == ExecutionStatus.FAILED: status = ExecutionStatus.FAILED if error_code is None: diff --git a/src/sqlbuild/executor/scenario/models.py b/src/sqlbuild/executor/scenario/models.py index 22c2fd7251..6764fe43ad 100644 --- a/src/sqlbuild/executor/scenario/models.py +++ b/src/sqlbuild/executor/scenario/models.py @@ -21,6 +21,7 @@ from sqlbuild.executor.scenario.classes.scenario_interrupts import ScenarioInterrupts from sqlbuild.executor.scenario.types import ScenarioLocalRunStatus, ScenarioSnapshotState from sqlbuild.executor.scheduling.types import ExecutionStatus +from sqlbuild.runtime.execution_limits.models import RunItemTimeouts @dataclass(frozen=True) @@ -136,6 +137,15 @@ class ScenarioLocalReplaySource: capture_dialect: str | None = None +@dataclass(frozen=True) +class ScenarioTestSettings: + """Cleanup, concurrency and time-limit settings for one warehouse scenario test run.""" + + retain: bool + concurrency: int = 1 + timeouts: RunItemTimeouts | None = None + + @dataclass(frozen=True) class ScenarioCaptureSettings: """Capture provenance and safety settings for one scenario snapshot run.""" diff --git a/src/sqlbuild/executor/testing/_helpers/comparison.py b/src/sqlbuild/executor/testing/_helpers/comparison.py new file mode 100644 index 0000000000..6e92e77f9e --- /dev/null +++ b/src/sqlbuild/executor/testing/_helpers/comparison.py @@ -0,0 +1,72 @@ +"""Run one unit test's comparison against its actual outputs and collect step results.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild._native import SqlTestPlanEntry, render_sql_test_comparison_sql +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.executor.testing._helpers.difference_samples import add_difference_samples +from sqlbuild.executor.testing.classes.sql_test_actual_outputs import SqlTestActualOutputs +from sqlbuild.executor.testing.models import StepResult +from sqlbuild.executor.testing.types import SqlTestOutcome + + +def compare_actual_outputs( + *, + test_entry: SqlTestPlanEntry, + adapter: BaseAdapter, + connection: Any, + inline_sql: str, + outputs: SqlTestActualOutputs, +) -> list[StepResult]: + """Store the actual outputs once where supported, then count and sample the differences.""" + + relations: dict[int, str] = outputs.materialize() + comparison_sql: str = inline_sql + if relations: + (comparison_sql,) = render_sql_test_comparison_sql( + (test_entry,), + set_difference_operator=adapter.render_set_difference_operator(), + sql_analysis_dialect=adapter.sql_analysis_dialect(), + actual_relations=relations, + ) + cursor: Any = adapter.execute(connection=connection, sql=comparison_sql) + rows: list[Any] = cursor.fetchall() + return add_difference_samples( + step_results=_build_step_results(rows), + test_entry=test_entry, + adapter=adapter, + connection=connection, + relations=relations, + ) + + +def _build_step_results(rows: list[Any]) -> list[StepResult]: + """Convert comparison query rows into per-model step results.""" + + step_results: list[StepResult] = [] + row: Any + for row in sorted(rows, key=lambda item: int(item[0])): + model_name: str = str(row[1]) + actual_count: int = int(row[2]) + expected_count: int = int(row[3]) + unexpected_count: int = int(row[4]) + missing_count: int = int(row[5]) + outcome: SqlTestOutcome + if unexpected_count == 0 and missing_count == 0 and actual_count == expected_count: + outcome = SqlTestOutcome.PASS + else: + outcome = SqlTestOutcome.FAIL + step_results.append( + StepResult( + model_name=model_name, + outcome=outcome, + actual_row_count=actual_count, + expected_row_count=expected_count, + mismatched_row_count=unexpected_count, + unexpected_row_count=unexpected_count, + missing_row_count=missing_count, + ) + ) + return step_results diff --git a/src/sqlbuild/executor/testing/_helpers/difference_samples.py b/src/sqlbuild/executor/testing/_helpers/difference_samples.py index 9d6f02f246..05c6a0b36b 100644 --- a/src/sqlbuild/executor/testing/_helpers/difference_samples.py +++ b/src/sqlbuild/executor/testing/_helpers/difference_samples.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections.abc import Mapping from dataclasses import replace from typing import Any, Final @@ -16,6 +17,7 @@ StepResult, ) from sqlbuild.executor.testing.types import SqlTestDifferenceDirection, SqlTestOutcome +from sqlbuild.runtime.execution_limits.exceptions import ItemTimeoutError _ROW_LIMIT: Final[int] = 3 _COLUMN_LIMIT: Final[int] = 12 @@ -30,18 +32,23 @@ def add_difference_samples( test_entry: SqlTestPlanEntry, adapter: BaseAdapter, connection: Any, + relations: Mapping[int, str] | None = None, ) -> list[StepResult]: """Attach bounded, redacted samples to failed expected-output steps.""" sampled_results: list[StepResult] = [] - expected_steps: tuple[ChainStep, ...] = tuple( - step for step in test_entry.chain if step.expected_cte_sql is not None + expected_steps: tuple[tuple[int, ChainStep], ...] = tuple( + (index, step) + for index, step in enumerate(test_entry.chain) + if step.expected_cte_sql is not None ) + stored: Mapping[int, str] = relations or {} for step_index, step_result in enumerate(step_results): if step_index >= len(expected_steps) or step_result.outcome != SqlTestOutcome.FAIL: sampled_results.append(step_result) continue - step: ChainStep = expected_steps[step_index] + chain_index, step = expected_steps[step_index] + actual_relation: str | None = stored.get(chain_index) unexpected_samples: tuple[SqlTestDifferenceSample, ...] = () missing_samples: tuple[SqlTestDifferenceSample, ...] = () if step_result.unexpected_row_count is not None and step_result.unexpected_row_count > 0: @@ -51,6 +58,7 @@ def add_difference_samples( direction=SqlTestDifferenceDirection.UNEXPECTED, adapter=adapter, connection=connection, + actual_relation=actual_relation, ) if step_result.missing_row_count is not None and step_result.missing_row_count > 0: missing_samples = _best_effort_difference_samples( @@ -59,6 +67,7 @@ def add_difference_samples( direction=SqlTestDifferenceDirection.MISSING, adapter=adapter, connection=connection, + actual_relation=actual_relation, ) sampled_results.append( replace( @@ -107,6 +116,7 @@ def _fetch_difference_samples( direction: SqlTestDifferenceDirection, adapter: BaseAdapter, connection: Any, + actual_relation: str | None, ) -> tuple[SqlTestDifferenceSample, ...]: sql: str = render_sql_test_difference_sample_sql( test_entry=test_entry, @@ -115,6 +125,7 @@ def _fetch_difference_samples( sample_limit=_ROW_LIMIT, set_difference_operator=adapter.render_set_difference_operator(), sql_analysis_dialect=adapter.sql_analysis_dialect(), + actual_relation=actual_relation, ) cursor: Any = adapter.execute(connection=connection, sql=sql) description: Any | None = getattr(cursor, "description", None) @@ -144,6 +155,7 @@ def _best_effort_difference_samples( direction: SqlTestDifferenceDirection, adapter: BaseAdapter, connection: Any, + actual_relation: str | None, ) -> tuple[SqlTestDifferenceSample, ...]: """Return samples when available without replacing a known comparison failure.""" @@ -154,7 +166,10 @@ def _best_effort_difference_samples( direction=direction, adapter=adapter, connection=connection, + actual_relation=actual_relation, ) + except ItemTimeoutError: + raise except Exception as error: log_debug_event( logger=_LOGGER, diff --git a/src/sqlbuild/executor/testing/_helpers/expected_column_probe.py b/src/sqlbuild/executor/testing/_helpers/expected_column_probe.py index 27df21d025..6eb8de39d1 100644 --- a/src/sqlbuild/executor/testing/_helpers/expected_column_probe.py +++ b/src/sqlbuild/executor/testing/_helpers/expected_column_probe.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections.abc import Mapping from typing import Any from sqlbuild._native import ( @@ -18,7 +19,11 @@ def missing_expected_columns_message( - *, test_entry: SqlTestPlanEntry, adapter: BaseAdapter, connection: Any + *, + test_entry: SqlTestPlanEntry, + adapter: BaseAdapter, + connection: Any, + relations: Mapping[int, str] | None = None, ) -> str | None: """Probe actual step columns after a failed comparison and describe missing ones.""" @@ -26,7 +31,11 @@ def missing_expected_columns_message( if step.expected_columns is None or step.expected_cte_sql is None: continue available_columns: frozenset[str] | None = _probe_actual_columns( - test_entry=test_entry, step_index=step_index, adapter=adapter, connection=connection + test_entry=test_entry, + step_index=step_index, + adapter=adapter, + connection=connection, + relations=relations or {}, ) if available_columns is None: continue @@ -41,13 +50,19 @@ def missing_expected_columns_message( def _probe_actual_columns( - *, test_entry: SqlTestPlanEntry, step_index: int, adapter: BaseAdapter, connection: Any + *, + test_entry: SqlTestPlanEntry, + step_index: int, + adapter: BaseAdapter, + connection: Any, + relations: Mapping[int, str], ) -> frozenset[str] | None: (sql,) = render_sql_test_comparison_sql( (test_entry,), set_difference_operator=adapter.render_set_difference_operator(), sql_analysis_dialect=adapter.sql_analysis_dialect(), probe_step_index=step_index, + actual_relations=dict(relations), ) if not sql: return None diff --git a/src/sqlbuild/executor/testing/_helpers/timeout_policy.py b/src/sqlbuild/executor/testing/_helpers/timeout_policy.py new file mode 100644 index 0000000000..9b81d8caf5 --- /dev/null +++ b/src/sqlbuild/executor/testing/_helpers/timeout_policy.py @@ -0,0 +1,19 @@ +"""The time limit one SQL unit test runs under.""" + +from __future__ import annotations + +from sqlbuild._native import SqlTestPlanEntry +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, RunItemTimeouts + + +def sql_test_timeout( + *, entry: SqlTestPlanEntry, timeouts: RunItemTimeouts | None +) -> ItemTimeoutPolicy | None: + """The CLI limit, else the test's TEST() timeout, else the target's or project's.""" + + location: str = entry.source_path.as_posix() if entry.source_path else entry.name + return (timeouts or RunItemTimeouts()).resolve( + declared=entry.timeout, + declared_seconds=entry.timeout_seconds, + declared_source=f"TEST() timeout in {location}", + ) diff --git a/src/sqlbuild/executor/testing/_helpers/timeout_result.py b/src/sqlbuild/executor/testing/_helpers/timeout_result.py new file mode 100644 index 0000000000..d11266f569 --- /dev/null +++ b/src/sqlbuild/executor/testing/_helpers/timeout_result.py @@ -0,0 +1,103 @@ +"""The result of a SQL unit test that ran out of time.""" + +from __future__ import annotations + +from sqlbuild._native import SqlTestPlanEntry +from sqlbuild.executor.testing.constants import ( + INLINED_MODEL_DISPLAY_LIMIT, + MINUTE_SUGGESTION_THRESHOLD_SECONDS, + SECONDS_PER_MINUTE, + SQL_TEST_TIMEOUT_CODE, +) +from sqlbuild.executor.testing.models import SqlTestExecutionResult, StepResult +from sqlbuild.executor.testing.types import SqlTestOutcome +from sqlbuild.runtime.execution_limits.main.timeout_report_lines import timeout_report_lines +from sqlbuild.runtime.execution_limits.models import ItemTimeoutReport + +_SCENARIO_ADVICE: tuple[str, ...] = ( + " 2. If this test checks behaviour across several models, consider a scenario", + " (tests/scenarios/). A unit test inlines every real model into one statement,", + " so the warehouse compiles and optimises the whole chain at once; a scenario", + " builds each model step by step, compiling each statement separately.", +) + + +def timed_out_result( + *, test_entry: SqlTestPlanEntry, report: ItemTimeoutReport +) -> SqlTestExecutionResult: + """A TIMEOUT result naming the limit, its source, the cancellation and how to proceed.""" + + message: str = timeout_message(test_entry=test_entry, report=report) + model_name: str = next( + (step.model_name for step in test_entry.chain if step.expected_cte_sql is not None), + test_entry.name, + ) + return SqlTestExecutionResult( + test_name=test_entry.name, + outcome=SqlTestOutcome.TIMEOUT, + source_path=test_entry.source_path, + block_index=test_entry.block_index, + parent_name=test_entry.parent_name, + case_name=test_entry.case_name, + case_index=test_entry.case_index, + case_fingerprint=test_entry.case_fingerprint, + parameter_schema=test_entry.parameter_schema, + parameter_values=test_entry.parameter_values, + step_results=( + StepResult( + model_name=model_name, + outcome=SqlTestOutcome.TIMEOUT, + error_code=SQL_TEST_TIMEOUT_CODE, + error_message=message, + ), + ), + error_code=SQL_TEST_TIMEOUT_CODE, + error_message=message, + real_model_count=test_entry.real_model_count, + timeout=report, + ) + + +def timeout_message(*, test_entry: SqlTestPlanEntry, report: ItemTimeoutReport) -> str: + """Key facts first, then the options in order: mock, consider a scenario, raise the limit.""" + + lines: list[str] = [ + *timeout_report_lines(report=report, label="SQL unit test"), + ] + inlined: list[str] = [ + step.model_name for step in test_entry.chain if step.expected_cte_sql is None + ] + raised: str = _raised_limit(report) + raise_lines: list[str] = [ + f'set TEST(timeout "{raised}") on this test, [settings] test_timeout = "{raised}"', + "in sqlbuild_project.toml (or sqlbuild_local.toml), or --test-timeout for one run.", + ] + if test_entry.real_model_count is None or not inlined: + lines.extend(["", f"To allow more time, {raise_lines[0]}", raise_lines[1]]) + return "\n".join(lines) + shown: list[str] = inlined[-INLINED_MODEL_DISPLAY_LIMIT:] + more: str = f" and {len(inlined) - len(shown)} more" if len(inlined) > len(shown) else "" + lines.extend( + [ + "", + f"This test inlines {len(inlined)} real " + f"{'model' if len(inlined) == 1 else 'models'} beyond the model under test into one " + "statement:", + f" {', '.join(shown)}{more}", + "", + "Options:", + " 1. Mock the inlined models closest to the model under test to cut the chain,", + f" for example __ref__{shown[-1]} AS (SELECT ...).", + *_SCENARIO_ADVICE, + f" 3. To allow more time, {raise_lines[0]}", + f" {raise_lines[1]}", + ] + ) + return "\n".join(lines) + + +def _raised_limit(report: ItemTimeoutReport) -> str: + seconds: float = max(report.elapsed_seconds * 2, 1.0) + if seconds >= MINUTE_SUGGESTION_THRESHOLD_SECONDS: + return f"{int(seconds // SECONDS_PER_MINUTE) + 1}m" + return f"{int(seconds) + 1}s" diff --git a/src/sqlbuild/executor/testing/classes/__init__.py b/src/sqlbuild/executor/testing/classes/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py b/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py new file mode 100644 index 0000000000..9b9714a4a3 --- /dev/null +++ b/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py @@ -0,0 +1,104 @@ +"""Evaluate a unit test's tested chain once and let every comparison read the stored output.""" + +from __future__ import annotations + +import logging +from types import TracebackType +from typing import Any +from uuid import uuid4 + +from sqlbuild._native import SqlTestPlanEntry, render_sql_test_comparison_sql +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.diagnostics.classes.diagnostic_record_redactor import DiagnosticRecordRedactor +from sqlbuild.diagnostics.main.log_debug_event import log_debug_event +from sqlbuild.runtime.execution_limits.exceptions import ( + ExecutionDurationLimitError, + ItemTimeoutError, +) +from sqlbuild.runtime.execution_limits.main.item_deadline_suspended import ( + item_deadline_suspended, +) + +_LOGGER: logging.Logger = logging.getLogger("sqlbuild.execution") +_OUTPUT_PREFIX: str = "__sqb_test_output" + + +class SqlTestActualOutputs: + """The actual outputs of one test's compared chain steps, stored once and dropped after.""" + + def __init__( + self, *, test_entry: SqlTestPlanEntry, adapter: BaseAdapter, connection: Any + ) -> None: + self._test_entry: SqlTestPlanEntry = test_entry + self._adapter: BaseAdapter = adapter + self._connection: Any = connection + self._created: list[str] = [] + self.relations: dict[int, str] = {} + + def __enter__(self) -> SqlTestActualOutputs: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self._release() + + def materialize(self) -> dict[int, str]: + """Store each compared step's output, or return no relations to compare inline.""" + + token: str = uuid4().hex[:12] + try: + for index in range(len(self._test_entry.chain)): + (sql,) = render_sql_test_comparison_sql( + (self._test_entry,), + set_difference_operator=self._adapter.render_set_difference_operator(), + sql_analysis_dialect=self._adapter.sql_analysis_dialect(), + output_step_index=index, + actual_relations=dict(self.relations), + ) + if not sql: + continue + relation: str | None = self._adapter.create_sql_test_output( + connection=self._connection, + sql=sql, + name=f"{_OUTPUT_PREFIX}_{token}_{index}", + ) + if relation is None: + self._release() + return {} + if relation.startswith(_OUTPUT_PREFIX): + self._created.append(relation) + self.relations[index] = relation + except (ItemTimeoutError, ExecutionDurationLimitError): + raise + except Exception as error: + log_debug_event( + logger=_LOGGER, + message="SQL test output could not be stored; comparing inline", + test_name=self._test_entry.name, + sqlbuild_error=DiagnosticRecordRedactor.text(str(error)), + ) + self._release() + return {} + return dict(self.relations) + + def _release(self) -> None: + created: list[str] = list(self._created) + self._created.clear() + self.relations.clear() + with item_deadline_suspended(): + for relation in created: + try: + self._adapter.drop_sql_test_output( + connection=self._connection, relation=relation + ) + except Exception as error: + log_debug_event( + logger=_LOGGER, + message="SQL test output cleanup failed", + test_name=self._test_entry.name, + sqlbuild_error=DiagnosticRecordRedactor.text(str(error)), + ) diff --git a/src/sqlbuild/executor/testing/constants.py b/src/sqlbuild/executor/testing/constants.py index 5cabd16b52..515a77c461 100644 --- a/src/sqlbuild/executor/testing/constants.py +++ b/src/sqlbuild/executor/testing/constants.py @@ -1,8 +1,14 @@ """SQL test executor diagnostic codes.""" +from sqlbuild.runtime.execution_limits.constants import ITEM_TIMEOUT_CODE + SQL_TEST_TOO_LARGE_CODE: str = "T001" SQL_TEST_EXECUTION_ERROR_CODE: str = "T002" SQL_TEST_ASSERTION_FAILED_CODE: str = "T003" SQL_TEST_ENTRY_MISSING_CODE: str = "T004" SQL_TEST_EXPECTED_COLUMNS_CODE: str = "T005" SQL_TEST_SIZE_BYTES: str = "bytes" +SQL_TEST_TIMEOUT_CODE: str = ITEM_TIMEOUT_CODE +INLINED_MODEL_DISPLAY_LIMIT: int = 8 +MINUTE_SUGGESTION_THRESHOLD_SECONDS: int = 120 +SECONDS_PER_MINUTE: int = 60 diff --git a/src/sqlbuild/executor/testing/main/_execute.py b/src/sqlbuild/executor/testing/main/_execute.py index 69951e0ae0..6b1751f928 100644 --- a/src/sqlbuild/executor/testing/main/_execute.py +++ b/src/sqlbuild/executor/testing/main/_execute.py @@ -2,15 +2,19 @@ from __future__ import annotations +from dataclasses import replace from typing import Any from sqlbuild._native import SqlTestPlanEntry, render_sql_test_comparison_sql from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.compiler.compile.exceptions import CompileInputError -from sqlbuild.executor.testing._helpers.difference_samples import add_difference_samples +from sqlbuild.executor.testing._helpers.comparison import compare_actual_outputs from sqlbuild.executor.testing._helpers.expected_column_probe import ( missing_expected_columns_message, ) +from sqlbuild.executor.testing._helpers.timeout_policy import sql_test_timeout +from sqlbuild.executor.testing._helpers.timeout_result import timed_out_result +from sqlbuild.executor.testing.classes.sql_test_actual_outputs import SqlTestActualOutputs from sqlbuild.executor.testing.constants import ( SQL_TEST_ASSERTION_FAILED_CODE, SQL_TEST_EXECUTION_ERROR_CODE, @@ -22,6 +26,10 @@ ) from sqlbuild.executor.testing.models import SqlTestExecutionResult, StepResult from sqlbuild.executor.testing.types import SqlTestOutcome +from sqlbuild.runtime.execution_limits.exceptions import ItemTimeoutError +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope +from sqlbuild.runtime.execution_limits.models import RunItemTimeouts +from sqlbuild.runtime.execution_limits.types import ItemKind from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.models import OperationAttributes @@ -32,6 +40,39 @@ def execute_sql_test( adapter: BaseAdapter, connection: Any, quality_scope: str = "standalone", + timeouts: RunItemTimeouts | None = None, +) -> SqlTestExecutionResult: + """Execute one SQL unit test within its time limit, cancelling warehouse work at expiry.""" + + with item_deadline_scope( + kind=ItemKind.TEST, + name=test_entry.name, + policy=sql_test_timeout(entry=test_entry, timeouts=timeouts), + adapter=adapter, + connection=connection, + ) as deadline: + try: + result: SqlTestExecutionResult = _execute_sql_test( + test_entry=test_entry, + adapter=adapter, + connection=connection, + quality_scope=quality_scope, + ) + except ItemTimeoutError: + if deadline is None: + raise + return timed_out_result(test_entry=test_entry, report=deadline.report()) + if deadline is not None and deadline.expired: + return timed_out_result(test_entry=test_entry, report=deadline.report()) + return replace(result, real_model_count=test_entry.real_model_count) + + +def _execute_sql_test( + *, + test_entry: SqlTestPlanEntry, + adapter: BaseAdapter, + connection: Any, + quality_scope: str, ) -> SqlTestExecutionResult: """Execute one SQL unit test and collect bounded diagnostics for failed comparisons.""" @@ -82,20 +123,33 @@ def execute_sql_test( metadata={"item_count": len(test_entry.chain) + len(test_entry.assertions)}, attributes=OperationAttributes(phase="assert", target_kind="sql_test", scope=quality_scope), ) as lifecycle: - try: - cursor: Any = adapter.execute(connection=connection, sql=comparison_sql) - rows: list[Any] = cursor.fetchall() - step_results: list[StepResult] = _build_step_results(rows) - step_results = add_difference_samples( - step_results=step_results, - test_entry=test_entry, - adapter=adapter, - connection=connection, - ) - except Exception as error: - expected_columns_message: str | None = missing_expected_columns_message( - test_entry=test_entry, adapter=adapter, connection=connection - ) + step_results: list[StepResult] = [] + failure: tuple[Exception, str | None] | None = None + with SqlTestActualOutputs( + test_entry=test_entry, adapter=adapter, connection=connection + ) as outputs: + try: + step_results = compare_actual_outputs( + test_entry=test_entry, + adapter=adapter, + connection=connection, + inline_sql=comparison_sql, + outputs=outputs, + ) + except ItemTimeoutError: + raise + except Exception as error: + failure = ( + error, + missing_expected_columns_message( + test_entry=test_entry, + adapter=adapter, + connection=connection, + relations=outputs.relations, + ), + ) + if failure is not None: + error, expected_columns_message = failure error_code: str = ( SQL_TEST_EXECUTION_ERROR_CODE if expected_columns_message is None @@ -160,33 +214,3 @@ def execute_sql_test( else None, error_message=error_message, ) - - -def _build_step_results(rows: list[Any]) -> list[StepResult]: - """Convert comparison query rows into per-model step results.""" - - step_results: list[StepResult] = [] - row: Any - for row in sorted(rows, key=lambda item: int(item[0])): - model_name: str = str(row[1]) - actual_count: int = int(row[2]) - expected_count: int = int(row[3]) - unexpected_count: int = int(row[4]) - missing_count: int = int(row[5]) - outcome: SqlTestOutcome - if unexpected_count == 0 and missing_count == 0 and actual_count == expected_count: - outcome = SqlTestOutcome.PASS - else: - outcome = SqlTestOutcome.FAIL - step_results.append( - StepResult( - model_name=model_name, - outcome=outcome, - actual_row_count=actual_count, - expected_row_count=expected_count, - mismatched_row_count=unexpected_count, - unexpected_row_count=unexpected_count, - missing_row_count=missing_count, - ) - ) - return step_results diff --git a/src/sqlbuild/executor/testing/models.py b/src/sqlbuild/executor/testing/models.py index 40247d1a60..45187c57b9 100644 --- a/src/sqlbuild/executor/testing/models.py +++ b/src/sqlbuild/executor/testing/models.py @@ -7,6 +7,7 @@ from sqlbuild._native import SqlTestParameterDeclaration, SqlValue from sqlbuild.executor.testing.types import SqlTestOutcome +from sqlbuild.runtime.execution_limits.models import ItemTimeoutReport @dataclass(frozen=True) @@ -62,3 +63,5 @@ class SqlTestExecutionResult: error_code: str | None = None error_help: str | None = None error_message: str | None = None + real_model_count: int | None = None + timeout: ItemTimeoutReport | None = None diff --git a/src/sqlbuild/executor/testing/types.py b/src/sqlbuild/executor/testing/types.py index b0dc160aaa..5fd39cbf76 100644 --- a/src/sqlbuild/executor/testing/types.py +++ b/src/sqlbuild/executor/testing/types.py @@ -9,6 +9,7 @@ class SqlTestOutcome(StrEnum): PASS = "pass" FAIL = "fail" ERROR = "error" + TIMEOUT = "timeout" class SqlTestDifferenceDirection(StrEnum): diff --git a/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py b/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py new file mode 100644 index 0000000000..074f4ff640 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py @@ -0,0 +1,30 @@ +"""Ask an adapter to cancel a connection's running statement and classify the answer.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild.runtime.execution_limits.types import CancellationOutcome + + +def cancel_connection(*, adapter: Any, connection: Any) -> tuple[CancellationOutcome, str | None]: + """The cancellation outcome for one connection and a detail when it is not confirmed.""" + + adapter_name: str = str(getattr(adapter, "adapter_name", type(adapter).__name__)) + try: + requested: bool = bool(adapter.interrupt_connection(connection)) + except Exception as error: + return CancellationOutcome.FAILED, f"{adapter_name} cancellation failed: {error}" + if not requested: + return ( + CancellationOutcome.UNSUPPORTED, + f"adapter '{adapter_name}' cannot cancel a running statement; the warehouse may " + "still be running it", + ) + if getattr(adapter, "statement_cancellation_is_asynchronous", False): + return ( + CancellationOutcome.UNKNOWN, + f"cancellation was requested; {adapter_name} stops work asynchronously, so its " + "final state is not confirmed", + ) + return CancellationOutcome.CONFIRMED, None diff --git a/src/sqlbuild/runtime/execution_limits/_helpers/context.py b/src/sqlbuild/runtime/execution_limits/_helpers/context.py index 31f32b4e13..e0a259cc6c 100644 --- a/src/sqlbuild/runtime/execution_limits/_helpers/context.py +++ b/src/sqlbuild/runtime/execution_limits/_helpers/context.py @@ -2,9 +2,19 @@ from contextvars import ContextVar +from sqlbuild.runtime.execution_limits.classes.active_connections import ActiveConnections +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline from sqlbuild.runtime.execution_limits.models import ExecutionDeadline current_deadline_context: ContextVar[ExecutionDeadline | None] = ContextVar( "sqlbuild_execution_deadline", default=None, ) +current_item_deadline_context: ContextVar[ItemDeadline | None] = ContextVar( + "sqlbuild_item_deadline", + default=None, +) +current_invocation_connections_context: ContextVar[ActiveConnections | None] = ContextVar( + "sqlbuild_invocation_connections", + default=None, +) diff --git a/src/sqlbuild/runtime/execution_limits/classes/__init__.py b/src/sqlbuild/runtime/execution_limits/classes/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/sqlbuild/runtime/execution_limits/classes/active_connections.py b/src/sqlbuild/runtime/execution_limits/classes/active_connections.py new file mode 100644 index 0000000000..d4f2bac42b --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/classes/active_connections.py @@ -0,0 +1,34 @@ +"""Connections running statements under a deadline, so the deadline can cancel them.""" + +from __future__ import annotations + +import threading +from typing import Any + + +class ActiveConnections: + """Adapters and connections with a statement in flight, counted per connection.""" + + def __init__(self) -> None: + self._connections: dict[int, tuple[Any, Any, int]] = {} + self._lock: threading.Lock = threading.Lock() + + def add(self, *, adapter: Any, connection: Any) -> None: + with self._lock: + _, _, count = self._connections.get(id(connection), (adapter, connection, 0)) + self._connections[id(connection)] = (adapter, connection, count + 1) + + def remove(self, *, connection: Any) -> None: + with self._lock: + entry: tuple[Any, Any, int] | None = self._connections.get(id(connection)) + if entry is None: + return + adapter, _, count = entry + if count <= 1: + del self._connections[id(connection)] + else: + self._connections[id(connection)] = (adapter, connection, count - 1) + + def snapshot(self) -> list[tuple[Any, Any]]: + with self._lock: + return [(adapter, connection) for adapter, connection, _ in self._connections.values()] diff --git a/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py b/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py new file mode 100644 index 0000000000..a782c5e14a --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py @@ -0,0 +1,114 @@ +"""A per-item deadline shared by the executing thread and its expiry watchdog.""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +from sqlbuild.runtime.execution_limits._helpers.cancellation import cancel_connection +from sqlbuild.runtime.execution_limits.classes.active_connections import ActiveConnections +from sqlbuild.runtime.execution_limits.constants import ( + EXPIRY_REPORT_WAIT_SECONDS, + OUTCOME_SEVERITY, +) +from sqlbuild.runtime.execution_limits.exceptions import ItemTimeoutError +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, ItemTimeoutReport +from sqlbuild.runtime.execution_limits.types import CancellationOutcome, ItemKind + + +class ItemDeadline: + """One test, scenario or model attempt's deadline, its statements and its cancellation.""" + + def __init__(self, *, kind: ItemKind, name: str, policy: ItemTimeoutPolicy) -> None: + self.kind: ItemKind = kind + self.name: str = name + self.policy: ItemTimeoutPolicy = policy + self.started: float = time.monotonic() + self.deadline: float = self.started + policy.seconds + self.expired: bool = False + self.connections: ActiveConnections = ActiveConnections() + self._active_statements: int = 0 + self._awaiting_confirmation: bool = False + self._cancellation: CancellationOutcome | None = None + self._cancellation_detail: str | None = None + self._query_ids: list[str] = [] + self._lock: threading.Lock = threading.Lock() + self._expiring: threading.Lock = threading.Lock() + + @property + def remaining_seconds(self) -> float: + return max(0.0, self.deadline - time.monotonic()) + + def enforce(self) -> None: + if not self.expired and self.remaining_seconds > 0: + return + self.expired = True + raise ItemTimeoutError(kind=self.kind, name=self.name, limit=self.policy.limit) + + def statement_started(self) -> None: + with self._lock: + self._active_statements += 1 + + def statement_finished(self, *, query_id: str | None) -> None: + with self._lock: + self._active_statements = max(0, self._active_statements - 1) + if query_id: + self._query_ids.append(query_id) + + def statement_submitted(self, *, query_id: str | None) -> None: + if not query_id: + return + with self._lock: + self._query_ids.append(query_id) + + def expire(self, *, fallback: tuple[Any, Any]) -> None: + """Mark the deadline expired and cancel every statement still running under it.""" + + with self._expiring: + self._cancel_running(fallback=fallback) + + def _cancel_running(self, *, fallback: tuple[Any, Any]) -> None: + targets: list[tuple[Any, Any]] = self.connections.snapshot() + with self._lock: + self.expired = True + if not targets and self._active_statements: + targets = [fallback] + if not targets: + return + outcome: CancellationOutcome = CancellationOutcome.CONFIRMED + details: list[str] = [] + for adapter, connection in targets: + target_outcome, detail = cancel_connection(adapter=adapter, connection=connection) + if OUTCOME_SEVERITY[target_outcome] > OUTCOME_SEVERITY[outcome]: + outcome = target_outcome + if detail is not None and detail not in details: + details.append(detail) + with self._lock: + self._cancellation = outcome + self._awaiting_confirmation = outcome == CancellationOutcome.CONFIRMED + self._cancellation_detail = "; ".join(details) or None + + def report(self) -> ItemTimeoutReport: + """What happened at expiry, once any cancellation in progress has finished.""" + + if self._expiring.acquire(timeout=EXPIRY_REPORT_WAIT_SECONDS): + self._expiring.release() + with self._lock: + outcome: CancellationOutcome = self._cancellation or CancellationOutcome.NOT_NEEDED + if self._awaiting_confirmation: + outcome = ( + CancellationOutcome.UNKNOWN + if self._active_statements + else CancellationOutcome.CONFIRMED + ) + return ItemTimeoutReport( + kind=self.kind, + name=self.name, + limit=self.policy.limit, + source=self.policy.source, + elapsed_seconds=max(0.0, time.monotonic() - self.started), + query_ids=tuple(dict.fromkeys(self._query_ids)), + cancellation=outcome, + cancellation_detail=self._cancellation_detail, + ) diff --git a/src/sqlbuild/runtime/execution_limits/constants.py b/src/sqlbuild/runtime/execution_limits/constants.py new file mode 100644 index 0000000000..252cd233ee --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/constants.py @@ -0,0 +1,21 @@ +"""Runtime execution-limit constants.""" + +from sqlbuild.runtime.execution_limits.types import CancellationOutcome + +OUTCOME_SEVERITY: dict[CancellationOutcome, int] = { + CancellationOutcome.NOT_NEEDED: 0, + CancellationOutcome.CONFIRMED: 1, + CancellationOutcome.UNKNOWN: 2, + CancellationOutcome.UNSUPPORTED: 3, + CancellationOutcome.FAILED: 4, +} +REPORTED_QUERY_ID_LIMIT: int = 5 +CANCELLATION_TEXT: dict[CancellationOutcome, str] = { + CancellationOutcome.CONFIRMED: "confirmed; the running statement stopped", + CancellationOutcome.UNKNOWN: "unknown", + CancellationOutcome.FAILED: "failed", + CancellationOutcome.UNSUPPORTED: "not supported", + CancellationOutcome.NOT_NEEDED: "not needed; no statement was running", +} +ITEM_TIMEOUT_CODE: str = "C417" +EXPIRY_REPORT_WAIT_SECONDS: float = 30.0 diff --git a/src/sqlbuild/runtime/execution_limits/exceptions.py b/src/sqlbuild/runtime/execution_limits/exceptions.py index adc141e991..bfaf02983e 100644 --- a/src/sqlbuild/runtime/execution_limits/exceptions.py +++ b/src/sqlbuild/runtime/execution_limits/exceptions.py @@ -1,5 +1,7 @@ """Runtime execution-limit exceptions.""" +from __future__ import annotations + class ExecutionDurationLimitError(Exception): """Raised when a build exceeds its configured maximum duration.""" @@ -25,3 +27,12 @@ def __init__( ) self.help = remediation super().__init__(self.message) + + +class ItemTimeoutError(Exception): + """Raised inside one test, scenario or model attempt once its own time limit has passed.""" + + code: str = "C417" + + def __init__(self, *, kind: str, name: str, limit: str) -> None: + super().__init__(f"{kind} '{name}' exceeded its {limit} time limit") diff --git a/src/sqlbuild/runtime/execution_limits/main/current_execution_deadline.py b/src/sqlbuild/runtime/execution_limits/main/current_execution_deadline.py index f7d86929aa..872b3726a8 100644 --- a/src/sqlbuild/runtime/execution_limits/main/current_execution_deadline.py +++ b/src/sqlbuild/runtime/execution_limits/main/current_execution_deadline.py @@ -1,10 +1,20 @@ """Current execution deadline lookup.""" -from sqlbuild.runtime.execution_limits._helpers.context import current_deadline_context +from sqlbuild.runtime.execution_limits._helpers.context import ( + current_deadline_context, + current_item_deadline_context, +) +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline from sqlbuild.runtime.execution_limits.models import ExecutionDeadline -def current_execution_deadline() -> ExecutionDeadline | None: - """Return the active build deadline, when configured.""" +def current_execution_deadline() -> ExecutionDeadline | ItemDeadline | None: + """Return the active deadline that expires first: the build's or the current item's.""" - return current_deadline_context.get() + build: ExecutionDeadline | None = current_deadline_context.get() + item: ItemDeadline | None = current_item_deadline_context.get() + if item is None: + return build + if build is None or item.remaining_seconds <= build.remaining_seconds: + return item + return build diff --git a/src/sqlbuild/runtime/execution_limits/main/current_item_deadline.py b/src/sqlbuild/runtime/execution_limits/main/current_item_deadline.py new file mode 100644 index 0000000000..f89000598f --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/current_item_deadline.py @@ -0,0 +1,10 @@ +"""Current per-item deadline lookup.""" + +from sqlbuild.runtime.execution_limits._helpers.context import current_item_deadline_context +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline + + +def current_item_deadline() -> ItemDeadline | None: + """The deadline of the test, scenario or model attempt running in this context, if any.""" + + return current_item_deadline_context.get() diff --git a/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py b/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py index 091aa06edc..e858e61b4d 100644 --- a/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py +++ b/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py @@ -1,12 +1,19 @@ """Current execution deadline enforcement.""" -from sqlbuild.runtime.execution_limits._helpers.context import current_deadline_context +from sqlbuild.runtime.execution_limits._helpers.context import ( + current_deadline_context, + current_item_deadline_context, +) +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline from sqlbuild.runtime.execution_limits.models import ExecutionDeadline def enforce_execution_deadline() -> None: - """Raise when the active build deadline has expired.""" + """Raise when the active build deadline or the current item's deadline has expired.""" deadline: ExecutionDeadline | None = current_deadline_context.get() if deadline is not None: deadline.enforce() + item: ItemDeadline | None = current_item_deadline_context.get() + if item is not None: + item.enforce() diff --git a/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py b/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py index 2e83f96b03..e4fc2ec605 100644 --- a/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py +++ b/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py @@ -2,12 +2,18 @@ from __future__ import annotations +import threading import time from collections.abc import Generator from contextlib import contextmanager from contextvars import Token -from sqlbuild.runtime.execution_limits._helpers.context import current_deadline_context +from sqlbuild.runtime.execution_limits._helpers.cancellation import cancel_connection +from sqlbuild.runtime.execution_limits._helpers.context import ( + current_deadline_context, + current_invocation_connections_context, +) +from sqlbuild.runtime.execution_limits.classes.active_connections import ActiveConnections from sqlbuild.runtime.execution_limits.exceptions import ExecutionDurationLimitError from sqlbuild.runtime.execution_limits.models import ExecutionDeadline @@ -31,7 +37,16 @@ def execution_deadline_scope( max_duration=max_duration, remediation=remediation, ) + active: ActiveConnections = ActiveConnections() token: Token[ExecutionDeadline | None] = current_deadline_context.set(deadline) + connections_token: Token[ActiveConnections | None] = current_invocation_connections_context.set( + active + ) + timer: threading.Timer = threading.Timer( + max_duration_seconds, _cancel_running, kwargs={"active": active} + ) + timer.daemon = True + timer.start() try: try: yield @@ -44,4 +59,11 @@ def execution_deadline_scope( else: deadline.enforce() finally: + timer.cancel() + current_invocation_connections_context.reset(connections_token) current_deadline_context.reset(token) + + +def _cancel_running(*, active: ActiveConnections) -> None: + for adapter, connection in active.snapshot(): + _ = cancel_connection(adapter=adapter, connection=connection) diff --git a/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py b/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py new file mode 100644 index 0000000000..98e7b9dcdc --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py @@ -0,0 +1,42 @@ +"""Per-item deadlines with a watchdog that cancels in-flight warehouse work at expiry.""" + +from __future__ import annotations + +import threading +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import Token +from typing import Any + +from sqlbuild.runtime.execution_limits._helpers.context import current_item_deadline_context +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy +from sqlbuild.runtime.execution_limits.types import ItemKind + + +@contextmanager +def item_deadline_scope( + *, + kind: ItemKind, + name: str, + policy: ItemTimeoutPolicy | None, + adapter: Any, + connection: Any, +) -> Generator[ItemDeadline | None]: + """Bound one item's statements by its policy and cancel what is running when it expires.""" + + if policy is None: + yield None + return + deadline: ItemDeadline = ItemDeadline(kind=kind, name=name, policy=policy) + token: Token[ItemDeadline | None] = current_item_deadline_context.set(deadline) + timer: threading.Timer = threading.Timer( + policy.seconds, deadline.expire, kwargs={"fallback": (adapter, connection)} + ) + timer.daemon = True + timer.start() + try: + yield deadline + finally: + timer.cancel() + current_item_deadline_context.reset(token) diff --git a/src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py b/src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py new file mode 100644 index 0000000000..003c466432 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py @@ -0,0 +1,21 @@ +"""Statements that must run after an item's deadline, such as cleanup.""" + +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import Token + +from sqlbuild.runtime.execution_limits._helpers.context import current_item_deadline_context +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline + + +@contextmanager +def item_deadline_suspended() -> Generator[None]: + """Run cleanup after an item's deadline without the deadline refusing its statements.""" + + token: Token[ItemDeadline | None] = current_item_deadline_context.set(None) + try: + yield + finally: + current_item_deadline_context.reset(token) diff --git a/src/sqlbuild/runtime/execution_limits/main/timeout_report_lines.py b/src/sqlbuild/runtime/execution_limits/main/timeout_report_lines.py new file mode 100644 index 0000000000..7eff83935f --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/timeout_report_lines.py @@ -0,0 +1,29 @@ +"""The shared facts of a per-item timeout message.""" + +from __future__ import annotations + +from sqlbuild.runtime.execution_limits.constants import ( + CANCELLATION_TEXT, + REPORTED_QUERY_ID_LIMIT, +) +from sqlbuild.runtime.execution_limits.models import ItemTimeoutReport + + +def timeout_report_lines(*, report: ItemTimeoutReport, label: str) -> list[str]: + """The headline, then indented limit source, cancellation outcome and query ID lines.""" + + cancellation: str = CANCELLATION_TEXT[report.cancellation] + if report.cancellation_detail: + cancellation = f"{cancellation} ({report.cancellation_detail})" + lines: list[str] = [ + f"{label} '{report.name}' timed out after {report.elapsed_seconds:.1f}s " + f"(limit {report.limit}).", + f" Limit set by: {report.source}", + f" Cancellation: {cancellation}", + ] + if report.query_ids: + shown: tuple[str, ...] = report.query_ids[-REPORTED_QUERY_ID_LIMIT:] + omitted: int = len(report.query_ids) - len(shown) + suffix: str = f" (+{omitted} earlier)" if omitted else "" + lines.append(f" Query IDs: {', '.join(shown)}{suffix}") + return lines diff --git a/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py b/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py new file mode 100644 index 0000000000..d23e9e2529 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py @@ -0,0 +1,34 @@ +"""Register a statement's connection with the deadlines that may need to cancel it.""" + +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from typing import Any + +from sqlbuild.runtime.execution_limits._helpers.context import ( + current_invocation_connections_context, + current_item_deadline_context, +) +from sqlbuild.runtime.execution_limits.classes.active_connections import ActiveConnections +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline + + +@contextmanager +def tracked_statement_connection(*, adapter: Any, connection: Any) -> Generator[None]: + """Record that a statement runs on `connection` so a deadline can cancel it there.""" + + registries: list[ActiveConnections] = [] + invocation: ActiveConnections | None = current_invocation_connections_context.get() + if invocation is not None: + registries.append(invocation) + deadline: ItemDeadline | None = current_item_deadline_context.get() + if deadline is not None: + registries.append(deadline.connections) + for registry in registries: + registry.add(adapter=adapter, connection=connection) + try: + yield + finally: + for registry in registries: + registry.remove(connection=connection) diff --git a/src/sqlbuild/runtime/execution_limits/models.py b/src/sqlbuild/runtime/execution_limits/models.py index cf85018d84..9da8dd6a20 100644 --- a/src/sqlbuild/runtime/execution_limits/models.py +++ b/src/sqlbuild/runtime/execution_limits/models.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from sqlbuild.runtime.execution_limits.exceptions import ExecutionDurationLimitError +from sqlbuild.runtime.execution_limits.types import CancellationOutcome, ItemKind @dataclass(frozen=True) @@ -29,3 +30,55 @@ def enforce(self) -> None: max_duration=self.max_duration, remediation=self.remediation, ) + + +@dataclass(frozen=True) +class ItemTimeoutPolicy: + """The resolved time limit of one test, scenario or model attempt and where it came from.""" + + limit: str + seconds: int + source: str + + +@dataclass(frozen=True) +class RunItemTimeouts: + """A run's per-item limit, and whether it overrides limits the items declare.""" + + timeout_limit: str | None = None + timeout_seconds: int | None = None + timeout_source: str | None = None + timeout_overrides_items: bool = False + + def resolve( + self, *, declared: str | None, declared_seconds: int | None, declared_source: str + ) -> ItemTimeoutPolicy | None: + """The CLI limit, else the item's declared limit, else the run's configured limit.""" + + if not self.timeout_overrides_items and declared is not None: + if declared_seconds is None: + return None + return ItemTimeoutPolicy( + limit=declared, seconds=declared_seconds, source=declared_source + ) + if self.timeout_limit is None or self.timeout_seconds is None: + return None + return ItemTimeoutPolicy( + limit=self.timeout_limit, + seconds=self.timeout_seconds, + source=self.timeout_source or "run settings", + ) + + +@dataclass(frozen=True) +class ItemTimeoutReport: + """What happened when one item ran out of time.""" + + kind: ItemKind + name: str + limit: str + source: str + elapsed_seconds: float + query_ids: tuple[str, ...] + cancellation: CancellationOutcome + cancellation_detail: str | None diff --git a/src/sqlbuild/runtime/execution_limits/types.py b/src/sqlbuild/runtime/execution_limits/types.py new file mode 100644 index 0000000000..13a9b00c94 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/types.py @@ -0,0 +1,23 @@ +"""Runtime execution-limit domain types.""" + +from __future__ import annotations + +from enum import StrEnum + + +class ItemKind(StrEnum): + """The kind of work one per-item deadline bounds.""" + + TEST = "test" + SCENARIO = "scenario" + MODEL = "model" + + +class CancellationOutcome(StrEnum): + """What SQLBuild knows about the warehouse work it cancelled at a deadline.""" + + CONFIRMED = "confirmed" + UNKNOWN = "unknown" + FAILED = "failed" + UNSUPPORTED = "unsupported" + NOT_NEEDED = "not_needed" diff --git a/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py b/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py index c19f5368ad..33ec8eff2f 100644 --- a/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py +++ b/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py @@ -12,6 +12,8 @@ from contextvars import ContextVar, Token from types import TracebackType +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_item_deadline import current_item_deadline from sqlbuild.runtime.observability._helpers.dispatcher import ( current_event_dispatcher, dispatcher_scope, @@ -82,6 +84,8 @@ def __init__( ) self._pending_failure: tuple[BaseException, str | None, str | None] | None = None self._statement_facts: tuple[str, str] | None = None + self._item_deadline: ItemDeadline | None = None + self._terminal_query_id: str | None = None @staticmethod @contextmanager @@ -120,6 +124,9 @@ def __enter__(self) -> StatementLifecycle: self._started_monotonic = self._clock() self._execution_owner_key = execution_owner_key self._publish(event_type="statement_started", payload=self._base_payload()) + self._item_deadline = current_item_deadline() + if self._item_deadline is not None: + self._item_deadline.statement_started() self._owner_token = _STATEMENT_LIFECYCLE_OWNER.set(self) self._monitor = StatementMonitor( on_submitted=lambda query_id: self.submitted(query_id=query_id), @@ -150,6 +157,8 @@ def submitted(self, *, query_id: str | None = None, job_id: str | None = None) - payload["query_id"] = query_id if job_id is not None: payload["job_id"] = job_id + if self._item_deadline is not None: + self._item_deadline.statement_submitted(query_id=query_id or job_id) self._publish(event_type="statement_submitted", payload=payload) def set_query_id_provider(self, provider: Callable[[], str | None]) -> None: @@ -270,6 +279,9 @@ def __exit__( self.failed(error=exc_value) finally: self._stop_monitor() + if self._item_deadline is not None: + self._item_deadline.statement_finished(query_id=self._terminal_query_id) + self._item_deadline = None if self._owner_token is not None: _STATEMENT_LIFECYCLE_OWNER.reset(self._owner_token) self._owner_token = None @@ -317,6 +329,8 @@ def _publish_terminal(self, *, event_type: str, payload: Mapping[str, JSONValue] if self._terminal: return self._terminal = True + terminal_id: JSONValue | None = payload.get("query_id") or payload.get("job_id") + self._terminal_query_id = terminal_id if isinstance(terminal_id, str) else None self._publish(event_type=event_type, payload=payload) def _defer_completed( diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py index ba30010c86..dc8a025b84 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py @@ -366,3 +366,54 @@ class DirectLogicExtractionE2ETestCase: overrides: dict[str, str] expected_exit_code: int expected_output_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class RealModelLimitE2ETestCase: + """One chain length and limit setting run through `sqb test`.""" + + description: str + local_files: dict[str, str] + expected_exit_code: int + expected_output_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class SingleReadComparisonE2ETestCase: + """One unit test whose comparison reads its stored actual output.""" + + description: str + test_sql: str + expected_exit_code: int + expected_output_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class UnitTestTimeoutE2ETestCase: + """One slow unit test run under a time limit.""" + + description: str + test_sql: str + command: tuple[str, ...] + expected_output_fragments: tuple[str, ...] + expected_absent_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class ModelTimeoutE2ETestCase: + """One slow model with a max_duration built next to an independent model.""" + + description: str + max_duration: str + expected_output_fragments: tuple[str, ...] + expected_tables: tuple[tuple[str, bool], ...] + + +@dataclass(frozen=True) +class ScenarioTimeoutE2ETestCase: + """One slow scenario run under a time limit.""" + + description: str + header_timeout: str + command: tuple[str, ...] + expected_output_fragments: tuple[str, ...] diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py index 22a10e811d..3012f029bd 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py @@ -1930,3 +1930,61 @@ def build_helper_reference_project_files( f'TEST (name "{test_name}");\n\nWITH\n{_HELPER_REFERENCE_TESTS[test_name]}\nSELECT 1\n' ) return files + + +COST_GUARD_PROJECT_TOML: str = """ +name = "shop" +adapter = "duckdb" +default_target = "dev" + +[connection] +database = "warehouse.duckdb" + +[targets.dev] +schema = "main" +""".strip() +COST_GUARD_RAW_ORDERS_SOURCE: str = ( + "sources:\n" + " - name: raw_orders\n description: Raw orders.\n" + " schema: main\n" + " table: raw_orders\n" + " columns:\n" + " - name: order_id\n type: INTEGER\n" + " - name: amount\n type: INTEGER\n" +) +SLOW_SCALAR_SQL: str = "(SELECT COUNT(*) FROM range(100000000000))" + + +def build_order_chain_files(*, chain_length: int) -> dict[str, str]: + """A source-fed chain of `chain_length` order models and one unit test of the last one.""" + + files: dict[str, str] = { + "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML, + "sources/raw_orders.yml": COST_GUARD_RAW_ORDERS_SOURCE, + "models/order_step_1.sql": ( + "MODEL (description 'Order step 1.', materialized table);\n\n" + 'SELECT order_id, amount FROM __source("raw_orders")\n' + ), + } + for step in range(2, chain_length + 1): + files[f"models/order_step_{step}.sql"] = ( + f"MODEL (description 'Order step {step}.', materialized table);\n\n" + f'SELECT order_id, amount + 1 AS amount FROM __ref("order_step_{step - 1}")\n' + ) + files["tests/unit/order_chain_totals.sql"] = ( + "TEST();\n\nWITH\n__source__raw_orders AS (\n SELECT 1 AS order_id, 5 AS amount\n),\n" + f"__expected__order_step_{chain_length} AS (\n" + f" SELECT 1 AS order_id, {4 + chain_length} AS amount\n)\n" + ) + return files + + +def stored_test_outputs(*, database_path: Path) -> list[str]: + """Persistent relations named like a stored unit-test output.""" + + with duckdb.connect(str(database_path), read_only=True) as connection: + rows: list[tuple[str]] = connection.execute( + "SELECT table_name FROM information_schema.tables " + "WHERE table_name LIKE '__sqb_test_output%'" + ).fetchall() + return [row[0] for row in rows] diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py new file mode 100644 index 0000000000..e93ae30a95 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py @@ -0,0 +1,352 @@ +"""E2E coverage for unit-test cost guards: real-model limit, single-read comparison, time limits.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.test._test_types import ( + ModelTimeoutE2ETestCase, + RealModelLimitE2ETestCase, + ScenarioTimeoutE2ETestCase, + SingleReadComparisonE2ETestCase, + UnitTestTimeoutE2ETestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.main.test.helpers import ( + COST_GUARD_PROJECT_TOML, + COST_GUARD_RAW_ORDERS_SOURCE, + SLOW_SCALAR_SQL, + build_order_chain_files, + stored_test_outputs, +) +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import ( + prepare_inline_project, + run_sqb, + table_exists, +) + +_ORDER_FILES: dict[str, str] = { + "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML, + "models/stg_orders.sql": ( + "MODEL (description 'Staged orders.', materialized table);\n\n" + "SELECT * FROM (VALUES (1, 100, 5)) AS t(order_id, customer_id, amount)\n" + ), + "models/order_totals.sql": ( + "MODEL (description 'Order totals.', materialized table);\n\n" + 'SELECT order_id, customer_id, amount * 2 AS total FROM __ref("stg_orders")\n' + ), +} +_SLOW_FILES: dict[str, str] = { + "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML, + "sources/raw_orders.yml": COST_GUARD_RAW_ORDERS_SOURCE, + "models/stg_orders.sql": ( + "MODEL (description 'Staged orders.', materialized table);\n\n" + f'SELECT order_id, {SLOW_SCALAR_SQL} AS scanned FROM __source("raw_orders")\n' + ), + "models/order_scans.sql": ( + "MODEL (description 'Order scans.', materialized table);\n\n" + 'SELECT order_id, scanned FROM __ref("stg_orders")\n' + ), + "tests/unit/order_ids_are_kept.sql": ( + "TEST();\n\nWITH\n__ref__stg_orders AS (\n SELECT 1 AS order_id, 0 AS scanned\n),\n" + "__expected__order_scans AS (\n SELECT 1 AS order_id\n)\n" + ), +} + + +@pytest.mark.parametrize( + "test_case", + [ + RealModelLimitE2ETestCase( + description="default limit rejects an eleven-model chain at compile time", + local_files={}, + expected_exit_code=1, + expected_output_fragments=( + "SQL unit test 'order_chain_totals' executes 11 real models (limit 10).", + "Model under test: order_step_11", + "order_step_1 order_step_11 -> order_step_10 -> order_step_9", + "1. Mock the direct inputs of the model under test to cut the chain:", + "__ref__order_step_10 AS (SELECT ...)", + "2. If this test checks behaviour across several models, consider a scenario", + "compiles and optimises the whole chain at once", + "3. Raise the limit in sqlbuild_project.toml (or sqlbuild_local.toml):", + "unit_test_max_real_models = 11", + ), + ), + RealModelLimitE2ETestCase( + description="local limit allows the chain and the row shows its real model count", + local_files={"sqlbuild_local.toml": "[settings]\nunit_test_max_real_models = 11\n"}, + expected_exit_code=0, + expected_output_fragments=(" 11 real models", "PASS=1 FAIL=0 TOTAL=1"), + ), + ], + ids=lambda case: case.description, +) +def test_given_long_unmocked_chain_when_testing_then_real_model_limit_applies( + tmp_path: Path, test_case: RealModelLimitE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files={**build_order_chain_files(chain_length=11), **test_case.local_files}, + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "test"), project_dir=project_dir + ) + output: str = result.stdout + result.stderr + + assert result.returncode == test_case.expected_exit_code, output + for fragment in test_case.expected_output_fragments: + assert fragment in output, output + + +@pytest.mark.parametrize( + "test_case", + [ + SingleReadComparisonE2ETestCase( + description="passing test", + test_sql=( + "TEST();\n\nWITH\n__ref__stg_orders AS (\n" + " SELECT 1 AS order_id, 100 AS customer_id, 5 AS amount\n),\n" + "__expected__order_totals AS (\n" + " SELECT 1 AS order_id, 100 AS customer_id, 10 AS total\n)\n" + ), + expected_exit_code=0, + expected_output_fragments=("PASS=1 FAIL=0 TOTAL=1", " 1 real model"), + ), + SingleReadComparisonE2ETestCase( + description="column subset test", + test_sql=( + "TEST();\n\nWITH\n__ref__stg_orders AS (\n" + " SELECT 1 AS order_id, 100 AS customer_id, 5 AS amount\n),\n" + "__expected__order_totals AS (\n SELECT 1 AS order_id, 10 AS total\n)\n" + ), + expected_exit_code=0, + expected_output_fragments=("PASS=1 FAIL=0 TOTAL=1",), + ), + SingleReadComparisonE2ETestCase( + description="failing test keeps counts and samples", + test_sql=( + "TEST();\n\nWITH\n__ref__stg_orders AS (\n" + " SELECT 1 AS order_id, 100 AS customer_id, 5 AS amount\n),\n" + "__expected__order_totals AS (\n SELECT 1 AS order_id, 11 AS total\n)\n" + ), + expected_exit_code=1, + expected_output_fragments=( + "PASS=0 FAIL=1 TOTAL=1", + "unexpected=1; missing=1", + "unexpected sample 1: order_id=1, total=10", + "missing sample 1: order_id=1, total=11", + "differing column total: actual=10, expected=11", + "[T003] test 'order_totals_doubles' failed for models: order_totals", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_unit_test_when_comparing_stored_output_then_results_and_samples_are_unchanged( + tmp_path: Path, test_case: SingleReadComparisonE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files={**_ORDER_FILES, "tests/unit/order_totals_doubles.sql": test_case.test_sql}, + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "--debug", "test"), project_dir=project_dir + ) + output: str = result.stdout + result.stderr + + assert result.returncode == test_case.expected_exit_code, output + for fragment in test_case.expected_output_fragments: + assert fragment in output, output + assert "phase=assert kind=CREATE" in output, output + assert output.index("kind=CREATE") < output.index("kind=WITH") < output.index("kind=DROP") + assert stored_test_outputs(database_path=project_dir / "warehouse.duckdb") == [] + + +@pytest.mark.parametrize( + "test_case", + [ + UnitTestTimeoutE2ETestCase( + description="inlined slow model times out and suggests a scenario", + test_sql=( + 'TEST(timeout "1s");\n\nWITH\n' + "__source__raw_orders AS (\n SELECT 1 AS order_id, 5 AS amount\n),\n" + "__expected__order_scans AS (\n SELECT 1 AS order_id\n)\n" + ), + command=("--no-color", "test"), + expected_output_fragments=( + "TIMEOUT", + "SQL unit test 'order_scans_slow' timed out after", + "(limit 1s).", + "Limit set by: TEST() timeout in tests/unit/order_scans_slow.sql", + "Cancellation: confirmed; the running statement stopped", + "This test inlines 1 real model beyond the model under test", + "__ref__stg_orders AS (SELECT ...)", + "2. If this test checks behaviour across several models, consider a scenario", + "3. To allow more time, set TEST(timeout", + "PASS=1 FAIL=0 TIMEOUT=1 TOTAL=2", + ), + expected_absent_fragments=(), + ), + UnitTestTimeoutE2ETestCase( + description="command line limit applies to the model under test alone", + test_sql=( + "TEST();\n\nWITH\n" + "__source__raw_orders AS (\n SELECT 1 AS order_id, 5 AS amount\n),\n" + "__expected__stg_orders AS (\n SELECT 1 AS order_id\n)\n" + ), + command=("--no-color", "test", "--test-timeout", "1s"), + expected_output_fragments=( + "SQL unit test 'order_scans_slow' timed out after", + "Limit set by: --test-timeout", + 'To allow more time, set TEST(timeout "', + "PASS=1 FAIL=0 TIMEOUT=1 TOTAL=2", + ), + expected_absent_fragments=("consider a scenario",), + ), + ], + ids=lambda case: case.description, +) +def test_given_slow_unit_test_when_limit_expires_then_test_times_out_and_run_continues( + tmp_path: Path, test_case: UnitTestTimeoutE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files={**_SLOW_FILES, "tests/unit/order_scans_slow.sql": test_case.test_sql}, + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=test_case.command, project_dir=project_dir + ) + output: str = result.stdout + result.stderr + + assert result.returncode == 1, output + for fragment in test_case.expected_output_fragments: + assert fragment in output, output + for fragment in test_case.expected_absent_fragments: + assert fragment not in output, output + + +@pytest.mark.parametrize( + "test_case", + [ + ModelTimeoutE2ETestCase( + description="header max_duration stops the slow model", + max_duration="1s", + expected_output_fragments=( + "TIMEOUT", + "[C417] Model 'order_scans' timed out after", + "Limit set by: max_duration in MODEL() header", + "Cancellation: confirmed; the running statement stopped", + "Skipped (1): order_scan_summary", + ), + expected_tables=( + ("order_scans", False), + ("order_scan_summary", False), + ("customers", True), + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_model_max_duration_when_building_then_slow_model_times_out_and_dependents_skip( + tmp_path: Path, test_case: ModelTimeoutE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files={ + "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML, + "models/order_scans.sql": ( + "MODEL (description 'Order scans.', materialized table, " + f'max_duration "{test_case.max_duration}");\n\n' + f"SELECT 1 AS order_id, {SLOW_SCALAR_SQL} AS scanned\n" + ), + "models/order_scan_summary.sql": ( + "MODEL (description 'Order scan summary.', materialized table);\n\n" + 'SELECT COUNT(*) AS orders FROM __ref("order_scans")\n' + ), + "models/customers.sql": ( + "MODEL (description 'Customers.', materialized table);\n\n" + "SELECT 100 AS customer_id\n" + ), + }, + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--concurrency", "2"), project_dir=project_dir + ) + output: str = result.stdout + result.stderr + + assert result.returncode == 1, output + for fragment in test_case.expected_output_fragments: + assert fragment in output, output + for table_name, expected_exists in test_case.expected_tables: + assert ( + table_exists(db_path=project_dir / "warehouse.duckdb", table_name=table_name) + == expected_exists + ), table_name + + +@pytest.mark.parametrize( + "test_case", + [ + ScenarioTimeoutE2ETestCase( + description="SCENARIO() timeout stops the slow step", + header_timeout="1s", + command=("--no-color", "scenario", "test"), + expected_output_fragments=( + "[C417]: Scenario 'order_scans_flow' timed out after", + "Limit set by: SCENARIO() timeout in tests/scenarios/order_scans_flow.sql", + "Cancellation: confirmed; the running statement stopped", + "PASS=0 FAIL=1 TOTAL=1", + ), + ), + ScenarioTimeoutE2ETestCase( + description="command line limit overrides the header", + header_timeout="10m", + command=("--no-color", "scenario", "test", "--scenario-timeout", "1s"), + expected_output_fragments=( + "Scenario 'order_scans_flow' timed out after", + "Limit set by: --scenario-timeout", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_slow_scenario_when_limit_expires_then_scenario_times_out( + tmp_path: Path, test_case: ScenarioTimeoutE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files={ + **_SLOW_FILES, + "tests/scenarios/order_scans_flow.sql": ( + 'SCENARIO (\n description "Order scans keep order ids",\n' + f' timeout "{test_case.header_timeout}"\n);\n\nWITH\n' + "__source__raw_orders AS (\n SELECT 1 AS order_id, 5 AS amount\n),\n" + "__expected__order_scans AS (\n SELECT 1 AS order_id\n)\n" + ), + }, + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=test_case.command, project_dir=project_dir + ) + output: str = result.stdout + result.stderr + + assert result.returncode == 1, output + for fragment in test_case.expected_output_fragments: + assert fragment in output, output + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/_test_types.py b/tests/unit/src/sqlbuild/adapters/_test_types.py index dc73592300..579168a736 100644 --- a/tests/unit/src/sqlbuild/adapters/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/_test_types.py @@ -164,3 +164,10 @@ class DuckDbRowCountEstimateTestCase: name: str expected_row_count: int | None expected_detail: str | None = None + + +@dataclass(frozen=True) +class StatementCancellationTestCase: + description: str + expected_limit_range: tuple[int, int] + expected_asynchronous: bool diff --git a/tests/unit/src/sqlbuild/adapters/helpers.py b/tests/unit/src/sqlbuild/adapters/helpers.py index c2f811c3f7..a0610885f5 100644 --- a/tests/unit/src/sqlbuild/adapters/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/helpers.py @@ -2,10 +2,13 @@ from __future__ import annotations +import threading from collections.abc import Callable, Iterator from dataclasses import dataclass from typing import Any +import pytest + class FetchedRows: """A cursor result that returns fixed rows.""" @@ -139,3 +142,122 @@ def execute_through_cursor(*, connection: MetadataConnection, sql: str) -> Metad """Adapter execute stub that runs SQL on a fresh fake cursor.""" return connection.cursor().execute(sql) + + +BLOCKING_SQL: str = "SELECT * FROM orders" + + +class BlockingCursor: + """A DB-API cursor whose query blocks until another thread cancels it.""" + + def __init__(self, *, executed: list[str], started: threading.Event) -> None: + self.executed: list[str] = executed + self.started: threading.Event = started + self.cancelled: threading.Event = threading.Event() + self.connection: object | None = None + self.description: object | None = None + self.rowcount: int = -1 + self._request_id: str = "11111111-1111-4111-8111-111111111111" + self.sfqid: str | None = None + self.statement_params: dict[str, str] | None = None + + def execute(self, sql: str, *args: Any, **kwargs: Any) -> BlockingCursor: + del args + self.executed.append(sql) + self.statement_params = kwargs.get("_statement_params") + responses: dict[bool, Callable[[], BlockingCursor]] = { + True: self._block_until_cancelled, + False: lambda: self, + } + return responses[sql == BLOCKING_SQL]() + + def _block_until_cancelled(self) -> BlockingCursor: + self.started.set() + assert self.cancelled.wait(timeout=5), "statement was not cancelled" + raise RuntimeError("statement cancelled") + + def cancel(self) -> None: + self.cancelled.set() + + def close(self) -> None: + return None + + +class CancellableRawConnection: + """A raw driver connection exposing each driver's cancel entry point.""" + + def __init__(self) -> None: + self.executed: list[str] = [] + self.cursors: list[BlockingCursor] = [] + self.cancel_requests: list[tuple[str, str]] = [] + self.started: threading.Event = threading.Event() + self._conn: CancellableRawConnection = self + + def cursor(self) -> BlockingCursor: + cursor: BlockingCursor = BlockingCursor(executed=self.executed, started=self.started) + cursor.connection = self + self.cursors.append(cursor) + return cursor + + def execute(self, sql: str) -> None: + self.executed.append(sql) + + def cancel_safe(self) -> None: + self._cancel_running() + + def cancel(self) -> None: + self._cancel_running() + + def _cancel_query(self, sql: str, request_id: str) -> dict[str, bool]: + self.cancel_requests.append((sql, request_id)) + self._cancel_running() + return {"success": True} + + def _cancel_running(self) -> None: + for cursor in self.cursors: + cursor.cancel() + + +def cancel_from_watchdog(*, adapter: Any, connection: Any, raw: CancellableRawConnection) -> None: + """Run a blocking query while another thread asks the adapter to cancel it.""" + + def watchdog() -> None: + assert raw.started.wait(timeout=5) + assert adapter.interrupt_connection(connection) is True + + thread: threading.Thread = threading.Thread(target=watchdog) + thread.start() + with pytest.raises(RuntimeError, match="statement cancelled"): + connection.execute(BLOCKING_SQL) + thread.join(timeout=5) + + +class CancellableJob: + """A BigQuery job whose result blocks until it is cancelled.""" + + def __init__(self) -> None: + self.job_id: str = "job-1" + self.cancelled: threading.Event = threading.Event() + self.waiting: threading.Event = threading.Event() + + def result(self) -> list[object]: + self.waiting.set() + assert self.cancelled.wait(timeout=5), "job was not cancelled" + raise RuntimeError("statement cancelled") + + def cancel(self) -> bool: + self.cancelled.set() + return True + + +class CancellableJobClient: + """A BigQuery client returning one cancellable job and recording each job config.""" + + def __init__(self) -> None: + self.job: CancellableJob = CancellableJob() + self.job_configs: list[Any] = [] + + def query(self, sql: str, **kwargs: Any) -> CancellableJob: + del sql + self.job_configs.append(kwargs.get("job_config")) + return self.job diff --git a/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py b/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py new file mode 100644 index 0000000000..ce05e6632e --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py @@ -0,0 +1,189 @@ +"""Each built-in adapter's warehouse statement limit and cancellation at an execution deadline.""" + +from __future__ import annotations + +import threading + +import pytest + +from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter +from sqlbuild.adapters.bigquery.classes.bigquery_connection import _BigQueryConnection +from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter +from sqlbuild.adapters.databricks.classes.databricks_connection import _DatabricksConnection +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter +from sqlbuild.adapters.postgres.classes.postgres_connection import _PostgresConnection +from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection +from sqlbuild.adapters.sqlserver.classes.sqlserver_adapter import SqlServerAdapter +from sqlbuild.adapters.sqlserver.classes.sqlserver_connection import _SqlServerConnection +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy +from sqlbuild.runtime.execution_limits.types import ItemKind +from tests.unit.src.sqlbuild.adapters._test_types import StatementCancellationTestCase +from tests.unit.src.sqlbuild.adapters.helpers import ( + CancellableJobClient, + CancellableRawConnection, + cancel_from_watchdog, +) + +_POLICY: ItemTimeoutPolicy = ItemTimeoutPolicy(limit="60s", seconds=60, source="test") +_REQUEST_ID: str = "11111111-1111-4111-8111-111111111111" + + +@pytest.mark.parametrize( + "test_case", + [ + StatementCancellationTestCase( + description="postgres statement_timeout in milliseconds", + expected_limit_range=(50_000, 60_000), + expected_asynchronous=False, + ) + ], + ids=lambda case: case.description, +) +def test_given_postgres_deadline_when_executing_then_statement_timeout_bounds_and_cancel_stops( + test_case: StatementCancellationTestCase, +) -> None: + raw: CancellableRawConnection = CancellableRawConnection() + connection: _PostgresConnection = _PostgresConnection(raw) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=None, connection=connection + ): + cancel_from_watchdog(adapter=PostgresAdapter(), connection=connection, raw=raw) + limit_statement: str = raw.executed[0] + raw.executed.clear() + _ = connection.execute("SELECT 1") + + assert limit_statement.startswith("SET statement_timeout = ") + low, high = test_case.expected_limit_range + assert low <= int(limit_statement.rsplit(" ", 1)[1]) <= high + assert raw.executed == ["RESET statement_timeout", "SELECT 1"] + assert PostgresAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous + + +@pytest.mark.parametrize( + "test_case", + [ + StatementCancellationTestCase( + description="databricks STATEMENT_TIMEOUT in seconds", + expected_limit_range=(50, 60), + expected_asynchronous=True, + ) + ], + ids=lambda case: case.description, +) +def test_given_databricks_deadline_when_executing_then_session_timeout_set_and_cursor_cancelled( + test_case: StatementCancellationTestCase, +) -> None: + raw: CancellableRawConnection = CancellableRawConnection() + connection: _DatabricksConnection = _DatabricksConnection(raw) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=None, connection=connection + ): + cancel_from_watchdog(adapter=DatabricksAdapter(), connection=connection, raw=raw) + + assert raw.executed[0].startswith("SET STATEMENT_TIMEOUT = ") + low, high = test_case.expected_limit_range + assert low <= int(raw.executed[0].rsplit(" ", 1)[1]) <= high + assert ( + DatabricksAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous + ) + + +@pytest.mark.parametrize( + "test_case", + [ + StatementCancellationTestCase( + description="sql server driver cancel", + expected_limit_range=(0, 0), + expected_asynchronous=False, + ) + ], + ids=lambda case: case.description, +) +def test_given_sqlserver_running_statement_when_interrupted_then_driver_cancel_is_called( + test_case: StatementCancellationTestCase, +) -> None: + raw: CancellableRawConnection = CancellableRawConnection() + connection: _SqlServerConnection = _SqlServerConnection(raw) + + cancel_from_watchdog(adapter=SqlServerAdapter(), connection=connection, raw=raw) + + assert raw.cursors[0].cancelled.is_set() + assert ( + SqlServerAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous + ) + + +@pytest.mark.parametrize( + "test_case", + [ + StatementCancellationTestCase( + description="snowflake STATEMENT_TIMEOUT_IN_SECONDS and cancel request", + expected_limit_range=(50, 60), + expected_asynchronous=False, + ) + ], + ids=lambda case: case.description, +) +def test_given_snowflake_running_statement_when_interrupted_then_cancel_request_names_it( + test_case: StatementCancellationTestCase, +) -> None: + raw: CancellableRawConnection = CancellableRawConnection() + connection: _SnowflakeConnection = _SnowflakeConnection(raw) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=None, connection=connection + ): + cancel_from_watchdog(adapter=SnowflakeAdapter(), connection=connection, raw=raw) + + assert raw.cancel_requests == [("SELECT * FROM orders", _REQUEST_ID)] + statement_params: dict[str, str] | None = raw.cursors[0].statement_params + assert statement_params is not None + low, high = test_case.expected_limit_range + assert low <= int(statement_params["STATEMENT_TIMEOUT_IN_SECONDS"]) <= high + assert ( + SnowflakeAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous + ) + + +@pytest.mark.parametrize( + "test_case", + [ + StatementCancellationTestCase( + description="bigquery job timeout and jobs.cancel", + expected_limit_range=(50_000, 60_000), + expected_asynchronous=True, + ) + ], + ids=lambda case: case.description, +) +def test_given_bigquery_deadline_when_executing_then_job_timeout_set_and_job_cancelled( + test_case: StatementCancellationTestCase, +) -> None: + client: CancellableJobClient = CancellableJobClient() + connection: _BigQueryConnection = _BigQueryConnection(client=client, location=None) + adapter: BigQueryAdapter = BigQueryAdapter() + + def watchdog() -> None: + assert client.job.waiting.wait(timeout=5) + assert adapter.interrupt_connection(connection) is True + + thread: threading.Thread = threading.Thread(target=watchdog) + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=None, connection=connection + ): + thread.start() + with pytest.raises(RuntimeError, match="statement cancelled"): + _ = connection.execute("SELECT * FROM orders") + thread.join(timeout=5) + + low, high = test_case.expected_limit_range + assert low <= int(client.job_configs[0].job_timeout_ms) <= high + assert BigQueryAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/cli/commands/main/execution/test_build_process_resources.py b/tests/unit/src/sqlbuild/cli/commands/main/execution/test_build_process_resources.py index b3284c2102..a23b1feaee 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/execution/test_build_process_resources.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/execution/test_build_process_resources.py @@ -70,14 +70,14 @@ def test_given_failed_debug_build_when_unwinding_then_reports_resources_and_part "test_case", [ BuildDurationLimitTestCase( - description="DuckDB target cannot guarantee statement cancellation", + description="DuckDB target rejects duration beyond its cancellable range", project_toml=( 'name = "shop"\nadapter = "duckdb"\ndefault_target = "dev"\n\n' - '[targets.dev.execution_limits]\nmax_duration = "30m"\n' + '[targets.dev.execution_limits]\nmax_duration = "30d"\n' ), - remediation="Use an adapter with safe statement cancellation.", - expected_code="C415", - expected_error_fragment="cannot cancel active statements safely", + remediation="Choose a duration supported by the configured warehouse.", + expected_code="C416", + expected_error_fragment="max_duration exceeds the largest duration supported", ), BuildDurationLimitTestCase( description="Snowflake target rejects duration beyond statement timeout range", diff --git a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py index 75f64d2ecc..2a276b5d77 100644 --- a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py +++ b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py @@ -26,6 +26,7 @@ ScenarioRunOptions, ScenarioRunResult, ScenarioSnapshotStateResult, + ScenarioTestSettings, ) from sqlbuild.executor.scenario.types import ScenarioLocalRunStatus, ScenarioSnapshotState from sqlbuild.executor.scheduling.types import ExecutionStatus @@ -118,7 +119,7 @@ def complete_scenario( connection_config={"database": "pipeline.duckdb"}, adapter=adapter, project_name="waffle_shop", - retain=False, + settings=ScenarioTestSettings(retain=False), on_scenario_start=lambda scenario: started_names.append(scenario.name), on_scenario_complete=complete_scenario, ) @@ -176,7 +177,7 @@ def test_given_warehouse_test_pipeline_when_connecting_then_publishes_connection connection_config={"database": "pipeline.duckdb"}, adapter=adapter, project_name="waffle_shop", - retain=False, + settings=ScenarioTestSettings(retain=False), ) assert tuple(event.event_type for event in lifecycle_events) == test_case.expected_event_types diff --git a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_testing.py b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_testing.py index f4521e26ec..8690416f6c 100644 --- a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_testing.py +++ b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_testing.py @@ -27,6 +27,7 @@ dispatcher_scope, invocation_scope, ) +from sqlbuild.runtime.execution_limits.models import RunItemTimeouts from tests.unit.src.sqlbuild.executor.pipeline._helpers._test_types import ( SqlTestConcurrencyTestCase, SqlTestFunctionPreflightTestCase, @@ -180,8 +181,7 @@ def test_given_parameterized_chain_when_running_then_one_operation_owns_statemen assert event_types == ( "resource_attempt_started", "operation_started", - "statement_started", - "statement_completed", + *("statement_started", "statement_completed") * 3, "operation_completed", "resource_attempt_completed", ) @@ -239,9 +239,13 @@ def test_given_concurrent_workers_when_running_tests_then_execution_overlaps_on_ connection_ids_lock: threading.Lock = threading.Lock() def execute_with_overlap( - *, test_entry: SqlTestPlanEntry, adapter: DuckDbAdapter, connection: object + *, + test_entry: SqlTestPlanEntry, + adapter: DuckDbAdapter, + connection: object, + timeouts: RunItemTimeouts | None, ) -> SqlTestExecutionResult: - del adapter + del adapter, timeouts with connection_ids_lock: connection_ids.append(id(connection)) overlap_barrier.wait(timeout=2) diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/__init__.py b/tests/unit/src/sqlbuild/runtime/execution_limits/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py new file mode 100644 index 0000000000..bfab8a9e88 --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py @@ -0,0 +1,54 @@ +from dataclasses import dataclass + +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy +from sqlbuild.runtime.execution_limits.types import CancellationOutcome, ItemKind +from tests.unit.src.sqlbuild.runtime.execution_limits.helpers import FakeCancellingAdapter + + +@dataclass(frozen=True) +class RunItemTimeoutsTestCase: + description: str + run_limit: str | None + run_seconds: int | None + overrides_items: bool + declared: str | None + declared_seconds: int | None + expected_policy: ItemTimeoutPolicy | None + + +@dataclass(frozen=True) +class DeadlineCancellationTestCase: + description: str + interrupt_result: bool + adapter_type: type[FakeCancellingAdapter] + asynchronous: bool + statement_query_id: str | None + expected_outcome: CancellationOutcome + + +@dataclass(frozen=True) +class IdleExpiryTestCase: + description: str + kind: ItemKind + expected_outcome: CancellationOutcome + expected_interrupted: bool + + +@dataclass(frozen=True) +class SuspendedCleanupTestCase: + description: str + expected_enforced: bool + + +@dataclass(frozen=True) +class EarliestDeadlineTestCase: + description: str + build_seconds: int + item_seconds: int + expected_item_is_current: bool + + +@dataclass(frozen=True) +class NoPolicyTestCase: + description: str + expected_deadline: None diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py b/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py new file mode 100644 index 0000000000..75f475fe57 --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py @@ -0,0 +1,51 @@ +"""Fakes for execution-limit tests.""" + +from __future__ import annotations + +import threading +from typing import Any, ClassVar + +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_item_deadline import current_item_deadline +from sqlbuild.runtime.execution_limits.main.tracked_statement_connection import ( + tracked_statement_connection, +) + + +class FakeCancellingAdapter: + """An adapter whose interrupt answers as configured and records that it was asked.""" + + adapter_name: ClassVar[str] = "fake" + + def __init__(self, *, result: bool, asynchronous: bool) -> None: + self.result: bool = result + self.statement_cancellation_is_asynchronous: bool = asynchronous + self.interrupted: threading.Event = threading.Event() + + def interrupt_connection(self, connection: Any) -> bool: + del connection + self.interrupted.set() + return self.result + + +class FakeRefusingAdapter(FakeCancellingAdapter): + """An adapter whose cancel request raises.""" + + def interrupt_connection(self, connection: Any) -> bool: + del connection + self.interrupted.set() + raise RuntimeError("cancel refused") + + +def run_statement_until_cancelled( + *, adapter: FakeCancellingAdapter, finished_query_id: str | None +) -> None: + """Run one fake statement under the current item deadline until the watchdog interrupts it.""" + + deadline: ItemDeadline | None = current_item_deadline() + assert deadline is not None + with tracked_statement_connection(adapter=adapter, connection=object()): + deadline.statement_started() + deadline.statement_submitted(query_id="query-1") + assert adapter.interrupted.wait(timeout=5) + deadline.statement_finished(query_id=finished_query_id) diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py b/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py new file mode 100644 index 0000000000..f7d8dd0b35 --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py @@ -0,0 +1,222 @@ +"""Per-item deadlines: expiry, watchdog cancellation and honest cancellation outcomes.""" + +from __future__ import annotations + +import time + +import pytest + +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.exceptions import ItemTimeoutError +from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( + current_execution_deadline, +) +from sqlbuild.runtime.execution_limits.main.enforce_execution_deadline import ( + enforce_execution_deadline, +) +from sqlbuild.runtime.execution_limits.main.execution_deadline_scope import ( + execution_deadline_scope, +) +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope +from sqlbuild.runtime.execution_limits.main.item_deadline_suspended import ( + item_deadline_suspended, +) +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, ItemTimeoutReport +from sqlbuild.runtime.execution_limits.types import CancellationOutcome, ItemKind +from tests.unit.src.sqlbuild.runtime.execution_limits._test_types import ( + DeadlineCancellationTestCase, + EarliestDeadlineTestCase, + IdleExpiryTestCase, + NoPolicyTestCase, + SuspendedCleanupTestCase, +) +from tests.unit.src.sqlbuild.runtime.execution_limits.helpers import ( + FakeCancellingAdapter, + FakeRefusingAdapter, + run_statement_until_cancelled, +) + +_ONE_SECOND: ItemTimeoutPolicy = ItemTimeoutPolicy(limit="1s", seconds=1, source="test") +_EXPIRED: ItemTimeoutPolicy = ItemTimeoutPolicy(limit="1s", seconds=0, source="test") + + +@pytest.mark.parametrize( + "test_case", + [ + DeadlineCancellationTestCase( + description="synchronous cancel stops the statement", + interrupt_result=True, + adapter_type=FakeCancellingAdapter, + asynchronous=False, + statement_query_id="query-1", + expected_outcome=CancellationOutcome.CONFIRMED, + ), + DeadlineCancellationTestCase( + description="asynchronous cancel is reported unknown", + interrupt_result=True, + adapter_type=FakeCancellingAdapter, + asynchronous=True, + statement_query_id="query-1", + expected_outcome=CancellationOutcome.UNKNOWN, + ), + DeadlineCancellationTestCase( + description="adapter without cancellation is reported unsupported", + interrupt_result=False, + adapter_type=FakeCancellingAdapter, + asynchronous=False, + statement_query_id=None, + expected_outcome=CancellationOutcome.UNSUPPORTED, + ), + DeadlineCancellationTestCase( + description="cancel request error is reported failed", + interrupt_result=True, + adapter_type=FakeRefusingAdapter, + asynchronous=False, + statement_query_id=None, + expected_outcome=CancellationOutcome.FAILED, + ), + ], + ids=lambda case: case.description, +) +def test_given_running_statement_when_item_deadline_expires_then_cancellation_outcome_is_reported( + test_case: DeadlineCancellationTestCase, +) -> None: + adapter: FakeCancellingAdapter = test_case.adapter_type( + result=test_case.interrupt_result, asynchronous=test_case.asynchronous + ) + + with item_deadline_scope( + kind=ItemKind.TEST, + name="orders_rollup", + policy=_ONE_SECOND, + adapter=adapter, + connection=None, + ) as deadline: + run_statement_until_cancelled( + adapter=adapter, finished_query_id=test_case.statement_query_id + ) + + assert isinstance(deadline, ItemDeadline) + assert deadline.expired + report: ItemTimeoutReport = deadline.report() + assert report.cancellation == test_case.expected_outcome + assert "query-1" in report.query_ids + assert report.source == "test" + + +@pytest.mark.parametrize( + "test_case", + [ + IdleExpiryTestCase( + description="model with no running statement", + kind=ItemKind.MODEL, + expected_outcome=CancellationOutcome.NOT_NEEDED, + expected_interrupted=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_no_running_statement_when_item_deadline_expires_then_cancellation_is_not_needed( + test_case: IdleExpiryTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with item_deadline_scope( + kind=test_case.kind, name="orders", policy=_EXPIRED, adapter=adapter, connection=None + ) as deadline: + time.sleep(0.05) + with pytest.raises(ItemTimeoutError): + enforce_execution_deadline() + + assert isinstance(deadline, ItemDeadline) + assert adapter.interrupted.is_set() == test_case.expected_interrupted + assert deadline.report().cancellation == test_case.expected_outcome + + +@pytest.mark.parametrize( + "test_case", + [SuspendedCleanupTestCase(description="cleanup after expiry", expected_enforced=True)], + ids=lambda case: case.description, +) +def test_given_expired_item_when_suspended_for_cleanup_then_statements_are_allowed( + test_case: SuspendedCleanupTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_EXPIRED, adapter=adapter, connection=None + ): + time.sleep(0.05) + with item_deadline_suspended(): + enforce_execution_deadline() + enforced: bool = True + + assert enforced == test_case.expected_enforced + + +@pytest.mark.parametrize( + "test_case", + [ + EarliestDeadlineTestCase( + description="item expires first", + build_seconds=3600, + item_seconds=1, + expected_item_is_current=True, + ), + EarliestDeadlineTestCase( + description="build expires first", + build_seconds=1, + item_seconds=3600, + expected_item_is_current=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_build_and_item_deadlines_when_reading_current_then_earliest_expiry_applies( + test_case: EarliestDeadlineTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with ( + execution_deadline_scope( + max_duration_seconds=test_case.build_seconds, + max_duration=f"{test_case.build_seconds}s", + target_name="dev", + remediation=None, + ), + item_deadline_scope( + kind=ItemKind.TEST, + name="orders", + policy=ItemTimeoutPolicy( + limit=f"{test_case.item_seconds}s", seconds=test_case.item_seconds, source="test" + ), + adapter=adapter, + connection=None, + ), + ): + current: object = current_execution_deadline() + + assert isinstance(current, ItemDeadline) == test_case.expected_item_is_current + + +@pytest.mark.parametrize( + "test_case", + [NoPolicyTestCase(description="no configured limit", expected_deadline=None)], + ids=lambda case: case.description, +) +def test_given_no_policy_when_entering_scope_then_no_deadline_applies( + test_case: NoPolicyTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=None, adapter=adapter, connection=None + ) as deadline: + current: object = current_execution_deadline() + + assert deadline is test_case.expected_deadline + assert current is test_case.expected_deadline + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/test_run_item_timeouts.py b/tests/unit/src/sqlbuild/runtime/execution_limits/test_run_item_timeouts.py new file mode 100644 index 0000000000..a368a1be5a --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/test_run_item_timeouts.py @@ -0,0 +1,88 @@ +"""Precedence of CLI, declared and configured per-item time limits.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, RunItemTimeouts +from tests.unit.src.sqlbuild.runtime.execution_limits._test_types import RunItemTimeoutsTestCase + + +@pytest.mark.parametrize( + "test_case", + [ + RunItemTimeoutsTestCase( + description="configured limit applies without a declared limit", + run_limit="30s", + run_seconds=30, + overrides_items=False, + declared=None, + declared_seconds=None, + expected_policy=ItemTimeoutPolicy( + limit="30s", seconds=30, source="[settings] test_timeout" + ), + ), + RunItemTimeoutsTestCase( + description="declared limit beats the configured limit", + run_limit="30s", + run_seconds=30, + overrides_items=False, + declared="2m", + declared_seconds=120, + expected_policy=ItemTimeoutPolicy( + limit="2m", seconds=120, source="TEST() timeout in tests/unit/orders.sql" + ), + ), + RunItemTimeoutsTestCase( + description="declared disabled turns the configured limit off", + run_limit="30s", + run_seconds=30, + overrides_items=False, + declared="disabled", + declared_seconds=None, + expected_policy=None, + ), + RunItemTimeoutsTestCase( + description="command line limit beats the declared limit", + run_limit="5s", + run_seconds=5, + overrides_items=True, + declared="2m", + declared_seconds=120, + expected_policy=ItemTimeoutPolicy( + limit="5s", seconds=5, source="[settings] test_timeout" + ), + ), + RunItemTimeoutsTestCase( + description="disabled run limit applies no limit", + run_limit="disabled", + run_seconds=None, + overrides_items=False, + declared=None, + declared_seconds=None, + expected_policy=None, + ), + ], + ids=lambda case: case.description, +) +def test_given_run_and_declared_limits_when_resolving_then_most_specific_applies( + test_case: RunItemTimeoutsTestCase, +) -> None: + timeouts: RunItemTimeouts = RunItemTimeouts( + timeout_limit=test_case.run_limit, + timeout_seconds=test_case.run_seconds, + timeout_source="[settings] test_timeout", + timeout_overrides_items=test_case.overrides_items, + ) + + policy: ItemTimeoutPolicy | None = timeouts.resolve( + declared=test_case.declared, + declared_seconds=test_case.declared_seconds, + declared_source="TEST() timeout in tests/unit/orders.sql", + ) + + assert policy == test_case.expected_policy + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From e332e10546cb2ebc9efe2f0db635a12dae3ba8ab Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:27:49 +0100 Subject: [PATCH 02/11] docs: document unit test cost guards and per-item time limits --- website/src/content/docs/docs/cli/build.mdx | 6 ++ .../src/content/docs/docs/cli/scenario.mdx | 1 + website/src/content/docs/docs/cli/test.mdx | 9 ++ .../docs/concepts/models/configuration.mdx | 40 +++++++++ .../docs/concepts/project-configuration.mdx | 72 +++++++++++++-- .../content/docs/docs/concepts/scenarios.mdx | 14 ++- .../content/docs/docs/concepts/testing.mdx | 87 +++++++++++++++++++ 7 files changed, 219 insertions(+), 10 deletions(-) diff --git a/website/src/content/docs/docs/cli/build.mdx b/website/src/content/docs/docs/cli/build.mdx index f572ec6dbe..296229fbd4 100644 --- a/website/src/content/docs/docs/cli/build.mdx +++ b/website/src/content/docs/docs/cli/build.mdx @@ -45,6 +45,7 @@ own option. `--version` works only directly after `sqb`, as in `sqb --version`: | `--defer-sources-to` | Read managed source data from another target | | `--fail-fast` | Stop on first failure and skip remaining nodes | | `--concurrency` | Number of worker connections (default: 1) | +| `--test-timeout ` | Time limit for each unit test the build runs, such as `30s`, or `disabled`; see [Time limits](/docs/concepts/testing#time-limits) | | `--verbose`, `-v` | Show lifecycle SQL inline after each model | | `--start-cursor` | Override the start cursor, typed from the selected incremental models' cursor type; fails with `S303` if the selection mixes timestamp and integer cursors. With no cursor model selected, an integer value is read as an integer cursor and anything else as an ISO timestamp | | `--end-cursor` | Override the end cursor, typed the same way as `--start-cursor` | @@ -141,6 +142,11 @@ When a model fails: - Staging/delta tables are retained for inspection - Failure details show the model name, failed phase, and error message +A model that exceeds its [`max_duration`](/docs/concepts/models/configuration#execution-time-limit) +fails with `C417`. The failure states the elapsed time, the limit and where it was set, whether +cancelling its running statement was confirmed, and the query IDs SQLBuild saw; downstream models are +skipped as for any other failure. + ## Fingerprints After a successful build, SQLBuild writes version identities to `_sqlbuild_fingerprints` in the target schema. These are used on subsequent runs to detect changes and skip unchanged work. See [Planning and Change Detection](/docs/concepts/planning) for details. diff --git a/website/src/content/docs/docs/cli/scenario.mdx b/website/src/content/docs/docs/cli/scenario.mdx index 209b687c1f..082f28b154 100644 --- a/website/src/content/docs/docs/cli/scenario.mdx +++ b/website/src/content/docs/docs/cli/scenario.mdx @@ -26,6 +26,7 @@ sqb scenario test [flags] | `--local` | Run locally against DuckDB using captured JSONL snapshots | | `--strict` | Treat missing/stale local snapshots as errors instead of skips | | `--concurrency N` | Run up to `N` scenarios at once (default: the project `concurrency` setting) | +| `--scenario-timeout ` | Time limit for each scenario in this warehouse run, such as `10m`, or `disabled`; overrides `SCENARIO(timeout ...)`, the target and `[settings] scenario_timeout`. Not applied with `--local` | | `--sync-snapshots` | Capture missing/stale snapshots before local run (requires `--local`) | | `--refresh` | Recapture all selected snapshots before local run (requires `--local`) | | `--force` | Bypass snapshot capture safety limits | diff --git a/website/src/content/docs/docs/cli/test.mdx b/website/src/content/docs/docs/cli/test.mdx index a5c21bee1c..7e79389134 100644 --- a/website/src/content/docs/docs/cli/test.mdx +++ b/website/src/content/docs/docs/cli/test.mdx @@ -25,6 +25,7 @@ Positional `SELECTOR` arguments work like `--select`, and both forms can be comb | `--select`, `-s` | Select tests by name (`order_status_rules` or `test:order_status_rules`), or the tests targeting specific models | | `--exclude` | Exclude tests by name, or the tests targeting specific models | | `--warehouse ` | Snowflake warehouse for this invocation; overrides the target's `query` [warehouse group](/docs/concepts/project-configuration#command-group-warehouses) and the connection warehouse | +| `--test-timeout ` | Time limit for each test in this run, such as `30s` or `2m`, or `disabled`; overrides `TEST(timeout ...)`, the target and `[settings] test_timeout` | Test names are globally unique, so a bare test name selects exactly that test. `test:` is the explicit form and accepts the same glob patterns as names, such as `test:orders_*`. A bare pattern @@ -42,6 +43,11 @@ declared parameter types and nullability, typed values, and a content fingerprin are strings in JSON so their scale is preserved. Each test check in JSON output reports its `duration_ms`. +Each test row shows how many real (unmocked) models the test executes, such as `3 real models`. A +test that runs past its [time limit](/docs/concepts/testing#time-limits) is reported as `TIMEOUT`, +counted separately in the summary (`PASS=4 FAIL=0 TIMEOUT=1 TOTAL=5`), and the run continues +with the remaining tests. + ## Examples ```bash @@ -56,4 +62,7 @@ sqb test order_status_rules # Run one case of a parameterized test sqb test test:order_status_rules --case open_case + +# Allow each test two minutes in this run +sqb test --test-timeout 2m ``` diff --git a/website/src/content/docs/docs/concepts/models/configuration.mdx b/website/src/content/docs/docs/concepts/models/configuration.mdx index 248f29979a..bd6d51ac94 100644 --- a/website/src/content/docs/docs/concepts/models/configuration.mdx +++ b/website/src/content/docs/docs/concepts/models/configuration.mdx @@ -38,6 +38,7 @@ which accepts any keys. | `sql_analysis` | Per-model SQL-analysis override | | `dynamic_columns` | Named pivot column families with `pivot_column`, `value_column`, `aggregate`, `type`, and `name_pattern` | | `audit_factories` | Audit factories that generate audit cases for this model; see [Audits](/docs/concepts/audits) | +| `max_duration` | Time limit for one build attempt of this model, or `"disabled"`; see [Execution time limit](#execution-time-limit) | ## Descriptions @@ -210,3 +211,42 @@ See [Snapshots](/docs/concepts/snapshots) for strategy combinations, defaults, a | `row_diff_tolerances` | Numeric comparison tolerances | | `row_diff_sample_rows` | Deterministic unique-key sample size; `0` disables inherited sampling | | `row_diff_sample_seed` | Integer seed used for deterministic key hashing | + +## Execution time limit + +`max_duration` bounds how long one build attempt of a model may run. It has no default: a model +without it, or with `max_duration "disabled"`, runs until it finishes. Set it in the usual layers, +with the `MODEL()` header taking precedence over the nearest path default, and the path default over +`[defaults]`: + +```toml +[defaults] +max_duration = "30m" + +[path_defaults."marts/order_history"] +max_duration = "2h" +``` + +```sql +MODEL ( + description 'Daily customer order counts.', + materialized table, + max_duration "10m" +); +``` + +The limit is a positive fixed duration such as `"90s"`, `"15m"` or `"1h30m"`. One deadline covers the +whole attempt, starting when the model begins executing: its hooks, statements, attached audits and +every microbatch. Time spent waiting for upstream models or for a free worker does not count. + +At the deadline SQLBuild cancels the model's running statements through the adapter (see the +[adapter table](/docs/concepts/project-configuration#duration-limit)), runs no further statements +for the attempt, and reports the model as timed out with the elapsed time, the limit, where it was +set, the cancellation outcome and the query IDs it saw. Statements that already completed are not +undone, but the attempt stops before any later promotion or swap, so a staged table is not +published. Models that depend on it are skipped as for any failure, and independent models continue. +A model that finishes successfully just after its limit keeps its successful result. + +The per-model limit is separate from the target's whole-build +[`max_duration`](/docs/concepts/project-configuration#duration-limit). When both apply, each +statement is bounded by whichever deadline comes first. diff --git a/website/src/content/docs/docs/concepts/project-configuration.mdx b/website/src/content/docs/docs/concepts/project-configuration.mdx index 14b25205ec..04d4516448 100644 --- a/website/src/content/docs/docs/concepts/project-configuration.mdx +++ b/website/src/content/docs/docs/concepts/project-configuration.mdx @@ -345,9 +345,9 @@ Both policies default to `false`. `sqb clone --from prod --to dev` requires `all ### Execution limits -Named targets can set optional limits for `sqb build`. Limits are useful for development targets -where an unexpectedly broad selector or long-running command should stop instead of consuming an -unbounded amount of warehouse compute. +Named targets can set optional limits for `sqb build`, and per-item time limits for unit tests and +scenarios. Limits are useful for development targets where an unexpectedly broad selector or +long-running command should stop instead of consuming an unbounded amount of warehouse compute. ```toml [targets.dev] @@ -358,6 +358,8 @@ schema = "dev_alice" [targets.dev.execution_limits] max_models = 125 max_duration = "45m" +test_timeout = "60s" +scenario_timeout = "10m" remediation = """ Review the expanded plan and confirm that the selected scope is intentional. Automated tools must request approval before changing this policy. @@ -380,17 +382,38 @@ key, and confirms that no warehouse changes were made. `max_duration` is a positive fixed duration such as `30m`, `2h`, or `90s`. It covers planning and warehouse-mutating build phases after target resolution. SQLBuild stops dispatching work at the -deadline and passes the remaining duration to each warehouse statement so in-flight work is -cancelled by the adapter. - -Safe in-flight cancellation is currently supported by the Snowflake adapter, whose maximum accepted -duration is seven days. SQLBuild rejects an unsupported duration or an adapter that cannot guarantee -active statement cancellation rather than presenting a limit that may leave warehouse work running. +deadline, bounds each warehouse statement by the remaining time where the warehouse supports a +statement limit, and cancels statements still running at the deadline. + +Every built-in adapter enforces the limit: + +| Adapter | Warehouse statement limit | Cancellation at the deadline | Longest accepted duration | +|---------|---------------------------|------------------------------|---------------------------| +| Snowflake | `STATEMENT_TIMEOUT_IN_SECONDS` per statement | Driver cancel request for the running query | 7 days | +| DuckDB, MotherDuck | none | Connection interrupt | about 24 days | +| PostgreSQL | `statement_timeout` set for each statement | Driver cancel (`cancel_safe`) | about 24 days | +| BigQuery | Job timeout (`job_timeout_ms`) on each query job | `jobs.cancel` (asynchronous) | about 24 days | +| Databricks | `STATEMENT_TIMEOUT` session setting, capped at 48 hours | Cursor cancel (asynchronous) | about 24 days | +| SQL Server | none | Driver cancel | about 24 days | + +BigQuery and Databricks stop cancelled work asynchronously, so SQLBuild reports their cancellation +as requested but unconfirmed. SQL Server has no per-statement limit; the driver's query timeout is +process-wide, so SQLBuild relies on cancelling the running statement instead. A project adapter that +cannot cancel a running statement is rejected for `max_duration` rather than presenting a limit that +may leave warehouse work running. Under a per-item timeout, such an adapter's items stop before their +next statement and report the cancellation as not supported. A duration failure may occur after earlier models have completed. Its error therefore distinguishes the timeout from a pre-execution model-limit failure and states that warehouse changes may already have occurred. +#### Unit test and scenario timeouts + +`test_timeout` and `scenario_timeout` set the per-test and per-scenario time limit on this target, +replacing `[settings] test_timeout` and `[settings] scenario_timeout`. They follow the +[timeout precedence](#unit-test-and-scenario-limits) described under Settings and accept a positive +fixed duration or `"disabled"`. + #### Custom remediation `remediation` is optional project-authored text. SQLBuild appends it to model-count, duration, and @@ -467,6 +490,7 @@ instead. - Microbatch: `microbatch_strategy`, `batch_size`, `batch_concurrency`, `max_microbatches` - Row diff: `row_diff_exclude_columns`, `row_diff_tolerances`, `row_diff_sample_rows`, `row_diff_sample_seed` +- Execution: `max_duration` (see [Execution time limit](/docs/concepts/models/configuration#execution-time-limit)) ## Constants @@ -571,6 +595,9 @@ auto_load_sources = true table_promotion_mode = "staged" default_audit_severity = "error" default_audit_run_scope = "delta_and_final" +unit_test_max_real_models = 10 +test_timeout = "30s" +scenario_timeout = "15m" ``` | Field | Default | Description | @@ -585,6 +612,33 @@ default_audit_run_scope = "delta_and_final" | `table_promotion_mode` | adapter default | `staged` (CTAS to staging, audit, then promote) or `immediate` (CTAS directly to target, then audit) | | `default_audit_severity` | `error` | Default severity for audits: `warn` or `error` | | `default_audit_run_scope` | `delta_and_final` | Default run scope for audits: `final` or `delta_and_final` | +| `unit_test_max_real_models` | `10` | Most real (unmocked) models one SQL unit test may execute; a larger test fails compilation. See [Unit test and scenario limits](#unit-test-and-scenario-limits). | +| `test_timeout` | `30s` | Time limit for each SQL unit test, or `"disabled"` | +| `scenario_timeout` | none | Time limit for each scenario, or `"disabled"` | + +### Unit test and scenario limits + +A SQL unit test runs the real SQL of every unmocked model between its mocks and the model under +test as one warehouse statement. `unit_test_max_real_models` stops that statement from silently +growing into a whole-chain integration query: a test that executes more real models fails +compilation and names every real model, the dependency path that pulls it in, and the inputs to +mock. See [Testing](/docs/concepts/testing#cost-guards) for an example. + +`test_timeout` bounds each test's warehouse work, measured from the test's first statement. The +limit that applies to a test is, from highest precedence: + +1. `--test-timeout` on `sqb test` or `sqb build` +2. `TEST(timeout "...")` in the test file +3. `[targets..execution_limits] test_timeout` +4. `[settings] test_timeout`, then the 30-second default + +`scenario_timeout` works the same way with `--scenario-timeout`, `SCENARIO(timeout "...")`, +`[targets..execution_limits] scenario_timeout` and `[settings] scenario_timeout`; it has +no default. At the limit, SQLBuild cancels the item's running statement, reports the item as +`TIMEOUT` with the elapsed time, the limit, where the limit was set, the cancellation outcome and +the query IDs it saw, and continues with independent items. Every value accepts a positive fixed +duration such as `"45s"` or `"10m"`, or `"disabled"`. `sqlbuild_local.toml` accepts all three +settings. ### Table promotion mode diff --git a/website/src/content/docs/docs/concepts/scenarios.mdx b/website/src/content/docs/docs/concepts/scenarios.mdx index 1641f44af2..56c0252873 100644 --- a/website/src/content/docs/docs/concepts/scenarios.mdx +++ b/website/src/content/docs/docs/concepts/scenarios.mdx @@ -18,6 +18,13 @@ Scenarios materialize physical relations in the warehouse and build the actual m Use unit tests for transformation logic. Use scenarios for coherent end-to-end validation. +Cost is a second reason to move a check into a scenario. A unit test inlines every real model in its +chain into one SQL statement, so the warehouse compiles and optimises the whole chain at once, and +that cost can grow much faster than the query's size. A scenario materialises each model step by +step, so each statement is compiled separately, as in a real build. SQLBuild limits how many real +models one unit test may execute (see [Cost guards](/docs/concepts/testing#cost-guards)); when a test +reaches that limit because it checks behaviour across several models, it may belong here. + ## Authoring scenarios Scenario files live under `tests/scenarios/` and use one file per scenario: @@ -116,12 +123,17 @@ available only to the model are not included. See ### SCENARIO() header -The header is metadata only: +The header holds metadata and an optional time limit: | Field | Description | |-------|-------------| | `description` | Required description of the behaviour the scenario proves ([Descriptions](/docs/concepts/models/configuration#descriptions)) | | `tags` | Optional list of string tags | +| `timeout` | Optional time limit for the scenario's warehouse run, such as `"10m"`, or `"disabled"`; overrides `[settings] scenario_timeout` and the target's `scenario_timeout` | + +At the limit SQLBuild cancels the scenario's running statement, still drops the relations the +scenario created unless `--retain` is set, and reports the scenario as failed with `C417`, the +elapsed time, the limit and where it was set. ## How scenarios execute diff --git a/website/src/content/docs/docs/concepts/testing.mdx b/website/src/content/docs/docs/concepts/testing.mdx index 7e74906495..a480c2d3fa 100644 --- a/website/src/content/docs/docs/concepts/testing.mdx +++ b/website/src/content/docs/docs/concepts/testing.mdx @@ -764,6 +764,93 @@ tests/ ... ``` +## Cost guards + +A unit test runs the real SQL of every unmocked model between its mocks and the model under test as +one warehouse statement. That is cheap for one or two models, but the warehouse compiles and +optimises the whole chain at once, and that cost can grow much faster than the query's size. +SQLBuild therefore bounds each unit test in two ways. + +### Real-model limit + +`[settings] unit_test_max_real_models` (default `10`, also accepted in `sqlbuild_local.toml`) caps +how many real models one unit test may execute, counting the model under test. A test over the limit +fails compilation before any warehouse connection opens: + +```text +error[S000]: Invalid SQL test fixtures: +- SQL unit test 'customer_totals_rollup' executes 3 real models (limit 2). + Test file: tests/unit/customer_totals_rollup.sql + Model under test: customer_totals + +Real models executed, and the dependency path that pulls each one in: + customer_totals model under test + order_totals customer_totals -> order_totals + stg_orders customer_totals -> order_totals -> stg_orders + +Options: + 1. Mock the direct inputs of the model under test to cut the chain: + __ref__order_totals AS (SELECT ...) + 2. If this test checks behaviour across several models, consider a scenario + (tests/scenarios/). A unit test inlines every real model into one statement, + so the warehouse compiles and optimises the whole chain at once; a scenario + builds each model step by step, compiling each statement separately. + 3. Raise the limit in sqlbuild_project.toml (or sqlbuild_local.toml): + [settings] + unit_test_max_real_models = 3 +``` + +`sqb test` shows each test's real-model count on its row, such as `3 real models`, so a chain that +is growing towards the limit is visible early. Direct macro, UDF and table function tests execute +no models and are not counted. + +### Time limits + +Each unit test has a time limit, `30s` by default, measured from its first warehouse statement. +Set it, from highest precedence, with `--test-timeout` on `sqb test` or `sqb build`, with +`TEST(timeout "45s")` in the test file, with `[targets..execution_limits] test_timeout`, or +with `[settings] test_timeout`. Each accepts a positive fixed duration or `"disabled"`. + +At the limit SQLBuild cancels the test's running statement, reports the test as `TIMEOUT` with the +elapsed time, the limit, where it was set and the cancellation outcome, and continues with the +remaining tests. When the test inlines real models beyond the model under test, the message also +lists them and the same options as the real-model limit: + +```text +SQL unit test 'order_scans_kept' timed out after 1.0s (limit 1s). + Limit set by: TEST() timeout in tests/unit/order_scans_kept.sql + Cancellation: confirmed; the running statement stopped + +This test inlines 1 real model beyond the model under test into one statement: + stg_orders + +Options: + 1. Mock the inlined models closest to the model under test to cut the chain, + for example __ref__stg_orders AS (SELECT ...). + 2. If this test checks behaviour across several models, consider a scenario + (tests/scenarios/). A unit test inlines every real model into one statement, + so the warehouse compiles and optimises the whole chain at once; a scenario + builds each model step by step, compiling each statement separately. + 3. To allow more time, set TEST(timeout "3s") on this test, [settings] test_timeout = "3s" + in sqlbuild_project.toml (or sqlbuild_local.toml), or --test-timeout for one run. +``` + +A test that runs only the model under test reports the timeout, the limit, its source and how to +change it, without the scenario advice. See +[Duration limit](/docs/concepts/project-configuration#duration-limit) for how each adapter cancels +a running statement. + +### Single-read comparison + +A test compares its actual output with the expected rows several ways: both row counts, rows +missing from each side, the expected-column check and, on failure, row samples. Where the warehouse +supports it, SQLBuild evaluates each compared model's actual output once into a temporary relation +and runs every comparison against it, so the chain is not recomputed for each read. Results, +messages and samples are unchanged. DuckDB, MotherDuck and PostgreSQL use a session temporary +table that SQLBuild drops after the test, including after a failure or timeout; Snowflake reads the +first query's persisted result through `RESULT_SCAN`. Other adapters compare the inline query as +before. `target/run/tests/` keeps the self-contained inline comparison query for each test. + ## Skipping SQL analysis Compile checks SQL test fixtures and assertions like model SQL. For SQL the parser cannot read, From c01743408b281cc3df5808a1d18912d5232b9931 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 07:58:44 +0100 Subject: [PATCH 03/11] feat: attach structured meta to audit cases and attachments Audit cases and SQL audit attachments accept an optional JSON mapping that is reported with each result but never reaches audit SQL, evaluation, thresholds or audit fingerprints. Results store it in a new nullable audit_meta_json column, added to existing result tables only when a result carries meta; audit_completed events and JSON check output include it as audit_meta. Non-JSON meta is a compile error that names the offending value. --- .../src/stage_capture/constants.rs | 2 + .../src/lifecycle_events/constants.rs | 4 +- .../model_files/_helpers/header_columns.rs | 1 + .../src/model_files/models.rs | 1 + .../src/header_metadata/_helpers/audits.rs | 18 ++ .../src/header_metadata/constants.rs | 3 +- .../header_metadata/main/json_meta_problem.rs | 49 ++++ .../src/header_metadata/main/mod.rs | 1 + .../src/header_metadata/models.rs | 2 + .../tests/test_parse_header_metadata.rs | 41 ++++ .../_helpers/authored_values/python_nodes.rs | 27 ++- .../_helpers/compile_records/constants.rs | 1 + .../_helpers/compiled_project/audit_loop.rs | 19 ++ .../_helpers/compiled_project/audit_views.rs | 15 ++ .../compiled_project/model_loop_contracts.rs | 3 +- .../_helpers/contract_records/constants.rs | 1 + .../_helpers/model_config/header_metadata.rs | 5 +- .../_helpers/planner_records/constants.rs | 1 + .../planner_strategy/audit_entries.rs | 1 + .../_helpers/records/getter_readers.rs | 2 + .../bindings/_helpers/result_output/checks.rs | 2 +- .../_helpers/rules/custom_host_columns.rs | 3 + .../references/docs/concepts/audits.md | 33 ++- src/sqlbuild/_native.pyi | 17 +- src/sqlbuild/compiler/compile/types.py | 4 + .../audit_results/_helpers/ddl_lock.py | 6 +- .../audit_results/_helpers/meta_column.py | 77 +++++++ .../executor/audit_results/_helpers/sql.py | 24 +- .../executor/audit_results/constants.py | 2 + .../executor/audit_results/main/_write.py | 23 +- src/sqlbuild/executor/audit_results/models.py | 11 + .../executor/auditing/_helpers/execution.py | 1 + .../auditing/_helpers/result_projection.py | 11 +- src/sqlbuild/executor/auditing/models.py | 1 + .../executor/run/_helpers/reuse/audit.py | 1 + src/sqlbuild/python_nodes/models.py | 2 + .../cli/commands/main/audit/_test_types.py | 33 +++ .../cli/commands/main/audit/helpers.py | 191 ++++++++++++++++ .../commands/main/audit/test_audit_meta.py | 210 ++++++++++++++++++ .../expected/header_contract.json | 2 +- .../expected/model_schema.json | 2 +- .../compiled_project/expected/parity.json | 2 +- .../expected/waffle_shop.json | 24 +- .../audit_results/main/_test_types.py | 15 ++ .../executor/audit_results/main/helpers.py | 27 +++ .../executor/audit_results/main/test_write.py | 92 +++++++- .../main/test_parse_audit_instance.py | 62 ++++++ .../audit_results/_helpers/test_goldens.py | 2 +- .../executor/auditing/main/helpers.py | 2 + .../sqlbuild/python_nodes/main/_test_types.py | 14 ++ .../python_nodes/main/test_audit_factory.py | 74 ++++++ .../src/content/docs/docs/concepts/audits.mdx | 32 ++- 52 files changed, 1149 insertions(+), 50 deletions(-) create mode 100644 crates/sqlbuild-model-config/src/header_metadata/main/json_meta_problem.rs create mode 100644 src/sqlbuild/executor/audit_results/_helpers/meta_column.py create mode 100644 tests/e2e/src/sqlbuild/cli/commands/main/audit/test_audit_meta.py diff --git a/crates/sqlbuild-cache/src/stage_capture/constants.rs b/crates/sqlbuild-cache/src/stage_capture/constants.rs index a3877a527a..816052a6f1 100644 --- a/crates/sqlbuild-cache/src/stage_capture/constants.rs +++ b/crates/sqlbuild-cache/src/stage_capture/constants.rs @@ -80,6 +80,7 @@ pub const PROJECTED_ATTRIBUTES: [(&str, &[&str]); 15] = [ "run_scope", "always_run", "description", + "meta", "declaration_usages", "name", ], @@ -109,6 +110,7 @@ pub const PROJECTED_ATTRIBUTES: [(&str, &[&str]); 15] = [ "run_scope", "always_run", "description", + "meta", ], ), ( diff --git a/crates/sqlbuild-core/src/lifecycle_events/constants.rs b/crates/sqlbuild-core/src/lifecycle_events/constants.rs index 8743401fb4..630457ce08 100644 --- a/crates/sqlbuild-core/src/lifecycle_events/constants.rs +++ b/crates/sqlbuild-core/src/lifecycle_events/constants.rs @@ -161,7 +161,7 @@ pub(crate) const NONNEGATIVE_INTEGER_PAYLOAD_FIELDS: [&str; 25] = [ ]; pub(crate) const FINITE_NUMBER_PAYLOAD_FIELDS: [&str; 1] = ["measured_value"]; pub(crate) const BOOLEAN_PAYLOAD_FIELDS: [&str; 2] = ["evidence_truncated", "reused"]; -pub(crate) const JSON_PAYLOAD_FIELDS: [&str; 2] = ["evidence", "thresholds"]; +pub(crate) const JSON_PAYLOAD_FIELDS: [&str; 3] = ["audit_meta", "evidence", "thresholds"]; pub const RESOURCE_SKIP_CODES: [&str; 4] = ["dependency", "explicit", "fan_in", "scheduler"]; pub const RESOURCE_SKIP_MODES: [&str; 2] = ["hard", "soft"]; pub const HOOK_PHASES: [&str; 2] = ["post_hooks", "pre_hooks"]; @@ -497,6 +497,7 @@ pub(crate) const LIFECYCLE_EVENT_CATALOG_V1: [CatalogEntry; 20] = [ "attachment_kind", "audit_definition_name", "audit_description", + "audit_meta", "audit_name", "binding_key", "definition_fingerprint", @@ -851,6 +852,7 @@ pub(crate) const LIFECYCLE_EVENT_CATALOG_V2: [CatalogEntry; 23] = [ "attachment_kind", "audit_definition_name", "audit_description", + "audit_meta", "audit_name", "binding_key", "definition_fingerprint", diff --git a/crates/sqlbuild-discovery/src/model_files/_helpers/header_columns.rs b/crates/sqlbuild-discovery/src/model_files/_helpers/header_columns.rs index 8991ccd664..c56446f348 100644 --- a/crates/sqlbuild-discovery/src/model_files/_helpers/header_columns.rs +++ b/crates/sqlbuild-discovery/src/model_files/_helpers/header_columns.rs @@ -50,5 +50,6 @@ fn native_audit(audit: &ParsedAudit>) -> NativeAudit { thresholds: audit.thresholds, minimum_samples: audit.minimum_samples.as_ref().map(NativeNode::authored), evidence_limit: audit.evidence_limit.as_ref().map(NativeNode::authored), + meta: audit.meta.as_ref().map(NativeNode::authored), } } diff --git a/crates/sqlbuild-discovery/src/model_files/models.rs b/crates/sqlbuild-discovery/src/model_files/models.rs index a42b7646d4..3bbc25da51 100644 --- a/crates/sqlbuild-discovery/src/model_files/models.rs +++ b/crates/sqlbuild-discovery/src/model_files/models.rs @@ -80,4 +80,5 @@ pub struct NativeAudit { pub thresholds: Option, pub minimum_samples: Option, pub evidence_limit: Option, + pub meta: Option, } diff --git a/crates/sqlbuild-model-config/src/header_metadata/_helpers/audits.rs b/crates/sqlbuild-model-config/src/header_metadata/_helpers/audits.rs index 39e58b373e..b89eb3456f 100644 --- a/crates/sqlbuild-model-config/src/header_metadata/_helpers/audits.rs +++ b/crates/sqlbuild-model-config/src/header_metadata/_helpers/audits.rs @@ -7,6 +7,7 @@ use crate::header_metadata::_helpers::text::{ }; use crate::header_metadata::_helpers::thresholds::thresholds; use crate::header_metadata::constants::{AUDIT_OPTION_KEYS, AUDIT_SEVERITIES}; +use crate::header_metadata::main::json_meta_problem::json_meta_problem; use crate::header_metadata::models::ParsedAudit; use crate::types::{AuthoredNode, NodeKind}; @@ -96,6 +97,7 @@ fn configured_audit( option_site, (definition, "evidence_limit"), )?; + let meta = optional_meta(entry(options, "meta"), option_site)?; let arguments: Vec<(N, N)> = options .iter() .filter(|(key, _)| !is_option_key(key)) @@ -112,9 +114,24 @@ fn configured_audit( thresholds, minimum_samples, evidence_limit, + meta, }) } +fn optional_meta( + node: Option<&N>, + site: Site<'_>, +) -> Result, ConfigError> { + match node { + None => Ok(None), + Some(value) if value.kind() == NodeKind::Null => Ok(None), + Some(value) => match json_meta_problem(value) { + None => Ok(Some(value.clone())), + Some(problem) => Err(site.error(&format!("'meta' {problem}"))), + }, + } +} + fn bare_audit(definition_name: N) -> ParsedAudit { ParsedAudit { definition_name, @@ -127,6 +144,7 @@ fn bare_audit(definition_name: N) -> ParsedAudit { thresholds: None, minimum_samples: None, evidence_limit: None, + meta: None, } } diff --git a/crates/sqlbuild-model-config/src/header_metadata/constants.rs b/crates/sqlbuild-model-config/src/header_metadata/constants.rs index 949cb97071..cc093ed623 100644 --- a/crates/sqlbuild-model-config/src/header_metadata/constants.rs +++ b/crates/sqlbuild-model-config/src/header_metadata/constants.rs @@ -6,7 +6,7 @@ pub const MODEL_COLUMN_KEYS: [&str; 5] = /// The column key a MODEL header may declare and a SCHEMA declaration may not. pub const MIGRATE_FROM_KEY: &str = "migrate_from"; /// Audit options that are not passed to the audit as arguments. -pub const AUDIT_OPTION_KEYS: [&str; 8] = [ +pub const AUDIT_OPTION_KEYS: [&str; 9] = [ "name", "description", "severity", @@ -15,6 +15,7 @@ pub const AUDIT_OPTION_KEYS: [&str; 8] = [ "thresholds", "minimum_samples", "evidence_limit", + "meta", ]; /// The `thresholds` keys of the warning and error bounds. pub const THRESHOLD_WARN_KEY: &str = "warn"; diff --git a/crates/sqlbuild-model-config/src/header_metadata/main/json_meta_problem.rs b/crates/sqlbuild-model-config/src/header_metadata/main/json_meta_problem.rs new file mode 100644 index 0000000000..09e3a0a67b --- /dev/null +++ b/crates/sqlbuild-model-config/src/header_metadata/main/json_meta_problem.rs @@ -0,0 +1,49 @@ +//! `json_meta_problem`: why an authored audit `meta` value is not a mapping of JSON values. + +use crate::types::{AuthoredNode, NodeKind}; + +/// The first reason `meta` is not a string-keyed mapping of JSON values, or `None` when it is. +pub fn json_meta_problem(meta: &N) -> Option { + if meta.kind() != NodeKind::Map { + return Some(format!( + "must be a mapping of JSON values, got {}", + meta.python_repr() + )); + } + mapping_problem(meta, "meta") + .map(|problem| format!("must map string keys to JSON values; {problem}")) +} + +fn mapping_problem(node: &N, path: &str) -> Option { + for (key, value) in node.entries() { + let Some(name) = key.text() else { + return Some(format!( + "key {} in {path} is not a string", + key.python_repr() + )); + }; + if let Some(problem) = value_problem(&value, &format!("{path}['{name}']")) { + return Some(problem); + } + } + None +} + +fn value_problem(node: &N, path: &str) -> Option { + match node.kind() { + NodeKind::Null | NodeKind::Bool(_) | NodeKind::Int { .. } | NodeKind::Str => None, + NodeKind::List | NodeKind::Tuple => node + .items() + .iter() + .enumerate() + .find_map(|(index, item)| value_problem(item, &format!("{path}[{index}]"))), + NodeKind::Map => mapping_problem(node, path), + NodeKind::Other => match node.number() { + Some(number) if number.is_finite() => None, + _ => Some(format!( + "{path} is not JSON-serialisable: {}", + node.python_repr() + )), + }, + } +} diff --git a/crates/sqlbuild-model-config/src/header_metadata/main/mod.rs b/crates/sqlbuild-model-config/src/header_metadata/main/mod.rs index d195d8f0d8..2baf67df15 100644 --- a/crates/sqlbuild-model-config/src/header_metadata/main/mod.rs +++ b/crates/sqlbuild-model-config/src/header_metadata/main/mod.rs @@ -2,6 +2,7 @@ pub mod audit_factory_references; pub mod check_measurement_contract; pub mod check_threshold_bound; pub mod check_threshold_pair; +pub mod json_meta_problem; pub mod merged_audit_identities; pub mod parse_audit_list; pub mod parse_dynamic_columns; diff --git a/crates/sqlbuild-model-config/src/header_metadata/models.rs b/crates/sqlbuild-model-config/src/header_metadata/models.rs index 9a1ea2e896..aa3a505fbc 100644 --- a/crates/sqlbuild-model-config/src/header_metadata/models.rs +++ b/crates/sqlbuild-model-config/src/header_metadata/models.rs @@ -17,6 +17,8 @@ pub struct ParsedAudit { pub thresholds: Option, pub minimum_samples: Option, pub evidence_limit: Option, + /// The authored mapping of JSON values, reported with results and never evaluated. + pub meta: Option, } /// One measurement threshold bound with its limits as Python floats. diff --git a/crates/sqlbuild-model-config/src/header_metadata/tests/test_parse_header_metadata.rs b/crates/sqlbuild-model-config/src/header_metadata/tests/test_parse_header_metadata.rs index 1d41b5ec28..d04effd6c7 100644 --- a/crates/sqlbuild-model-config/src/header_metadata/tests/test_parse_header_metadata.rs +++ b/crates/sqlbuild-model-config/src/header_metadata/tests/test_parse_header_metadata.rs @@ -198,6 +198,47 @@ fn given_header_metadata_when_parsing_then_python_parses_and_errors_result() { "/project/models/orders.sql model audit 'a' 'evidence_limit' must be a non-negative integer", ), }, + HeaderMetadataTestCase { + description: "JSON meta is an option, not an argument", + columns: Value::Null, + audits: Value::List(vec![map(&[( + "row_count", + map(&[ + ("minimum", Value::Int(1)), + ( + "meta", + map(&[ + ("region", Value::Str("emea")), + ("weight", Value::Float), + ("tiers", Value::List(vec![Value::Str("gold"), Value::Null])), + ]), + ), + ]), + )])]), + expected_summary: Ok(&["row_count:minimum"]), + }, + HeaderMetadataTestCase { + description: "a meta date fails with its path", + columns: Value::Null, + audits: Value::List(vec![map(&[( + "a", + map(&[( + "meta", + map(&[("window", map(&[("since", Value::Date("2020, 1, 1"))]))]), + )]), + )])]), + expected_summary: Err( + "/project/models/orders.sql model audit 'a' 'meta' must map string keys to JSON values; meta['window']['since'] is not JSON-serialisable: datetime.date(2020, 1, 1)", + ), + }, + HeaderMetadataTestCase { + description: "a non-mapping meta fails", + columns: Value::Null, + audits: Value::List(vec![map(&[("a", map(&[("meta", Value::Str("emea"))]))])]), + expected_summary: Err( + "/project/models/orders.sql model audit 'a' 'meta' must be a mapping of JSON values, got 'emea'", + ), + }, HeaderMetadataTestCase { description: "a boolean count fails", columns: Value::Null, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/authored_values/python_nodes.rs b/crates/sqlbuild-python/src/bindings/_helpers/authored_values/python_nodes.rs index dafe233301..688f801583 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/authored_values/python_nodes.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/authored_values/python_nodes.rs @@ -10,9 +10,11 @@ use pyo3::types::{ use pyo3::types::{PyTupleMethods, PyTypeMethods}; use pyo3::{pyfunction, wrap_pyfunction}; use sqlbuild_core::identifier_text::main::is_lowercase_identifier::is_lowercase_identifier; +use sqlbuild_model_config::header_metadata::main::json_meta_problem::json_meta_problem; use crate::bindings::_helpers::boundary::keyword_request::Request; use crate::bindings::_helpers::compile_driver::python_calls::imported; +use crate::bindings::_helpers::model_config::authored_nodes::PyNode; use crate::bindings::_helpers::native_host::facade::FacadeHost; use crate::bindings::_helpers::native_host::host; use crate::bindings::_helpers::native_host::user_code::UserCodeHost; @@ -143,16 +145,17 @@ fn valid_count(value: &Bound<'_, PyAny>) -> PyResult { Ok(!below(value, 0)?) } -/// Reject an audit case without a lowercase name, a definition, or valid optional counts. +/// Reject an audit case without a lowercase name, a definition, valid counts, or JSON meta. #[pyfunction] -#[pyo3(signature = (*, name, definition, minimum_samples, evidence_limit))] -fn validate_audit_case( - name: &Bound<'_, PyAny>, - definition: &Bound<'_, PyAny>, - minimum_samples: &Bound<'_, PyAny>, - evidence_limit: &Bound<'_, PyAny>, -) -> Result<(), PyErr> { - let py: Python<'_> = name.py(); +#[pyo3(signature = (**request))] +fn validate_audit_case(request: Option<&Bound<'_, PyDict>>) -> Result<(), PyErr> { + let request = Request::new("validate_audit_case", request)?; + let py: Python<'_> = request.py(); + let name: &Bound<'_, PyAny> = &request.get("name")?; + let definition: &Bound<'_, PyAny> = &request.get("definition")?; + let minimum_samples: &Bound<'_, PyAny> = &request.get("minimum_samples")?; + let evidence_limit: &Bound<'_, PyAny> = &request.get("evidence_limit")?; + let meta: Option> = request.optional("meta")?; let Ok(text) = name.downcast::() else { return Err(name_type_error(name)?); }; @@ -168,6 +171,12 @@ fn validate_audit_case( if !valid_count(evidence_limit)? { return Err(shared_input_error(py, AUDIT_CASE_EVIDENCE_LIMIT)); } + if let Some(meta) = meta.as_ref().filter(|meta| !meta.is_none()) + && let Some(problem) = json_meta_problem(&PyNode(meta.clone())) + { + let message: String = format!("AuditCase '{text}' meta {problem}"); + return Err(shared_input_error(py, &message)); + } Ok(()) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs index f6cf74c7a1..4de06e17e0 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_records/constants.rs @@ -1485,6 +1485,7 @@ const R_COMPILED_AUDIT: RecordSpec = RecordSpec { FieldSpec::new("run_scope", FieldDefault::None), FieldSpec::new("always_run", FieldDefault::False), FieldSpec::new("description", FieldDefault::None), + FieldSpec::new("meta", FieldDefault::None), ], }; const R_COMPILED_SQL_SCENARIO: RecordSpec = RecordSpec { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_loop.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_loop.rs index 78eecf3de5..20869e5c8f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_loop.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_loop.rs @@ -39,6 +39,8 @@ use crate::bindings::_helpers::model_config::authored_nodes::PyNode; use crate::bindings::_helpers::model_config::header_metadata::thresholds; use crate::bindings::_helpers::model_config::model_validator::NativeModelValidator; use crate::bindings::_helpers::records::fields::field as native_field; +use crate::bindings::_helpers::result_output::integration_values::json_value; +use sqlbuild_core::json::models::JsonValue; use sqlbuild_model_config::errors::ConfigError; use sqlbuild_model_config::header_metadata::main::parse_measurement_policy::parse_measurement_policy; @@ -341,9 +343,25 @@ impl<'py> AuditScope<'_, 'py> { usages: PyTuple::new(py, usages)?.unbind(), name: None, description: None, + meta: None, }) } + /// The instance's `meta` as JSON, or `None`; a value JSON cannot carry fails the compile. + fn meta(&self, instance: &Bound<'py, PyAny>, label: &str) -> PyResult> { + let meta: Bound<'py, PyAny> = native_field(instance, "meta")?; + if meta.is_none() { + return Ok(None); + } + match json_value(&meta)? { + Some(value) => Ok(Some(value)), + None => { + let message: String = format!("{label}: 'meta' must be a mapping of JSON values"); + Err(compile_error(self.context.classes, &message, None)) + } + } + } + /// One attached generic audit instance, as `build_attached_audit_input` builds it. fn attached( &self, @@ -477,6 +495,7 @@ impl<'py> AuditScope<'_, 'py> { usages: PyTuple::new(py, usages)?.unbind(), name: native_field(instance, "name")?.extract()?, description: native_field(instance, "description")?.extract()?, + meta: self.meta(instance, &label)?, }) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_views.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_views.rs index 62e96eb2a8..98af5da878 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_views.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/audit_views.rs @@ -14,6 +14,8 @@ use crate::bindings::_helpers::compiled_project::model_input_view::{ use crate::bindings::_helpers::compiled_project::model_loop_host::ReferenceRow; use crate::bindings::_helpers::compiled_project::project::NativeCompiledProject; use crate::bindings::_helpers::records::fields::field as native_field; +use crate::bindings::_helpers::scopes::scope_report_objects::json_to_python; +use sqlbuild_core::json::models::JsonValue; const AUDIT_RESOURCE: &str = "audit"; const AUDIT_CLASSES: &str = "audits"; @@ -63,6 +65,8 @@ pub(crate) struct AttachedAudit { pub(crate) usages: Py, pub(crate) name: Option, pub(crate) description: Option, + /// The authored mapping of JSON values reported with results, or `None`. + pub(crate) meta: Option, } impl AttachedAudit { @@ -276,6 +280,16 @@ macro_rules! audit_getters { Ok(audit.attached.description.clone()) }) } + + #[getter] + pub(crate) fn meta(&self, py: Python<'_>) -> PyResult> { + read_audit(&self.owner, self.index, py, |audit, py| { + match &audit.attached.meta { + Some(meta) => json_to_python(py, meta.clone(), false), + None => Ok(py.None()), + } + }) + } } }; } @@ -438,6 +452,7 @@ impl NativeCompiledAudit { .to_owned() .into_any(), self.description(py)?.into_bound_py_any(py)?, + self.meta(py)?.into_bound(py), ]) } } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/model_loop_contracts.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/model_loop_contracts.rs index 843076350c..009cc5c581 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/model_loop_contracts.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/model_loop_contracts.rs @@ -37,7 +37,7 @@ const LOCATION_FIELD: &str = "location"; const ACCEPTED_VALUES_AUDIT: &str = "accepted_values"; const CONTRACT_KEY: &str = "contract"; /// `AuditCase` fields copied onto each generated `SchemaAuditInstance` under the same name. -const COPIED_CASE_FIELDS: [&str; 8] = [ +const COPIED_CASE_FIELDS: [&str; 9] = [ "name", "severity", "run_scope", @@ -46,6 +46,7 @@ const COPIED_CASE_FIELDS: [&str; 8] = [ "thresholds", "minimum_samples", "evidence_limit", + "meta", ]; const PIVOT_COLUMN_FIELD: &str = "pivot_column"; const VALUE_COLUMN_FIELD: &str = "value_column"; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs index 2654ff7f31..8030fa854d 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/contract_records/constants.rs @@ -569,6 +569,7 @@ const R_SCHEMA_AUDIT_INSTANCE: RecordSpec = RecordSpec { FieldSpec::new("thresholds", FieldDefault::None), FieldSpec::new("minimum_samples", FieldDefault::None), FieldSpec::new("evidence_limit", FieldDefault::None), + FieldSpec::new("meta", FieldDefault::None), FieldSpec { name: "location", default: FieldDefault::None, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs b/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs index 250844aba9..881b358b83 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs @@ -35,7 +35,7 @@ const SCHEMA_COLUMN_FIELDS: [&str; 8] = [ "migrate_from", ]; /// `SchemaAuditInstance` fields in declaration order. -const AUDIT_INSTANCE_FIELDS: [&str; 11] = [ +const AUDIT_INSTANCE_FIELDS: [&str; 12] = [ "definition_name", "arguments", "name", @@ -46,6 +46,7 @@ const AUDIT_INSTANCE_FIELDS: [&str; 11] = [ "thresholds", "minimum_samples", "evidence_limit", + "meta", "location", ]; /// `MeasurementThresholds` fields in declaration order. @@ -324,6 +325,7 @@ fn native_audit_instance<'py>( thresholds(py, classes, audit.thresholds.as_ref())?, count(audit.minimum_samples.as_ref())?, count(audit.evidence_limit.as_ref())?, + count(audit.meta.as_ref())?, location.clone(), ], ) @@ -565,6 +567,7 @@ fn audit_instance<'py>( thresholds(py, classes, audit.thresholds.as_ref())?, optional(py, audit.minimum_samples.as_ref()), optional(py, audit.evidence_limit.as_ref()), + optional(py, audit.meta.as_ref()), location.clone(), ], ) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs index 90d4e7c6cf..888d7e6765 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_records/constants.rs @@ -772,6 +772,7 @@ const R_AUDIT_PLAN_ENTRY: RecordSpec = RecordSpec { FieldSpec::new("always_run", FieldDefault::False), FieldSpec::new("description", FieldDefault::None), FieldSpec::new("reads_outside_target_lineage", FieldDefault::False), + FieldSpec::new("meta", FieldDefault::None), ], }; const R_CHAIN_STEP: RecordSpec = RecordSpec { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs index 0f72547a56..5a6cd73896 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/planner_strategy/audit_entries.rs @@ -144,6 +144,7 @@ fn plan_audit<'py>( "attached_column_name", "always_run", "description", + "meta", ] { values.set_item(field, native_field(&audit, field)?)?; } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/records/getter_readers.rs b/crates/sqlbuild-python/src/bindings/_helpers/records/getter_readers.rs index 8b2d3f11e8..0af5b3562f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/records/getter_readers.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/records/getter_readers.rs @@ -328,6 +328,7 @@ pub(crate) fn read_native_audit_input<'py>( "description" => this .description(py) .and_then(|value| value.into_bound_py_any(py)), + "meta" => this.meta(py).and_then(|value| value.into_bound_py_any(py)), "declaration_usages" => this .declaration_usages(py) .and_then(|value| value.into_bound_py_any(py)), @@ -405,6 +406,7 @@ pub(crate) fn read_native_compiled_audit<'py>( "description" => this .description(py) .and_then(|value| value.into_bound_py_any(py)), + "meta" => this.meta(py).and_then(|value| value.into_bound_py_any(py)), "key" => this.key(py).and_then(|value| value.into_bound_py_any(py)), "scope_deps" => this .scope_deps(py) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/result_output/checks.rs b/crates/sqlbuild-python/src/bindings/_helpers/result_output/checks.rs index 3a29d5e392..3d3b5d7582 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/result_output/checks.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/result_output/checks.rs @@ -292,7 +292,7 @@ pub(crate) fn audit_checks<'py>( entries.push(("name", native_field(&result, "audit_name")?)); entries.extend(fields( &result, - &["audit_definition_name", "audit_description"], + &["audit_definition_name", "audit_description", "audit_meta"], )?); entries.push(("check_id", any(py, audit_check_id(&result)?)?)); entries.push(("passed", any(py, audit_passed(&result)?)?)); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs index 559ebfa4f5..0101fba536 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs @@ -38,6 +38,7 @@ struct PayloadAudit { thresholds: Option<(Option, Option)>, minimum_samples: Option, evidence_limit: Option, + meta: Option, } /// A threshold bound with its limits as `f64` bits, so every value survives exactly. @@ -127,6 +128,7 @@ impl PayloadAudit { }), minimum_samples: audit.minimum_samples.as_ref().map(PayloadValue::new), evidence_limit: audit.evidence_limit.as_ref().map(PayloadValue::new), + meta: audit.meta.as_ref().map(PayloadValue::new), } } @@ -145,6 +147,7 @@ impl PayloadAudit { }), minimum_samples: self.minimum_samples.map(PayloadValue::authored), evidence_limit: self.evidence_limit.map(PayloadValue::authored), + meta: self.meta.map(PayloadValue::authored), } } } diff --git a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/audits.md b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/audits.md index 6486305610..46f21a727d 100644 --- a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/audits.md +++ b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/audits.md @@ -11,6 +11,7 @@ Online: https://sqlbuild.com/docs/concepts/audits/ - How audits work - Measurement audits - Audit factories +- Audit metadata - Result history - Built-in audits - Custom generic audits @@ -106,12 +107,42 @@ def order_quality(): Attach it with `MODEL (audit_factories [order_quality])`. Generated cases compile to the same audit contract as directly authored instances. +## Audit metadata + +Use `meta` to attach structured context, such as the segment a case covers or why its condition +excludes some rows, so dashboards can group and filter results without parsing descriptions: + +```python +AuditCase( + name="valid_order_rate__emea", + definition="valid_order_rate", + arguments={"condition": "order_id IS NOT NULL AND region = 'emea'"}, + thresholds=MeasurementThresholds(error=below(99)), + meta={ + "region": "emea", + "exclusion": "orders before 2020", + "exclusion_reason": "imported without order IDs", + }, +) +``` + +SQL attachments accept the same option, for example +`valid_order_rate (condition "order_id IS NOT NULL", meta (region "emea"), ...)`. + +`meta` must be a mapping with string keys and JSON values; anything else is a compile error that +names the offending value. It is never passed to the audit SQL and never affects evaluation, +thresholds, or audit fingerprints, so changing only `meta` reruns nothing that a `description` +change would not. It is recorded with each result in the `audit_meta_json` column of +`_sqlbuild_audit_results`, as `audit_meta` in the `audit_completed` lifecycle event, and as +`audit_meta` on audit checks in JSON output. + ## Result history Native warehouse adapters best-effort append confirmed audit outcomes to `_sqlbuild_audit_results`. Rows are immutable and use deterministic IDs, so retrying the same result is idempotent. Projection failure is reported separately and does not change the audit outcome or -command exit code. Lifecycle sinks can also consume the corresponding `audit_completed` fact, which +command exit code. A results table created before `audit_meta_json` existed gains that nullable +column the first time a result with `meta` is recorded. Lifecycle sinks can also consume the corresponding `audit_completed` fact, which is published as each audit finishes. ## Built-in audits diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index a1e9eb7203..26335e717b 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -1727,6 +1727,8 @@ class NativeAuditInput: @property def description(self) -> str | None: ... @property + def meta(self) -> dict[str, object] | None: ... + @property def declaration_usages(self) -> tuple[UsageRecord, ...]: ... @property def name(self) -> str | None: ... @@ -1780,6 +1782,8 @@ class NativeCompiledAudit: def always_run(self) -> bool: ... @property def description(self) -> str | None: ... + @property + def meta(self) -> dict[str, object] | None: ... class NativeProjectView: @property @@ -5462,7 +5466,12 @@ def retry_policy_exceptions( ) -> tuple[type[BaseException], ...]: ... def retry_delay_seconds(*, retry_policy: RetryPolicy, retry_index: int) -> float: ... def validate_audit_case( - *, name: str, definition: str, minimum_samples: int | None, evidence_limit: int | None + *, + name: str, + definition: str, + minimum_samples: int | None, + evidence_limit: int | None, + meta: object = ..., ) -> None: ... def normalize_hook_reads( value: SqlResourceRef | tuple[SqlResourceRef, ...] | list[SqlResourceRef], @@ -14205,6 +14214,7 @@ class SchemaAuditInstance(NativeRecord): thresholds: MeasurementThresholds | None minimum_samples: int | None evidence_limit: int | None + meta: dict[str, object] | None location: SourceLocation | None def __init__( self, @@ -14218,6 +14228,7 @@ class SchemaAuditInstance(NativeRecord): thresholds: MeasurementThresholds | None = ..., minimum_samples: int | None = ..., evidence_limit: int | None = ..., + meta: dict[str, object] | None = ..., location: SourceLocation | None = ..., ) -> None: ... @@ -15935,6 +15946,7 @@ class AuditPlanEntry(NativeRecord): always_run: bool description: str | None reads_outside_target_lineage: bool + meta: dict[str, object] | None def __init__( self, key: CompiledObjectKey, @@ -15962,6 +15974,7 @@ class AuditPlanEntry(NativeRecord): always_run: bool = ..., description: str | None = ..., reads_outside_target_lineage: bool = ..., + meta: dict[str, object] | None = ..., ) -> None: ... class ChainStep(NativeRecord): @@ -17357,6 +17370,7 @@ class CompiledAudit(NativeRecord): run_scope: str | None always_run: bool description: str | None + meta: dict[str, object] | None def __init__( self, key: CompiledObjectKey, @@ -17381,6 +17395,7 @@ class CompiledAudit(NativeRecord): run_scope: str | None = ..., always_run: bool = ..., description: str | None = ..., + meta: dict[str, object] | None = ..., ) -> None: ... class CompiledSqlScenario(NativeRecord): diff --git a/src/sqlbuild/compiler/compile/types.py b/src/sqlbuild/compiler/compile/types.py index f2f1b9d938..a093f3faa9 100644 --- a/src/sqlbuild/compiler/compile/types.py +++ b/src/sqlbuild/compiler/compile/types.py @@ -206,6 +206,8 @@ def run_scope(self) -> str | None: ... def always_run(self) -> bool: ... @property def description(self) -> str | None: ... + @property + def meta(self) -> dict[str, object] | None: ... class CompiledProjectView(Protocol): @@ -434,6 +436,8 @@ def always_run(self) -> bool: ... @property def description(self) -> str | None: ... @property + def meta(self) -> dict[str, object] | None: ... + @property def declaration_usages(self) -> tuple[UsageRecord, ...]: ... @property def name(self) -> str | None: ... diff --git a/src/sqlbuild/executor/audit_results/_helpers/ddl_lock.py b/src/sqlbuild/executor/audit_results/_helpers/ddl_lock.py index 95987e0e29..74c637fbf3 100644 --- a/src/sqlbuild/executor/audit_results/_helpers/ddl_lock.py +++ b/src/sqlbuild/executor/audit_results/_helpers/ddl_lock.py @@ -8,8 +8,8 @@ _audit_result_ddl_lock: Lock = Lock() -def run_with_audit_result_ddl_lock(callback: Callable[[], None]) -> None: - """Run audit-result DDL under a process-local lock.""" +def run_with_audit_result_ddl_lock[T](callback: Callable[[], T]) -> T: + """Run audit-result DDL under a process-local lock and return its result.""" with _audit_result_ddl_lock: - callback() + return callback() diff --git a/src/sqlbuild/executor/audit_results/_helpers/meta_column.py b/src/sqlbuild/executor/audit_results/_helpers/meta_column.py new file mode 100644 index 0000000000..859b89424a --- /dev/null +++ b/src/sqlbuild/executor/audit_results/_helpers/meta_column.py @@ -0,0 +1,77 @@ +"""Additive evolution of the audit result table's nullable meta column.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.contract.types import AdapterExecute, FrameworkType +from sqlbuild.executor.audit_results._helpers.sql import build_qualified_table_name +from sqlbuild.executor.audit_results.constants import ( + AUDIT_RESULTS_TABLE_NAME, + COLUMN_AUDIT_META_JSON, +) +from sqlbuild.executor.audit_results.models import AuditResultColumnEvolution + +_LOGGER: logging.Logger = logging.getLogger("sqlbuild.execution") + + +def ensure_meta_column( + *, + connection: Any, + execute: AdapterExecute[Any, Any], + database: str | None, + schema: str, + evolution: AuditResultColumnEvolution, + render_qualified_name: Callable[..., str | None], + render_framework_type: Callable[[FrameworkType], str], +) -> bool: + """Add the meta column to a table created before it existed; report whether it exists.""" + + if _meta_column_present( + connection=connection, database=database, schema=schema, evolution=evolution + ): + return True + statements: tuple[str, ...] = evolution.render_add_columns( + destination=build_qualified_table_name( + database=database, + schema=schema, + render_qualified_name=render_qualified_name, + ), + columns=( + ColumnInfo( + name=COLUMN_AUDIT_META_JSON, + type=render_framework_type(FrameworkType.STRING), + ), + ), + ) + try: + for statement in statements: + _ = execute(connection=connection, sql=statement) + except Exception as error: + if _meta_column_present( + connection=connection, database=database, schema=schema, evolution=evolution + ): + return True + _LOGGER.warning( + "Audit result meta was not stored because column '%s' could not be added (%s)", + COLUMN_AUDIT_META_JSON, + error, + ) + return False + return True + + +def _meta_column_present( + *, + connection: Any, + database: str | None, + schema: str, + evolution: AuditResultColumnEvolution, +) -> bool: + columns: tuple[ColumnInfo, ...] = evolution.read_columns( + connection=connection, database=database, schema=schema, name=AUDIT_RESULTS_TABLE_NAME + ) + return any(column.name.lower() == COLUMN_AUDIT_META_JSON for column in columns) diff --git a/src/sqlbuild/executor/audit_results/_helpers/sql.py b/src/sqlbuild/executor/audit_results/_helpers/sql.py index cd79175624..f9bab89dee 100644 --- a/src/sqlbuild/executor/audit_results/_helpers/sql.py +++ b/src/sqlbuild/executor/audit_results/_helpers/sql.py @@ -17,6 +17,7 @@ COLUMN_ATTACHMENT_KIND, COLUMN_AUDIT_DEFINITION_NAME, COLUMN_AUDIT_DESCRIPTION, + COLUMN_AUDIT_META_JSON, COLUMN_AUDIT_NAME, COLUMN_BINDING_KEY, COLUMN_DEFINITION_FINGERPRINT, @@ -110,14 +111,26 @@ def build_insert_sql( schema: str, records: Sequence[AuditResultRecord], render_qualified_name: Callable[..., str | None], + include_meta: bool = True, ) -> str: if not records: raise ExecutorInputError("audit result insert requires at least one record") qualified_name: str = build_qualified_table_name( database=database, schema=schema, render_qualified_name=render_qualified_name ) - rows: str = ", ".join(_record_values(record) for record in records) - return f"INSERT INTO {qualified_name} ({', '.join(AUDIT_RESULT_COLUMNS)}) VALUES {rows}" + columns: tuple[str, ...] = insert_columns(records=records, include_meta=include_meta) + rows: str = ", ".join(_record_values(record=record, columns=columns) for record in records) + return f"INSERT INTO {qualified_name} ({', '.join(columns)}) VALUES {rows}" + + +def insert_columns( + *, records: Sequence[AuditResultRecord], include_meta: bool = True +) -> tuple[str, ...]: + """Return inserted columns; the additive meta column is named only when a record has meta.""" + + if include_meta and any(record.audit_meta_json is not None for record in records): + return AUDIT_RESULT_COLUMNS + return tuple(column for column in AUDIT_RESULT_COLUMNS if column != COLUMN_AUDIT_META_JSON) def build_read_latest_sql( @@ -150,7 +163,7 @@ def build_read_latest_sql( ) -def _record_values(record: AuditResultRecord) -> str: +def _record_values(*, record: AuditResultRecord, columns: tuple[str, ...]) -> str: values: dict[str, object | None] = { COLUMN_RESULT_ID: record.result_id, COLUMN_SCHEMA_VERSION: record.schema_version, @@ -191,12 +204,11 @@ def _record_values(record: AuditResultRecord) -> str: COLUMN_SQL_DIGEST: record.sql_digest, COLUMN_METADATA_JSON: record.metadata_json, COLUMN_REUSED: _text_boolean(record.reused), + COLUMN_AUDIT_META_JSON: record.audit_meta_json, } return ( "(" - + ", ".join( - _column_literal(column=column, value=values[column]) for column in AUDIT_RESULT_COLUMNS - ) + + ", ".join(_column_literal(column=column, value=values[column]) for column in columns) + ")" ) diff --git a/src/sqlbuild/executor/audit_results/constants.py b/src/sqlbuild/executor/audit_results/constants.py index 551379eaaf..f6fd51bb91 100644 --- a/src/sqlbuild/executor/audit_results/constants.py +++ b/src/sqlbuild/executor/audit_results/constants.py @@ -44,6 +44,7 @@ COLUMN_SQL_DIGEST: str = "sql_digest" COLUMN_METADATA_JSON: str = "metadata_json" COLUMN_REUSED: str = "reused" +COLUMN_AUDIT_META_JSON: str = "audit_meta_json" AUDIT_RESULT_COLUMNS: tuple[str, ...] = ( COLUMN_RESULT_ID, @@ -85,6 +86,7 @@ COLUMN_SQL_DIGEST, COLUMN_METADATA_JSON, COLUMN_REUSED, + COLUMN_AUDIT_META_JSON, ) AUDIT_RESULT_COLUMN_TYPES: dict[str, StateSqlValueType] = { diff --git a/src/sqlbuild/executor/audit_results/main/_write.py b/src/sqlbuild/executor/audit_results/main/_write.py index a41bcce832..1480e77b9c 100644 --- a/src/sqlbuild/executor/audit_results/main/_write.py +++ b/src/sqlbuild/executor/audit_results/main/_write.py @@ -7,9 +7,10 @@ from sqlbuild.adapter.contract.types import AdapterExecute, FrameworkType from sqlbuild.executor.audit_results._helpers.ddl_lock import run_with_audit_result_ddl_lock +from sqlbuild.executor.audit_results._helpers.meta_column import ensure_meta_column from sqlbuild.executor.audit_results._helpers.sql import build_create_table_sql, build_insert_sql from sqlbuild.executor.audit_results.exceptions import AuditResultStorageError -from sqlbuild.executor.audit_results.models import AuditResultRecord +from sqlbuild.executor.audit_results.models import AuditResultColumnEvolution, AuditResultRecord def write_audit_result_records( @@ -23,8 +24,9 @@ def write_audit_result_records( render_framework_type: Callable[[FrameworkType], str], render_create_table_sql: Callable[..., str] | None = None, render_create_index_sqls: Callable[..., tuple[str, ...]] | None = None, + column_evolution: AuditResultColumnEvolution | None = None, ) -> None: - """Create storage if needed and append all records in one insert.""" + """Create storage if needed, add the meta column when records carry meta, and append.""" if not records: return @@ -39,19 +41,32 @@ def write_audit_result_records( render_framework_type=render_framework_type, ) ) + needs_meta: bool = any(record.audit_meta_json is not None for record in records) - def initialize_audit_result_table() -> None: + def initialize_audit_result_table() -> bool: _ = execute(connection=connection, sql=create_sql) if render_create_index_sqls is not None: for index_sql in render_create_index_sqls(database=database, schema=schema): _ = execute(connection=connection, sql=index_sql) + if not needs_meta or column_evolution is None: + return True + return ensure_meta_column( + connection=connection, + execute=execute, + database=database, + schema=schema, + evolution=column_evolution, + render_qualified_name=render_qualified_name, + render_framework_type=render_framework_type, + ) - _ = run_with_audit_result_ddl_lock(initialize_audit_result_table) + meta_available: bool = run_with_audit_result_ddl_lock(initialize_audit_result_table) insert_sql: str = build_insert_sql( database=database, schema=schema, records=records, render_qualified_name=render_qualified_name, + include_meta=meta_available, ) _ = execute(connection=connection, sql=insert_sql) except AuditResultStorageError: diff --git a/src/sqlbuild/executor/audit_results/models.py b/src/sqlbuild/executor/audit_results/models.py index a5c7166c9a..230845cc9f 100644 --- a/src/sqlbuild/executor/audit_results/models.py +++ b/src/sqlbuild/executor/audit_results/models.py @@ -2,9 +2,11 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from datetime import datetime +from sqlbuild.adapter.contract.models import ColumnInfo from sqlbuild.executor.audit_results._helpers.identity import ( build_audit_result_id as build_audit_result_id, ) @@ -53,3 +55,12 @@ class AuditResultRecord: metadata_json: str | None reused: bool audit_description: str | None = None + audit_meta_json: str | None = None + + +@dataclass(frozen=True) +class AuditResultColumnEvolution: + """Adapter callbacks that add result columns missing from a table created earlier.""" + + read_columns: Callable[..., tuple[ColumnInfo, ...]] + render_add_columns: Callable[..., tuple[str, ...]] diff --git a/src/sqlbuild/executor/auditing/_helpers/execution.py b/src/sqlbuild/executor/auditing/_helpers/execution.py index edfe1f4fd4..5ad6cdb1ce 100644 --- a/src/sqlbuild/executor/auditing/_helpers/execution.py +++ b/src/sqlbuild/executor/auditing/_helpers/execution.py @@ -258,6 +258,7 @@ def _base_result( audit_name=audit.name, audit_definition_name=audit.definition_name, audit_description=audit.description, + audit_meta=audit.meta, attachment_kind=audit.attachment_kind, severity=audit.severity, outcome=outcome, diff --git a/src/sqlbuild/executor/auditing/_helpers/result_projection.py b/src/sqlbuild/executor/auditing/_helpers/result_projection.py index 5e1ca99d0f..0c350d4400 100644 --- a/src/sqlbuild/executor/auditing/_helpers/result_projection.py +++ b/src/sqlbuild/executor/auditing/_helpers/result_projection.py @@ -26,7 +26,11 @@ from sqlbuild.executor.audit_results.constants import AUDIT_RESULT_SCHEMA_VERSION from sqlbuild.executor.audit_results.exceptions import AuditResultStorageError from sqlbuild.executor.audit_results.main._write import write_audit_result_records -from sqlbuild.executor.audit_results.models import AuditResultRecord, build_audit_result_id +from sqlbuild.executor.audit_results.models import ( + AuditResultColumnEvolution, + AuditResultRecord, + build_audit_result_id, +) from sqlbuild.executor.auditing.models import AuditExecutionResult, AuditResultProjection from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.main.current_execution_identity import ( @@ -258,6 +262,9 @@ def project_audit_result_batch_impl( render_framework_type=adapter.render_framework_type, render_create_table_sql=adapter.render_create_audit_result_table_sql, render_create_index_sqls=adapter.render_create_audit_result_index_sqls, + column_evolution=AuditResultColumnEvolution( + read_columns=adapter.get_columns, render_add_columns=adapter.render_add_columns + ), ) except AuditResultStorageError as error: _LOGGER.warning( @@ -316,6 +323,7 @@ def _build_record( # noqa: PLR0913 audit_name=result.audit_name, audit_definition_name=result.audit_definition_name, audit_description=result.audit_description, + audit_meta_json=_json_or_none(result.audit_meta), binding_key=audit_id.binding_key, definition_fingerprint=audit_id.definition_fingerprint, execution_fingerprint=audit_id.execution_fingerprint, @@ -407,6 +415,7 @@ def _publish_audit_completed(record: AuditResultRecord) -> None: "evidence_sql": record.evidence_sql, "execution_error": record.execution_error, "audit_description": record.audit_description, + "audit_meta": json.loads(record.audit_meta_json) if record.audit_meta_json else None, } payload.update({key: value for key, value in optional.items() if value is not None}) OperationLifecycle.publish_audit_completed(payload=payload) diff --git a/src/sqlbuild/executor/auditing/models.py b/src/sqlbuild/executor/auditing/models.py index 27c26b1b5a..6c07bef9a4 100644 --- a/src/sqlbuild/executor/auditing/models.py +++ b/src/sqlbuild/executor/auditing/models.py @@ -72,3 +72,4 @@ class AuditExecutionResult: evidence_sql: str | None = None audit_description: str | None = None execution_error: str | None = None + audit_meta: Mapping[str, object] | None = None diff --git a/src/sqlbuild/executor/run/_helpers/reuse/audit.py b/src/sqlbuild/executor/run/_helpers/reuse/audit.py index 5f5b0ee20d..d39d362cb9 100644 --- a/src/sqlbuild/executor/run/_helpers/reuse/audit.py +++ b/src/sqlbuild/executor/run/_helpers/reuse/audit.py @@ -54,6 +54,7 @@ def reused_final_audit_results_by_binding_key( reused=True, audit_definition_name=audit.definition_name, audit_description=audit.description, + audit_meta=audit.meta, evaluation_mode=audit.evaluation_mode, sample_unit=audit.sample_unit, minimum_samples=audit.minimum_samples, diff --git a/src/sqlbuild/python_nodes/models.py b/src/sqlbuild/python_nodes/models.py index 2f5c61e697..6ecf319de8 100644 --- a/src/sqlbuild/python_nodes/models.py +++ b/src/sqlbuild/python_nodes/models.py @@ -176,6 +176,7 @@ class AuditCase: thresholds: MeasurementThresholds | None = None minimum_samples: int | None = None evidence_limit: int | None = None + meta: dict[str, object] | None = None def __post_init__(self) -> None: validate_audit_case( @@ -183,4 +184,5 @@ def __post_init__(self) -> None: definition=self.definition, minimum_samples=self.minimum_samples, evidence_limit=self.evidence_limit, + meta=self.meta, ) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/audit/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/audit/_test_types.py index 3aa603d3f1..309b8d22e1 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/audit/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/audit/_test_types.py @@ -2,7 +2,9 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass, field +from pathlib import Path @dataclass(frozen=True) @@ -51,3 +53,34 @@ class InvalidDrainTimeoutE2ETestCase: description: str shutdown_timeout_toml: str expected_error: str + + +@dataclass(frozen=True) +class AuditMetaRoundTripE2ETestCase: + """Declared audit meta and where each audit's meta must reappear unchanged.""" + + description: str + prepare_results_table: Callable[[Path], None] + expected_meta: dict[str, object] + + +@dataclass(frozen=True) +class AuditMetaChangeE2ETestCase: + """One header edit whose rebuild is compared with an equivalent description-only edit.""" + + description: str + original_header_fragment: str + meta_only_fragment: str + description_only_fragment: str + expected_meta_after_change: dict[str, object] + expected_plan_reasons: dict[str, str] + + +@dataclass(frozen=True) +class InvalidAuditMetaE2ETestCase: + """One project file whose audit meta must fail compilation with a clear message.""" + + description: str + relative_path: str + contents: str + expected_fragments: tuple[str, ...] diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/audit/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/audit/helpers.py index d44567130e..f610101374 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/audit/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/audit/helpers.py @@ -3,8 +3,20 @@ from __future__ import annotations import json +import subprocess from pathlib import Path from textwrap import dedent +from typing import Any, cast + +import duckdb + +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.executor.audit_results.constants import COLUMN_AUDIT_META_JSON +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import ( + prepare_inline_project, + query_duckdb, + run_sqb, +) ORDER_COLUMNS: tuple[str, ...] = tuple(f"order_attribute_{index:02d}" for index in range(12)) @@ -94,3 +106,182 @@ def events_of_type(events: list[dict[str, object]], event_type: str) -> list[dic """Return events with one lifecycle type in their existing order.""" return list(filter(lambda event: event["event_type"] == event_type, events)) + + +AUDIT_META_ORDERS_HEADER: str = ( + 'description "Valid order rate",\n' + ' meta (region "emea", tiers ["gold", "silver"], weight 1.5, flags (late false)),' +) + + +def audit_meta_project_files( + *, measurement_options: str = AUDIT_META_ORDERS_HEADER +) -> dict[str, str]: + """Return a DuckDB project whose SQL, column and factory audits each declare meta.""" + + return { + "sqlbuild_project.toml": dedent( + """ + name = "audit_meta" + adapter = "duckdb" + + [connection] + database = "audit_meta.duckdb" + """ + ).lstrip(), + "audits/generic/valid_order_rate.sql": dedent( + """ + AUDIT ( + evaluation measurement, + value valid_rate, + sample_count total_rows, + sample_unit rows + ); + + MEASURE ( + SELECT + COUNT(*) AS total_rows, + 100.0 * AVG(CASE WHEN @condition THEN 1 ELSE 0 END) AS valid_rate + FROM @relation + ); + """ + ).lstrip(), + "audits/generic/expression_is_true.sql": ( + "AUDIT ();\n\nSELECT * FROM @relation WHERE NOT (@expression)\n" + ), + "python/factories/order_quality.py": dedent( + """ + from sqlbuild.audits import AuditCase, AuditSeverity, audit_factory + + + @audit_factory + def order_quality(): + return [ + AuditCase( + name="positive_amount", + definition="expression_is_true", + arguments={"expression": "amount > 0"}, + severity=AuditSeverity.ERROR, + meta={"segment": "core", "reference": "orders-quality", "rank": 1}, + ) + ] + """ + ).lstrip(), + "models/orders.sql": ( + "MODEL (\n" + ' description "Test model orders.",\n' + " materialized table,\n" + " audit_factories [order_quality],\n" + " columns (\n" + ' order_id (audits [not_null (meta (owner "fulfillment"))]),\n' + " ),\n" + " audits [\n" + " valid_order_rate (\n" + " condition \"region = 'emea'\",\n" + " thresholds (error (below 50)),\n" + f" {measurement_options}\n" + " ),\n" + ' expression_is_true (name "known_region", expression "region IS NOT NULL"),\n' + " ],\n" + ");\n\n" + "SELECT * FROM (VALUES (1, 'emea', 10), (2, 'emea', 20), (3, 'apac', 30))" + " AS orders(order_id, region, amount)\n" + ), + } + + +AUDIT_META_DATABASE: str = "audit_meta.duckdb" + + +def prepare_audit_meta_project(*, tmp_path: Path, measurement_options: str) -> Path: + """Write the audit meta project with a sink that records audit lifecycle events.""" + + return prepare_inline_project( + tmp_path=tmp_path, + project_name="audit_meta", + repo_files={ + **audit_meta_project_files(measurement_options=measurement_options), + "sinks/orders_events.py": recording_sink_source( + event_kinds="{LifecycleEventKind.AUDIT}", delay_seconds=0 + ), + }, + ) + + +def create_legacy_audit_results_table(project_dir: Path) -> None: + """Create the audit results table as it was before the meta column existed.""" + + create_sql: str = DuckDbAdapter().render_create_audit_result_table_sql( + database=None, schema="main" + ) + with duckdb.connect(str(project_dir / AUDIT_META_DATABASE)) as connection: + connection.execute(create_sql.replace(f", {COLUMN_AUDIT_META_JSON} VARCHAR", "")) + + +def leave_audit_results_table_absent(project_dir: Path) -> None: + """Leave the audit results table for the build to create.""" + + del project_dir + + +def stored_audit_meta(project_dir: Path) -> dict[str, object]: + """Map each stored audit name to its decoded meta, `None` when absent.""" + + rows: list[tuple[Any, ...]] = query_duckdb( + db_path=project_dir / AUDIT_META_DATABASE, + sql=( + "SELECT audit_name, COALESCE(audit_meta_json, 'null') " + "FROM main._sqlbuild_audit_results ORDER BY occurred_at" + ), + ) + return {str(name): json.loads(str(meta)) for name, meta in rows} + + +def stored_audit_identities(project_dir: Path) -> dict[str, tuple[object, ...]]: + """Map each audit's latest stored result to its identity, outcome and measurement.""" + + rows: list[tuple[Any, ...]] = query_duckdb( + db_path=project_dir / AUDIT_META_DATABASE, + sql=( + "SELECT audit_name, binding_key, definition_fingerprint, execution_fingerprint, " + "outcome, measured_value, sample_count " + "FROM main._sqlbuild_audit_results ORDER BY occurred_at" + ), + ) + return {str(row[0]): tuple(row[1:]) for row in rows} + + +def audit_completed_meta(event_path: Path) -> dict[str, object]: + """Map each audit_completed event's audit name to its `audit_meta` payload field.""" + + events: list[dict[str, object]] = events_of_type( + read_sequenced_events(event_path), "audit_completed" + ) + payloads: list[dict[str, object]] = [ + cast(dict[str, object], event["payload"]) for event in events + ] + return {str(payload["audit_name"]): payload.get("audit_meta") for payload in payloads} + + +def json_check_meta(stdout: str) -> dict[str, object]: + """Map each JSON-output check name to its `audit_meta` field.""" + + return {str(check["name"]): check.get("audit_meta") for check in json.loads(stdout)["checks"]} + + +def json_asset_statuses(stdout: str) -> dict[str, str]: + """Map each JSON-output asset name to its status.""" + + return {str(asset["name"]): str(asset["status"]) for asset in json.loads(stdout)["assets"]} + + +def plan_reasons(project_dir: Path) -> dict[str, str]: + """Map each planned model to its plan reason from `sqb plan --json`.""" + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "plan", "--json"), project_dir=project_dir + ) + assert result.returncode == 0, result.stdout + result.stderr + return { + str(model["name"]): str(model["reason"]) for model in json.loads(result.stdout)["models"] + } diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/audit/test_audit_meta.py b/tests/e2e/src/sqlbuild/cli/commands/main/audit/test_audit_meta.py new file mode 100644 index 0000000000..9244bc4d7a --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/audit/test_audit_meta.py @@ -0,0 +1,210 @@ +"""E2E coverage for structured audit meta through results, lifecycle events and JSON output.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.audit._test_types import ( + AuditMetaChangeE2ETestCase, + AuditMetaRoundTripE2ETestCase, + InvalidAuditMetaE2ETestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.main.audit.helpers import ( + AUDIT_META_ORDERS_HEADER, + audit_completed_meta, + create_legacy_audit_results_table, + json_asset_statuses, + json_check_meta, + leave_audit_results_table_absent, + plan_reasons, + prepare_audit_meta_project, + stored_audit_identities, + stored_audit_meta, +) +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import run_sqb + +_DECLARED_META: dict[str, object] = { + "valid_order_rate": { + "region": "emea", + "tiers": ["gold", "silver"], + "weight": 1.5, + "flags": {"late": False}, + }, + "positive_amount": {"segment": "core", "reference": "orders-quality", "rank": 1}, + "not_null": {"owner": "fulfillment"}, + "known_region": None, +} +_INVALID_FACTORY: str = """from datetime import date +from sqlbuild.audits import AuditCase, audit_factory + + +@audit_factory +def order_quality(): + return [ + AuditCase( + name="positive_amount", + definition="expression_is_true", + arguments={"expression": "amount > 0"}, + meta={"since": date(2020, 1, 1)}, + ) + ] +""" +_INVALID_HEADER: str = """MODEL (description "Test model orders.", materialized table, + audits [expression_is_true (expression "amount > 0", meta "emea")]); + +SELECT 1 AS amount +""" + + +@pytest.mark.parametrize( + "test_case", + ( + AuditMetaRoundTripE2ETestCase( + "fresh results table", leave_audit_results_table_absent, _DECLARED_META + ), + AuditMetaRoundTripE2ETestCase( + "results table without meta column", create_legacy_audit_results_table, _DECLARED_META + ), + ), + ids=lambda case: case.description, +) +def test_given_audit_meta_when_building_then_meta_reaches_results_events_and_json( + test_case: AuditMetaRoundTripE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_audit_meta_project( + tmp_path=tmp_path, measurement_options=AUDIT_META_ORDERS_HEADER + ) + test_case.prepare_results_table(project_dir) + event_path: Path = tmp_path / "events.jsonl" + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--json"), + project_dir=project_dir, + env={"ORDERS_EVENT_PATH": str(event_path)}, + ) + + assert result.returncode == 0, result.stdout + result.stderr + assert stored_audit_meta(project_dir) == test_case.expected_meta + assert audit_completed_meta(event_path) == test_case.expected_meta + assert json_check_meta(result.stdout) == test_case.expected_meta + + +@pytest.mark.parametrize( + "test_case", + ( + AuditMetaChangeE2ETestCase( + "meta-only edit matches a description-only edit", + original_header_fragment=AUDIT_META_ORDERS_HEADER, + meta_only_fragment=AUDIT_META_ORDERS_HEADER.replace('"emea"', '"europe"'), + description_only_fragment=AUDIT_META_ORDERS_HEADER.replace( + "Valid order rate", "Valid order share" + ), + expected_meta_after_change={ + "region": "europe", + "tiers": ["gold", "silver"], + "weight": 1.5, + "flags": {"late": False}, + }, + expected_plan_reasons={"orders": "no_change"}, + ), + ), + ids=lambda case: case.description, +) +def test_given_meta_only_change_when_rebuilding_then_outcomes_and_reruns_match_description_change( + test_case: AuditMetaChangeE2ETestCase, + tmp_path: Path, +) -> None: + statuses: dict[str, dict[str, str]] = {} + baselines: dict[str, dict[str, tuple[object, ...]]] = {} + rebuilt: dict[str, dict[str, tuple[object, ...]]] = {} + project_dirs: dict[str, Path] = {} + reasons: dict[str, dict[str, str]] = {} + for variant, fragment in ( + ("meta", test_case.meta_only_fragment), + ("description", test_case.description_only_fragment), + ): + project_dir: Path = prepare_audit_meta_project( + tmp_path=tmp_path / variant, measurement_options=test_case.original_header_fragment + ) + first: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--json"), + project_dir=project_dir, + env={"ORDERS_EVENT_PATH": str(tmp_path / f"{variant}-first.jsonl")}, + ) + assert first.returncode == 0, first.stdout + first.stderr + baselines[variant] = stored_audit_identities(project_dir) + model_path: Path = project_dir / "models" / "orders.sql" + model_path.write_text( + model_path.read_text(encoding="utf-8").replace( + test_case.original_header_fragment, fragment + ), + encoding="utf-8", + ) + reasons[variant] = plan_reasons(project_dir) + second: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--json"), + project_dir=project_dir, + env={"ORDERS_EVENT_PATH": str(tmp_path / f"{variant}-second.jsonl")}, + ) + assert second.returncode == 0, second.stdout + second.stderr + statuses[variant] = json_asset_statuses(second.stdout) + rebuilt[variant] = stored_audit_identities(project_dir) + project_dirs[variant] = project_dir + + assert reasons["meta"] == reasons["description"] == test_case.expected_plan_reasons + assert statuses["meta"] == statuses["description"] + assert rebuilt["meta"] == baselines["meta"] + assert rebuilt["description"] == baselines["description"] + assert ( + stored_audit_meta(project_dirs["meta"])["valid_order_rate"] + == test_case.expected_meta_after_change + ) + + +@pytest.mark.parametrize( + "test_case", + ( + InvalidAuditMetaE2ETestCase( + "factory meta with a date", + "python/factories/order_quality.py", + _INVALID_FACTORY, + ( + "AuditCase 'positive_amount' meta must map string keys to JSON values", + "meta['since'] is not JSON-serialisable", + ), + ), + InvalidAuditMetaE2ETestCase( + "header meta that is not a mapping", + "models/orders.sql", + _INVALID_HEADER, + ("'meta' must be a mapping of JSON values",), + ), + ), + ids=lambda case: case.description, +) +def test_given_non_json_meta_when_compiling_then_compile_error_names_the_value( + test_case: InvalidAuditMetaE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_audit_meta_project( + tmp_path=tmp_path, measurement_options=AUDIT_META_ORDERS_HEADER + ) + (project_dir / test_case.relative_path).write_text(test_case.contents, encoding="utf-8") + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "compile"), project_dir=project_dir + ) + + output: str = result.stdout + result.stderr + assert result.returncode != 0, output + assert [fragment in output for fragment in test_case.expected_fragments] == [True] * len( + test_case.expected_fragments + ), output + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/header_contract.json b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/header_contract.json index d0c8c20319..f4fca57ba9 100644 --- a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/header_contract.json +++ b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/header_contract.json @@ -52,7 +52,7 @@ "query_sql": "\"SELECT 'PLACED' AS status, 3 AS total\"", "references": "()", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='order_statuses', model_schema=None, description=None, type_enforcement=True, meta={}, columns=(SchemaColumn(name='status', type='VARCHAR', nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='accepted_values', arguments={'values': ('PLACED', 'SHIPPED')}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/order_statuses.sql'), line=3, column=12, end_line=3, end_column=18), migrate_from=None), SchemaColumn(name='total', type='INTEGER', nullable=None, description=None, meta={}, audits=(), location=SourceLocation(path=PosixPath('models/order_statuses.sql'), line=3, column=59, end_line=3, end_column=64), migrate_from=None)), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'total > 0'}, name='positive_total', description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), tags=())", + "schema_entry": "SchemaModelEntry(name='order_statuses', model_schema=None, description=None, type_enforcement=True, meta={}, columns=(SchemaColumn(name='status', type='VARCHAR', nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='accepted_values', arguments={'values': ('PLACED', 'SHIPPED')}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/order_statuses.sql'), line=3, column=12, end_line=3, end_column=18), migrate_from=None), SchemaColumn(name='total', type='INTEGER', nullable=None, description=None, meta={}, audits=(), location=SourceLocation(path=PosixPath('models/order_statuses.sql'), line=3, column=59, end_line=3, end_column=64), migrate_from=None)), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'total > 0'}, name='positive_total', description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT @enum(\"order_status\").PLACED AS status, 3 AS total', expanded_sql=\"SELECT 'PLACED' AS status, 3 AS total\", passes=((ExpansionSpan(source_start=7, source_end=35, output_start=7, output_end=15),), ()))", "sql_validation_enabled": "True" } diff --git a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/model_schema.json b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/model_schema.json index ea77bd3b77..a3d8b6491a 100644 --- a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/model_schema.json +++ b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/model_schema.json @@ -16,7 +16,7 @@ "query_sql": "\"SELECT 1 AS order_id, 'PLACED' AS status, 'x' AS note\"", "references": "()", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='orders', model_schema='order_shape', description='Order shape.', type_enforcement=True, meta={}, columns=(SchemaColumn(name='order_id', type='INTEGER', nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('schemas/orders.sql'), line=4, column=12, end_line=4, end_column=20), migrate_from=None), SchemaColumn(name='status', type='VARCHAR', nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='accepted_values', arguments={'values': ('PLACED', 'SHIPPED')}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('schemas/orders.sql'), line=4, column=53, end_line=4, end_column=59), migrate_from=None), SchemaColumn(name='note', type='VARCHAR', nullable=None, description=None, meta={}, audits=(), location=SourceLocation(path=PosixPath('models/orders.sql'), line=4, column=40, end_line=4, end_column=44), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='orders', model_schema='order_shape', description='Order shape.', type_enforcement=True, meta={}, columns=(SchemaColumn(name='order_id', type='INTEGER', nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('schemas/orders.sql'), line=4, column=12, end_line=4, end_column=20), migrate_from=None), SchemaColumn(name='status', type='VARCHAR', nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='accepted_values', arguments={'values': ('PLACED', 'SHIPPED')}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('schemas/orders.sql'), line=4, column=53, end_line=4, end_column=59), migrate_from=None), SchemaColumn(name='note', type='VARCHAR', nullable=None, description=None, meta={}, audits=(), location=SourceLocation(path=PosixPath('models/orders.sql'), line=4, column=40, end_line=4, end_column=44), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql=\"SELECT 1 AS order_id, 'PLACED' AS status, 'x' AS note\", expanded_sql=\"SELECT 1 AS order_id, 'PLACED' AS status, 'x' AS note\", passes=((), ()))", "sql_validation_enabled": "True" }, diff --git a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/parity.json b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/parity.json index 11a257d31b..1b0b842896 100644 --- a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/parity.json +++ b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/parity.json @@ -20,7 +20,7 @@ "query_sql": "\"SELECT 1 AS order_id, 'x' AS note\"", "references": "()", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='described_columns', model_schema=None, description=None, type_enforcement=True, meta={}, columns=(SchemaColumn(name='order_id', type='INTEGER', nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/marts/described_columns.sql'), line=2, column=12, end_line=2, end_column=20), migrate_from=None), SchemaColumn(name='note', type=None, nullable=None, description='Free text.', meta={}, audits=(), location=SourceLocation(path=PosixPath('models/marts/described_columns.sql'), line=2, column=80, end_line=2, end_column=84), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='described_columns', model_schema=None, description=None, type_enforcement=True, meta={}, columns=(SchemaColumn(name='order_id', type='INTEGER', nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/marts/described_columns.sql'), line=2, column=12, end_line=2, end_column=20), migrate_from=None), SchemaColumn(name='note', type=None, nullable=None, description='Free text.', meta={}, audits=(), location=SourceLocation(path=PosixPath('models/marts/described_columns.sql'), line=2, column=80, end_line=2, end_column=84), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql=\"SELECT 1 AS order_id, 'x' AS note\", expanded_sql=\"SELECT 1 AS order_id, 'x' AS note\", passes=((), ()))", "sql_validation_enabled": "True" }, diff --git a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/waffle_shop.json b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/waffle_shop.json index 66efea9cb8..ff1ffdcce1 100644 --- a/tests/integration/src/sqlbuild/compiler/compiled_project/expected/waffle_shop.json +++ b/tests/integration/src/sqlbuild/compiler/compiled_project/expected/waffle_shop.json @@ -14,7 +14,7 @@ "query_sql": "'SELECT\\n customer_id,\\n MAX(ordered_at) AS last_ordered_at,\\n MAX(order_status) AS latest_order_status,\\n COUNT(*) AS total_orders,\\n SUM(line_total_cents) AS total_revenue_cents\\nFROM __ref(\"fact_orders\")\\nGROUP BY customer_id'", "references": "(CompileSqlReference(ref_kind=, ref_name='fact_orders', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='customer_status_snapshot', model_schema=None, description='Merge-based customer snapshot with timestamp cursor and explicit backfill policies.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='customer_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/intermediate/customer_status_snapshot.sql'), line=16, column=5, end_line=16, end_column=16), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='customer_status_snapshot', model_schema=None, description='Merge-based customer snapshot with timestamp cursor and explicit backfill policies.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='customer_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/intermediate/customer_status_snapshot.sql'), line=16, column=5, end_line=16, end_column=16), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n customer_id,\\n MAX(ordered_at) AS last_ordered_at,\\n MAX(order_status) AS latest_order_status,\\n COUNT(*) AS total_orders,\\n SUM(line_total_cents) AS total_revenue_cents\\nFROM __ref(\"fact_orders\")\\nGROUP BY customer_id', expanded_sql='SELECT\\n customer_id,\\n MAX(ordered_at) AS last_ordered_at,\\n MAX(order_status) AS latest_order_status,\\n COUNT(*) AS total_orders,\\n SUM(line_total_cents) AS total_revenue_cents\\nFROM __ref(\"fact_orders\")\\nGROUP BY customer_id', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -31,7 +31,7 @@ "query_sql": "'SELECT\\n order_id,\\n customer_id,\\n order_status,\\n ordered_at,\\n line_total_cents\\nFROM __ref(\"fact_orders\")'", "references": "(CompileSqlReference(ref_kind=, ref_name='fact_orders', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='order_status_index', model_schema=None, description='Integer-cursor incremental projection over fact orders.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/intermediate/order_status_index.sql'), line=11, column=5, end_line=11, end_column=13), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='order_status_index', model_schema=None, description='Integer-cursor incremental projection over fact orders.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/intermediate/order_status_index.sql'), line=11, column=5, end_line=11, end_column=13), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n order_id,\\n customer_id,\\n order_status,\\n ordered_at,\\n line_total_cents\\nFROM __ref(\"fact_orders\")', expanded_sql='SELECT\\n order_id,\\n customer_id,\\n order_status,\\n ordered_at,\\n line_total_cents\\nFROM __ref(\"fact_orders\")', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -48,7 +48,7 @@ "query_sql": "'SELECT\\n DATE_TRUNC(\\'day\\', activity_hour) AS activity_day,\\n SUM(orders_placed) AS orders_placed,\\n SUM(waffles_ordered) AS waffles_ordered,\\n SUM(revenue_cents) AS revenue_cents\\nFROM __ref(\"hourly_order_activity\")\\nGROUP BY DATE_TRUNC(\\'day\\', activity_hour)'", "references": "(CompileSqlReference(ref_kind=, ref_name='hourly_order_activity', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='daily_activity_rollup', model_schema=None, description='Downstream daily microbatch rollup with a wider batch size than its hourly upstream.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='activity_day', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/marts/daily_activity_rollup.sql'), line=18, column=5, end_line=18, end_column=17), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'orders_placed >= 0'}, name='daily_orders_placed_is_non_negative', description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), tags=())", + "schema_entry": "SchemaModelEntry(name='daily_activity_rollup', model_schema=None, description='Downstream daily microbatch rollup with a wider batch size than its hourly upstream.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='activity_day', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/marts/daily_activity_rollup.sql'), line=18, column=5, end_line=18, end_column=17), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'orders_placed >= 0'}, name='daily_orders_placed_is_non_negative', description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n @timestamp_trunc(\\'day\\', \\'activity_hour\\') AS activity_day,\\n SUM(orders_placed) AS orders_placed,\\n SUM(waffles_ordered) AS waffles_ordered,\\n SUM(revenue_cents) AS revenue_cents\\nFROM __ref(\"hourly_order_activity\")\\nGROUP BY @timestamp_trunc(\\'day\\', \\'activity_hour\\')', expanded_sql='SELECT\\n DATE_TRUNC(\\'day\\', activity_hour) AS activity_day,\\n SUM(orders_placed) AS orders_placed,\\n SUM(waffles_ordered) AS waffles_ordered,\\n SUM(revenue_cents) AS revenue_cents\\nFROM __ref(\"hourly_order_activity\")\\nGROUP BY DATE_TRUNC(\\'day\\', activity_hour)', passes=((), (ExpansionSpan(source_start=9, source_end=49, output_start=9, output_end=41), ExpansionSpan(source_start=232, source_end=272, output_start=224, output_end=256))))", "sql_validation_enabled": "True" }, @@ -65,7 +65,7 @@ "query_sql": "'SELECT\\n CAST(o.ordered_at AS DATE) AS order_date,\\n COUNT(DISTINCT o.order_id) AS order_count,\\n SUM(o.quantity) AS waffles_ordered,\\n COUNT(DISTINCT o.customer_id) AS unique_customers\\nFROM __ref(\"stg_orders\") o\\nWHERE CAST(o.ordered_at AS DATE) >= CAST(@@@partition_start AS DATE)\\n AND CAST(o.ordered_at AS DATE) < CAST(@@@partition_end AS DATE)\\nGROUP BY CAST(o.ordered_at AS DATE)'", "references": "(CompileSqlReference(ref_kind=, ref_name='stg_orders', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='daily_order_partitioned', model_schema=None, description='Partition-tracked daily order summary using custom materialization.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_date', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/marts/daily_order_partitioned.sql'), line=16, column=5, end_line=16, end_column=15), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'waffles_ordered > 0'}, name='waffles_ordered_is_positive', description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), tags=())", + "schema_entry": "SchemaModelEntry(name='daily_order_partitioned', model_schema=None, description='Partition-tracked daily order summary using custom materialization.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_date', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/marts/daily_order_partitioned.sql'), line=16, column=5, end_line=16, end_column=15), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'waffles_ordered > 0'}, name='waffles_ordered_is_positive', description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n CAST(o.ordered_at AS DATE) AS order_date,\\n COUNT(DISTINCT o.order_id) AS order_count,\\n SUM(o.quantity) AS waffles_ordered,\\n COUNT(DISTINCT o.customer_id) AS unique_customers\\nFROM __ref(\"stg_orders\") o\\nWHERE CAST(o.ordered_at AS DATE) >= CAST(@@@partition_start AS DATE)\\n AND CAST(o.ordered_at AS DATE) < CAST(@@@partition_end AS DATE)\\nGROUP BY CAST(o.ordered_at AS DATE)', expanded_sql='SELECT\\n CAST(o.ordered_at AS DATE) AS order_date,\\n COUNT(DISTINCT o.order_id) AS order_count,\\n SUM(o.quantity) AS waffles_ordered,\\n COUNT(DISTINCT o.customer_id) AS unique_customers\\nFROM __ref(\"stg_orders\") o\\nWHERE CAST(o.ordered_at AS DATE) >= CAST(@@@partition_start AS DATE)\\n AND CAST(o.ordered_at AS DATE) < CAST(@@@partition_end AS DATE)\\nGROUP BY CAST(o.ordered_at AS DATE)', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -82,7 +82,7 @@ "query_sql": "'SELECT\\n CAST(o.ordered_at AS DATE) AS revenue_date,\\n COUNT(DISTINCT o.order_id) AS order_count,\\n SUM(o.quantity) AS waffles_sold,\\n SUM(p.amount_cents) AS total_revenue_cents,\\n SUM(p.amount_cents) / COUNT(DISTINCT o.order_id) AS avg_order_value_cents\\nFROM __ref(\"stg_orders\") o\\nINNER JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id AND p.payment_status = \\'success\\'\\nGROUP BY CAST(o.ordered_at AS DATE)'", "references": "(CompileSqlReference(ref_kind=, ref_name='stg_orders', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='stg_payments', ref_package=None, call_argument_count=None))", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='daily_revenue', model_schema=None, description='Daily revenue aggregation.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='revenue_date', type=None, nullable=False, description=None, meta={}, audits=(), location=SourceLocation(path=PosixPath('models/marts/daily_revenue.sql'), line=12, column=5, end_line=12, end_column=17), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'total_revenue_cents >= 0'}, name='revenue_is_non_negative', description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), tags=())", + "schema_entry": "SchemaModelEntry(name='daily_revenue', model_schema=None, description='Daily revenue aggregation.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='revenue_date', type=None, nullable=False, description=None, meta={}, audits=(), location=SourceLocation(path=PosixPath('models/marts/daily_revenue.sql'), line=12, column=5, end_line=12, end_column=17), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'total_revenue_cents >= 0'}, name='revenue_is_non_negative', description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n CAST(o.ordered_at AS DATE) AS revenue_date,\\n COUNT(DISTINCT o.order_id) AS order_count,\\n SUM(o.quantity) AS waffles_sold,\\n SUM(p.amount_cents) AS total_revenue_cents,\\n SUM(p.amount_cents) / COUNT(DISTINCT o.order_id) AS avg_order_value_cents\\nFROM __ref(\"stg_orders\") o\\nINNER JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id AND p.payment_status = \\'success\\'\\nGROUP BY CAST(o.ordered_at AS DATE)', expanded_sql='SELECT\\n CAST(o.ordered_at AS DATE) AS revenue_date,\\n COUNT(DISTINCT o.order_id) AS order_count,\\n SUM(o.quantity) AS waffles_sold,\\n SUM(p.amount_cents) AS total_revenue_cents,\\n SUM(p.amount_cents) / COUNT(DISTINCT o.order_id) AS avg_order_value_cents\\nFROM __ref(\"stg_orders\") o\\nINNER JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id AND p.payment_status = \\'success\\'\\nGROUP BY CAST(o.ordered_at AS DATE)', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -99,7 +99,7 @@ "query_sql": "'SELECT\\n c.customer_id,\\n c.first_name,\\n c.last_name,\\n c.email,\\n c.created_at AS customer_since,\\n COUNT(DISTINCT o.order_id) AS lifetime_orders,\\n COALESCE(SUM(p.amount_cents), 0) AS lifetime_spend_cents\\nFROM __ref(\"stg_customers\") c\\nLEFT JOIN __ref(\"stg_orders\") o ON c.customer_id = o.customer_id\\nLEFT JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id AND p.payment_status = \\'success\\'\\nGROUP BY c.customer_id, c.first_name, c.last_name, c.email, c.created_at'", "references": "(CompileSqlReference(ref_kind=, ref_name='stg_customers', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='stg_orders', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='stg_payments', ref_package=None, call_argument_count=None))", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='dim_customers', model_schema=None, description='Customer dimension with lifetime metrics.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='customer_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/marts/dim_customers.sql'), line=6, column=5, end_line=6, end_column=16), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='dim_customers', model_schema=None, description='Customer dimension with lifetime metrics.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='customer_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/marts/dim_customers.sql'), line=6, column=5, end_line=6, end_column=16), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n c.customer_id,\\n c.first_name,\\n c.last_name,\\n c.email,\\n c.created_at AS customer_since,\\n COUNT(DISTINCT o.order_id) AS lifetime_orders,\\n COALESCE(SUM(p.amount_cents), 0) AS lifetime_spend_cents\\nFROM __ref(\"stg_customers\") c\\nLEFT JOIN __ref(\"stg_orders\") o ON c.customer_id = o.customer_id\\nLEFT JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id AND p.payment_status = \\'success\\'\\nGROUP BY c.customer_id, c.first_name, c.last_name, c.email, c.created_at', expanded_sql='SELECT\\n c.customer_id,\\n c.first_name,\\n c.last_name,\\n c.email,\\n c.created_at AS customer_since,\\n COUNT(DISTINCT o.order_id) AS lifetime_orders,\\n COALESCE(SUM(p.amount_cents), 0) AS lifetime_spend_cents\\nFROM __ref(\"stg_customers\") c\\nLEFT JOIN __ref(\"stg_orders\") o ON c.customer_id = o.customer_id\\nLEFT JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id AND p.payment_status = \\'success\\'\\nGROUP BY c.customer_id, c.first_name, c.last_name, c.email, c.created_at', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -116,7 +116,7 @@ "query_sql": "'SELECT\\n o.order_id,\\n o.customer_id,\\n o.waffle_type_id,\\n w.waffle_name,\\n w.category AS waffle_category,\\n o.quantity,\\n w.price_cents * o.quantity AS line_total_cents,\\n o.ordered_at,\\n o.status AS order_status,\\n __udf(\"is_completed_order\")(o.status) AS is_completed_order,\\n __udf(\"is_completed_order_py\")(o.status) AS is_completed_order_py,\\n p.payment_method,\\n p.payment_status,\\n p.amount_cents AS payment_amount_cents\\nFROM __ref(\"stg_orders\") o\\nLEFT JOIN __seed(\"waffle_types\") w ON o.waffle_type_id = w.waffle_type_id\\nLEFT JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id'", "references": "(CompileSqlReference(ref_kind=, ref_name='is_completed_order', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='is_completed_order_py', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='stg_orders', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='waffle_types', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='stg_payments', ref_package=None, call_argument_count=None))", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='fact_orders', model_schema=None, description='Order fact table with waffle and payment details.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/marts/fact_orders.sql'), line=6, column=5, end_line=6, end_column=13), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='fact_orders', model_schema=None, description='Order fact table with waffle and payment details.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/marts/fact_orders.sql'), line=6, column=5, end_line=6, end_column=13), migrate_from=None),), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n o.order_id,\\n o.customer_id,\\n o.waffle_type_id,\\n w.waffle_name,\\n w.category AS waffle_category,\\n o.quantity,\\n @line_total_cents(\"w.price_cents\", \"o.quantity\") AS line_total_cents,\\n o.ordered_at,\\n o.status AS order_status,\\n __udf(\"is_completed_order\")(o.status) AS is_completed_order,\\n __udf(\"is_completed_order_py\")(o.status) AS is_completed_order_py,\\n p.payment_method,\\n p.payment_status,\\n p.amount_cents AS payment_amount_cents\\nFROM __ref(\"stg_orders\") o\\nLEFT JOIN __seed(\"waffle_types\") w ON o.waffle_type_id = w.waffle_type_id\\nLEFT JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id', expanded_sql='SELECT\\n o.order_id,\\n o.customer_id,\\n o.waffle_type_id,\\n w.waffle_name,\\n w.category AS waffle_category,\\n o.quantity,\\n w.price_cents * o.quantity AS line_total_cents,\\n o.ordered_at,\\n o.status AS order_status,\\n __udf(\"is_completed_order\")(o.status) AS is_completed_order,\\n __udf(\"is_completed_order_py\")(o.status) AS is_completed_order_py,\\n p.payment_method,\\n p.payment_status,\\n p.amount_cents AS payment_amount_cents\\nFROM __ref(\"stg_orders\") o\\nLEFT JOIN __seed(\"waffle_types\") w ON o.waffle_type_id = w.waffle_type_id\\nLEFT JOIN __ref(\"stg_payments\") p ON o.order_id = p.order_id', passes=((), (ExpansionSpan(source_start=124, source_end=172, output_start=124, output_end=150),)))", "sql_validation_enabled": "True" }, @@ -133,7 +133,7 @@ "query_sql": "'SELECT\\n h.activity_hour,\\n h.orders_placed,\\n d.orders_placed AS day_orders_placed,\\n h.waffles_ordered,\\n h.revenue_cents\\nFROM __ref(\"hourly_order_activity\") h\\nINNER JOIN __ref(\"daily_activity_rollup\") d\\n ON DATE_TRUNC(\\'day\\', h.activity_hour) = d.activity_day'", "references": "(CompileSqlReference(ref_kind=, ref_name='hourly_order_activity', ref_package=None, call_argument_count=None), CompileSqlReference(ref_kind=, ref_name='daily_activity_rollup', ref_package=None, call_argument_count=None))", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='hourly_activity_with_daily_context', model_schema=None, description='Hourly downstream microbatch model that depends on a coarser-grain daily upstream.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='activity_hour', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/marts/hourly_activity_with_daily_context.sql'), line=17, column=5, end_line=17, end_column=18), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'day_orders_placed >= orders_placed'}, name='day_orders_cover_hourly_orders', description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), tags=())", + "schema_entry": "SchemaModelEntry(name='hourly_activity_with_daily_context', model_schema=None, description='Hourly downstream microbatch model that depends on a coarser-grain daily upstream.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='activity_hour', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/marts/hourly_activity_with_daily_context.sql'), line=17, column=5, end_line=17, end_column=18), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'day_orders_placed >= orders_placed'}, name='day_orders_cover_hourly_orders', description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n h.activity_hour,\\n h.orders_placed,\\n d.orders_placed AS day_orders_placed,\\n h.waffles_ordered,\\n h.revenue_cents\\nFROM __ref(\"hourly_order_activity\") h\\nINNER JOIN __ref(\"daily_activity_rollup\") d\\n ON @timestamp_trunc(\\'day\\', \\'h.activity_hour\\') = d.activity_day', expanded_sql='SELECT\\n h.activity_hour,\\n h.orders_placed,\\n d.orders_placed AS day_orders_placed,\\n h.waffles_ordered,\\n h.revenue_cents\\nFROM __ref(\"hourly_order_activity\") h\\nINNER JOIN __ref(\"daily_activity_rollup\") d\\n ON DATE_TRUNC(\\'day\\', h.activity_hour) = d.activity_day', passes=((), (ExpansionSpan(source_start=211, source_end=253, output_start=211, output_end=245),)))", "sql_validation_enabled": "True" }, @@ -150,7 +150,7 @@ "query_sql": "'SELECT\\n DATE_TRUNC(\\'hour\\', o.ordered_at) AS activity_hour,\\n COUNT(*) AS orders_placed,\\n SUM(o.quantity) AS waffles_ordered,\\n SUM(o.line_total_cents) AS revenue_cents\\nFROM __ref(\"fact_orders\") o\\nGROUP BY DATE_TRUNC(\\'hour\\', o.ordered_at)'", "references": "(CompileSqlReference(ref_kind=, ref_name='fact_orders', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='hourly_order_activity', model_schema=None, description='Hourly order activity aggregated via microbatch incremental.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='activity_hour', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/marts/hourly_order_activity.sql'), line=17, column=5, end_line=17, end_column=18), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'orders_placed >= 0'}, name='orders_placed_is_non_negative', description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), tags=())", + "schema_entry": "SchemaModelEntry(name='hourly_order_activity', model_schema=None, description='Hourly order activity aggregated via microbatch incremental.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='activity_hour', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/marts/hourly_order_activity.sql'), line=17, column=5, end_line=17, end_column=18), migrate_from=None),), dynamic_columns=(), audits=(SchemaAuditInstance(definition_name='expression_is_true', arguments={'expression': 'orders_placed >= 0'}, name='orders_placed_is_non_negative', description=None, severity=None, run_scope='delta_and_final', always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n @timestamp_trunc(\\'hour\\', \\'o.ordered_at\\') AS activity_hour,\\n COUNT(*) AS orders_placed,\\n SUM(o.quantity) AS waffles_ordered,\\n SUM(o.line_total_cents) AS revenue_cents\\nFROM __ref(\"fact_orders\") o\\nGROUP BY @timestamp_trunc(\\'hour\\', \\'o.ordered_at\\')', expanded_sql='SELECT\\n DATE_TRUNC(\\'hour\\', o.ordered_at) AS activity_hour,\\n COUNT(*) AS orders_placed,\\n SUM(o.quantity) AS waffles_ordered,\\n SUM(o.line_total_cents) AS revenue_cents\\nFROM __ref(\"fact_orders\") o\\nGROUP BY DATE_TRUNC(\\'hour\\', o.ordered_at)', passes=((), (ExpansionSpan(source_start=9, source_end=49, output_start=9, output_end=41), ExpansionSpan(source_start=215, source_end=255, output_start=207, output_end=239))))", "sql_validation_enabled": "True" }, @@ -184,7 +184,7 @@ "query_sql": "'SELECT\\n id AS customer_id,\\n first_name,\\n last_name,\\n email,\\n created_at\\nFROM __source(\"raw_customers\")'", "references": "(CompileSqlReference(ref_kind=, ref_name='raw_customers', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='stg_customers', model_schema=None, description='Cleaned customer records.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='customer_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/staging/stg_customers.sql'), line=6, column=5, end_line=6, end_column=16), migrate_from=None), SchemaColumn(name='email', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/staging/stg_customers.sql'), line=7, column=5, end_line=7, end_column=10), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='stg_customers', model_schema=None, description='Cleaned customer records.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='customer_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/staging/stg_customers.sql'), line=6, column=5, end_line=6, end_column=16), migrate_from=None), SchemaColumn(name='email', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/staging/stg_customers.sql'), line=7, column=5, end_line=7, end_column=10), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n id AS customer_id,\\n first_name,\\n last_name,\\n email,\\n created_at\\nFROM __source(\"raw_customers\")', expanded_sql='SELECT\\n id AS customer_id,\\n first_name,\\n last_name,\\n email,\\n created_at\\nFROM __source(\"raw_customers\")', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -201,7 +201,7 @@ "query_sql": "'SELECT\\n id AS order_id,\\n customer_id,\\n waffle_type_id,\\n quantity,\\n ordered_at,\\n status\\nFROM __source(\"raw_orders\")'", "references": "(CompileSqlReference(ref_kind=, ref_name='raw_orders', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='stg_orders', model_schema=None, description='Cleaned order records.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=6, column=5, end_line=6, end_column=13), migrate_from=None), SchemaColumn(name='customer_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=7, column=5, end_line=7, end_column=16), migrate_from=None), SchemaColumn(name='waffle_type_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='relationships', arguments={'to': '__seed(\"waffle_types\")', 'field': 'waffle_type_id'}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=8, column=5, end_line=8, end_column=19), migrate_from=None), SchemaColumn(name='status', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='accepted_values', arguments={'values': ['placed', 'preparing', 'ready', 'completed', 'cancelled']}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None),), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=11, column=5, end_line=11, end_column=11), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='stg_orders', model_schema=None, description='Cleaned order records.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='order_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=6, column=5, end_line=6, end_column=13), migrate_from=None), SchemaColumn(name='customer_id', type=None, nullable=False, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=7, column=5, end_line=7, end_column=16), migrate_from=None), SchemaColumn(name='waffle_type_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='relationships', arguments={'to': '__seed(\"waffle_types\")', 'field': 'waffle_type_id'}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=8, column=5, end_line=8, end_column=19), migrate_from=None), SchemaColumn(name='status', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='accepted_values', arguments={'values': ['placed', 'preparing', 'ready', 'completed', 'cancelled']}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None),), location=SourceLocation(path=PosixPath('models/staging/stg_orders.sql'), line=11, column=5, end_line=11, end_column=11), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n id AS order_id,\\n customer_id,\\n waffle_type_id,\\n quantity,\\n ordered_at,\\n status\\nFROM __source(\"raw_orders\")', expanded_sql='SELECT\\n id AS order_id,\\n customer_id,\\n waffle_type_id,\\n quantity,\\n ordered_at,\\n status\\nFROM __source(\"raw_orders\")', passes=((), ()))", "sql_validation_enabled": "True" }, @@ -218,7 +218,7 @@ "query_sql": "'SELECT\\n id AS payment_id,\\n order_id,\\n amount_cents,\\n payment_method,\\n paid_at,\\n status AS payment_status\\nFROM __source(\"raw_payments\")'", "references": "(CompileSqlReference(ref_kind=, ref_name='raw_payments', ref_package=None, call_argument_count=None),)", "rejected_sql_analysis_opt_out": "None", - "schema_entry": "SchemaModelEntry(name='stg_payments', model_schema=None, description='Cleaned payment records.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='payment_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/staging/stg_payments.sql'), line=6, column=5, end_line=6, end_column=15), migrate_from=None), SchemaColumn(name='order_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None), SchemaAuditInstance(definition_name='relationships', arguments={'to': '__source(\"raw_orders\")', 'field': 'id'}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None)), location=SourceLocation(path=PosixPath('models/staging/stg_payments.sql'), line=7, column=5, end_line=7, end_column=13), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", + "schema_entry": "SchemaModelEntry(name='stg_payments', model_schema=None, description='Cleaned payment records.', type_enforcement=None, meta={}, columns=(SchemaColumn(name='payment_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='unique', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/staging/stg_payments.sql'), line=6, column=5, end_line=6, end_column=15), migrate_from=None), SchemaColumn(name='order_id', type=None, nullable=None, description=None, meta={}, audits=(SchemaAuditInstance(definition_name='not_null', arguments={}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None), SchemaAuditInstance(definition_name='relationships', arguments={'to': '__source(\"raw_orders\")', 'field': 'id'}, name=None, description=None, severity=None, run_scope=None, always_run=False, thresholds=None, minimum_samples=None, evidence_limit=None, meta=None)), location=SourceLocation(path=PosixPath('models/staging/stg_payments.sql'), line=7, column=5, end_line=7, end_column=13), migrate_from=None)), dynamic_columns=(), audits=(), tags=())", "sql_expansion": "CompiledSqlExpansion(authored_sql='SELECT\\n id AS payment_id,\\n order_id,\\n amount_cents,\\n payment_method,\\n paid_at,\\n status AS payment_status\\nFROM __source(\"raw_payments\")', expanded_sql='SELECT\\n id AS payment_id,\\n order_id,\\n amount_cents,\\n payment_method,\\n paid_at,\\n status AS payment_status\\nFROM __source(\"raw_payments\")', passes=((), ()))", "sql_validation_enabled": "True" } diff --git a/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py b/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py index 7e39314372..81526a0636 100644 --- a/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py +++ b/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py @@ -1,6 +1,10 @@ """Dataclass-backed native audit result integration cases.""" +from collections.abc import Callable from dataclasses import dataclass +from typing import Any + +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter @dataclass(frozen=True) @@ -9,3 +13,14 @@ class AuditResultWriteTestCase: description: str expected_row_count: int + + +@dataclass(frozen=True) +class AuditResultMetaWriteTestCase: + """One audit result write against a table that may predate the meta column.""" + + description: str + prepare: Callable[[DuckDbAdapter, Any], None] + audit_meta_json: str | None + expected_meta_column: bool + expected_meta: tuple[str | None, ...] diff --git a/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py b/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py index 84e3622d65..3065b0e0dd 100644 --- a/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py +++ b/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py @@ -1,9 +1,36 @@ """Helpers for native audit result integration tests.""" from datetime import UTC, datetime +from typing import Any +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.executor.audit_results.constants import AUDIT_RESULT_COLUMNS, COLUMN_AUDIT_META_JSON from sqlbuild.executor.audit_results.models import AuditResultRecord +LEGACY_AUDIT_RESULT_COLUMNS: tuple[str, ...] = tuple( + filter(lambda column: column != COLUMN_AUDIT_META_JSON, AUDIT_RESULT_COLUMNS) +) + + +def legacy_results_table_sql(adapter: DuckDbAdapter) -> str: + """Return results-table DDL as it was before the meta column was added.""" + + return adapter.render_create_audit_result_table_sql(database=None, schema="main").replace( + f", {COLUMN_AUDIT_META_JSON} VARCHAR", "" + ) + + +def create_legacy_results_table(adapter: DuckDbAdapter, connection: Any) -> None: + """Create the results table without the meta column.""" + + connection.execute(legacy_results_table_sql(adapter)) + + +def leave_results_table_absent(adapter: DuckDbAdapter, connection: Any) -> None: + """Leave the results table for the writer to create.""" + + del adapter, connection + def build_record() -> AuditResultRecord: return AuditResultRecord( diff --git a/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py b/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py index f78de9e5ff..9329fe7c00 100644 --- a/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py +++ b/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py @@ -1,18 +1,28 @@ """DuckDB integration tests for append-only audit result history.""" +from dataclasses import replace from pathlib import Path from typing import Any import pytest from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter -from sqlbuild.executor.audit_results.constants import AUDIT_RESULT_COLUMNS +from sqlbuild.executor.audit_results.constants import ( + AUDIT_RESULT_COLUMNS, + COLUMN_AUDIT_META_JSON, +) from sqlbuild.executor.audit_results.main._write import write_audit_result_records -from sqlbuild.executor.audit_results.models import AuditResultRecord +from sqlbuild.executor.audit_results.models import AuditResultColumnEvolution, AuditResultRecord from tests.integration.src.sqlbuild.executor.audit_results.main._test_types import ( + AuditResultMetaWriteTestCase, AuditResultWriteTestCase, ) -from tests.integration.src.sqlbuild.executor.audit_results.main.helpers import build_record +from tests.integration.src.sqlbuild.executor.audit_results.main.helpers import ( + LEGACY_AUDIT_RESULT_COLUMNS, + build_record, + create_legacy_results_table, + leave_results_table_absent, +) @pytest.mark.parametrize( @@ -94,5 +104,81 @@ def test_given_repeated_result_records_when_writing_then_auto_creates_and_append record.sql_digest, record.metadata_json, "false", + None, ) assert rows == [expected] * test_case.expected_row_count + + +@pytest.mark.parametrize( + "test_case", + ( + AuditResultMetaWriteTestCase( + description="legacy table gains the meta column when meta is written", + prepare=create_legacy_results_table, + audit_meta_json='{"region":"emea"}', + expected_meta_column=True, + expected_meta=('{"region":"emea"}',), + ), + AuditResultMetaWriteTestCase( + description="legacy table is left unchanged without meta", + prepare=create_legacy_results_table, + audit_meta_json=None, + expected_meta_column=False, + expected_meta=(None,), + ), + AuditResultMetaWriteTestCase( + description="new table stores meta", + prepare=leave_results_table_absent, + audit_meta_json='{"region":"emea"}', + expected_meta_column=True, + expected_meta=('{"region":"emea"}',), + ), + ), + ids=lambda case: case.description, +) +def test_given_results_table_generation_when_writing_meta_then_column_is_added_additively( + test_case: AuditResultMetaWriteTestCase, + tmp_path: Path, +) -> None: + adapter: DuckDbAdapter = DuckDbAdapter() + connection: Any = adapter.connect({"database": str(tmp_path / "audit_results.duckdb")}) + record: AuditResultRecord = replace(build_record(), audit_meta_json=test_case.audit_meta_json) + try: + test_case.prepare(adapter, connection) + write_audit_result_records( + connection=connection, + execute=adapter.execute, + database=None, + schema="main", + records=(record,), + render_qualified_name=adapter.render_qualified_name, + render_framework_type=adapter.render_framework_type, + render_create_table_sql=adapter.render_create_audit_result_table_sql, + render_create_index_sqls=adapter.render_create_audit_result_index_sqls, + column_evolution=AuditResultColumnEvolution( + read_columns=adapter.get_columns, render_add_columns=adapter.render_add_columns + ), + ) + columns: tuple[str, ...] = tuple( + column.name + for column in adapter.get_columns( + connection=connection, database=None, schema="main", name="_sqlbuild_audit_results" + ) + ) + stored_meta: tuple[str | None, ...] = tuple( + row[0] + for row in connection.execute( + f"SELECT json_extract_string(to_json(results), '$.{COLUMN_AUDIT_META_JSON}') " + "FROM main._sqlbuild_audit_results AS results" + ).fetchall() + ) + finally: + adapter.close(connection) + + assert (COLUMN_AUDIT_META_JSON in columns) is test_case.expected_meta_column + assert columns[: len(LEGACY_AUDIT_RESULT_COLUMNS)] == LEGACY_AUDIT_RESULT_COLUMNS + assert stored_meta == test_case.expected_meta + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instance.py b/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instance.py index 2ff4dfb4f9..c0c373ac63 100644 --- a/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instance.py +++ b/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instance.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import date from pathlib import Path import pytest @@ -174,3 +175,64 @@ def test_given_invalid_measurement_policy_options_when_parsing_then_clear_errors label="model orders", error_class=ValueError, ) + + +@pytest.mark.parametrize( + "test_case", + [ + ParseAuditInstanceTestCase( + description="meta is a reserved option kept unchanged", + raw_audit={ + "expression_is_true": { + "expression": "amount > 0", + "meta": {"region": "emea", "tiers": ["gold", 2], "weight": 1.5}, + } + }, + expected_definition_name="expression_is_true", + expected_always_run=False, + expected_argument_keys=("expression",), + ) + ], + ids=lambda case: case.description, +) +def test_given_audit_meta_when_parsing_then_meta_is_not_an_argument( + test_case: ParseAuditInstanceTestCase, +) -> None: + result: SchemaAuditInstance = parse_audit_instance( + raw_audit=test_case.raw_audit, + file_path=Path("models/orders.sql"), + label="model orders", + error_class=ValueError, + ) + + assert result.definition_name == test_case.expected_definition_name + assert tuple(sorted(result.arguments)) == test_case.expected_argument_keys + assert result.meta == {"region": "emea", "tiers": ["gold", 2], "weight": 1.5} + + +@pytest.mark.parametrize( + "test_case", + [ + ParseAuditInstanceErrorTestCase( + description="meta must be a mapping", + raw_audit={"not_null": {"meta": "emea"}}, + expected_error_fragment="audit 'not_null' 'meta' must be a mapping of JSON values", + ), + ParseAuditInstanceErrorTestCase( + description="meta values must be JSON", + raw_audit={"not_null": {"meta": {"since": date(2020, 1, 1)}}}, + expected_error_fragment=r"'meta' must map string keys to JSON values; meta\['since'\]", + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_audit_meta_when_parsing_then_compile_error_names_the_value( + test_case: ParseAuditInstanceErrorTestCase, +) -> None: + with pytest.raises(ValueError, match=test_case.expected_error_fragment): + parse_audit_instance( + raw_audit=test_case.raw_audit, + file_path=Path("models/orders.sql"), + label="model orders", + error_class=ValueError, + ) diff --git a/tests/unit/src/sqlbuild/executor/audit_results/_helpers/test_goldens.py b/tests/unit/src/sqlbuild/executor/audit_results/_helpers/test_goldens.py index 48b5e26f41..68dd03c486 100644 --- a/tests/unit/src/sqlbuild/executor/audit_results/_helpers/test_goldens.py +++ b/tests/unit/src/sqlbuild/executor/audit_results/_helpers/test_goldens.py @@ -2,5 +2,5 @@ # ruff: noqa: E501 -CREATE_TABLE_SQL_GOLDEN: str = "CREATE TABLE IF NOT EXISTS analytics._sqlbuild_audit_results (result_id VARCHAR NOT NULL, schema_version BIGINT NOT NULL, occurred_at TIMESTAMP NOT NULL, invocation_id VARCHAR NOT NULL, run_id VARCHAR NOT NULL, audit_name VARCHAR NOT NULL, audit_definition_name VARCHAR NOT NULL, audit_description VARCHAR, binding_key VARCHAR NOT NULL, definition_fingerprint VARCHAR NOT NULL, execution_fingerprint VARCHAR NOT NULL, evaluation_mode VARCHAR NOT NULL, run_scope_phase VARCHAR NOT NULL, attachment_kind VARCHAR NOT NULL, attached_target_kind VARCHAR, attached_target_name VARCHAR, attached_column_name VARCHAR, target_database VARCHAR, target_schema VARCHAR, target_name VARCHAR, severity VARCHAR NOT NULL, outcome VARCHAR NOT NULL, execution_error VARCHAR, violation_count BIGINT, measured_value VARCHAR, sample_count BIGINT, sample_unit VARCHAR, minimum_samples BIGINT, thresholds_json VARCHAR, evidence_json VARCHAR, evidence_count BIGINT, evidence_truncated VARCHAR, evidence_error VARCHAR, measurement_sql VARCHAR, evidence_sql VARCHAR, executed_sql VARCHAR, sql_digest VARCHAR, metadata_json VARCHAR, reused VARCHAR NOT NULL)" +CREATE_TABLE_SQL_GOLDEN: str = "CREATE TABLE IF NOT EXISTS analytics._sqlbuild_audit_results (result_id VARCHAR NOT NULL, schema_version BIGINT NOT NULL, occurred_at TIMESTAMP NOT NULL, invocation_id VARCHAR NOT NULL, run_id VARCHAR NOT NULL, audit_name VARCHAR NOT NULL, audit_definition_name VARCHAR NOT NULL, audit_description VARCHAR, binding_key VARCHAR NOT NULL, definition_fingerprint VARCHAR NOT NULL, execution_fingerprint VARCHAR NOT NULL, evaluation_mode VARCHAR NOT NULL, run_scope_phase VARCHAR NOT NULL, attachment_kind VARCHAR NOT NULL, attached_target_kind VARCHAR, attached_target_name VARCHAR, attached_column_name VARCHAR, target_database VARCHAR, target_schema VARCHAR, target_name VARCHAR, severity VARCHAR NOT NULL, outcome VARCHAR NOT NULL, execution_error VARCHAR, violation_count BIGINT, measured_value VARCHAR, sample_count BIGINT, sample_unit VARCHAR, minimum_samples BIGINT, thresholds_json VARCHAR, evidence_json VARCHAR, evidence_count BIGINT, evidence_truncated VARCHAR, evidence_error VARCHAR, measurement_sql VARCHAR, evidence_sql VARCHAR, executed_sql VARCHAR, sql_digest VARCHAR, metadata_json VARCHAR, reused VARCHAR NOT NULL, audit_meta_json VARCHAR)" INSERT_SQL_GOLDEN: str = "INSERT INTO analytics._sqlbuild_audit_results (result_id, schema_version, occurred_at, invocation_id, run_id, audit_name, audit_definition_name, audit_description, binding_key, definition_fingerprint, execution_fingerprint, evaluation_mode, run_scope_phase, attachment_kind, attached_target_kind, attached_target_name, attached_column_name, target_database, target_schema, target_name, severity, outcome, execution_error, violation_count, measured_value, sample_count, sample_unit, minimum_samples, thresholds_json, evidence_json, evidence_count, evidence_truncated, evidence_error, measurement_sql, evidence_sql, executed_sql, sql_digest, metadata_json, reused) VALUES ('result-1', CAST(1 AS BIGINT), '2026-01-02T03:04:05+00:00', 'invocation-1', 'run-1', 'valid_orders', 'dq_column_rate', 'Valid order percentage', 'orders.valid_orders', 'definition-fingerprint', 'execution-fingerprint', 'measurement', 'final', 'direct', 'model', 'orders', NULL, NULL, 'analytics', 'orders', 'warn', 'warn', NULL, CAST(NULL AS BIGINT), '99.25', CAST(12 AS BIGINT), 'rows', CAST(10 AS BIGINT), '{\"warn\":{\"below\":99.9}}', '[{\"note\":\"it''s invalid\"}]', CAST(1 AS BIGINT), 'false', NULL, 'SELECT AVG(is_valid) AS value FROM orders WHERE note = ''quoted''', NULL, 'SELECT AVG(is_valid) AS value FROM orders WHERE note = ''quoted''', 'sql-digest', NULL, 'false'), ('result-2', CAST(1 AS BIGINT), '2026-01-02T03:04:05+00:00', 'invocation-1', 'run-1', 'valid_orders', 'dq_column_rate', 'Valid order percentage', 'orders.valid_orders', 'definition-fingerprint', 'execution-fingerprint', 'measurement', 'final', 'direct', 'model', 'orders', NULL, NULL, 'analytics', 'orders', 'warn', 'warn', NULL, CAST(NULL AS BIGINT), NULL, CAST(12 AS BIGINT), 'rows', CAST(10 AS BIGINT), '{\"warn\":{\"below\":99.9}}', '[{\"note\":\"it''s invalid\"}]', CAST(1 AS BIGINT), 'false', NULL, 'SELECT AVG(is_valid) AS value FROM orders WHERE note = ''quoted''', NULL, 'SELECT AVG(is_valid) AS value FROM orders WHERE note = ''quoted''', 'sql-digest', NULL, 'false')" diff --git a/tests/unit/src/sqlbuild/executor/auditing/main/helpers.py b/tests/unit/src/sqlbuild/executor/auditing/main/helpers.py index 7db523e4b5..c46c3bee22 100644 --- a/tests/unit/src/sqlbuild/executor/auditing/main/helpers.py +++ b/tests/unit/src/sqlbuild/executor/auditing/main/helpers.py @@ -74,6 +74,8 @@ def writer_adapter() -> Any: render_framework_type=lambda *args, **kwargs: "", render_create_audit_result_table_sql=lambda *args, **kwargs: "", render_create_audit_result_index_sqls=lambda *args, **kwargs: (), + get_columns=lambda *args, **kwargs: (), + render_add_columns=lambda *args, **kwargs: (), ), ) diff --git a/tests/unit/src/sqlbuild/python_nodes/main/_test_types.py b/tests/unit/src/sqlbuild/python_nodes/main/_test_types.py index fb7bbf10b5..b38ab648a9 100644 --- a/tests/unit/src/sqlbuild/python_nodes/main/_test_types.py +++ b/tests/unit/src/sqlbuild/python_nodes/main/_test_types.py @@ -28,6 +28,20 @@ class InvalidAuditCaseNameTestCase: expected_error_fragment: str +@dataclass(frozen=True) +class AuditCaseMetaTestCase: + description: str + meta: dict[str, object] | None + expected_meta_json: str | None + + +@dataclass(frozen=True) +class InvalidAuditCaseMetaTestCase: + description: str + meta: object + expected_error_fragment: str + + @dataclass(frozen=True) class RetryDelayTestCase: description: str diff --git a/tests/unit/src/sqlbuild/python_nodes/main/test_audit_factory.py b/tests/unit/src/sqlbuild/python_nodes/main/test_audit_factory.py index ac848f03a9..b21c8e7225 100644 --- a/tests/unit/src/sqlbuild/python_nodes/main/test_audit_factory.py +++ b/tests/unit/src/sqlbuild/python_nodes/main/test_audit_factory.py @@ -2,6 +2,10 @@ from __future__ import annotations +import json +from datetime import date +from typing import cast + import pytest from sqlbuild.audits import ( @@ -17,8 +21,10 @@ from sqlbuild.errors.contracts.exceptions import SharedInputError from sqlbuild.python_nodes.models import AuditFactoryDefinition from tests.unit.src.sqlbuild.python_nodes.main._test_types import ( + AuditCaseMetaTestCase, AuditDecoratorTestCase, AuditThresholdHelperTestCase, + InvalidAuditCaseMetaTestCase, InvalidAuditCaseNameTestCase, ) @@ -81,5 +87,73 @@ def test_given_invalid_name_when_constructing_audit_case_then_input_error_is_rai AuditCase(name=test_case.name, definition="expression_is_true") +@pytest.mark.parametrize( + "test_case", + [ + AuditCaseMetaTestCase("absent", None, "null"), + AuditCaseMetaTestCase("empty", {}, "{}"), + AuditCaseMetaTestCase( + "nested JSON values", + { + "region": "emea", + "priority": 2, + "weight": 0.5, + "active": True, + "owner": None, + "channels": ["web", ("store", 3)], + "exclusion": {"reason": "before launch", "since": "2020-01-01"}, + }, + '{"region": "emea", "priority": 2, "weight": 0.5, "active": true, "owner": null, ' + '"channels": ["web", ["store", 3]], ' + '"exclusion": {"reason": "before launch", "since": "2020-01-01"}}', + ), + ], + ids=lambda case: case.description, +) +def test_given_json_meta_when_constructing_audit_case_then_meta_is_kept_unchanged( + test_case: AuditCaseMetaTestCase, +) -> None: + case: AuditCase = AuditCase( + name="valid_orders", definition="expression_is_true", meta=test_case.meta + ) + + assert case.meta is test_case.meta + assert json.dumps(case.meta) == test_case.expected_meta_json + + +@pytest.mark.parametrize( + "test_case", + [ + InvalidAuditCaseMetaTestCase( + "not a mapping", ["emea"], "meta must be a mapping of JSON values" + ), + InvalidAuditCaseMetaTestCase( + "non-string key", {1: "emea"}, "key 1 in meta is not a string" + ), + InvalidAuditCaseMetaTestCase( + "nested date", {"window": {"since": date(2020, 1, 1)}}, r"meta\['window'\]\['since'\]" + ), + InvalidAuditCaseMetaTestCase( + "non-finite number", {"weight": float("nan")}, "is not JSON-serialisable: nan" + ), + InvalidAuditCaseMetaTestCase( + "set in list", {"regions": [{"emea"}]}, r"meta\['regions'\]\[0\]" + ), + ], + ids=lambda case: case.description, +) +def test_given_non_json_meta_when_constructing_audit_case_then_input_error_names_the_value( + test_case: InvalidAuditCaseMetaTestCase, +) -> None: + with pytest.raises(SharedInputError, match=test_case.expected_error_fragment) as error: + AuditCase( + name="valid_orders", + definition="expression_is_true", + meta=cast("dict[str, object]", test_case.meta), + ) + + assert "AuditCase 'valid_orders'" in str(error.value) + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/website/src/content/docs/docs/concepts/audits.mdx b/website/src/content/docs/docs/concepts/audits.mdx index a764da23ca..0dc17ff775 100644 --- a/website/src/content/docs/docs/concepts/audits.mdx +++ b/website/src/content/docs/docs/concepts/audits.mdx @@ -88,12 +88,42 @@ def order_quality(): Attach it with `MODEL (audit_factories [order_quality])`. Generated cases compile to the same audit contract as directly authored instances. +## Audit metadata + +Use `meta` to attach structured context, such as the segment a case covers or why its condition +excludes some rows, so dashboards can group and filter results without parsing descriptions: + +```python +AuditCase( + name="valid_order_rate__emea", + definition="valid_order_rate", + arguments={"condition": "order_id IS NOT NULL AND region = 'emea'"}, + thresholds=MeasurementThresholds(error=below(99)), + meta={ + "region": "emea", + "exclusion": "orders before 2020", + "exclusion_reason": "imported without order IDs", + }, +) +``` + +SQL attachments accept the same option, for example +`valid_order_rate (condition "order_id IS NOT NULL", meta (region "emea"), ...)`. + +`meta` must be a mapping with string keys and JSON values; anything else is a compile error that +names the offending value. It is never passed to the audit SQL and never affects evaluation, +thresholds, or audit fingerprints, so changing only `meta` reruns nothing that a `description` +change would not. It is recorded with each result in the `audit_meta_json` column of +`_sqlbuild_audit_results`, as `audit_meta` in the `audit_completed` lifecycle event, and as +`audit_meta` on audit checks in JSON output. + ## Result history Native warehouse adapters best-effort append confirmed audit outcomes to `_sqlbuild_audit_results`. Rows are immutable and use deterministic IDs, so retrying the same result is idempotent. Projection failure is reported separately and does not change the audit outcome or -command exit code. Lifecycle sinks can also consume the corresponding `audit_completed` fact, which +command exit code. A results table created before `audit_meta_json` existed gains that nullable +column the first time a result with `meta` is recorded. Lifecycle sinks can also consume the corresponding `audit_completed` fact, which is published as each audit finishes. ## Built-in audits From d33e8ff20a814dea0866d441b16f0cc853d13897 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:50:13 +0100 Subject: [PATCH 04/11] fix: keep unit-test output storage optional for project adapters StrictAdapter now provides create_sql_test_output and drop_sql_test_output defaults that store nothing, so the executor keeps the inline comparison query and project adapters need no new methods. Built-in adapters that support single-read comparison override them. --- .../adapter/contract/classes/base_adapter.py | 11 --- .../contract/classes/strict_adapter.py | 13 +-- .../bigquery/classes/bigquery_adapter.py | 11 --- .../databricks/classes/databricks_adapter.py | 11 --- .../sqlserver/classes/sqlserver_adapter.py | 11 --- .../classes/strict_adapter/_test_types.py | 15 ++++ .../classes/strict_adapter/helpers.py | 22 +++++ .../strict_adapter/test_strict_adapter.py | 87 +++++++++++++++++++ .../content/docs/docs/concepts/adapters.mdx | 2 + 9 files changed, 133 insertions(+), 50 deletions(-) diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index 2ea80590f8..e3eb73e660 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -156,17 +156,6 @@ def interrupt_connection(self, connection: Any) -> bool: del connection return False - def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: - """A relation holding a unit test's evaluated output, or `None` to compare inline.""" - - del connection, sql, name - return None - - def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: - """Remove a relation `create_sql_test_output` returned.""" - - del connection, relation - def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: """Return whether ``list_relations`` reports every relation in this scope.""" diff --git a/src/sqlbuild/adapter/contract/classes/strict_adapter.py b/src/sqlbuild/adapter/contract/classes/strict_adapter.py index ebddef75ae..ac1bf07a68 100644 --- a/src/sqlbuild/adapter/contract/classes/strict_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/strict_adapter.py @@ -105,15 +105,16 @@ def interrupt_connection(self, connection: Any) -> bool: """Cancel a connection's running statement from another thread; False if unsupported.""" ... - @abstractmethod def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: - """Evaluate a unit test's actual output once and return a relation reading it.""" - ... + """Optional: a relation holding a unit test's output, or `None` to compare inline.""" + + del connection, sql, name + return None - @abstractmethod def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: - """Remove a relation `create_sql_test_output` returned.""" - ... + """Optional: remove a relation `create_sql_test_output` returned.""" + + del connection, relation @abstractmethod def lists_relations_authoritatively(self, *, database: str | None, schema: str) -> bool: diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index e8e55487c9..fa50a792f6 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -1508,17 +1508,6 @@ def close(self, connection: _BigQueryConnection) -> None: def supports_transactions(self) -> bool: return False - def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: - """Keep unit-test comparisons inline; this adapter stores no test outputs.""" - - del connection, sql, name - return None - - def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: - """Nothing to drop: no test outputs are stored.""" - - del connection, relation - def interrupt_connection(self, connection: Any) -> bool: """Request cancellation of the connection's jobs; BigQuery stops them asynchronously.""" diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index 164f32b272..91ff016fd2 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -1166,17 +1166,6 @@ def render_apply_historical_check_snapshot_changes( invalidate_hard_deletes=invalidate_hard_deletes, ) - def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: - """Keep unit-test comparisons inline; this adapter stores no test outputs.""" - - del connection, sql, name - return None - - def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: - """Nothing to drop: no test outputs are stored.""" - - del connection, relation - def interrupt_connection(self, connection: Any) -> bool: """Request cancellation of the connection's running statements.""" diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index c7628b588b..87482677b8 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -2799,17 +2799,6 @@ def supports_table_functions(self) -> bool: def supports_unqualified_function_fingerprints(self) -> bool: return False - def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: - """Keep unit-test comparisons inline; this adapter stores no test outputs.""" - - del connection, sql, name - return None - - def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: - """Nothing to drop: no test outputs are stored.""" - - del connection, relation - def interrupt_connection(self, connection: Any) -> bool: """Cancel the connection's running statement through the driver.""" diff --git a/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/_test_types.py b/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/_test_types.py index 58cdb56f82..0bb976dcc0 100644 --- a/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/_test_types.py +++ b/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/_test_types.py @@ -18,3 +18,18 @@ class TableFreshnessMetadataUnsupportedTestCase: description: str expected_supports_metadata: bool expected_error_fragment: str + + +@dataclass(frozen=True) +class OptionalStrictMethodTestCase: + description: str + expected_optional_methods: frozenset[str] + expected_output_relation: None + + +@dataclass(frozen=True) +class StrictAdapterInlineSqlTestTestCase: + description: str + expected_total: int + expected_outcome: str + expected_statement_count: int diff --git a/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/helpers.py b/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/helpers.py index 9c3f7649fe..899eff0e2d 100644 --- a/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/helpers.py +++ b/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/helpers.py @@ -9,6 +9,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter from sqlbuild.adapter.discovery.main.builtins import builtin_adapter_classes +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter def first_class_adapter_contract_violations() -> tuple[str, ...]: @@ -86,3 +87,24 @@ def _skip_method_owner(owners: dict[str, type[object]], cls: type[object]) -> No _METHOD_OWNER_COLLECTORS: MappingProxyType[bool, Callable[..., None]] = MappingProxyType( {False: _skip_method_owner, True: _collect_method_owner} ) + + +def delegating_strict_adapter_class() -> type[StrictAdapter]: + """A project-style `StrictAdapter` that implements only the required methods, via DuckDB.""" + + inner: DuckDbAdapter = DuckDbAdapter() + required: frozenset[str] = cast(Any, StrictAdapter).__abstractmethods__ + namespace: dict[str, object] = { + name: _delegate(inner=inner, name=name) for name in sorted(required) + } + namespace["adapter_name"] = "project_duckdb" + namespace["sql_analysis_dialect_name"] = "duckdb" + return cast(type[StrictAdapter], type("ProjectDuckDbAdapter", (StrictAdapter,), namespace)) + + +def _delegate(*, inner: DuckDbAdapter, name: str) -> Callable[..., object]: + def method(self: StrictAdapter, *args: object, **kwargs: object) -> object: + del self + return getattr(inner, name)(*args, **kwargs) + + return method diff --git a/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/test_strict_adapter.py b/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/test_strict_adapter.py index b2c86971cf..f1951d4d43 100644 --- a/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/test_strict_adapter.py +++ b/tests/unit/src/sqlbuild/adapter/contract/classes/strict_adapter/test_strict_adapter.py @@ -2,20 +2,34 @@ import inspect from collections import defaultdict +from pathlib import Path from typing import Any, cast import pytest +from sqlbuild._native import ( + ChainStep, + CompiledObjectKey, + CompiledResourceType, + PlanOutput, + SqlTestPlanEntry, +) from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.executor.pipeline._helpers.testing import run_test_pipeline +from sqlbuild.executor.testing.models import SqlTestExecutionResult +from sqlbuild.observability import EventDispatcher, LifecycleEvent, dispatcher_scope from tests.unit.src.sqlbuild.adapter.contract.classes.strict_adapter._test_types import ( FirstClassAdapterImplementationContractTestCase, + OptionalStrictMethodTestCase, StrictAdapterContractTestCase, + StrictAdapterInlineSqlTestTestCase, TableFreshnessMetadataUnsupportedTestCase, ) from tests.unit.src.sqlbuild.adapter.contract.classes.strict_adapter.helpers import ( + delegating_strict_adapter_class, first_class_adapter_contract_violations, ) @@ -65,6 +79,79 @@ def test_given_first_class_adapters_when_checking_contract_then_no_method_comes_ assert violations == test_case.expected_violations +@pytest.mark.parametrize( + "test_case", + [ + OptionalStrictMethodTestCase( + description="unit-test output storage is optional and defaults to inline", + expected_optional_methods=frozenset({"create_sql_test_output", "drop_sql_test_output"}), + expected_output_relation=None, + ) + ], + ids=lambda case: case.description, +) +def test_given_strict_adapter_when_reading_optional_methods_then_defaults_keep_inline_comparison( + test_case: OptionalStrictMethodTestCase, +) -> None: + adapter: StrictAdapter = delegating_strict_adapter_class()() + abstract_methods: frozenset[str] = cast(Any, StrictAdapter).__abstractmethods__ + + relation: str | None = adapter.create_sql_test_output( + connection=None, sql="SELECT 1", name="orders_output" + ) + adapter.drop_sql_test_output(connection=None, relation="orders_output") + + assert test_case.expected_optional_methods.isdisjoint(abstract_methods) + assert test_case.expected_optional_methods <= frozenset(StrictAdapter.__dict__) + assert relation is test_case.expected_output_relation + + +@pytest.mark.parametrize( + "test_case", + [ + StrictAdapterInlineSqlTestTestCase( + description="project strict adapter without output storage compares inline", + expected_total=1, + expected_outcome="pass", + expected_statement_count=1, + ) + ], + ids=lambda case: case.description, +) +def test_given_strict_adapter_without_optional_methods_when_running_unit_test_then_inline_passes( + tmp_path: Path, test_case: StrictAdapterInlineSqlTestTestCase +) -> None: + adapter: StrictAdapter = delegating_strict_adapter_class()() + test_entry: SqlTestPlanEntry = SqlTestPlanEntry( + key=CompiledObjectKey(resource_type=CompiledResourceType.SQL_TEST, name="order_totals"), + name="order_totals", + chain=( + ChainStep( + model_name="orders", + resolved_sql="SELECT 1 AS order_id, 10 AS total", + expected_cte_sql="SELECT 1 AS order_id, 10 AS total", + ), + ), + ) + events: list[LifecycleEvent] = [] + dispatcher: EventDispatcher = EventDispatcher() + dispatcher.subscribe_lifecycle(subscriber=events.append, accepts_opaque=False) + + with dispatcher_scope(dispatcher): + results: tuple[SqlTestExecutionResult, ...] = run_test_pipeline( + plan=PlanOutput(test_entries=(test_entry,)), + connection_config={"database": str(tmp_path / "orders.duckdb")}, + adapter=cast(Any, adapter), + run_id="strict-run", + ) + + assert len(results) == test_case.expected_total + assert results[0].outcome == test_case.expected_outcome, results[0].error_message + assert [event.event_type for event in events].count( + "statement_started" + ) == test_case.expected_statement_count + + @pytest.mark.parametrize( "test_case", [ diff --git a/website/src/content/docs/docs/concepts/adapters.mdx b/website/src/content/docs/docs/concepts/adapters.mdx index bdbd328538..d65f2b3a77 100644 --- a/website/src/content/docs/docs/concepts/adapters.mdx +++ b/website/src/content/docs/docs/concepts/adapters.mdx @@ -94,6 +94,8 @@ Set `sql_analysis_dialect_name` to the SQL dialect name that matches your engine For full control with no inherited defaults, subclass `StrictAdapter` instead. Every method is abstract and must be implemented explicitly. SQLBuild raises a clear error listing any unimplemented methods. +A few `StrictAdapter` methods are optional and have a documented default. `create_sql_test_output` and `drop_sql_test_output` let an adapter store a unit test's actual output once so every comparison reads it (see [Single-read comparison](/docs/concepts/testing#single-read-comparison)). By default `create_sql_test_output` returns `None`, and SQLBuild runs the test's inline comparison query instead, so adapters that do not implement them keep working unchanged. + ### Discovery rules - SQLBuild discovers all `.py` files under `adapters/` recursively (excluding `__init__.py` and files starting with `_`) From c59123a3bad3d72746016945fab12196bba40278 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:50:14 +0100 Subject: [PATCH 05/11] feat: count timed-out models and tests separately in build and test summaries The sqb build footer reports TIMEOUT= apart from FAIL=, and build JSON adds summary.timeout_count, with failure_count excluding timeouts. sqb test JSON adds timeout_count likewise. A timeout still fails the run: the exit code is non-zero and the JSON status is failed. --- .../src/build_aggregate/constants.rs | 2 + .../main/aggregate_build_counts.rs | 10 +++- .../src/build_aggregate/models.rs | 1 + .../src/build_progress/_helpers/footer.rs | 19 ++++--- .../main/format_build_footer.rs | 29 +++++----- .../_helpers/result_output/aggregation.rs | 1 + .../_helpers/result_output/documents.rs | 8 ++- src/sqlbuild/executor/build/models.py | 1 + src/sqlbuild/executor/pipeline/main/run.py | 2 +- .../test_build_no_tests_no_audits.py | 3 + .../cli/commands/main/build/test_build.py | 1 + .../cli/commands/main/test/_test_types.py | 9 +++ .../cli/commands/main/test/helpers.py | 20 +++++++ .../commands/main/test/test_cost_guards.py | 56 +++++++++++++------ .../test_execution_protocol_v1.py | 1 + website/src/content/docs/docs/cli/build.mdx | 6 +- website/src/content/docs/docs/cli/test.mdx | 5 +- 17 files changed, 129 insertions(+), 45 deletions(-) diff --git a/crates/sqlbuild-core/src/build_aggregate/constants.rs b/crates/sqlbuild-core/src/build_aggregate/constants.rs index 0ac789c3c3..ec787e5d19 100644 --- a/crates/sqlbuild-core/src/build_aggregate/constants.rs +++ b/crates/sqlbuild-core/src/build_aggregate/constants.rs @@ -8,3 +8,5 @@ pub(crate) const AUDIT_WARN: &str = "warn"; pub(crate) const AUDIT_ERROR: &str = "error"; pub(crate) const AUDIT_INSUFFICIENT: &str = "insufficient"; pub(crate) const SOURCE_FRESHNESS_BLOCKED_CODE: &str = "B005"; +pub(crate) const TEST_TIMEOUT: &str = "timeout"; +pub(crate) const ITEM_TIMEOUT_CODE: &str = "C417"; diff --git a/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs b/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs index f749904e5b..a849b886d6 100644 --- a/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs +++ b/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs @@ -1,18 +1,19 @@ //! Aggregate counts and overall status of one build. use crate::build_aggregate::constants::{ - AUDIT_ERROR, AUDIT_INSUFFICIENT, AUDIT_WARN, FAILED, SKIPPED, SOURCE_FRESHNESS_BLOCKED_CODE, - SUCCESS, TEST_PASS, + AUDIT_ERROR, AUDIT_INSUFFICIENT, AUDIT_WARN, FAILED, ITEM_TIMEOUT_CODE, SKIPPED, + SOURCE_FRESHNESS_BLOCKED_CODE, SUCCESS, TEST_PASS, TEST_TIMEOUT, }; use crate::build_aggregate::models::{BuildCounts, BuildResultFacts}; -/// Success, failure, skipped, warning and insufficient counts, and whether the build failed. +/// Success, failure, timeout, skipped, warning and insufficient counts, and whether it failed. #[must_use] pub fn aggregate_build_counts(facts: &BuildResultFacts) -> BuildCounts { let mut counts: BuildCounts = BuildCounts::default(); for model in &facts.models { match model.status.as_str() { SUCCESS => counts.success += 1, + FAILED if model.error_code.as_deref() == Some(ITEM_TIMEOUT_CODE) => counts.timeout += 1, FAILED => counts.failure += 1, SKIPPED => { counts.skipped += 1; @@ -42,6 +43,8 @@ pub fn aggregate_build_counts(facts: &BuildResultFacts) -> BuildCounts { for outcome in &facts.test_outcomes { if outcome == TEST_PASS { counts.success += 1; + } else if outcome == TEST_TIMEOUT { + counts.timeout += 1; } else { counts.failure += 1; } @@ -58,6 +61,7 @@ pub fn aggregate_build_counts(facts: &BuildResultFacts) -> BuildCounts { counts.warning += matching(&facts.source_audit_outcomes, AUDIT_WARN) + matching(&facts.end_audit_outcomes, AUDIT_WARN); counts.failed = counts.failure > 0 + || counts.timeout > 0 || facts .end_audit_outcomes .iter() diff --git a/crates/sqlbuild-core/src/build_aggregate/models.rs b/crates/sqlbuild-core/src/build_aggregate/models.rs index 2ceed85105..066374563b 100644 --- a/crates/sqlbuild-core/src/build_aggregate/models.rs +++ b/crates/sqlbuild-core/src/build_aggregate/models.rs @@ -27,6 +27,7 @@ pub struct BuildResultFacts { pub struct BuildCounts { pub success: usize, pub failure: usize, + pub timeout: usize, pub skipped: usize, pub warning: usize, pub insufficient: usize, diff --git a/crates/sqlbuild-core/src/build_progress/_helpers/footer.rs b/crates/sqlbuild-core/src/build_progress/_helpers/footer.rs index f3601560cb..a9ac63ec88 100644 --- a/crates/sqlbuild-core/src/build_progress/_helpers/footer.rs +++ b/crates/sqlbuild-core/src/build_progress/_helpers/footer.rs @@ -4,17 +4,18 @@ use crate::build_progress::_helpers::error_rows::failure_error_block; use crate::build_progress::_helpers::text::pad; use crate::build_progress::constants::{ DELTA_SUFFIX, ERROR, FAILED, PASS, PYTHON_NODE_NAME_WIDTH, SKIPPED, SKIPPED_ROLLUP_CAP, - SUCCESS, WARN, + SUCCESS, TIMEOUT_CODE, TIMEOUT_OUTCOME, WARN, }; use crate::build_progress::models::{AuditResultText, BuildResultText, ErrorText, PythonNodeText}; use crate::cli_text::models::CliStyle; -/// Pass, warn, fail and skip counts of the footer. +/// Pass, warn, fail, timeout and skip counts of the footer. #[derive(Default)] pub(crate) struct FooterCounts { pub(crate) pass: i64, pub(crate) warn: i64, pub(crate) fail: i64, + pub(crate) timeout: i64, pub(crate) skip: i64, } @@ -42,7 +43,11 @@ fn build_audits(result: &BuildResultText) -> impl Iterator FooterCounts { let mut counts: FooterCounts = FooterCounts::default(); for model in &result.models { - counts.status(&model.status, !model.warnings.is_empty()); + if model.status == FAILED && model.error.code.as_deref() == Some(TIMEOUT_CODE) { + counts.timeout += 1; + } else { + counts.status(&model.status, !model.warnings.is_empty()); + } } for audit in build_audits(result) { match audit.outcome.as_str() { @@ -63,10 +68,10 @@ pub(crate) fn count_results(result: &BuildResultText, python: &[PythonNodeText]) counts.status(&load.status, !load.warnings.is_empty()); } for test in &result.tests { - if test.outcome == PASS { - counts.pass += 1; - } else { - counts.fail += 1; + match test.outcome.as_str() { + PASS => counts.pass += 1, + TIMEOUT_OUTCOME => counts.timeout += 1, + _ => counts.fail += 1, } } for node in python { diff --git a/crates/sqlbuild-core/src/build_progress/main/format_build_footer.rs b/crates/sqlbuild-core/src/build_progress/main/format_build_footer.rs index 97255700f9..c4b17a9bd6 100644 --- a/crates/sqlbuild-core/src/build_progress/main/format_build_footer.rs +++ b/crates/sqlbuild-core/src/build_progress/main/format_build_footer.rs @@ -27,19 +27,22 @@ pub fn format_build_footer( ) -> String { let counts: FooterCounts = count_results(result, python); let insufficient: i64 = result.insufficient_count; - let total: i64 = counts.pass + counts.warn + counts.fail + counts.skip; - let summary: String = format_summary_footer( - style, - &[ - count("PASS", counts.pass), - count("WARN", counts.warn), - count("FAIL", counts.fail), - count("INSUFFICIENT", insufficient), - count("SKIP", counts.skip), - count("TOTAL", total + insufficient), - ], - Some(&format!("{elapsed:.2}s")), - ); + let total: i64 = counts.pass + counts.warn + counts.fail + counts.timeout + counts.skip; + let mut summary_counts: Vec = vec![ + count("PASS", counts.pass), + count("WARN", counts.warn), + count("FAIL", counts.fail), + ]; + if counts.timeout > 0 { + summary_counts.push(count("TIMEOUT", counts.timeout)); + } + summary_counts.extend([ + count("INSUFFICIENT", insufficient), + count("SKIP", counts.skip), + count("TOTAL", total + insufficient), + ]); + let summary: String = + format_summary_footer(style, &summary_counts, Some(&format!("{elapsed:.2}s"))); let python_failed: bool = python.iter().any(|node| node.status == FAILED); let (state, label) = if result.failed || python_failed { (CompletionState::Fail, "Completed with errors") diff --git a/crates/sqlbuild-python/src/bindings/_helpers/result_output/aggregation.rs b/crates/sqlbuild-python/src/bindings/_helpers/result_output/aggregation.rs index ebc31c43f0..77c916c459 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/result_output/aggregation.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/result_output/aggregation.rs @@ -76,6 +76,7 @@ pub(crate) fn aggregated_build_result<'py>( keywords.extend([ ("success_count", counts.success.into_bound_py_any(py)?), ("failure_count", counts.failure.into_bound_py_any(py)?), + ("timeout_count", counts.timeout.into_bound_py_any(py)?), ("skipped_count", counts.skipped.into_bound_py_any(py)?), ("warning_count", counts.warning.into_bound_py_any(py)?), ( diff --git a/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs b/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs index d9d92598d3..e36dc3c57f 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs @@ -26,6 +26,7 @@ const SUCCESS: &str = "success"; const SKIPPED: &str = "skipped"; const ERROR_OUTCOME: &str = "error"; const PASS_OUTCOME: &str = "pass"; +const TIMEOUT_OUTCOME: &str = "timeout"; const WARN_OUTCOME: &str = "warn"; const BUILD_OPTIONS: (&[&str], &[&str]) = ( &["result", "plan"], @@ -347,6 +348,7 @@ fn build_json<'py>( "failure_count", count(&counted, "failure_count")? + python_fail + check_fail, ), + ("timeout_count", count(&counted, "timeout_count")?), ( "skipped_count", count(&counted, "skipped_count")? + python_skipped, @@ -494,17 +496,19 @@ pub(crate) fn format_test_execution_json<'py>( let projection: Projection<'py> = Projection::installed(py)?; let results: Vec> = projection.terminal_results(&results)?; let pass_count: usize = count_status(&results, "outcome", PASS_OUTCOME)?; - let fail_count: usize = results.len() - pass_count; + let timeout_count: usize = count_status(&results, "outcome", TIMEOUT_OUTCOME)?; + let fail_count: usize = results.len() - pass_count - timeout_count; emitted( &projection, Document { command: "test".to_owned(), - status: outcome_status(fail_count), + status: outcome_status(fail_count + timeout_count), summary: summary( py, &[ ("pass_count", pass_count), ("fail_count", fail_count), + ("timeout_count", timeout_count), ("total_count", results.len()), ], )?, diff --git a/src/sqlbuild/executor/build/models.py b/src/sqlbuild/executor/build/models.py index c56c637503..818b7d9eaf 100644 --- a/src/sqlbuild/executor/build/models.py +++ b/src/sqlbuild/executor/build/models.py @@ -230,6 +230,7 @@ class BuildExecutionResult: end_audit_results: tuple[AuditExecutionResult, ...] = field(default_factory=tuple) success_count: int = 0 failure_count: int = 0 + timeout_count: int = 0 skipped_count: int = 0 warning_count: int = 0 insufficient_count: int = 0 diff --git a/src/sqlbuild/executor/pipeline/main/run.py b/src/sqlbuild/executor/pipeline/main/run.py index eba798e2f5..01a639d716 100644 --- a/src/sqlbuild/executor/pipeline/main/run.py +++ b/src/sqlbuild/executor/pipeline/main/run.py @@ -137,7 +137,7 @@ def run_build_pipeline( "run_kind": "build", "duration_ms": int((time.monotonic() - started) * 1000), "succeeded_count": result.success_count, - "failed_count": result.failure_count, + "failed_count": result.failure_count + result.timeout_count, "skipped_count": result.skipped_count, }, ) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/no_tests_no_audits/test_build_no_tests_no_audits.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/no_tests_no_audits/test_build_no_tests_no_audits.py index 78727ef17e..3100cd5be9 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/no_tests_no_audits/test_build_no_tests_no_audits.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/no_tests_no_audits/test_build_no_tests_no_audits.py @@ -572,6 +572,7 @@ def test_given_task_selector_with_json_output_when_running_run_then_json_include assert payload["summary"] == { "success_count": 1, "failure_count": 0, + "timeout_count": 0, "skipped_count": 0, "warning_count": 0, "python_check_pass_count": 0, @@ -641,6 +642,7 @@ def test_given_failing_task_selector_when_running_run_then_command_fails( assert payload["summary"] == { "success_count": 0, "failure_count": 1, + "timeout_count": 0, "skipped_count": 0, "warning_count": 0, "python_check_pass_count": 0, @@ -1095,6 +1097,7 @@ def test_given_skip_and_asset_selection_with_json_when_running_run_then_json_rec assert payload["summary"] == { "success_count": 1, "failure_count": 0, + "timeout_count": 0, "skipped_count": 1, "warning_count": 0, "python_check_pass_count": 0, diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_build.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_build.py index 0e80b749db..57a5c0b7bc 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_build.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_build.py @@ -1294,6 +1294,7 @@ def test_given_python_asset_with_json_output_when_building_then_json_includes_py assert payload["summary"] == { "success_count": 1, "failure_count": 0, + "timeout_count": 0, "skipped_count": 0, "warning_count": 0, "python_check_pass_count": 0, diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py index dc8a025b84..113355dba5 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py @@ -417,3 +417,12 @@ class ScenarioTimeoutE2ETestCase: header_timeout: str command: tuple[str, ...] expected_output_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class ModelTimeoutJsonE2ETestCase: + """One slow model's build reported as JSON.""" + + description: str + expected_status: str + expected_summary_counts: dict[str, int] diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py index 3012f029bd..35d5e955cb 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py @@ -1988,3 +1988,23 @@ def stored_test_outputs(*, database_path: Path) -> list[str]: "WHERE table_name LIKE '__sqb_test_output%'" ).fetchall() return [row[0] for row in rows] + + +def slow_model_files(*, max_duration: str) -> dict[str, str]: + """A slow model with `max_duration`, a model reading it and an independent model.""" + + return { + "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML, + "models/order_scans.sql": ( + "MODEL (description 'Order scans.', materialized table, " + f'max_duration "{max_duration}");\n\n' + f"SELECT 1 AS order_id, {SLOW_SCALAR_SQL} AS scanned\n" + ), + "models/order_scan_summary.sql": ( + "MODEL (description 'Order scan summary.', materialized table);\n\n" + 'SELECT COUNT(*) AS orders FROM __ref("order_scans")\n' + ), + "models/customers.sql": ( + "MODEL (description 'Customers.', materialized table);\n\nSELECT 100 AS customer_id\n" + ), + } diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py index e93ae30a95..4820b2fc76 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py @@ -2,13 +2,16 @@ from __future__ import annotations +import json import subprocess from pathlib import Path +from typing import Any import pytest from tests.e2e.src.sqlbuild.cli.commands.main.test._test_types import ( ModelTimeoutE2ETestCase, + ModelTimeoutJsonE2ETestCase, RealModelLimitE2ETestCase, ScenarioTimeoutE2ETestCase, SingleReadComparisonE2ETestCase, @@ -19,6 +22,7 @@ COST_GUARD_RAW_ORDERS_SOURCE, SLOW_SCALAR_SQL, build_order_chain_files, + slow_model_files, stored_test_outputs, ) from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import ( @@ -246,6 +250,7 @@ def test_given_slow_unit_test_when_limit_expires_then_test_times_out_and_run_con "Limit set by: max_duration in MODEL() header", "Cancellation: confirmed; the running statement stopped", "Skipped (1): order_scan_summary", + "FAIL=0 TIMEOUT=1 INSUFFICIENT=0 SKIP=1 TOTAL=3", ), expected_tables=( ("order_scans", False), @@ -262,22 +267,7 @@ def test_given_model_max_duration_when_building_then_slow_model_times_out_and_de project_dir: Path = prepare_inline_project( tmp_path=tmp_path, project_name="shop", - repo_files={ - "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML, - "models/order_scans.sql": ( - "MODEL (description 'Order scans.', materialized table, " - f'max_duration "{test_case.max_duration}");\n\n' - f"SELECT 1 AS order_id, {SLOW_SCALAR_SQL} AS scanned\n" - ), - "models/order_scan_summary.sql": ( - "MODEL (description 'Order scan summary.', materialized table);\n\n" - 'SELECT COUNT(*) AS orders FROM __ref("order_scans")\n' - ), - "models/customers.sql": ( - "MODEL (description 'Customers.', materialized table);\n\n" - "SELECT 100 AS customer_id\n" - ), - }, + repo_files=slow_model_files(max_duration=test_case.max_duration), ) result: subprocess.CompletedProcess[str] = run_sqb( @@ -295,6 +285,40 @@ def test_given_model_max_duration_when_building_then_slow_model_times_out_and_de ), table_name +@pytest.mark.parametrize( + "test_case", + [ + ModelTimeoutJsonE2ETestCase( + description="build JSON counts the timeout apart from failures", + expected_status="failed", + expected_summary_counts={ + "success_count": 1, + "failure_count": 0, + "timeout_count": 1, + "skipped_count": 1, + }, + ), + ], + ids=lambda case: case.description, +) +def test_given_model_max_duration_when_building_as_json_then_summary_counts_timeout( + tmp_path: Path, test_case: ModelTimeoutJsonE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, project_name="shop", repo_files=slow_model_files(max_duration="1s") + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--json"), project_dir=project_dir + ) + payload: dict[str, Any] = json.loads(result.stdout) + + assert result.returncode == 1, result.stderr + assert payload["status"] == test_case.expected_status + for key, expected_count in test_case.expected_summary_counts.items(): + assert payload["summary"][key] == expected_count, payload["summary"] + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/unit/src/sqlbuild/cli/output/main/build_execution_json/test_execution_protocol_v1.py b/tests/unit/src/sqlbuild/cli/output/main/build_execution_json/test_execution_protocol_v1.py index 56b15cc7ce..f2a8d05065 100644 --- a/tests/unit/src/sqlbuild/cli/output/main/build_execution_json/test_execution_protocol_v1.py +++ b/tests/unit/src/sqlbuild/cli/output/main/build_execution_json/test_execution_protocol_v1.py @@ -90,6 +90,7 @@ def test_given_microbatch_limit_result_when_formatting_execution_json_then_metad expected_summary={ "success_count": 1, "failure_count": 0, + "timeout_count": 0, "skipped_count": 0, "warning_count": 0, "python_check_pass_count": 0, diff --git a/website/src/content/docs/docs/cli/build.mdx b/website/src/content/docs/docs/cli/build.mdx index 296229fbd4..6e10ac7b4b 100644 --- a/website/src/content/docs/docs/cli/build.mdx +++ b/website/src/content/docs/docs/cli/build.mdx @@ -145,7 +145,11 @@ When a model fails: A model that exceeds its [`max_duration`](/docs/concepts/models/configuration#execution-time-limit) fails with `C417`. The failure states the elapsed time, the limit and where it was set, whether cancelling its running statement was confirmed, and the query IDs SQLBuild saw; downstream models are -skipped as for any other failure. +skipped as for any other failure. Its row shows `TIMEOUT`, and timed-out models and unit tests are +counted under a separate `TIMEOUT=` in the summary line +(`PASS=1 WARN=0 FAIL=0 TIMEOUT=1 INSUFFICIENT=0 SKIP=1 TOTAL=3`) and under `timeout_count` in +`--json` output, not in `FAIL=` or `failure_count`. A timeout is still a failure for the exit code: +the build exits non-zero and its JSON `status` is `failed`. ## Fingerprints diff --git a/website/src/content/docs/docs/cli/test.mdx b/website/src/content/docs/docs/cli/test.mdx index 7e79389134..862689e4b6 100644 --- a/website/src/content/docs/docs/cli/test.mdx +++ b/website/src/content/docs/docs/cli/test.mdx @@ -45,8 +45,9 @@ are strings in JSON so their scale is preserved. Each test check in JSON output Each test row shows how many real (unmocked) models the test executes, such as `3 real models`. A test that runs past its [time limit](/docs/concepts/testing#time-limits) is reported as `TIMEOUT`, -counted separately in the summary (`PASS=4 FAIL=0 TIMEOUT=1 TOTAL=5`), and the run continues -with the remaining tests. +counted separately in the summary (`PASS=4 FAIL=0 TIMEOUT=1 TOTAL=5`) and in the `--json` +summary's `timeout_count`, and the run continues with the remaining tests. A timeout still fails the +run: `sqb test` exits non-zero. ## Examples From db20d09a6bdbcc364d28240744cb3a0c9a07893b Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 10:05:31 +0100 Subject: [PATCH 06/11] fix: keep JSON failure counts inclusive of timeouts Build JSON summary.failure_count and test JSON summary.fail_count count timed-out items again, as before; timeout_count is an additive subset. The human-readable footers still show FAIL= and TIMEOUT= separately, and the CLI docs explain the difference. --- .../main/aggregate_build_counts.rs | 12 +++-- .../src/build_aggregate/models.rs | 2 + .../_helpers/result_output/documents.rs | 4 +- src/sqlbuild/executor/pipeline/main/run.py | 2 +- .../commands/main/test/test_cost_guards.py | 47 ++++++++++++++++++- website/src/content/docs/docs/cli/build.mdx | 11 +++-- website/src/content/docs/docs/cli/test.mdx | 7 +-- 7 files changed, 67 insertions(+), 18 deletions(-) diff --git a/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs b/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs index a849b886d6..68db45ec51 100644 --- a/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs +++ b/crates/sqlbuild-core/src/build_aggregate/main/aggregate_build_counts.rs @@ -13,8 +13,12 @@ pub fn aggregate_build_counts(facts: &BuildResultFacts) -> BuildCounts { for model in &facts.models { match model.status.as_str() { SUCCESS => counts.success += 1, - FAILED if model.error_code.as_deref() == Some(ITEM_TIMEOUT_CODE) => counts.timeout += 1, - FAILED => counts.failure += 1, + FAILED => { + counts.failure += 1; + if model.error_code.as_deref() == Some(ITEM_TIMEOUT_CODE) { + counts.timeout += 1; + } + } SKIPPED => { counts.skipped += 1; if model.error_code.as_deref() == Some(SOURCE_FRESHNESS_BLOCKED_CODE) { @@ -43,10 +47,9 @@ pub fn aggregate_build_counts(facts: &BuildResultFacts) -> BuildCounts { for outcome in &facts.test_outcomes { if outcome == TEST_PASS { counts.success += 1; - } else if outcome == TEST_TIMEOUT { - counts.timeout += 1; } else { counts.failure += 1; + counts.timeout += usize::from(outcome == TEST_TIMEOUT); } } let model_audits = facts.models.iter().flat_map(|model| &model.audit_outcomes); @@ -61,7 +64,6 @@ pub fn aggregate_build_counts(facts: &BuildResultFacts) -> BuildCounts { counts.warning += matching(&facts.source_audit_outcomes, AUDIT_WARN) + matching(&facts.end_audit_outcomes, AUDIT_WARN); counts.failed = counts.failure > 0 - || counts.timeout > 0 || facts .end_audit_outcomes .iter() diff --git a/crates/sqlbuild-core/src/build_aggregate/models.rs b/crates/sqlbuild-core/src/build_aggregate/models.rs index 066374563b..850c068bf4 100644 --- a/crates/sqlbuild-core/src/build_aggregate/models.rs +++ b/crates/sqlbuild-core/src/build_aggregate/models.rs @@ -26,7 +26,9 @@ pub struct BuildResultFacts { #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct BuildCounts { pub success: usize, + /// Failed items, timed-out items included. pub failure: usize, + /// Timed-out models and tests, also counted in `failure`. pub timeout: usize, pub skipped: usize, pub warning: usize, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs b/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs index e36dc3c57f..52ef4e8003 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/result_output/documents.rs @@ -497,12 +497,12 @@ pub(crate) fn format_test_execution_json<'py>( let results: Vec> = projection.terminal_results(&results)?; let pass_count: usize = count_status(&results, "outcome", PASS_OUTCOME)?; let timeout_count: usize = count_status(&results, "outcome", TIMEOUT_OUTCOME)?; - let fail_count: usize = results.len() - pass_count - timeout_count; + let fail_count: usize = results.len() - pass_count; emitted( &projection, Document { command: "test".to_owned(), - status: outcome_status(fail_count + timeout_count), + status: outcome_status(fail_count), summary: summary( py, &[ diff --git a/src/sqlbuild/executor/pipeline/main/run.py b/src/sqlbuild/executor/pipeline/main/run.py index 01a639d716..eba798e2f5 100644 --- a/src/sqlbuild/executor/pipeline/main/run.py +++ b/src/sqlbuild/executor/pipeline/main/run.py @@ -137,7 +137,7 @@ def run_build_pipeline( "run_kind": "build", "duration_ms": int((time.monotonic() - started) * 1000), "succeeded_count": result.success_count, - "failed_count": result.failure_count + result.timeout_count, + "failed_count": result.failure_count, "skipped_count": result.skipped_count, }, ) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py index 4820b2fc76..a86a4ae2c8 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py @@ -289,11 +289,11 @@ def test_given_model_max_duration_when_building_then_slow_model_times_out_and_de "test_case", [ ModelTimeoutJsonE2ETestCase( - description="build JSON counts the timeout apart from failures", + description="build JSON counts the timeout within failures", expected_status="failed", expected_summary_counts={ "success_count": 1, - "failure_count": 0, + "failure_count": 1, "timeout_count": 1, "skipped_count": 1, }, @@ -319,6 +319,49 @@ def test_given_model_max_duration_when_building_as_json_then_summary_counts_time assert payload["summary"][key] == expected_count, payload["summary"] +@pytest.mark.parametrize( + "test_case", + [ + ModelTimeoutJsonE2ETestCase( + description="test JSON counts the timeout within failures", + expected_status="failed", + expected_summary_counts={ + "pass_count": 1, + "fail_count": 1, + "timeout_count": 1, + "total_count": 2, + }, + ), + ], + ids=lambda case: case.description, +) +def test_given_slow_unit_test_when_testing_as_json_then_summary_counts_timeout( + tmp_path: Path, test_case: ModelTimeoutJsonE2ETestCase +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files={ + **_SLOW_FILES, + "tests/unit/order_scans_slow.sql": ( + 'TEST(timeout "1s");\n\nWITH\n' + "__source__raw_orders AS (\n SELECT 1 AS order_id, 5 AS amount\n),\n" + "__expected__order_scans AS (\n SELECT 1 AS order_id\n)\n" + ), + }, + ) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "test", "--json"), project_dir=project_dir + ) + payload: dict[str, Any] = json.loads(result.stdout) + + assert result.returncode == 1, result.stderr + assert payload["status"] == test_case.expected_status + for key, expected_count in test_case.expected_summary_counts.items(): + assert payload["summary"][key] == expected_count, payload["summary"] + + @pytest.mark.parametrize( "test_case", [ diff --git a/website/src/content/docs/docs/cli/build.mdx b/website/src/content/docs/docs/cli/build.mdx index 6e10ac7b4b..9826556dc6 100644 --- a/website/src/content/docs/docs/cli/build.mdx +++ b/website/src/content/docs/docs/cli/build.mdx @@ -145,11 +145,12 @@ When a model fails: A model that exceeds its [`max_duration`](/docs/concepts/models/configuration#execution-time-limit) fails with `C417`. The failure states the elapsed time, the limit and where it was set, whether cancelling its running statement was confirmed, and the query IDs SQLBuild saw; downstream models are -skipped as for any other failure. Its row shows `TIMEOUT`, and timed-out models and unit tests are -counted under a separate `TIMEOUT=` in the summary line -(`PASS=1 WARN=0 FAIL=0 TIMEOUT=1 INSUFFICIENT=0 SKIP=1 TOTAL=3`) and under `timeout_count` in -`--json` output, not in `FAIL=` or `failure_count`. A timeout is still a failure for the exit code: -the build exits non-zero and its JSON `status` is `failed`. +skipped as for any other failure. Its row shows `TIMEOUT`. In the summary line, timed-out models and unit tests are counted under +`TIMEOUT=` and excluded from `FAIL=` +(`PASS=1 WARN=0 FAIL=0 TIMEOUT=1 INSUFFICIENT=0 SKIP=1 TOTAL=3`). In `--json` output, +`summary.failure_count` still includes them, and `summary.timeout_count` reports how many of those +failures were timeouts. A timeout is a failure for the exit code: the build exits non-zero and its +JSON `status` is `failed`. ## Fingerprints diff --git a/website/src/content/docs/docs/cli/test.mdx b/website/src/content/docs/docs/cli/test.mdx index 862689e4b6..e3aa46347d 100644 --- a/website/src/content/docs/docs/cli/test.mdx +++ b/website/src/content/docs/docs/cli/test.mdx @@ -45,9 +45,10 @@ are strings in JSON so their scale is preserved. Each test check in JSON output Each test row shows how many real (unmocked) models the test executes, such as `3 real models`. A test that runs past its [time limit](/docs/concepts/testing#time-limits) is reported as `TIMEOUT`, -counted separately in the summary (`PASS=4 FAIL=0 TIMEOUT=1 TOTAL=5`) and in the `--json` -summary's `timeout_count`, and the run continues with the remaining tests. A timeout still fails the -run: `sqb test` exits non-zero. +counted under `TIMEOUT=` and excluded from `FAIL=` in the summary line +(`PASS=4 FAIL=0 TIMEOUT=1 TOTAL=5`), and the run continues with the remaining tests. In `--json` +output, `summary.fail_count` still includes timed-out tests, and `summary.timeout_count` reports how +many of them timed out. A timeout still fails the run: `sqb test` exits non-zero. ## Examples From 1c78bba9a1338a3399b8316090c825cdc2170dc7 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 10:19:19 +0100 Subject: [PATCH 07/11] fix: write audit results without meta when the result columns cannot be read A failing column read while checking for the additive meta column now logs a warning and writes the records without meta instead of failing the write. Remove the unused Python build_read_latest_sql builder. --- .../audit_results/_helpers/meta_column.py | 32 ++++++++++-- .../executor/audit_results/_helpers/sql.py | 30 ----------- .../audit_results/main/_test_types.py | 9 ++++ .../executor/audit_results/main/helpers.py | 8 +++ .../executor/audit_results/main/test_write.py | 52 +++++++++++++++++++ 5 files changed, 97 insertions(+), 34 deletions(-) diff --git a/src/sqlbuild/executor/audit_results/_helpers/meta_column.py b/src/sqlbuild/executor/audit_results/_helpers/meta_column.py index 859b89424a..17cdbddc8e 100644 --- a/src/sqlbuild/executor/audit_results/_helpers/meta_column.py +++ b/src/sqlbuild/executor/audit_results/_helpers/meta_column.py @@ -30,9 +30,18 @@ def ensure_meta_column( ) -> bool: """Add the meta column to a table created before it existed; report whether it exists.""" - if _meta_column_present( - connection=connection, database=database, schema=schema, evolution=evolution - ): + try: + present: bool = _meta_column_present( + connection=connection, database=database, schema=schema, evolution=evolution + ) + except Exception as error: + _LOGGER.warning( + "Audit result meta was not stored because the columns of the result table could " + "not be read (%s)", + error, + ) + return False + if present: return True statements: tuple[str, ...] = evolution.render_add_columns( destination=build_qualified_table_name( @@ -51,7 +60,7 @@ def ensure_meta_column( for statement in statements: _ = execute(connection=connection, sql=statement) except Exception as error: - if _meta_column_present( + if _meta_column_added_concurrently( connection=connection, database=database, schema=schema, evolution=evolution ): return True @@ -64,6 +73,21 @@ def ensure_meta_column( return True +def _meta_column_added_concurrently( + *, + connection: Any, + database: str | None, + schema: str, + evolution: AuditResultColumnEvolution, +) -> bool: + try: + return _meta_column_present( + connection=connection, database=database, schema=schema, evolution=evolution + ) + except Exception: + return False + + def _meta_column_present( *, connection: Any, diff --git a/src/sqlbuild/executor/audit_results/_helpers/sql.py b/src/sqlbuild/executor/audit_results/_helpers/sql.py index f9bab89dee..6caec57aa4 100644 --- a/src/sqlbuild/executor/audit_results/_helpers/sql.py +++ b/src/sqlbuild/executor/audit_results/_helpers/sql.py @@ -133,36 +133,6 @@ def insert_columns( return tuple(column for column in AUDIT_RESULT_COLUMNS if column != COLUMN_AUDIT_META_JSON) -def build_read_latest_sql( - *, - database: str | None, - schema: str, - audit_name: str, - binding_key: str, - execution_fingerprint: str, - run_scope_phase: str, - limit: int, - render_qualified_name: Callable[..., str | None], -) -> str: - qualified_name: str = build_qualified_table_name( - database=database, schema=schema, render_qualified_name=render_qualified_name - ) - predicates: tuple[tuple[str, str], ...] = ( - (COLUMN_AUDIT_NAME, audit_name), - (COLUMN_BINDING_KEY, binding_key), - (COLUMN_EXECUTION_FINGERPRINT, execution_fingerprint), - (COLUMN_RUN_SCOPE_PHASE, run_scope_phase), - ) - where_sql: str = " AND ".join( - f"{column} = {_column_literal(column=column, value=value)}" for column, value in predicates - ) - return ( - f"SELECT {', '.join(AUDIT_RESULT_COLUMNS)} FROM {qualified_name} " - f"WHERE {where_sql} ORDER BY {COLUMN_OCCURRED_AT} DESC, {COLUMN_RESULT_ID} DESC " - f"LIMIT {limit}" - ) - - def _record_values(*, record: AuditResultRecord, columns: tuple[str, ...]) -> str: values: dict[str, object | None] = { COLUMN_RESULT_ID: record.result_id, diff --git a/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py b/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py index 81526a0636..6273aa144e 100644 --- a/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py +++ b/tests/integration/src/sqlbuild/executor/audit_results/main/_test_types.py @@ -24,3 +24,12 @@ class AuditResultMetaWriteTestCase: audit_meta_json: str | None expected_meta_column: bool expected_meta: tuple[str | None, ...] + + +@dataclass(frozen=True) +class AuditResultMetaUnreadableColumnsTestCase: + """One meta write whose column read fails.""" + + description: str + expected_row_count: int + expected_meta_column: bool diff --git a/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py b/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py index 3065b0e0dd..b7b58bb404 100644 --- a/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py +++ b/tests/integration/src/sqlbuild/executor/audit_results/main/helpers.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime from typing import Any +from sqlbuild.adapter.contract.models import ColumnInfo from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.executor.audit_results.constants import AUDIT_RESULT_COLUMNS, COLUMN_AUDIT_META_JSON from sqlbuild.executor.audit_results.models import AuditResultRecord @@ -74,3 +75,10 @@ def build_record() -> AuditResultRecord: metadata_json='{"owner":"data"}', reused=False, ) + + +def failing_column_read(**kwargs: object) -> tuple[ColumnInfo, ...]: + """A column read that fails, as a warehouse metadata outage would.""" + + del kwargs + raise RuntimeError("metadata service unavailable") diff --git a/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py b/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py index 9329fe7c00..268a9f0a72 100644 --- a/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py +++ b/tests/integration/src/sqlbuild/executor/audit_results/main/test_write.py @@ -14,6 +14,7 @@ from sqlbuild.executor.audit_results.main._write import write_audit_result_records from sqlbuild.executor.audit_results.models import AuditResultColumnEvolution, AuditResultRecord from tests.integration.src.sqlbuild.executor.audit_results.main._test_types import ( + AuditResultMetaUnreadableColumnsTestCase, AuditResultMetaWriteTestCase, AuditResultWriteTestCase, ) @@ -21,6 +22,7 @@ LEGACY_AUDIT_RESULT_COLUMNS, build_record, create_legacy_results_table, + failing_column_read, leave_results_table_absent, ) @@ -180,5 +182,55 @@ def test_given_results_table_generation_when_writing_meta_then_column_is_added_a assert stored_meta == test_case.expected_meta +@pytest.mark.parametrize( + "test_case", + ( + AuditResultMetaUnreadableColumnsTestCase( + description="legacy table whose columns cannot be read keeps the record without meta", + expected_row_count=1, + expected_meta_column=False, + ), + ), + ids=lambda case: case.description, +) +def test_given_unreadable_columns_when_writing_meta_then_record_is_written_without_meta( + test_case: AuditResultMetaUnreadableColumnsTestCase, + tmp_path: Path, +) -> None: + adapter: DuckDbAdapter = DuckDbAdapter() + connection: Any = adapter.connect({"database": str(tmp_path / "audit_results.duckdb")}) + record: AuditResultRecord = replace(build_record(), audit_meta_json='{"region":"emea"}') + try: + create_legacy_results_table(adapter, connection) + write_audit_result_records( + connection=connection, + execute=adapter.execute, + database=None, + schema="main", + records=(record,), + render_qualified_name=adapter.render_qualified_name, + render_framework_type=adapter.render_framework_type, + render_create_table_sql=adapter.render_create_audit_result_table_sql, + render_create_index_sqls=adapter.render_create_audit_result_index_sqls, + column_evolution=AuditResultColumnEvolution( + read_columns=failing_column_read, render_add_columns=adapter.render_add_columns + ), + ) + row_count: int = connection.execute( + "SELECT COUNT(*) FROM main._sqlbuild_audit_results" + ).fetchone()[0] + columns: tuple[str, ...] = tuple( + column.name + for column in adapter.get_columns( + connection=connection, database=None, schema="main", name="_sqlbuild_audit_results" + ) + ) + finally: + adapter.close(connection) + + assert row_count == test_case.expected_row_count + assert (COLUMN_AUDIT_META_JSON in columns) is test_case.expected_meta_column + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) From 340dadbcdfe63f1581eca06c4ba73d79f2d036f0 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 10:55:26 +0100 Subject: [PATCH 08/11] fix: roll back, recover and bound only running time when an item times out Rollback and unit-test cleanup now run with every deadline suspended, so a model or scenario whose limit expires inside a transaction rolls it back before the timeout propagates. A connection whose statement was interrupted is probed before reuse by the next item, on build, test and scenario workers alike, and replaced when it is no longer usable. Item limits now count running time from the first warehouse statement and stop while cleanup or a queued BigQuery job waits. Snowflake relies on its statement limit, which excludes queueing, with a client backstop after a grace margin; BigQuery items are bounded by cancellation because a job timeout counts from job creation. Nested limits apply the earliest one, cancellation is issued only while the same statement is still registered, and work that finished before expiry keeps its result. PostgreSQL unit tests store their outputs in an ON COMMIT DROP table of one transaction bounded by SET LOCAL statement_timeout, so no setting outlives a pooled transaction, and a stored output that can no longer be read falls back to the inline comparison. --- fensu.toml | 6 + .../adapter/contract/classes/base_adapter.py | 1 + .../adapter/contract/classes/connection.py | 8 +- .../bigquery/classes/bigquery_connection.py | 34 ++- src/sqlbuild/adapters/bigquery/constants.py | 2 + .../postgres/classes/postgres_adapter.py | 28 ++- .../postgres/classes/postgres_connection.py | 50 ++++- src/sqlbuild/adapters/postgres/constants.py | 1 + .../snowflake/classes/snowflake_adapter.py | 1 + .../audit_results/_helpers/meta_column.py | 3 +- .../executor/build/classes/build_scheduler.py | 26 ++- .../executor/pipeline/_helpers/scenario.py | 18 +- .../executor/pipeline/_helpers/testing.py | 14 +- .../scenario/_helpers/execution/run.py | 8 +- .../executor/scheduling/classes/__init__.py | 0 .../classes/recovering_connection_pool.py | 19 ++ .../executor/testing/_helpers/comparison.py | 80 ++++++- .../classes/sql_test_actual_outputs.py | 16 +- .../executor/testing/main/_execute.py | 2 +- .../execution_limits/_helpers/cancellation.py | 2 + .../execution_limits/_helpers/interrupted.py | 7 + .../classes/active_connections.py | 9 +- .../classes/connection_recovery.py | 66 ++++++ .../classes/interrupted_connections.py | 26 +++ .../execution_limits/classes/item_deadline.py | 196 ++++++++++++++---- .../runtime/execution_limits/constants.py | 6 + .../main/current_build_deadline.py | 10 + .../main/enforce_execution_deadline.py | 10 +- .../main/execution_deadline_scope.py | 5 +- .../main/execution_deadlines_suspended.py | 38 ++++ .../main/item_deadline_scope.py | 26 ++- .../main/item_deadline_suspended.py | 21 -- .../main/tracked_statement_connection.py | 2 +- .../cli/commands/main/test/_test_types.py | 11 + .../cli/commands/main/test/helpers.py | 36 ++++ .../commands/main/test/test_cost_guards.py | 64 ++++++ .../unit/src/sqlbuild/adapters/_test_types.py | 10 + tests/unit/src/sqlbuild/adapters/helpers.py | 87 +++++++- .../sqlbuild/adapters/postgres/_test_types.py | 27 +++ .../postgres/test_deadline_transactions.py | 158 ++++++++++++++ .../adapters/test_statement_cancellation.py | 73 ++++++- .../runtime/execution_limits/_test_types.py | 17 ++ .../runtime/execution_limits/helpers.py | 9 + .../test_item_deadline_scope.py | 111 +++++++++- 44 files changed, 1194 insertions(+), 150 deletions(-) create mode 100644 src/sqlbuild/executor/scheduling/classes/__init__.py create mode 100644 src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py create mode 100644 src/sqlbuild/runtime/execution_limits/_helpers/interrupted.py create mode 100644 src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py create mode 100644 src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/current_build_deadline.py create mode 100644 src/sqlbuild/runtime/execution_limits/main/execution_deadlines_suspended.py delete mode 100644 src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py create mode 100644 tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py diff --git a/fensu.toml b/fensu.toml index 356ddab89e..1f7570c574 100644 --- a/fensu.toml +++ b/fensu.toml @@ -1323,3 +1323,9 @@ pending_budgets = [ "crates/sqlbuild-python/src/bindings/_helpers/state_rows/models.rs=3", "crates/sqlbuild-python/src/bindings/models.rs=4", ] + +[[targets.python.rule_exceptions]] +rule = "FFS120" +path = "src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py" +symbols = ["RecoveringConnectionPool.put"] +reason = "The override keeps queue.Queue.put's positional protocol, which every pool consumer calls." diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index e3eb73e660..a089443596 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -112,6 +112,7 @@ class BaseAdapter(RetentionAdapterMixin, StrictAdapter): allows_implicit_managed_write_schema: ClassVar[bool] = False execution_duration_limit_seconds: ClassVar[int | None] = None statement_cancellation_is_asynchronous: ClassVar[bool] = False + statement_limit_excludes_queueing: ClassVar[bool] = False metadata_inspection_concurrency: ClassVar[int] = 1 supports_session_warehouse: ClassVar[bool] = False _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( diff --git a/src/sqlbuild/adapter/contract/classes/connection.py b/src/sqlbuild/adapter/contract/classes/connection.py index 656a050917..56afb59516 100644 --- a/src/sqlbuild/adapter/contract/classes/connection.py +++ b/src/sqlbuild/adapter/contract/classes/connection.py @@ -16,6 +16,9 @@ from sqlbuild.runtime.execution_limits.main.enforce_execution_deadline import ( enforce_execution_deadline, ) +from sqlbuild.runtime.execution_limits.main.execution_deadlines_suspended import ( + execution_deadlines_suspended, +) from sqlbuild.runtime.execution_limits.main.tracked_statement_connection import ( tracked_statement_connection, ) @@ -42,7 +45,7 @@ def execute(self, *, connection: Any, sql: str) -> Any: result: Any = self._execute(connection=connection, sql=sql) except Exception as error: try: - enforce_execution_deadline() + enforce_execution_deadline(after_failure=True) except (ExecutionDurationLimitError, ItemTimeoutError) as deadline_error: raise deadline_error from error attributed_error: Exception = attribute_failed_sql(error=error, sql=sql) @@ -102,7 +105,8 @@ def transaction(self, connection: Any) -> Generator[None]: self.commit(connection) except BaseException: open_transactions.discard(id(connection)) - self.rollback(connection) + with execution_deadlines_suspended(): + self.rollback(connection) raise def supports_transactions(self) -> bool: diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py index 4b352a0d7b..55e09cbbb0 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py @@ -2,15 +2,22 @@ import math import threading +import time from typing import Any from sqlbuild.adapters.bigquery._helpers.statement_telemetry import affected_rows from sqlbuild.adapters.bigquery.classes.bigquery_cursor import _BigQueryCursor -from sqlbuild.adapters.bigquery.constants import SELECT_STATEMENT_TYPE +from sqlbuild.adapters.bigquery.constants import ( + PENDING_JOB_POLL_SECONDS, + PENDING_JOB_STATE, + SELECT_STATEMENT_TYPE, +) from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.main.current_build_deadline import current_build_deadline from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( current_execution_deadline, ) +from sqlbuild.runtime.execution_limits.main.current_item_deadline import current_item_deadline from sqlbuild.runtime.execution_limits.models import ExecutionDeadline from sqlbuild.runtime.observability.classes.statement_lifecycle import StatementLifecycle @@ -43,17 +50,21 @@ def execute(self, sql: str) -> _BigQueryCursor: return cursor def query_job(self, sql: str) -> Any: + """Submit a job; only the whole-build limit becomes a job timeout, as it counts queueing.""" + deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() - if deadline is None: + if deadline is not None: + deadline.enforce() + build: ExecutionDeadline | None = current_build_deadline() + if build is None: return self.client.query(sql, location=self.location) - deadline.enforce() from google.cloud import bigquery return self.client.query( sql, location=self.location, job_config=bigquery.QueryJobConfig( - job_timeout_ms=max(1, math.ceil(deadline.remaining_seconds * 1000)) + job_timeout_ms=max(1, math.ceil(build.remaining_seconds * 1000)) ), ) @@ -63,6 +74,7 @@ def job_result(self, job: Any) -> Any: with self._jobs_lock: self._active_jobs[id(job)] = job try: + _wait_while_pending(job=job, deadline=current_item_deadline()) return job.result() finally: with self._jobs_lock: @@ -82,6 +94,20 @@ def close(self) -> None: close() +def _wait_while_pending(*, job: Any, deadline: ItemDeadline | None) -> None: + """Stop an item's clock while its job waits in a queue, so only running time counts.""" + + if deadline is None or getattr(job, "state", None) != PENDING_JOB_STATE: + return + deadline.hold() + try: + while getattr(job, "state", None) == PENDING_JOB_STATE: + time.sleep(PENDING_JOB_POLL_SECONDS) + job.reload() + finally: + deadline.release() + + def _job_id(*, job: Any) -> str | None: value: object | None = getattr(job, "job_id", None) return value if isinstance(value, str) and value else None diff --git a/src/sqlbuild/adapters/bigquery/constants.py b/src/sqlbuild/adapters/bigquery/constants.py index 86105f9128..d3024011fb 100644 --- a/src/sqlbuild/adapters/bigquery/constants.py +++ b/src/sqlbuild/adapters/bigquery/constants.py @@ -32,3 +32,5 @@ ) MAX_CLONE_REFUSAL_ERROR_CHAIN: int = 8 NATIVE_TABLE_TYPE: str = "TABLE" +PENDING_JOB_STATE: str = "PENDING" +PENDING_JOB_POLL_SECONDS: float = 1.0 diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 26d4744c89..70348d455c 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -110,6 +110,9 @@ ) from sqlbuild.compiler.compile.types import FunctionLanguage from sqlbuild.diagnostics.main.log_sql import log_sql +from sqlbuild.runtime.execution_limits.main.execution_deadlines_suspended import ( + execution_deadlines_suspended, +) from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS @@ -143,15 +146,29 @@ def get_columns_for_relations( ) def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: - """Store a unit test's actual output once in a session temporary table.""" + """Store a unit test's output in a temporary table of one transaction, for poolers.""" - _ = self.execute(connection=connection, sql=f"CREATE TEMPORARY TABLE {name} AS {sql}") + if connection.claim_sql_test_transaction(): + _ = self.execute(connection=connection, sql="BEGIN") + try: + _ = self.execute( + connection=connection, + sql=f"CREATE TEMPORARY TABLE {name} ON COMMIT DROP AS {sql}", + ) + except BaseException: + with execution_deadlines_suspended(): + self.drop_sql_test_output(connection=connection, relation=name) + raise return name def drop_sql_test_output(self, *, connection: Any, relation: str) -> None: - """Drop the temporary table `create_sql_test_output` created.""" + """End the test's transaction, which drops its outputs, or drop one inside a caller's.""" - _ = self.execute(connection=connection, sql=f"DROP TABLE IF EXISTS {relation}") + if connection.release_sql_test_transaction(): + _ = self.execute(connection=connection, sql="ROLLBACK") + return + if connection.in_transaction(): + _ = self.execute(connection=connection, sql=f"DROP TABLE IF EXISTS {relation}") def interrupt_connection(self, connection: Any) -> bool: """Cancel a connection's running statement from another thread; False if unsupported.""" @@ -781,7 +798,8 @@ def view_definition_matches( ] ) finally: - connection.execute("ROLLBACK") + with execution_deadlines_suspended(): + connection.execute("ROLLBACK") return same_view_definition(definition=definition, sql=probe) def render_relation_revokes( diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_connection.py b/src/sqlbuild/adapters/postgres/classes/postgres_connection.py index eb61f542f7..e3a00275d8 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_connection.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_connection.py @@ -4,6 +4,7 @@ from typing import Any from sqlbuild.adapter.contract.classes.observed_cursor import ObservedCursor +from sqlbuild.adapters.postgres.constants import TRANSACTION_OPEN_STATUS from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline from sqlbuild.runtime.execution_limits.main.current_execution_deadline import ( current_execution_deadline, @@ -16,7 +17,9 @@ class _PostgresConnection: def __init__(self, raw_connection: Any) -> None: self.raw_connection: Any = raw_connection - self._statement_timeout_set: bool = False + self._local_timeout_set: bool = False + self._local_timeout_deadline: int | None = None + self._sql_test_transaction: bool = False def execute(self, sql: str) -> Any: self._apply_statement_timeout() @@ -31,15 +34,46 @@ def close(self) -> None: self.raw_connection.close() def _apply_statement_timeout(self) -> None: - """Bound the next statement by the remaining deadline, or restore the session default.""" + """Bound a statement in a transaction with SET LOCAL; never set it for a pooled session.""" deadline: ExecutionDeadline | ItemDeadline | None = current_execution_deadline() + if deadline is not None: + deadline.enforce() + if not self.in_transaction(): + self._local_timeout_set = False + self._local_timeout_deadline = None + return if deadline is None: - if self._statement_timeout_set: - self.raw_connection.execute("RESET statement_timeout") - self._statement_timeout_set = False + if self._local_timeout_set: + self.raw_connection.execute("SET LOCAL statement_timeout TO DEFAULT") + self._local_timeout_set = False + self._local_timeout_deadline = None + return + if self._local_timeout_set and self._local_timeout_deadline == id(deadline): return - deadline.enforce() milliseconds: int = max(1, math.ceil(deadline.remaining_seconds * 1000)) - self.raw_connection.execute(f"SET statement_timeout = {milliseconds}") - self._statement_timeout_set = True + self.raw_connection.execute(f"SET LOCAL statement_timeout = {milliseconds}") + self._local_timeout_set = True + self._local_timeout_deadline = id(deadline) + + def claim_sql_test_transaction(self) -> bool: + """Whether a unit test must open its own transaction, which it then owns.""" + + if self._sql_test_transaction or self.in_transaction(): + return False + self._sql_test_transaction = True + return True + + def release_sql_test_transaction(self) -> bool: + """Whether a unit test owned the open transaction, which it must now end.""" + + owned: bool = self._sql_test_transaction + self._sql_test_transaction = False + return owned + + def in_transaction(self) -> bool: + """Whether a healthy transaction is open on the server session.""" + + info: object = getattr(self.raw_connection, "info", None) + status: object = getattr(info, "transaction_status", None) + return getattr(status, "name", None) == TRANSACTION_OPEN_STATUS diff --git a/src/sqlbuild/adapters/postgres/constants.py b/src/sqlbuild/adapters/postgres/constants.py index d7d43f61b3..8910dbb44a 100644 --- a/src/sqlbuild/adapters/postgres/constants.py +++ b/src/sqlbuild/adapters/postgres/constants.py @@ -2,3 +2,4 @@ TABLE_FUNCTION_RETURN_TYPE: str = "TABLE" ROW_COUNT_RELATION_KINDS: frozenset[str] = frozenset({"r", "m"}) +TRANSACTION_OPEN_STATUS: str = "INTRANS" diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 98145c003f..05a67ee557 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -184,6 +184,7 @@ class SnowflakeAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter): backslash_escape_quotes=frozenset({"'"}), line_comment_prefixes=frozenset({"--", "//"}) ) execution_duration_limit_seconds: ClassVar[int | None] = MAX_STATEMENT_TIMEOUT_SECONDS + statement_limit_excludes_queueing: ClassVar[bool] = True metadata_inspection_concurrency: ClassVar[int] = _METADATA_INSPECTION_CONCURRENCY exact_column_inspection_limit: ClassVar[int] = _EXACT_COLUMN_INSPECTION_LIMIT diff --git a/src/sqlbuild/executor/audit_results/_helpers/meta_column.py b/src/sqlbuild/executor/audit_results/_helpers/meta_column.py index 17cdbddc8e..2a27aaf5a3 100644 --- a/src/sqlbuild/executor/audit_results/_helpers/meta_column.py +++ b/src/sqlbuild/executor/audit_results/_helpers/meta_column.py @@ -84,7 +84,8 @@ def _meta_column_added_concurrently( return _meta_column_present( connection=connection, database=database, schema=schema, evolution=evolution ) - except Exception: + except Exception as error: + _LOGGER.debug("Audit result columns could not be re-read after adding meta: %s", error) return False diff --git a/src/sqlbuild/executor/build/classes/build_scheduler.py b/src/sqlbuild/executor/build/classes/build_scheduler.py index b3d85ed1f4..a9e28234f1 100644 --- a/src/sqlbuild/executor/build/classes/build_scheduler.py +++ b/src/sqlbuild/executor/build/classes/build_scheduler.py @@ -107,6 +107,7 @@ ModelMaterializationContext, ) from sqlbuild.executor.run.types import MicrobatchBatchExecutor +from sqlbuild.executor.scheduling.classes.recovering_connection_pool import RecoveringConnectionPool from sqlbuild.executor.scheduling.main._run_worker import run_worker_with_completion from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.executor.seed.constants import SEED_ENTRY_MISSING_CODE @@ -123,6 +124,7 @@ from sqlbuild.microbatches.types import MicrobatchEventStore from sqlbuild.provider.main.runtime import ProviderContainer from sqlbuild.runtime.contracts.types import ExecutionResourceKind, NodeStartCallback +from sqlbuild.runtime.execution_limits.classes.connection_recovery import ConnectionRecovery from sqlbuild.runtime.execution_limits.main.enforce_execution_deadline import ( enforce_execution_deadline, ) @@ -250,7 +252,12 @@ def __init__( self._failed_sources: set[str] = set() self._completion_queue: queue.Queue[NodeCompletion] = queue.Queue() - self._connection_pool: queue.Queue[Any] = queue.Queue() + self._connection_recovery: ConnectionRecovery = ConnectionRecovery( + adapter=adapter, connect=lambda: adapter.connect(dict(connection_config)) + ) + self._connection_pool: queue.Queue[Any] = RecoveringConnectionPool( + recovery=self._connection_recovery + ) self._in_degree: dict[CompiledObjectKey, int] = {} self._selected_execution_keys: frozenset[CompiledObjectKey] = frozenset( @@ -276,10 +283,13 @@ def run( self._compute_in_degrees() self._seed_ready_queue() - if self._max_concurrency == 1: - self._run_serial() - else: - self._run_concurrent() + try: + if self._max_concurrency == 1: + self._run_serial() + else: + self._run_concurrent() + finally: + self._close_replacement_connections() enforce_execution_deadline() self._skip_remaining() @@ -334,6 +344,11 @@ def _provision_direct_microbatch_state(self) -> None: ): self._adapter.execute(connection=self._scheduler_connection, sql=index_sql) + def _close_replacement_connections(self) -> None: + """Close connections opened to replace broken ones; the originals belong to the caller.""" + + self._connection_recovery.close_replacements() + def _init_connection_pool(self) -> None: connection: Any for connection in self._connections: @@ -405,6 +420,7 @@ def _run_serial(self) -> None: | SqlTestExecutionResult | LoadExecutionResult ) = self._execute_node(key=key, connection=connection) + connection = self._connection_recovery.usable(connection) self._handle_completion(key=key, result=result) finally: self._connection_pool.put(connection) diff --git a/src/sqlbuild/executor/pipeline/_helpers/scenario.py b/src/sqlbuild/executor/pipeline/_helpers/scenario.py index e4d884f4d6..918f030374 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/scenario.py +++ b/src/sqlbuild/executor/pipeline/_helpers/scenario.py @@ -66,8 +66,10 @@ ScenarioTestSettings, ) from sqlbuild.executor.scenario.types import ScenarioLocalRunStatus, ScenarioSnapshotState +from sqlbuild.executor.scheduling.classes.recovering_connection_pool import RecoveringConnectionPool from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.runtime.contracts.models import ConnectionHooks +from sqlbuild.runtime.execution_limits.classes.connection_recovery import ConnectionRecovery from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy from sqlbuild.runtime.execution_limits.types import ItemKind @@ -190,7 +192,11 @@ def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: retain=retain, options=options, ) - if deadline is not None and deadline.expired: + if ( + deadline is not None + and deadline.expired + and result.status != ExecutionStatus.SUCCESS + ): return timed_out_scenario_result(result=result, report=deadline.report()) return result finally: @@ -238,7 +244,7 @@ def _sweep_interrupted_scenarios( interrupts: ScenarioInterrupts, adapter: BaseAdapter, connection_config: dict[str, object], - connection_pool: queue.SimpleQueue[Any], + connection_pool: queue.Queue[Any], ) -> None: """Drop started scenarios' relations once workers stop, on a connection still valid.""" @@ -401,7 +407,7 @@ def _scenario_target_connections( connection_hooks: ConnectionHooks | None, connection_count: int, abandoned: threading.Event | None = None, -) -> Iterator[queue.SimpleQueue[Any]]: +) -> Iterator[queue.Queue[Any]]: hooks: ConnectionHooks = connection_hooks if connection_hooks is not None else ConnectionHooks() if hooks.on_connection_start is not None: hooks.on_connection_start(connection_count) @@ -421,7 +427,10 @@ def _scenario_target_connections( raise if hooks.on_connection_complete is not None: hooks.on_connection_complete(connection_count, elapsed_seconds=time.monotonic() - start) - connection_pool: queue.SimpleQueue[Any] = queue.SimpleQueue() + recovery: ConnectionRecovery = ConnectionRecovery( + adapter=adapter, connect=lambda: adapter.connect(dict(connection_config)) + ) + connection_pool: queue.Queue[Any] = RecoveringConnectionPool(recovery=recovery) connection: Any for connection in connections: connection_pool.put(connection) @@ -433,6 +442,7 @@ def _scenario_target_connections( raise finally: if abandoned is None or not abandoned.is_set(): + recovery.close_replacements() close_connections(adapter=adapter, connections=connections, active_error=active_error) diff --git a/src/sqlbuild/executor/pipeline/_helpers/testing.py b/src/sqlbuild/executor/pipeline/_helpers/testing.py index be3e33156e..a1eb9145ff 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/testing.py +++ b/src/sqlbuild/executor/pipeline/_helpers/testing.py @@ -31,11 +31,13 @@ open_worker_connections, ) from sqlbuild.executor.pipeline.models import TestPipelineCallbacks +from sqlbuild.executor.scheduling.classes.recovering_connection_pool import RecoveringConnectionPool from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.executor.testing.constants import SQL_TEST_EXECUTION_ERROR_CODE from sqlbuild.executor.testing.main._execute import execute_sql_test from sqlbuild.executor.testing.models import SqlTestExecutionResult, StepResult from sqlbuild.executor.testing.types import SqlTestOutcome +from sqlbuild.runtime.execution_limits.classes.connection_recovery import ConnectionRecovery from sqlbuild.runtime.execution_limits.models import RunItemTimeouts from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.classes.resource_attempt_lifecycle import ( @@ -94,6 +96,9 @@ def run_test_pipeline( worker_count, elapsed_seconds=time.monotonic() - start ) active_error: BaseException | None = None + recovery: ConnectionRecovery = ConnectionRecovery( + adapter=adapter, connect=lambda: adapter.connect(dict(connection_config)) + ) try: preflight_start: float = time.monotonic() if resolved_callbacks.on_progress is not None: @@ -125,12 +130,14 @@ def run_test_pipeline( on_test_complete=resolved_callbacks.on_test_complete, run_id=canonical_run_id, timeouts=timeouts or RunItemTimeouts(), + recovery=recovery, ) return results except BaseException as error: active_error = error raise finally: + recovery.close_replacements() close_connections(adapter=adapter, connections=connections, active_error=active_error) @@ -145,14 +152,16 @@ def _run_tests( on_test_complete: Callable[[SqlTestExecutionResult], None] | None, run_id: str, timeouts: RunItemTimeouts, + recovery: ConnectionRecovery, ) -> tuple[SqlTestExecutionResult, ...]: if worker_count == 1: serial_results: list[SqlTestExecutionResult] = [] + serial_connection: Any = connections[0] for entry in entries: result: SqlTestExecutionResult = _execute_test_entry( entry=entry, adapter=adapter, - connection=connections[0], + connection=serial_connection, missing_functions=missing_functions_by_test.get(entry.key, ()), on_test_start=on_test_start, before_test_start=None, @@ -160,11 +169,12 @@ def _run_tests( timeouts=timeouts, ) serial_results.append(result) + serial_connection = recovery.usable(serial_connection) if on_test_complete is not None: on_test_complete(result) return tuple(serial_results) - connection_pool: queue.Queue[Any] = queue.Queue() + connection_pool: queue.Queue[Any] = RecoveringConnectionPool(recovery=recovery) for connection in connections: connection_pool.put(connection) start_coordinator: _AuthoredStartCoordinator = _AuthoredStartCoordinator() diff --git a/src/sqlbuild/executor/scenario/_helpers/execution/run.py b/src/sqlbuild/executor/scenario/_helpers/execution/run.py index d1ab571421..2ced4b0bae 100644 --- a/src/sqlbuild/executor/scenario/_helpers/execution/run.py +++ b/src/sqlbuild/executor/scenario/_helpers/execution/run.py @@ -36,8 +36,8 @@ ScenarioStepResults, ) from sqlbuild.executor.scheduling.types import ExecutionStatus -from sqlbuild.runtime.execution_limits.main.item_deadline_suspended import ( - item_deadline_suspended, +from sqlbuild.runtime.execution_limits.main.execution_deadlines_suspended import ( + execution_deadlines_suspended, ) from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle @@ -69,7 +69,7 @@ def execute_scenario_run_steps( if isinstance(error, (KeyboardInterrupt, ScenarioStopRequested)): options.interrupts.announce_stop() if not retain: - with item_deadline_suspended(): + with execution_deadlines_suspended(): cleanup_interrupted_scenarios( scenario_plans=(scenario_plan,), adapter=adapter, @@ -247,7 +247,7 @@ def _finish_scenario( if interrupts.stop_requested.is_set(): interrupts.announce_stop() if not retain and not interrupts.abandoned.is_set(): - with item_deadline_suspended(): + with execution_deadlines_suspended(): cleanup_result = execute_scenario_cleanup( scenario_plan=scenario_plan, adapter=adapter, diff --git a/src/sqlbuild/executor/scheduling/classes/__init__.py b/src/sqlbuild/executor/scheduling/classes/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py b/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py new file mode 100644 index 0000000000..2885e65bad --- /dev/null +++ b/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py @@ -0,0 +1,19 @@ +"""A worker connection pool that checks connections a deadline interrupted before reuse.""" + +from __future__ import annotations + +import queue +from typing import Any + +from sqlbuild.runtime.execution_limits.classes.connection_recovery import ConnectionRecovery + + +class RecoveringConnectionPool(queue.Queue[Any]): + """Returns each connection to the pool healthy, replacing one a timeout left unusable.""" + + def __init__(self, *, recovery: ConnectionRecovery) -> None: + super().__init__() + self._recovery: ConnectionRecovery = recovery + + def put(self, item: Any, block: bool = True, timeout: float | None = None) -> None: + super().put(self._recovery.usable(item), block=block, timeout=timeout) diff --git a/src/sqlbuild/executor/testing/_helpers/comparison.py b/src/sqlbuild/executor/testing/_helpers/comparison.py index 6e92e77f9e..3819e32a54 100644 --- a/src/sqlbuild/executor/testing/_helpers/comparison.py +++ b/src/sqlbuild/executor/testing/_helpers/comparison.py @@ -2,15 +2,22 @@ from __future__ import annotations +import logging +from collections.abc import Callable from typing import Any from sqlbuild._native import SqlTestPlanEntry, render_sql_test_comparison_sql from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.contract.types import RelationReadStatus +from sqlbuild.diagnostics.classes.diagnostic_record_redactor import DiagnosticRecordRedactor +from sqlbuild.diagnostics.main.log_debug_event import log_debug_event from sqlbuild.executor.testing._helpers.difference_samples import add_difference_samples from sqlbuild.executor.testing.classes.sql_test_actual_outputs import SqlTestActualOutputs from sqlbuild.executor.testing.models import StepResult from sqlbuild.executor.testing.types import SqlTestOutcome +_LOGGER: logging.Logger = logging.getLogger("sqlbuild.execution") + def compare_actual_outputs( *, @@ -23,14 +30,55 @@ def compare_actual_outputs( """Store the actual outputs once where supported, then count and sample the differences.""" relations: dict[int, str] = outputs.materialize() - comparison_sql: str = inline_sql - if relations: - (comparison_sql,) = render_sql_test_comparison_sql( - (test_entry,), - set_difference_operator=adapter.render_set_difference_operator(), - sql_analysis_dialect=adapter.sql_analysis_dialect(), - actual_relations=relations, + if not relations: + return _compare( + test_entry=test_entry, + adapter=adapter, + connection=connection, + comparison_sql=inline_sql, + relations={}, + ) + (stored_sql,) = render_sql_test_comparison_sql( + (test_entry,), + set_difference_operator=adapter.render_set_difference_operator(), + sql_analysis_dialect=adapter.sql_analysis_dialect(), + actual_relations=relations, + ) + try: + return _compare( + test_entry=test_entry, + adapter=adapter, + connection=connection, + comparison_sql=stored_sql, + relations=relations, ) + except Exception as error: + if not _stored_output_missing(adapter=adapter, error=error): + raise + log_debug_event( + logger=_LOGGER, + message="SQL test output was no longer readable; comparing inline", + test_name=test_entry.name, + sqlbuild_error=DiagnosticRecordRedactor.text(str(error)), + ) + outputs.release() + return _compare( + test_entry=test_entry, + adapter=adapter, + connection=connection, + comparison_sql=inline_sql, + relations={}, + ) + + +def _compare( + *, + test_entry: SqlTestPlanEntry, + adapter: BaseAdapter, + connection: Any, + comparison_sql: str, + relations: dict[int, str], +) -> list[StepResult]: cursor: Any = adapter.execute(connection=connection, sql=comparison_sql) rows: list[Any] = cursor.fetchall() return add_difference_samples( @@ -42,6 +90,24 @@ def compare_actual_outputs( ) +def _stored_output_missing(*, adapter: BaseAdapter, error: BaseException) -> bool: + """Whether a read failed because a stored output relation no longer exists.""" + + classify: Callable[[BaseException], RelationReadStatus | None] | None = getattr( + adapter, "_classify_relation_not_found", None + ) + if classify is None: + return False + current: BaseException | None = error + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + if classify(current) is RelationReadStatus.MISSING: + return True + current = current.__cause__ or current.__context__ + return False + + def _build_step_results(rows: list[Any]) -> list[StepResult]: """Convert comparison query rows into per-model step results.""" diff --git a/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py b/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py index 9b9714a4a3..1dc461346a 100644 --- a/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py +++ b/src/sqlbuild/executor/testing/classes/sql_test_actual_outputs.py @@ -15,8 +15,8 @@ ExecutionDurationLimitError, ItemTimeoutError, ) -from sqlbuild.runtime.execution_limits.main.item_deadline_suspended import ( - item_deadline_suspended, +from sqlbuild.runtime.execution_limits.main.execution_deadlines_suspended import ( + execution_deadlines_suspended, ) _LOGGER: logging.Logger = logging.getLogger("sqlbuild.execution") @@ -44,7 +44,7 @@ def __exit__( exc_value: BaseException | None, traceback: TracebackType | None, ) -> None: - self._release() + self.release() def materialize(self) -> dict[int, str]: """Store each compared step's output, or return no relations to compare inline.""" @@ -67,7 +67,7 @@ def materialize(self) -> dict[int, str]: name=f"{_OUTPUT_PREFIX}_{token}_{index}", ) if relation is None: - self._release() + self.release() return {} if relation.startswith(_OUTPUT_PREFIX): self._created.append(relation) @@ -81,15 +81,17 @@ def materialize(self) -> dict[int, str]: test_name=self._test_entry.name, sqlbuild_error=DiagnosticRecordRedactor.text(str(error)), ) - self._release() + self.release() return {} return dict(self.relations) - def _release(self) -> None: + def release(self) -> None: + """Drop every stored output; later comparisons run inline.""" + created: list[str] = list(self._created) self._created.clear() self.relations.clear() - with item_deadline_suspended(): + with execution_deadlines_suspended(): for relation in created: try: self._adapter.drop_sql_test_output( diff --git a/src/sqlbuild/executor/testing/main/_execute.py b/src/sqlbuild/executor/testing/main/_execute.py index 6b1751f928..b9988d86bf 100644 --- a/src/sqlbuild/executor/testing/main/_execute.py +++ b/src/sqlbuild/executor/testing/main/_execute.py @@ -62,7 +62,7 @@ def execute_sql_test( if deadline is None: raise return timed_out_result(test_entry=test_entry, report=deadline.report()) - if deadline is not None and deadline.expired: + if deadline is not None and deadline.expired and result.outcome == SqlTestOutcome.ERROR: return timed_out_result(test_entry=test_entry, report=deadline.report()) return replace(result, real_model_count=test_entry.real_model_count) diff --git a/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py b/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py index 074f4ff640..d96decc0e8 100644 --- a/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py +++ b/src/sqlbuild/runtime/execution_limits/_helpers/cancellation.py @@ -4,6 +4,7 @@ from typing import Any +from sqlbuild.runtime.execution_limits._helpers.interrupted import interrupted_connections from sqlbuild.runtime.execution_limits.types import CancellationOutcome @@ -11,6 +12,7 @@ def cancel_connection(*, adapter: Any, connection: Any) -> tuple[CancellationOut """The cancellation outcome for one connection and a detail when it is not confirmed.""" adapter_name: str = str(getattr(adapter, "adapter_name", type(adapter).__name__)) + interrupted_connections.mark(connection) try: requested: bool = bool(adapter.interrupt_connection(connection)) except Exception as error: diff --git a/src/sqlbuild/runtime/execution_limits/_helpers/interrupted.py b/src/sqlbuild/runtime/execution_limits/_helpers/interrupted.py new file mode 100644 index 0000000000..b45df9f99b --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/_helpers/interrupted.py @@ -0,0 +1,7 @@ +"""The process-wide record of connections a deadline interrupted.""" + +from sqlbuild.runtime.execution_limits.classes.interrupted_connections import ( + InterruptedConnections, +) + +interrupted_connections: InterruptedConnections = InterruptedConnections() diff --git a/src/sqlbuild/runtime/execution_limits/classes/active_connections.py b/src/sqlbuild/runtime/execution_limits/classes/active_connections.py index d4f2bac42b..89c387535d 100644 --- a/src/sqlbuild/runtime/execution_limits/classes/active_connections.py +++ b/src/sqlbuild/runtime/execution_limits/classes/active_connections.py @@ -3,6 +3,7 @@ from __future__ import annotations import threading +from collections.abc import Callable from typing import Any @@ -29,6 +30,10 @@ def remove(self, *, connection: Any) -> None: else: self._connections[id(connection)] = (adapter, connection, count - 1) - def snapshot(self) -> list[tuple[Any, Any]]: + def cancel_running[ResultT](self, *, cancel: Callable[[Any, Any], ResultT]) -> list[ResultT]: + """Cancel each still-registered statement, holding the lock so none can be replaced.""" + with self._lock: - return [(adapter, connection) for adapter, connection, _ in self._connections.values()] + return [ + cancel(adapter, connection) for adapter, connection, _ in self._connections.values() + ] diff --git a/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py b/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py new file mode 100644 index 0000000000..6574227011 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py @@ -0,0 +1,66 @@ +"""Check connections a deadline interrupted before reuse, replacing those left unusable.""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Callable +from typing import Any + +from sqlbuild.runtime.execution_limits._helpers.interrupted import interrupted_connections +from sqlbuild.runtime.execution_limits.constants import CONNECTION_PROBE_SQL +from sqlbuild.runtime.execution_limits.main.execution_deadlines_suspended import ( + execution_deadlines_suspended, +) + +_LOGGER: logging.Logger = logging.getLogger("sqlbuild.execution") + + +class ConnectionRecovery: + """Probe each interrupted connection, roll it back if needed, and replace it if still broken.""" + + def __init__(self, *, adapter: Any, connect: Callable[[], Any]) -> None: + self._adapter: Any = adapter + self._connect: Callable[[], Any] = connect + self._replacements: list[Any] = [] + self._lock: threading.Lock = threading.Lock() + + def usable(self, connection: Any) -> Any: + """`connection` when it is healthy, otherwise a newly opened replacement.""" + + if not interrupted_connections.take(connection): + return connection + with execution_deadlines_suspended(): + if self._responds(connection): + return connection + try: + self._adapter.rollback(connection) + except Exception as error: + _LOGGER.debug("rollback after an interrupted statement failed: %s", error) + if self._responds(connection): + return connection + _LOGGER.warning("A connection left unusable by a timed-out statement was replaced") + replacement: Any = self._connect() + with self._lock: + self._replacements.append(replacement) + return replacement + + def close_replacements(self) -> None: + """Close the connections this recovery opened; the originals are closed by their owner.""" + + with self._lock: + replacements: list[Any] = list(self._replacements) + self._replacements.clear() + for replacement in replacements: + try: + self._adapter.close(replacement) + except Exception as error: + _LOGGER.debug("closing a replacement connection failed: %s", error) + + def _responds(self, connection: Any) -> bool: + try: + _ = self._adapter.execute(connection=connection, sql=CONNECTION_PROBE_SQL) + except Exception as error: + _LOGGER.debug("connection probe after an interrupted statement failed: %s", error) + return False + return True diff --git a/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py b/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py new file mode 100644 index 0000000000..4e11723474 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py @@ -0,0 +1,26 @@ +"""Connections whose running statement a deadline interrupted, awaiting a health check.""" + +from __future__ import annotations + +import threading +from typing import Any + + +class InterruptedConnections: + """Identities of interrupted connections, each cleared when it is checked.""" + + def __init__(self) -> None: + self._connections: set[int] = set() + self._lock: threading.Lock = threading.Lock() + + def mark(self, connection: Any) -> None: + with self._lock: + self._connections.add(id(connection)) + + def take(self, connection: Any) -> bool: + """Whether `connection` was interrupted since it was last checked, clearing the mark.""" + + with self._lock: + interrupted: bool = id(connection) in self._connections + self._connections.discard(id(connection)) + return interrupted diff --git a/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py b/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py index a782c5e14a..9c06ca6f2a 100644 --- a/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py +++ b/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py @@ -4,6 +4,7 @@ import threading import time +from collections.abc import Iterator from typing import Any from sqlbuild.runtime.execution_limits._helpers.cancellation import cancel_connection @@ -11,6 +12,9 @@ from sqlbuild.runtime.execution_limits.constants import ( EXPIRY_REPORT_WAIT_SECONDS, OUTCOME_SEVERITY, + SERVER_LIMIT_TOLERANCE_SECONDS, + WATCHDOG_MAXIMUM_POLL_SECONDS, + WATCHDOG_MINIMUM_POLL_SECONDS, ) from sqlbuild.runtime.execution_limits.exceptions import ItemTimeoutError from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, ItemTimeoutReport @@ -18,76 +22,144 @@ class ItemDeadline: - """One test, scenario or model attempt's deadline, its statements and its cancellation.""" - - def __init__(self, *, kind: ItemKind, name: str, policy: ItemTimeoutPolicy) -> None: + """One test, scenario or model attempt's running-time budget and its cancellation.""" + + def __init__( + self, + *, + kind: ItemKind, + name: str, + policy: ItemTimeoutPolicy, + parent: ItemDeadline | None, + backstop_grace_seconds: float, + fallback: tuple[Any, Any], + ) -> None: self.kind: ItemKind = kind self.name: str = name self.policy: ItemTimeoutPolicy = policy - self.started: float = time.monotonic() - self.deadline: float = self.started + policy.seconds + self.parent: ItemDeadline | None = parent self.expired: bool = False self.connections: ActiveConnections = ActiveConnections() + self._grace: float = backstop_grace_seconds + self._fallback: tuple[Any, Any] = fallback + self._spent: float = 0.0 + self._running_since: float | None = None + self._clock_started: bool = False + self._holds: int = 0 + self._closed: bool = False self._active_statements: int = 0 self._awaiting_confirmation: bool = False self._cancellation: CancellationOutcome | None = None self._cancellation_detail: str | None = None self._query_ids: list[str] = [] - self._lock: threading.Lock = threading.Lock() + self._lock: threading.RLock = threading.RLock() self._expiring: threading.Lock = threading.Lock() + self._stop: threading.Event = threading.Event() + self._watchdog: threading.Thread | None = None + + def chain(self) -> Iterator[ItemDeadline]: + """This deadline and every enclosing one, innermost first.""" + + deadline: ItemDeadline | None = self + while deadline is not None: + yield deadline + deadline = deadline.parent @property def remaining_seconds(self) -> float: - return max(0.0, self.deadline - time.monotonic()) + """Running time left before the earliest limit in the chain, ignoring backstop grace.""" + + return max(0.0, min(deadline._own_remaining(grace=0.0) for deadline in self.chain())) + + def running_seconds(self) -> float: + with self._lock: + running: float = self._spent + if self._running_since is not None: + running += time.monotonic() - self._running_since + return running def enforce(self) -> None: - if not self.expired and self.remaining_seconds > 0: - return - self.expired = True - raise ItemTimeoutError(kind=self.kind, name=self.name, limit=self.policy.limit) + """Refuse new work once any deadline in the chain is past its limit and grace.""" + + for deadline in self.chain(): + deadline._raise_when_past(grace=deadline._grace) + + def enforce_after_failure(self) -> None: + """Treat a failed statement as a timeout once its running time reached a limit.""" + + for deadline in self.chain(): + deadline._raise_when_past(grace=-SERVER_LIMIT_TOLERANCE_SECONDS) def statement_started(self) -> None: - with self._lock: - self._active_statements += 1 + """Count a statement; the first one starts the running-time clock and the watchdog.""" + + for deadline in self.chain(): + deadline._start_statement() def statement_finished(self, *, query_id: str | None) -> None: - with self._lock: - self._active_statements = max(0, self._active_statements - 1) - if query_id: - self._query_ids.append(query_id) + for deadline in self.chain(): + with deadline._lock: + deadline._active_statements = max(0, deadline._active_statements - 1) + if query_id: + deadline._query_ids.append(query_id) def statement_submitted(self, *, query_id: str | None) -> None: if not query_id: return + for deadline in self.chain(): + with deadline._lock: + deadline._query_ids.append(query_id) + + def hold(self) -> None: + """Stop the clock, for warehouse queueing or cleanup that must not count.""" + + for deadline in self.chain(): + with deadline._lock: + deadline._holds += 1 + deadline._pause_clock() + + def release(self) -> None: + for deadline in self.chain(): + with deadline._lock: + deadline._holds = max(0, deadline._holds - 1) + deadline._resume_clock() + + def close(self) -> None: + """Stop the clock and the watchdog when the item's scope ends.""" + with self._lock: - self._query_ids.append(query_id) + self._closed = True + self._pause_clock() + self._stop.set() - def expire(self, *, fallback: tuple[Any, Any]) -> None: + def expire(self) -> None: """Mark the deadline expired and cancel every statement still running under it.""" with self._expiring: - self._cancel_running(fallback=fallback) - - def _cancel_running(self, *, fallback: tuple[Any, Any]) -> None: - targets: list[tuple[Any, Any]] = self.connections.snapshot() - with self._lock: self.expired = True - if not targets and self._active_statements: - targets = [fallback] - if not targets: - return - outcome: CancellationOutcome = CancellationOutcome.CONFIRMED - details: list[str] = [] - for adapter, connection in targets: - target_outcome, detail = cancel_connection(adapter=adapter, connection=connection) - if OUTCOME_SEVERITY[target_outcome] > OUTCOME_SEVERITY[outcome]: - outcome = target_outcome - if detail is not None and detail not in details: - details.append(detail) - with self._lock: - self._cancellation = outcome - self._awaiting_confirmation = outcome == CancellationOutcome.CONFIRMED - self._cancellation_detail = "; ".join(details) or None + outcomes: list[tuple[CancellationOutcome, str | None]] = ( + self.connections.cancel_running( + cancel=lambda adapter, connection: cancel_connection( + adapter=adapter, connection=connection + ) + ) + ) + with self._lock: + if not outcomes and self._active_statements: + adapter, connection = self._fallback + outcomes = [cancel_connection(adapter=adapter, connection=connection)] + if not outcomes: + return + outcome: CancellationOutcome = CancellationOutcome.CONFIRMED + details: list[str] = [] + for target_outcome, detail in outcomes: + if OUTCOME_SEVERITY[target_outcome] > OUTCOME_SEVERITY[outcome]: + outcome = target_outcome + if detail is not None and detail not in details: + details.append(detail) + self._cancellation = outcome + self._awaiting_confirmation = outcome == CancellationOutcome.CONFIRMED + self._cancellation_detail = "; ".join(details) or None def report(self) -> ItemTimeoutReport: """What happened at expiry, once any cancellation in progress has finished.""" @@ -107,8 +179,50 @@ def report(self) -> ItemTimeoutReport: name=self.name, limit=self.policy.limit, source=self.policy.source, - elapsed_seconds=max(0.0, time.monotonic() - self.started), + elapsed_seconds=self.running_seconds(), query_ids=tuple(dict.fromkeys(self._query_ids)), cancellation=outcome, cancellation_detail=self._cancellation_detail, ) + + def _own_remaining(self, *, grace: float) -> float: + return self.policy.seconds + grace - self.running_seconds() + + def _raise_when_past(self, *, grace: float) -> None: + if not self.expired and (not self._clock_started or self._own_remaining(grace=grace) > 0): + return + self.expired = True + raise ItemTimeoutError(kind=self.kind, name=self.name, limit=self.policy.limit) + + def _start_statement(self) -> None: + with self._lock: + self._active_statements += 1 + if self._clock_started or self._closed: + return + self._clock_started = True + self._resume_clock() + self._watchdog = threading.Thread( + target=self._watch, name=f"sqlbuild-deadline-{self.name}", daemon=True + ) + self._watchdog.start() + + def _pause_clock(self) -> None: + if self._running_since is not None: + self._spent += time.monotonic() - self._running_since + self._running_since = None + + def _resume_clock(self) -> None: + if self._clock_started and not self._closed and not self._holds: + self._running_since = self._running_since or time.monotonic() + + def _watch(self) -> None: + while True: + remaining: float = self._own_remaining(grace=self._grace) + if remaining <= 0: + self.expire() + return + poll: float = min( + WATCHDOG_MAXIMUM_POLL_SECONDS, max(WATCHDOG_MINIMUM_POLL_SECONDS, remaining) + ) + if self._stop.wait(timeout=poll): + return diff --git a/src/sqlbuild/runtime/execution_limits/constants.py b/src/sqlbuild/runtime/execution_limits/constants.py index 252cd233ee..7e91c30427 100644 --- a/src/sqlbuild/runtime/execution_limits/constants.py +++ b/src/sqlbuild/runtime/execution_limits/constants.py @@ -19,3 +19,9 @@ } ITEM_TIMEOUT_CODE: str = "C417" EXPIRY_REPORT_WAIT_SECONDS: float = 30.0 +BACKSTOP_MINIMUM_GRACE_SECONDS: float = 30.0 +BACKSTOP_GRACE_FRACTION: float = 0.5 +WATCHDOG_MINIMUM_POLL_SECONDS: float = 0.02 +WATCHDOG_MAXIMUM_POLL_SECONDS: float = 0.5 +SERVER_LIMIT_TOLERANCE_SECONDS: float = 0.05 +CONNECTION_PROBE_SQL: str = "SELECT 1" diff --git a/src/sqlbuild/runtime/execution_limits/main/current_build_deadline.py b/src/sqlbuild/runtime/execution_limits/main/current_build_deadline.py new file mode 100644 index 0000000000..87b97d423a --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/current_build_deadline.py @@ -0,0 +1,10 @@ +"""Current whole-build deadline lookup.""" + +from sqlbuild.runtime.execution_limits._helpers.context import current_deadline_context +from sqlbuild.runtime.execution_limits.models import ExecutionDeadline + + +def current_build_deadline() -> ExecutionDeadline | None: + """The target's whole-build `max_duration` deadline, when one is running.""" + + return current_deadline_context.get() diff --git a/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py b/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py index e858e61b4d..e7f8c38deb 100644 --- a/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py +++ b/src/sqlbuild/runtime/execution_limits/main/enforce_execution_deadline.py @@ -8,12 +8,16 @@ from sqlbuild.runtime.execution_limits.models import ExecutionDeadline -def enforce_execution_deadline() -> None: - """Raise when the active build deadline or the current item's deadline has expired.""" +def enforce_execution_deadline(*, after_failure: bool = False) -> None: + """Raise past a deadline; `after_failure` reports a warehouse-stopped statement as a timeout.""" deadline: ExecutionDeadline | None = current_deadline_context.get() if deadline is not None: deadline.enforce() item: ItemDeadline | None = current_item_deadline_context.get() - if item is not None: + if item is None: + return + if after_failure: + item.enforce_after_failure() + else: item.enforce() diff --git a/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py b/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py index e4fc2ec605..af578e143e 100644 --- a/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py +++ b/src/sqlbuild/runtime/execution_limits/main/execution_deadline_scope.py @@ -65,5 +65,6 @@ def execution_deadline_scope( def _cancel_running(*, active: ActiveConnections) -> None: - for adapter, connection in active.snapshot(): - _ = cancel_connection(adapter=adapter, connection=connection) + _ = active.cancel_running( + cancel=lambda adapter, connection: cancel_connection(adapter=adapter, connection=connection) + ) diff --git a/src/sqlbuild/runtime/execution_limits/main/execution_deadlines_suspended.py b/src/sqlbuild/runtime/execution_limits/main/execution_deadlines_suspended.py new file mode 100644 index 0000000000..fca7e32002 --- /dev/null +++ b/src/sqlbuild/runtime/execution_limits/main/execution_deadlines_suspended.py @@ -0,0 +1,38 @@ +"""Statements that must run whatever the deadlines say, such as rollback and cleanup.""" + +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import Token + +from sqlbuild.runtime.execution_limits._helpers.context import ( + current_deadline_context, + current_invocation_connections_context, + current_item_deadline_context, +) +from sqlbuild.runtime.execution_limits.classes.active_connections import ActiveConnections +from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.models import ExecutionDeadline + + +@contextmanager +def execution_deadlines_suspended() -> Generator[None]: + """Run statements no deadline refuses or cancels, with item clocks stopped meanwhile.""" + + item: ItemDeadline | None = current_item_deadline_context.get() + if item is not None: + item.hold() + item_token: Token[ItemDeadline | None] = current_item_deadline_context.set(None) + build_token: Token[ExecutionDeadline | None] = current_deadline_context.set(None) + connections_token: Token[ActiveConnections | None] = current_invocation_connections_context.set( + None + ) + try: + yield + finally: + current_invocation_connections_context.reset(connections_token) + current_deadline_context.reset(build_token) + current_item_deadline_context.reset(item_token) + if item is not None: + item.release() diff --git a/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py b/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py index 98e7b9dcdc..bd1079ce20 100644 --- a/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py +++ b/src/sqlbuild/runtime/execution_limits/main/item_deadline_scope.py @@ -2,7 +2,6 @@ from __future__ import annotations -import threading from collections.abc import Generator from contextlib import contextmanager from contextvars import Token @@ -10,6 +9,10 @@ from sqlbuild.runtime.execution_limits._helpers.context import current_item_deadline_context from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline +from sqlbuild.runtime.execution_limits.constants import ( + BACKSTOP_GRACE_FRACTION, + BACKSTOP_MINIMUM_GRACE_SECONDS, +) from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy from sqlbuild.runtime.execution_limits.types import ItemKind @@ -23,20 +26,25 @@ def item_deadline_scope( adapter: Any, connection: Any, ) -> Generator[ItemDeadline | None]: - """Bound one item's statements by its policy and cancel what is running when it expires.""" + """Bound one item's running time by its policy and cancel what is running when it expires.""" if policy is None: yield None return - deadline: ItemDeadline = ItemDeadline(kind=kind, name=name, policy=policy) - token: Token[ItemDeadline | None] = current_item_deadline_context.set(deadline) - timer: threading.Timer = threading.Timer( - policy.seconds, deadline.expire, kwargs={"fallback": (adapter, connection)} + grace: float = 0.0 + if getattr(adapter, "statement_limit_excludes_queueing", False): + grace = max(BACKSTOP_MINIMUM_GRACE_SECONDS, policy.seconds * BACKSTOP_GRACE_FRACTION) + deadline: ItemDeadline = ItemDeadline( + kind=kind, + name=name, + policy=policy, + parent=current_item_deadline_context.get(), + backstop_grace_seconds=grace, + fallback=(adapter, connection), ) - timer.daemon = True - timer.start() + token: Token[ItemDeadline | None] = current_item_deadline_context.set(deadline) try: yield deadline finally: - timer.cancel() + deadline.close() current_item_deadline_context.reset(token) diff --git a/src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py b/src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py deleted file mode 100644 index 003c466432..0000000000 --- a/src/sqlbuild/runtime/execution_limits/main/item_deadline_suspended.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Statements that must run after an item's deadline, such as cleanup.""" - -from __future__ import annotations - -from collections.abc import Generator -from contextlib import contextmanager -from contextvars import Token - -from sqlbuild.runtime.execution_limits._helpers.context import current_item_deadline_context -from sqlbuild.runtime.execution_limits.classes.item_deadline import ItemDeadline - - -@contextmanager -def item_deadline_suspended() -> Generator[None]: - """Run cleanup after an item's deadline without the deadline refusing its statements.""" - - token: Token[ItemDeadline | None] = current_item_deadline_context.set(None) - try: - yield - finally: - current_item_deadline_context.reset(token) diff --git a/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py b/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py index d23e9e2529..4d54c446df 100644 --- a/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py +++ b/src/sqlbuild/runtime/execution_limits/main/tracked_statement_connection.py @@ -24,7 +24,7 @@ def tracked_statement_connection(*, adapter: Any, connection: Any) -> Generator[ registries.append(invocation) deadline: ItemDeadline | None = current_item_deadline_context.get() if deadline is not None: - registries.append(deadline.connections) + registries.extend(enclosing.connections for enclosing in deadline.chain()) for registry in registries: registry.add(adapter=adapter, connection=connection) try: diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py index 113355dba5..b395a42855 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/_test_types.py @@ -426,3 +426,14 @@ class ModelTimeoutJsonE2ETestCase: description: str expected_status: str expected_summary_counts: dict[str, int] + + +@dataclass(frozen=True) +class TransactionTimeoutE2ETestCase: + """A model whose limit expires inside its promotion transaction.""" + + description: str + max_duration: str + expected_output_fragments: tuple[str, ...] + expected_published_rows: tuple[tuple[int, int], ...] + expected_later_model_rows: tuple[tuple[int], ...] diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py index 35d5e955cb..9028627451 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py @@ -1227,6 +1227,8 @@ def build_deep_shared_missing_mock_project_files(*, layers: int) -> dict[str, st 'adapter = "duckdb"\n\n' "[connection]\n" 'database = "deep_demo.duckdb"\n\n' + "[settings]\n" + f"unit_test_max_real_models = {2 * layers + 2}\n\n" "[defaults]\n" 'materialized = "table"\n' ), @@ -2008,3 +2010,37 @@ def slow_model_files(*, max_duration: str) -> dict[str, str]: "MODEL (description 'Customers.', materialized table);\n\nSELECT 100 AS customer_id\n" ), } + + +SLOW_SWAP_ADAPTER: str = ( + "from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter\n\n\n" + "class DuckDbSlowSwapAdapter(DuckDbAdapter):\n" + " adapter_name = 'duckdb_slow_swap'\n\n" + " def render_swap(self, *, left, right):\n" + " first, *rest = super().render_swap(left=left, right=right)\n" + " if 'order_totals' not in left:\n" + " return (first, *rest)\n" + " return (first, 'SELECT COUNT(*) FROM range(100000000000)', *rest)\n" +) + + +def slow_promotion_files( + *, database: Path, max_duration: str, order_total: int, customer_id: int +) -> dict[str, str]: + """A model whose promotion swap is slow on a rebuild, and an independent model after it.""" + + return { + "sqlbuild_project.toml": COST_GUARD_PROJECT_TOML.replace( + 'adapter = "duckdb"', 'adapter = "duckdb_slow_swap"' + ).replace('database = "warehouse.duckdb"', f'database = "{database}"'), + "adapters/duckdb_slow_swap.py": SLOW_SWAP_ADAPTER, + "models/order_totals.sql": ( + "MODEL (description 'Order totals.', materialized table, " + f'max_duration "{max_duration}");\n\n' + f"SELECT 1 AS order_id, {order_total} AS total\n" + ), + "models/z_customers.sql": ( + "MODEL (description 'Customers.', materialized table);\n\n" + f"SELECT {customer_id} AS customer_id\n" + ), + } diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py index a86a4ae2c8..1c840f8858 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_cost_guards.py @@ -15,6 +15,7 @@ RealModelLimitE2ETestCase, ScenarioTimeoutE2ETestCase, SingleReadComparisonE2ETestCase, + TransactionTimeoutE2ETestCase, UnitTestTimeoutE2ETestCase, ) from tests.e2e.src.sqlbuild.cli.commands.main.test.helpers import ( @@ -23,10 +24,12 @@ SLOW_SCALAR_SQL, build_order_chain_files, slow_model_files, + slow_promotion_files, stored_test_outputs, ) from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import ( prepare_inline_project, + query_duckdb, run_sqb, table_exists, ) @@ -415,5 +418,66 @@ def test_given_slow_scenario_when_limit_expires_then_scenario_times_out( assert fragment in output, output +@pytest.mark.parametrize( + "test_case", + [ + TransactionTimeoutE2ETestCase( + description="expiry inside the promotion swap keeps the published table", + max_duration="2s", + expected_output_fragments=( + "[C417] Model 'order_totals' timed out after", + "Limit set by: max_duration in MODEL() header", + "z_customers OK", + "PASS=1 WARN=0 FAIL=0 TIMEOUT=1", + ), + expected_published_rows=((1, 10),), + expected_later_model_rows=((200,),), + ), + ], + ids=lambda case: case.description, +) +def test_given_model_limit_expiring_in_promotion_when_building_then_published_table_is_kept( + tmp_path: Path, test_case: TransactionTimeoutE2ETestCase +) -> None: + database: Path = tmp_path / "shop" / "warehouse.duckdb" + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="shop", + repo_files=slow_promotion_files( + database=database, + max_duration=test_case.max_duration, + order_total=10, + customer_id=100, + ), + ) + first: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--concurrency", "1"), project_dir=project_dir + ) + assert first.returncode == 0, first.stdout + first.stderr + changed_files: dict[str, str] = slow_promotion_files( + database=database, max_duration=test_case.max_duration, order_total=20, customer_id=200 + ) + relative_path: str + for relative_path in ("models/order_totals.sql", "models/z_customers.sql"): + (project_dir / relative_path).write_text(changed_files[relative_path]) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build", "--concurrency", "1"), project_dir=project_dir + ) + output: str = result.stdout + result.stderr + + assert result.returncode == 1, output + for fragment in test_case.expected_output_fragments: + assert fragment in output, output + assert ( + tuple(query_duckdb(db_path=database, sql="SELECT order_id, total FROM main.order_totals")) + == test_case.expected_published_rows + ) + assert ( + tuple(query_duckdb(db_path=database, sql="SELECT customer_id FROM main.z_customers")) + == test_case.expected_later_model_rows + ) + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/_test_types.py b/tests/unit/src/sqlbuild/adapters/_test_types.py index 579168a736..b54f5f8247 100644 --- a/tests/unit/src/sqlbuild/adapters/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/_test_types.py @@ -171,3 +171,13 @@ class StatementCancellationTestCase: description: str expected_limit_range: tuple[int, int] expected_asynchronous: bool + + +@dataclass(frozen=True) +class BigQueryQueuedJobTestCase: + """One BigQuery job that waits in a queue before it runs.""" + + description: str + pending_polls: int + poll_seconds: float + expected_maximum_running_seconds: float diff --git a/tests/unit/src/sqlbuild/adapters/helpers.py b/tests/unit/src/sqlbuild/adapters/helpers.py index a0610885f5..e5d33e4c46 100644 --- a/tests/unit/src/sqlbuild/adapters/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/helpers.py @@ -3,6 +3,7 @@ from __future__ import annotations import threading +import time from collections.abc import Callable, Iterator from dataclasses import dataclass from typing import Any @@ -147,11 +148,40 @@ def execute_through_cursor(*, connection: MetadataConnection, sql: str) -> Metad BLOCKING_SQL: str = "SELECT * FROM orders" +@dataclass +class TransactionStatus: + """A psycopg-style transaction status, named like `psycopg.pq.TransactionStatus`.""" + + name: str = "IDLE" + + +class TransactionInfo: + """Tracks whether the fake connection has an open transaction, as psycopg reports it.""" + + def __init__(self) -> None: + self.transaction_status: TransactionStatus = TransactionStatus() + + def track(self, sql: str) -> None: + self.transaction_status = _TRANSACTION_STATUS_AFTER.get( + sql.strip().upper(), self.transaction_status + ) + + +_TRANSACTION_STATUS_AFTER: dict[str, TransactionStatus] = { + "BEGIN": TransactionStatus(name="INTRANS"), + "COMMIT": TransactionStatus(), + "ROLLBACK": TransactionStatus(), +} + + class BlockingCursor: """A DB-API cursor whose query blocks until another thread cancels it.""" - def __init__(self, *, executed: list[str], started: threading.Event) -> None: + def __init__( + self, *, executed: list[str], started: threading.Event, info: TransactionInfo + ) -> None: self.executed: list[str] = executed + self.info: TransactionInfo = info self.started: threading.Event = started self.cancelled: threading.Event = threading.Event() self.connection: object | None = None @@ -164,12 +194,13 @@ def __init__(self, *, executed: list[str], started: threading.Event) -> None: def execute(self, sql: str, *args: Any, **kwargs: Any) -> BlockingCursor: del args self.executed.append(sql) + self.info.track(sql) self.statement_params = kwargs.get("_statement_params") responses: dict[bool, Callable[[], BlockingCursor]] = { True: self._block_until_cancelled, False: lambda: self, } - return responses[sql == BLOCKING_SQL]() + return responses[sql.endswith(BLOCKING_SQL)]() def _block_until_cancelled(self) -> BlockingCursor: self.started.set() @@ -187,6 +218,7 @@ class CancellableRawConnection: """A raw driver connection exposing each driver's cancel entry point.""" def __init__(self) -> None: + self.info: TransactionInfo = TransactionInfo() self.executed: list[str] = [] self.cursors: list[BlockingCursor] = [] self.cancel_requests: list[tuple[str, str]] = [] @@ -194,13 +226,16 @@ def __init__(self) -> None: self._conn: CancellableRawConnection = self def cursor(self) -> BlockingCursor: - cursor: BlockingCursor = BlockingCursor(executed=self.executed, started=self.started) + cursor: BlockingCursor = BlockingCursor( + executed=self.executed, started=self.started, info=self.info + ) cursor.connection = self self.cursors.append(cursor) return cursor def execute(self, sql: str) -> None: self.executed.append(sql) + self.info.track(sql) def cancel_safe(self) -> None: self._cancel_running() @@ -250,14 +285,56 @@ def cancel(self) -> bool: return True +class QueuedJob: + """A BigQuery job that waits in a queue for a few polls before it runs and completes.""" + + def __init__(self, *, pending_polls: int, poll_seconds: float) -> None: + self.job_id: str = "job-queued" + self.state: str = "PENDING" + self.statement_type: str | None = None + self.num_dml_affected_rows: int | None = None + self._pending_polls: int = pending_polls + self._poll_seconds: float = poll_seconds + + def reload(self) -> None: + time.sleep(self._poll_seconds) + self._pending_polls -= 1 + self.state = ("RUNNING", "PENDING")[self._pending_polls > 0] + + def result(self) -> list[object]: + self.state = "DONE" + return [] + + def cancel(self) -> bool: + return True + + class CancellableJobClient: """A BigQuery client returning one cancellable job and recording each job config.""" def __init__(self) -> None: - self.job: CancellableJob = CancellableJob() + self.job: CancellableJob | QueuedJob = CancellableJob() self.job_configs: list[Any] = [] - def query(self, sql: str, **kwargs: Any) -> CancellableJob: + def query(self, sql: str, **kwargs: Any) -> CancellableJob | QueuedJob: del sql self.job_configs.append(kwargs.get("job_config")) return self.job + + +def without_limits(statements: list[str]) -> tuple[str, ...]: + """The statements a fake connection ran, without per-transaction statement limits.""" + + return tuple(filter(lambda sql: not sql.startswith("SET LOCAL statement_timeout"), statements)) + + +def expire_after_start(*, deadline: Any, raw: CancellableRawConnection) -> threading.Thread: + """Expire `deadline` from another thread once the fake connection's blocking query starts.""" + + def watchdog() -> None: + assert raw.started.wait(timeout=5) + deadline.expire() + + thread: threading.Thread = threading.Thread(target=watchdog) + thread.start() + return thread diff --git a/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py b/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py index 13ae7d8cae..5e302373e0 100644 --- a/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py @@ -125,3 +125,30 @@ class PostgresDependentViewRebindTestCase: description: str catalog_rows: tuple[tuple[object, ...], ...] expected_statements: tuple[str, ...] + + +@dataclass(frozen=True) +class DeadlineTransactionTestCase: + """One transaction whose item deadline expires before a later statement.""" + + description: str + expected_statements: tuple[str, ...] + + +@dataclass(frozen=True) +class SqlTestOutputTransactionTestCase: + """One unit test output stored and released inside the test's own transaction.""" + + description: str + setup_statements: tuple[str, ...] + expected_statements: tuple[str, ...] + expected_in_transaction: bool + + +@dataclass(frozen=True) +class FailedSqlTestOutputTestCase: + """One unit test output whose creation fails inside the test's own transaction.""" + + description: str + expected_statements: tuple[str, ...] + expected_in_transaction: bool diff --git a/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py b/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py new file mode 100644 index 0000000000..56340f3f58 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py @@ -0,0 +1,158 @@ +"""Postgres transactions and unit test outputs under per-item deadlines.""" + +from __future__ import annotations + +import threading + +import pytest + +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter +from sqlbuild.adapters.postgres.classes.postgres_connection import _PostgresConnection +from sqlbuild.runtime.execution_limits.exceptions import ItemTimeoutError +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope +from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy +from sqlbuild.runtime.execution_limits.types import ItemKind +from tests.unit.src.sqlbuild.adapters.helpers import ( + CancellableRawConnection, + expire_after_start, + without_limits, +) +from tests.unit.src.sqlbuild.adapters.postgres._test_types import ( + DeadlineTransactionTestCase, + FailedSqlTestOutputTestCase, + SqlTestOutputTransactionTestCase, +) + +_POLICY: ItemTimeoutPolicy = ItemTimeoutPolicy(limit="60s", seconds=60, source="test") +_OUTPUT: str = "__sqb_test_output_abc_0" + + +@pytest.mark.parametrize( + "test_case", + [ + DeadlineTransactionTestCase( + description="expiry mid-promotion rolls back before the timeout propagates", + expected_statements=( + "BEGIN", + "ALTER TABLE orders RENAME TO orders__prev", + "ROLLBACK", + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_expired_model_deadline_in_transaction_when_failing_then_rollback_runs_first( + test_case: DeadlineTransactionTestCase, +) -> None: + adapter: PostgresAdapter = PostgresAdapter() + raw: CancellableRawConnection = CancellableRawConnection() + connection: _PostgresConnection = _PostgresConnection(raw) + + with item_deadline_scope( + kind=ItemKind.MODEL, name="orders", policy=_POLICY, adapter=adapter, connection=connection + ) as deadline: + assert deadline is not None + with pytest.raises(ItemTimeoutError), adapter.transaction(connection): + _ = adapter.execute( + connection=connection, sql="ALTER TABLE orders RENAME TO orders__prev" + ) + deadline.expire() + _ = adapter.execute( + connection=connection, sql="ALTER TABLE orders__next RENAME TO orders" + ) + + assert without_limits(raw.executed) == test_case.expected_statements + assert not any(sql.startswith("SET statement_timeout") for sql in raw.executed) + assert not connection.in_transaction() + + +@pytest.mark.parametrize( + "test_case", + [ + SqlTestOutputTransactionTestCase( + description="test opens and ends its own transaction", + setup_statements=(), + expected_in_transaction=False, + expected_statements=( + "BEGIN", + f"CREATE TEMPORARY TABLE {_OUTPUT} ON COMMIT DROP AS SELECT 1 AS order_id", + f"SELECT * FROM {_OUTPUT}", + "ROLLBACK", + ), + ), + SqlTestOutputTransactionTestCase( + description="caller transaction is kept and the output dropped inside it", + setup_statements=("BEGIN",), + expected_in_transaction=True, + expected_statements=( + "BEGIN", + f"CREATE TEMPORARY TABLE {_OUTPUT} ON COMMIT DROP AS SELECT 1 AS order_id", + f"SELECT * FROM {_OUTPUT}", + f"DROP TABLE IF EXISTS {_OUTPUT}", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_sql_test_output_when_stored_and_released_then_one_transaction_holds_it( + test_case: SqlTestOutputTransactionTestCase, +) -> None: + adapter: PostgresAdapter = PostgresAdapter() + raw: CancellableRawConnection = CancellableRawConnection() + connection: _PostgresConnection = _PostgresConnection(raw) + for statement in test_case.setup_statements: + _ = adapter.execute(connection=connection, sql=statement) + + stored: str | None = adapter.create_sql_test_output( + connection=connection, sql="SELECT 1 AS order_id", name=_OUTPUT + ) + assert stored is not None + relation: str = stored + _ = adapter.execute(connection=connection, sql=f"SELECT * FROM {relation}") + adapter.drop_sql_test_output(connection=connection, relation=relation) + + assert without_limits(raw.executed) == test_case.expected_statements + assert connection.in_transaction() is test_case.expected_in_transaction + + +@pytest.mark.parametrize( + "test_case", + [ + FailedSqlTestOutputTestCase( + description="failed output creation ends the test's transaction", + expected_statements=("BEGIN", f"CREATE TEMPORARY TABLE {_OUTPUT} ON COMMIT DROP AS "), + expected_in_transaction=False, + ) + ], + ids=lambda case: case.description, +) +def test_given_failing_output_query_when_storing_then_transaction_is_rolled_back( + test_case: FailedSqlTestOutputTestCase, +) -> None: + adapter: PostgresAdapter = PostgresAdapter() + raw: CancellableRawConnection = CancellableRawConnection() + connection: _PostgresConnection = _PostgresConnection(raw) + failing_query: str = "SELECT * FROM orders" + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=adapter, connection=connection + ) as deadline: + assert deadline is not None + watchdog: threading.Thread = expire_after_start(deadline=deadline, raw=raw) + with pytest.raises(ItemTimeoutError): + _ = adapter.create_sql_test_output( + connection=connection, sql=failing_query, name=_OUTPUT + ) + watchdog.join(timeout=5) + + statements: tuple[str, ...] = without_limits(raw.executed) + assert statements[:2] == ( + test_case.expected_statements[0], + test_case.expected_statements[1] + failing_query, + ) + assert statements[-1] == "ROLLBACK" + assert connection.in_transaction() is test_case.expected_in_transaction + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py b/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py index ce05e6632e..65a472cfdd 100644 --- a/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py +++ b/tests/unit/src/sqlbuild/adapters/test_statement_cancellation.py @@ -16,13 +16,22 @@ from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection from sqlbuild.adapters.sqlserver.classes.sqlserver_adapter import SqlServerAdapter from sqlbuild.adapters.sqlserver.classes.sqlserver_connection import _SqlServerConnection +from sqlbuild.runtime.execution_limits.main.execution_deadline_scope import ( + execution_deadline_scope, +) from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy from sqlbuild.runtime.execution_limits.types import ItemKind -from tests.unit.src.sqlbuild.adapters._test_types import StatementCancellationTestCase +from tests.unit.src.sqlbuild.adapters._test_types import ( + BigQueryQueuedJobTestCase, + StatementCancellationTestCase, +) from tests.unit.src.sqlbuild.adapters.helpers import ( + BLOCKING_SQL, + CancellableJob, CancellableJobClient, CancellableRawConnection, + QueuedJob, cancel_from_watchdog, ) @@ -51,14 +60,22 @@ def test_given_postgres_deadline_when_executing_then_statement_timeout_bounds_an kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=None, connection=connection ): cancel_from_watchdog(adapter=PostgresAdapter(), connection=connection, raw=raw) - limit_statement: str = raw.executed[0] + outside_transaction: list[str] = list(raw.executed) + raw.executed.clear() + _ = connection.execute("BEGIN") + _ = connection.execute("SELECT 1") + _ = connection.execute("ROLLBACK") + in_transaction: list[str] = list(raw.executed) raw.executed.clear() _ = connection.execute("SELECT 1") - assert limit_statement.startswith("SET statement_timeout = ") + assert outside_transaction == [BLOCKING_SQL] + assert in_transaction[0] == "BEGIN" + assert in_transaction[1].startswith("SET LOCAL statement_timeout = ") low, high = test_case.expected_limit_range - assert low <= int(limit_statement.rsplit(" ", 1)[1]) <= high - assert raw.executed == ["RESET statement_timeout", "SELECT 1"] + assert low <= int(in_transaction[1].rsplit(" ", 1)[1]) <= high + assert in_transaction[2:] == ["SELECT 1", "ROLLBACK"] + assert raw.executed == ["SELECT 1"] assert PostgresAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous @@ -166,9 +183,11 @@ def test_given_bigquery_deadline_when_executing_then_job_timeout_set_and_job_can client: CancellableJobClient = CancellableJobClient() connection: _BigQueryConnection = _BigQueryConnection(client=client, location=None) adapter: BigQueryAdapter = BigQueryAdapter() + blocking_job: CancellableJob = CancellableJob() + client.job = blocking_job def watchdog() -> None: - assert client.job.waiting.wait(timeout=5) + assert blocking_job.waiting.wait(timeout=5) assert adapter.interrupt_connection(connection) is True thread: threading.Thread = threading.Thread(target=watchdog) @@ -179,11 +198,51 @@ def watchdog() -> None: with pytest.raises(RuntimeError, match="statement cancelled"): _ = connection.execute("SELECT * FROM orders") thread.join(timeout=5) + client.job = QueuedJob(pending_polls=0, poll_seconds=0.0) + with execution_deadline_scope( + max_duration_seconds=60, max_duration="60s", target_name="dev", remediation=None + ): + _ = connection.execute("SELECT * FROM orders") low, high = test_case.expected_limit_range - assert low <= int(client.job_configs[0].job_timeout_ms) <= high + assert client.job_configs[0] is None + assert low <= int(client.job_configs[1].job_timeout_ms) <= high assert BigQueryAdapter.statement_cancellation_is_asynchronous == test_case.expected_asynchronous +@pytest.mark.parametrize( + "test_case", + [ + BigQueryQueuedJobTestCase( + description="queued job time is excluded from the item running time", + pending_polls=8, + poll_seconds=0.05, + expected_maximum_running_seconds=0.2, + ) + ], + ids=lambda case: case.description, +) +def test_given_queued_bigquery_job_when_waiting_then_item_clock_is_stopped( + test_case: BigQueryQueuedJobTestCase, monkeypatch: pytest.MonkeyPatch +) -> None: + client: CancellableJobClient = CancellableJobClient() + client.job = QueuedJob(pending_polls=test_case.pending_polls, poll_seconds=0.0) + connection: _BigQueryConnection = _BigQueryConnection(client=client, location=None) + monkeypatch.setattr( + "sqlbuild.adapters.bigquery.classes.bigquery_connection.PENDING_JOB_POLL_SECONDS", + test_case.poll_seconds, + ) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=None, connection=connection + ) as deadline: + _ = connection.execute("SELECT * FROM orders") + assert deadline is not None + running: float = deadline.running_seconds() + + assert client.job_configs == [None] + assert running < test_case.expected_maximum_running_seconds + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py index bfab8a9e88..308f29b005 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py @@ -30,6 +30,8 @@ class DeadlineCancellationTestCase: class IdleExpiryTestCase: description: str kind: ItemKind + statements_run: int + expected_expired: bool expected_outcome: CancellationOutcome expected_interrupted: bool @@ -52,3 +54,18 @@ class EarliestDeadlineTestCase: class NoPolicyTestCase: description: str expected_deadline: None + + +@dataclass(frozen=True) +class NestedDeadlineTestCase: + description: str + outer_seconds: int + inner_seconds: int + expected_expired_name: str + + +@dataclass(frozen=True) +class HeldClockTestCase: + description: str + held_seconds: float + expected_maximum_running_seconds: float diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py b/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py index 75f475fe57..1ae300b0bf 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py @@ -49,3 +49,12 @@ def run_statement_until_cancelled( deadline.statement_submitted(query_id="query-1") assert adapter.interrupted.wait(timeout=5) deadline.statement_finished(query_id=finished_query_id) + + +def run_finished_statement() -> None: + """Run one fake statement to completion under the current item deadline.""" + + deadline: ItemDeadline | None = current_item_deadline() + assert deadline is not None + deadline.statement_started() + deadline.statement_finished(query_id=None) diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py b/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py index f7d8dd0b35..7f6aad8eec 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextlib import time import pytest @@ -17,22 +18,25 @@ from sqlbuild.runtime.execution_limits.main.execution_deadline_scope import ( execution_deadline_scope, ) -from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope -from sqlbuild.runtime.execution_limits.main.item_deadline_suspended import ( - item_deadline_suspended, +from sqlbuild.runtime.execution_limits.main.execution_deadlines_suspended import ( + execution_deadlines_suspended, ) +from sqlbuild.runtime.execution_limits.main.item_deadline_scope import item_deadline_scope from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy, ItemTimeoutReport from sqlbuild.runtime.execution_limits.types import CancellationOutcome, ItemKind from tests.unit.src.sqlbuild.runtime.execution_limits._test_types import ( DeadlineCancellationTestCase, EarliestDeadlineTestCase, + HeldClockTestCase, IdleExpiryTestCase, + NestedDeadlineTestCase, NoPolicyTestCase, SuspendedCleanupTestCase, ) from tests.unit.src.sqlbuild.runtime.execution_limits.helpers import ( FakeCancellingAdapter, FakeRefusingAdapter, + run_finished_statement, run_statement_until_cancelled, ) @@ -108,31 +112,119 @@ def test_given_running_statement_when_item_deadline_expires_then_cancellation_ou "test_case", [ IdleExpiryTestCase( - description="model with no running statement", + description="model past its limit between statements", kind=ItemKind.MODEL, + statements_run=1, + expected_expired=True, + expected_outcome=CancellationOutcome.NOT_NEEDED, + expected_interrupted=False, + ), + IdleExpiryTestCase( + description="rendering before the first statement does not count", + kind=ItemKind.TEST, + statements_run=0, + expected_expired=False, expected_outcome=CancellationOutcome.NOT_NEEDED, expected_interrupted=False, ), ], ids=lambda case: case.description, ) -def test_given_no_running_statement_when_item_deadline_expires_then_cancellation_is_not_needed( +def test_given_no_running_statement_when_limit_passes_then_only_started_items_expire( test_case: IdleExpiryTestCase, ) -> None: adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) with item_deadline_scope( - kind=test_case.kind, name="orders", policy=_EXPIRED, adapter=adapter, connection=None + kind=test_case.kind, name="orders", policy=_ONE_SECOND, adapter=adapter, connection=None ) as deadline: - time.sleep(0.05) - with pytest.raises(ItemTimeoutError): + for _ in range(test_case.statements_run): + run_finished_statement() + time.sleep(1.1) + with contextlib.suppress(ItemTimeoutError): enforce_execution_deadline() assert isinstance(deadline, ItemDeadline) + assert deadline.expired == test_case.expected_expired assert adapter.interrupted.is_set() == test_case.expected_interrupted assert deadline.report().cancellation == test_case.expected_outcome +@pytest.mark.parametrize( + "test_case", + [ + NestedDeadlineTestCase( + description="enclosing model limit cancels a test statement inside it", + outer_seconds=1, + inner_seconds=3600, + expected_expired_name="orders_model", + ), + ], + ids=lambda case: case.description, +) +def test_given_nested_deadlines_when_enclosing_one_expires_then_inner_statement_is_cancelled( + test_case: NestedDeadlineTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with item_deadline_scope( + kind=ItemKind.MODEL, + name="orders_model", + policy=ItemTimeoutPolicy( + limit=f"{test_case.outer_seconds}s", seconds=test_case.outer_seconds, source="test" + ), + adapter=adapter, + connection=None, + ) as outer: + with item_deadline_scope( + kind=ItemKind.TEST, + name="orders_test", + policy=ItemTimeoutPolicy( + limit=f"{test_case.inner_seconds}s", seconds=test_case.inner_seconds, source="test" + ), + adapter=adapter, + connection=None, + ) as inner: + assert inner is not None + assert inner.remaining_seconds <= test_case.outer_seconds + run_statement_until_cancelled(adapter=adapter, finished_query_id=None) + with pytest.raises(ItemTimeoutError) as raised: + enforce_execution_deadline() + + assert outer is not None + assert outer.expired + assert test_case.expected_expired_name in str(raised.value) + + +@pytest.mark.parametrize( + "test_case", + [ + HeldClockTestCase( + description="queued or cleanup time is not running time", + held_seconds=0.3, + expected_maximum_running_seconds=0.2, + ) + ], + ids=lambda case: case.description, +) +def test_given_held_clock_when_time_passes_then_running_time_excludes_it( + test_case: HeldClockTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_ONE_SECOND, adapter=adapter, connection=None + ) as deadline: + assert deadline is not None + run_finished_statement() + deadline.hold() + time.sleep(test_case.held_seconds) + deadline.release() + running: float = deadline.running_seconds() + + assert running < test_case.expected_maximum_running_seconds + + @pytest.mark.parametrize( "test_case", [SuspendedCleanupTestCase(description="cleanup after expiry", expected_enforced=True)], @@ -146,8 +238,9 @@ def test_given_expired_item_when_suspended_for_cleanup_then_statements_are_allow with item_deadline_scope( kind=ItemKind.TEST, name="orders", policy=_EXPIRED, adapter=adapter, connection=None ): + run_finished_statement() time.sleep(0.05) - with item_deadline_suspended(): + with execution_deadlines_suspended(): enforce_execution_deadline() enforced: bool = True From f2ed3f3e6e6e259c8b74492d8621420068e5a7be Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 10:55:31 +0100 Subject: [PATCH 09/11] docs: describe per-item running-time limits and upgrade notes Add the per-adapter enforcement matrix for item limits, the PostgreSQL transaction-scoped statement limit and single-read transaction, and upgrade notes for the new real-model compile check and for whole-build max_duration being accepted on every built-in adapter. --- .../docs/concepts/project-configuration.mdx | 31 ++++++++++++++++++- .../content/docs/docs/concepts/testing.mdx | 17 +++++++--- 2 files changed, 43 insertions(+), 5 deletions(-) diff --git a/website/src/content/docs/docs/concepts/project-configuration.mdx b/website/src/content/docs/docs/concepts/project-configuration.mdx index 04d4516448..2f8cd9fe14 100644 --- a/website/src/content/docs/docs/concepts/project-configuration.mdx +++ b/website/src/content/docs/docs/concepts/project-configuration.mdx @@ -391,7 +391,7 @@ Every built-in adapter enforces the limit: |---------|---------------------------|------------------------------|---------------------------| | Snowflake | `STATEMENT_TIMEOUT_IN_SECONDS` per statement | Driver cancel request for the running query | 7 days | | DuckDB, MotherDuck | none | Connection interrupt | about 24 days | -| PostgreSQL | `statement_timeout` set for each statement | Driver cancel (`cancel_safe`) | about 24 days | +| PostgreSQL | `SET LOCAL statement_timeout` inside transactions | Driver cancel (`cancel_safe`) | about 24 days | | BigQuery | Job timeout (`job_timeout_ms`) on each query job | `jobs.cancel` (asynchronous) | about 24 days | | Databricks | `STATEMENT_TIMEOUT` session setting, capped at 48 hours | Cursor cancel (asynchronous) | about 24 days | | SQL Server | none | Driver cancel | about 24 days | @@ -403,10 +403,39 @@ cannot cancel a running statement is rejected for `max_duration` rather than pre may leave warehouse work running. Under a per-item timeout, such an adapter's items stop before their next statement and report the cancellation as not supported. +PostgreSQL never sets `statement_timeout` for the whole session, so a limit cannot leak to another +client through a transaction-mode connection pooler; statements outside a transaction are bounded by +cancellation alone. + A duration failure may occur after earlier models have completed. Its error therefore distinguishes the timeout from a pre-execution model-limit failure and states that warehouse changes may already have occurred. +Upgrade note: earlier releases accepted `max_duration` only on Snowflake and rejected it on other +adapters. It is now accepted on every built-in adapter, so a target that sets it for a non-Snowflake +warehouse is enforced instead of failing at startup. + +#### Per-item time limits + +Unit test, scenario and model limits (`test_timeout`, `scenario_timeout` and a model's +`max_duration`) count running time only. The clock starts at the item's first warehouse statement, +so rendering and setup before it do not count, and it stops during rollback and cleanup, which always +run even after the limit. A timeout is reported only when the limit passed before the item's work +finished. Where a warehouse reports queueing separately, queued time is excluded: + +| Adapter | How the item limit is enforced | Queued time | +|---------|--------------------------------|-------------| +| Snowflake | `STATEMENT_TIMEOUT_IN_SECONDS` per statement, with a client-side backstop after a grace margin of half the limit (at least 30 seconds) | Excluded by the warehouse limit; the backstop counts it | +| DuckDB, MotherDuck | Connection interrupt at the limit | Not applicable | +| PostgreSQL | `SET LOCAL statement_timeout` inside transactions; cancellation at the limit | Lock waits count | +| BigQuery | `jobs.cancel` at the limit; no job timeout, because a job timeout counts from job creation | Excluded: the clock stops while a job is pending | +| Databricks | `STATEMENT_TIMEOUT` session setting and cursor cancel | May count | +| SQL Server | Driver cancel at the limit | Lock waits count | + +When a model contains unit tests or nested work with its own limit, the earliest limit applies and +cancels whatever is running under it. After a timeout interrupted a statement, SQLBuild checks the +connection before giving it to the next item and replaces it if it is no longer usable. + #### Unit test and scenario timeouts `test_timeout` and `scenario_timeout` set the per-test and per-scenario time limit on this target, diff --git a/website/src/content/docs/docs/concepts/testing.mdx b/website/src/content/docs/docs/concepts/testing.mdx index a480c2d3fa..0040e51570 100644 --- a/website/src/content/docs/docs/concepts/testing.mdx +++ b/website/src/content/docs/docs/concepts/testing.mdx @@ -800,6 +800,10 @@ Options: unit_test_max_real_models = 3 ``` +Upgrade note: the limit is new, so after upgrading, a selection containing a unit test that executes +more than 10 real models fails to compile (`sqb test`, and `sqb build` when it runs tests) until the +test is mocked or the limit is raised. + `sqb test` shows each test's real-model count on its row, such as `3 real models`, so a chain that is growing towards the limit is visible early. Direct macro, UDF and table function tests execute no models and are not counted. @@ -807,6 +811,9 @@ no models and are not counted. ### Time limits Each unit test has a time limit, `30s` by default, measured from its first warehouse statement. +Rendering, warehouse queueing where the warehouse reports it, and cleanup after the comparison do not +count; see [per-item time limits](/docs/concepts/project-configuration#per-item-time-limits) for +each adapter. Set it, from highest precedence, with `--test-timeout` on `sqb test` or `sqb build`, with `TEST(timeout "45s")` in the test file, with `[targets..execution_limits] test_timeout`, or with `[settings] test_timeout`. Each accepts a positive fixed duration or `"disabled"`. @@ -846,10 +853,12 @@ A test compares its actual output with the expected rows several ways: both row missing from each side, the expected-column check and, on failure, row samples. Where the warehouse supports it, SQLBuild evaluates each compared model's actual output once into a temporary relation and runs every comparison against it, so the chain is not recomputed for each read. Results, -messages and samples are unchanged. DuckDB, MotherDuck and PostgreSQL use a session temporary -table that SQLBuild drops after the test, including after a failure or timeout; Snowflake reads the -first query's persisted result through `RESULT_SCAN`. Other adapters compare the inline query as -before. `target/run/tests/` keeps the self-contained inline comparison query for each test. +messages and samples are unchanged. DuckDB and MotherDuck use a session temporary table that +SQLBuild drops after the test, including after a failure or timeout. PostgreSQL runs the test in +one transaction with an `ON COMMIT DROP` temporary table that the transaction's end removes, so the +reads stay on one server session behind a transaction-mode connection pooler. Snowflake reads the +first query's persisted result through `RESULT_SCAN`. If a stored output can no longer be read, +the test falls back to the inline comparison. Other adapters compare the inline query as before. `target/run/tests/` keeps the self-contained inline comparison query for each test. ## Skipping SQL analysis From f9c243de58ebe8982cde245eb33babbdab3705fe Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 11:03:19 +0100 Subject: [PATCH 10/11] fix: keep pool slots, test transactions and shared clocks correct after timeouts A connection that can neither be recovered nor replaced after a timeout is now returned to its pool marked broken, so later work fails with the driver's error instead of waiting forever for a lost slot; each later return retries the replacement, and microbatch subworker slots are always released. A PostgreSQL unit test claims its transaction only after BEGIN succeeds, so a refused BEGIN can no longer make a later release roll back a caller's transaction. A queued BigQuery job pauses an item's clock only while every statement under that deadline is queued, so one queued microbatch worker no longer pauses a model whose other workers are running. --- .../bigquery/classes/bigquery_connection.py | 4 +- .../postgres/classes/postgres_adapter.py | 3 +- .../postgres/classes/postgres_connection.py | 12 +-- .../executor/build/classes/build_scheduler.py | 6 +- .../classes/recovering_connection_pool.py | 6 +- .../classes/connection_recovery.py | 27 +++++- .../execution_limits/classes/item_deadline.py | 37 +++++++- .../sqlbuild/adapters/postgres/_test_types.py | 9 ++ .../postgres/test_deadline_transactions.py | 35 ++++++++ .../scheduling/classes/_test_types.py | 12 +++ .../executor/scheduling/classes/helpers.py | 56 ++++++++++++ .../test_recovering_connection_pool.py | 86 +++++++++++++++++++ .../runtime/execution_limits/_test_types.py | 10 +++ .../test_item_deadline_scope.py | 48 +++++++++++ 14 files changed, 333 insertions(+), 18 deletions(-) create mode 100644 tests/unit/src/sqlbuild/executor/scheduling/classes/_test_types.py create mode 100644 tests/unit/src/sqlbuild/executor/scheduling/classes/helpers.py create mode 100644 tests/unit/src/sqlbuild/executor/scheduling/classes/test_recovering_connection_pool.py diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py index 55e09cbbb0..890fdee67f 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py @@ -99,13 +99,13 @@ def _wait_while_pending(*, job: Any, deadline: ItemDeadline | None) -> None: if deadline is None or getattr(job, "state", None) != PENDING_JOB_STATE: return - deadline.hold() + deadline.statement_queued() try: while getattr(job, "state", None) == PENDING_JOB_STATE: time.sleep(PENDING_JOB_POLL_SECONDS) job.reload() finally: - deadline.release() + deadline.statement_dequeued() def _job_id(*, job: Any) -> str | None: diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 70348d455c..0d7ab58d6e 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -148,8 +148,9 @@ def get_columns_for_relations( def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None: """Store a unit test's output in a temporary table of one transaction, for poolers.""" - if connection.claim_sql_test_transaction(): + if connection.needs_sql_test_transaction(): _ = self.execute(connection=connection, sql="BEGIN") + connection.claim_sql_test_transaction() try: _ = self.execute( connection=connection, diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_connection.py b/src/sqlbuild/adapters/postgres/classes/postgres_connection.py index e3a00275d8..58d29c6823 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_connection.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_connection.py @@ -56,13 +56,15 @@ def _apply_statement_timeout(self) -> None: self._local_timeout_set = True self._local_timeout_deadline = id(deadline) - def claim_sql_test_transaction(self) -> bool: - """Whether a unit test must open its own transaction, which it then owns.""" + def needs_sql_test_transaction(self) -> bool: + """Whether a unit test must open its own transaction before storing outputs.""" + + return not self._sql_test_transaction and not self.in_transaction() + + def claim_sql_test_transaction(self) -> None: + """Record that the transaction a unit test just opened is its own to end.""" - if self._sql_test_transaction or self.in_transaction(): - return False self._sql_test_transaction = True - return True def release_sql_test_transaction(self) -> bool: """Whether a unit test owned the open transaction, which it must now end.""" diff --git a/src/sqlbuild/executor/build/classes/build_scheduler.py b/src/sqlbuild/executor/build/classes/build_scheduler.py index a9e28234f1..2c9515e434 100644 --- a/src/sqlbuild/executor/build/classes/build_scheduler.py +++ b/src/sqlbuild/executor/build/classes/build_scheduler.py @@ -680,8 +680,10 @@ def _run_microbatch_batch_worker( try: return execute(batch, connection) finally: - self._connection_pool.put(connection) - self._release_microbatch_subworker() + try: + self._connection_pool.put(connection) + finally: + self._release_microbatch_subworker() def _reserve_microbatch_subworker(self) -> bool: with self._microbatch_coordinator_lock: diff --git a/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py b/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py index 2885e65bad..ca6e3304e8 100644 --- a/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py +++ b/src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py @@ -16,4 +16,8 @@ def __init__(self, *, recovery: ConnectionRecovery) -> None: self._recovery: ConnectionRecovery = recovery def put(self, item: Any, block: bool = True, timeout: float | None = None) -> None: - super().put(self._recovery.usable(item), block=block, timeout=timeout) + returned: Any = item + try: + returned = self._recovery.usable(item) + finally: + super().put(returned, block=block, timeout=timeout) diff --git a/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py b/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py index 6574227011..578aa01202 100644 --- a/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py +++ b/src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py @@ -23,24 +23,41 @@ def __init__(self, *, adapter: Any, connect: Callable[[], Any]) -> None: self._adapter: Any = adapter self._connect: Callable[[], Any] = connect self._replacements: list[Any] = [] + self._broken: set[int] = set() self._lock: threading.Lock = threading.Lock() def usable(self, connection: Any) -> Any: - """`connection` when it is healthy, otherwise a newly opened replacement.""" + """`connection` if healthy, else a replacement, else the original kept marked broken.""" - if not interrupted_connections.take(connection): + interrupted: bool = interrupted_connections.take(connection) + with self._lock: + broken: bool = id(connection) in self._broken + if not interrupted and not broken: return connection with execution_deadlines_suspended(): if self._responds(connection): + self._forget_broken(connection) return connection try: self._adapter.rollback(connection) except Exception as error: _LOGGER.debug("rollback after an interrupted statement failed: %s", error) if self._responds(connection): + self._forget_broken(connection) return connection + try: + replacement: Any = self._connect() + except Exception as error: + with self._lock: + self._broken.add(id(connection)) + _LOGGER.warning( + "A connection left unusable by a timed-out statement could not be replaced (%s); " + "work that uses it will fail", + error, + ) + return connection _LOGGER.warning("A connection left unusable by a timed-out statement was replaced") - replacement: Any = self._connect() + self._forget_broken(connection) with self._lock: self._replacements.append(replacement) return replacement @@ -57,6 +74,10 @@ def close_replacements(self) -> None: except Exception as error: _LOGGER.debug("closing a replacement connection failed: %s", error) + def _forget_broken(self, connection: Any) -> None: + with self._lock: + self._broken.discard(id(connection)) + def _responds(self, connection: Any) -> bool: try: _ = self._adapter.execute(connection=connection, sql=CONNECTION_PROBE_SQL) diff --git a/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py b/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py index 9c06ca6f2a..b0b49a6b70 100644 --- a/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py +++ b/src/sqlbuild/runtime/execution_limits/classes/item_deadline.py @@ -46,6 +46,7 @@ def __init__( self._running_since: float | None = None self._clock_started: bool = False self._holds: int = 0 + self._queued_statements: int = 0 self._closed: bool = False self._active_statements: int = 0 self._awaiting_confirmation: bool = False @@ -102,6 +103,7 @@ def statement_finished(self, *, query_id: str | None) -> None: deadline._active_statements = max(0, deadline._active_statements - 1) if query_id: deadline._query_ids.append(query_id) + deadline._update_clock() def statement_submitted(self, *, query_id: str | None) -> None: if not query_id: @@ -111,18 +113,32 @@ def statement_submitted(self, *, query_id: str | None) -> None: deadline._query_ids.append(query_id) def hold(self) -> None: - """Stop the clock, for warehouse queueing or cleanup that must not count.""" + """Stop the clock for cleanup while no statement runs under this deadline.""" for deadline in self.chain(): with deadline._lock: deadline._holds += 1 - deadline._pause_clock() + deadline._update_clock() def release(self) -> None: for deadline in self.chain(): with deadline._lock: deadline._holds = max(0, deadline._holds - 1) - deadline._resume_clock() + deadline._update_clock() + + def statement_queued(self) -> None: + """Mark a running statement as waiting in a warehouse queue; only queued work pauses.""" + + for deadline in self.chain(): + with deadline._lock: + deadline._queued_statements += 1 + deadline._update_clock() + + def statement_dequeued(self) -> None: + for deadline in self.chain(): + with deadline._lock: + deadline._queued_statements = max(0, deadline._queued_statements - 1) + deadline._update_clock() def close(self) -> None: """Stop the clock and the watchdog when the item's scope ends.""" @@ -198,6 +214,7 @@ def _start_statement(self) -> None: with self._lock: self._active_statements += 1 if self._clock_started or self._closed: + self._update_clock() return self._clock_started = True self._resume_clock() @@ -212,9 +229,21 @@ def _pause_clock(self) -> None: self._running_since = None def _resume_clock(self) -> None: - if self._clock_started and not self._closed and not self._holds: + if self._clock_started and not self._closed and not self._paused(): self._running_since = self._running_since or time.monotonic() + def _paused(self) -> bool: + """Paused while held or queued, unless another statement is running meanwhile.""" + + held: bool = bool(self._holds or self._queued_statements) + return held and self._active_statements <= self._queued_statements + + def _update_clock(self) -> None: + if self._paused(): + self._pause_clock() + else: + self._resume_clock() + def _watch(self) -> None: while True: remaining: float = self._own_remaining(grace=self._grace) diff --git a/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py b/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py index 5e302373e0..5813cfbbab 100644 --- a/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/postgres/_test_types.py @@ -152,3 +152,12 @@ class FailedSqlTestOutputTestCase: description: str expected_statements: tuple[str, ...] expected_in_transaction: bool + + +@dataclass(frozen=True) +class RefusedBeginTestCase: + """A unit test whose BEGIN the deadline refuses, followed by work in a caller transaction.""" + + description: str + expected_statements: tuple[str, ...] + expected_in_transaction: bool diff --git a/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py b/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py index 56340f3f58..235a85b47e 100644 --- a/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py +++ b/tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py @@ -20,6 +20,7 @@ from tests.unit.src.sqlbuild.adapters.postgres._test_types import ( DeadlineTransactionTestCase, FailedSqlTestOutputTestCase, + RefusedBeginTestCase, SqlTestOutputTransactionTestCase, ) @@ -154,5 +155,39 @@ def test_given_failing_output_query_when_storing_then_transaction_is_rolled_back assert connection.in_transaction() is test_case.expected_in_transaction +@pytest.mark.parametrize( + "test_case", + [ + RefusedBeginTestCase( + description="a refused BEGIN leaves no ownership that could end a caller transaction", + expected_statements=("BEGIN", f"DROP TABLE IF EXISTS {_OUTPUT}"), + expected_in_transaction=True, + ) + ], + ids=lambda case: case.description, +) +def test_given_refused_begin_when_storing_output_then_later_release_keeps_caller_transaction( + test_case: RefusedBeginTestCase, +) -> None: + adapter: PostgresAdapter = PostgresAdapter() + raw: CancellableRawConnection = CancellableRawConnection() + connection: _PostgresConnection = _PostgresConnection(raw) + + with item_deadline_scope( + kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=adapter, connection=connection + ) as deadline: + assert deadline is not None + deadline.expire() + with pytest.raises(ItemTimeoutError): + _ = adapter.create_sql_test_output( + connection=connection, sql="SELECT 1 AS order_id", name=_OUTPUT + ) + _ = adapter.execute(connection=connection, sql="BEGIN") + adapter.drop_sql_test_output(connection=connection, relation=_OUTPUT) + + assert without_limits(raw.executed) == test_case.expected_statements + assert connection.in_transaction() is test_case.expected_in_transaction + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/executor/scheduling/classes/_test_types.py b/tests/unit/src/sqlbuild/executor/scheduling/classes/_test_types.py new file mode 100644 index 0000000000..8125446930 --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/scheduling/classes/_test_types.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class UnreplaceableConnectionTestCase: + """A timed-out item whose connection can neither be recovered nor replaced.""" + + description: str + later_items: int + expected_completions: tuple[tuple[str, str], ...] + expected_reconnect_attempts: int + expected_maximum_seconds: float diff --git a/tests/unit/src/sqlbuild/executor/scheduling/classes/helpers.py b/tests/unit/src/sqlbuild/executor/scheduling/classes/helpers.py new file mode 100644 index 0000000000..e18022c9d4 --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/scheduling/classes/helpers.py @@ -0,0 +1,56 @@ +"""Fakes for the recovering connection pool tests.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild.runtime.execution_limits._helpers.interrupted import interrupted_connections + + +class LostConnectionAdapter: + """An adapter whose connection dropped: every statement fails and reconnecting fails too.""" + + adapter_name: str = "fake" + + def __init__(self) -> None: + self.reconnect_attempts: int = 0 + + def execute(self, *, connection: Any, sql: str) -> None: + del connection, sql + raise ConnectionError("server closed the connection") + + def rollback(self, connection: Any) -> None: + del connection + raise ConnectionError("server closed the connection") + + def connect(self) -> Any: + self.reconnect_attempts += 1 + raise ConnectionError("network is unreachable") + + def close(self, connection: Any) -> None: + del connection + + +def timed_out_execute(connection: Any) -> str: + """An item whose statement a deadline interrupted on its connection.""" + + interrupted_connections.mark(connection) + return "timed out" + + +def later_execute(adapter: LostConnectionAdapter) -> Any: + """A later item that runs a statement on whatever connection the pool hands it.""" + + def execute(connection: Any) -> str: + adapter.execute(connection=connection, sql="SELECT 1") + return "ok" + + return execute + + +def success_completion(*, key: str, result: str) -> tuple[str, str]: + return key, result + + +def failure_completion(*, key: str, error: BaseException) -> tuple[str, str]: + return key, f"failed: {error}" diff --git a/tests/unit/src/sqlbuild/executor/scheduling/classes/test_recovering_connection_pool.py b/tests/unit/src/sqlbuild/executor/scheduling/classes/test_recovering_connection_pool.py new file mode 100644 index 0000000000..ca3ba2bb4c --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/scheduling/classes/test_recovering_connection_pool.py @@ -0,0 +1,86 @@ +"""A pool whose timed-out connection cannot be replaced keeps its slot and fails later work.""" + +from __future__ import annotations + +import queue +import time +from typing import Any + +import pytest + +from sqlbuild.executor.scheduling.classes.recovering_connection_pool import ( + RecoveringConnectionPool, +) +from sqlbuild.executor.scheduling.main._run_worker import run_worker_with_completion +from sqlbuild.runtime.execution_limits.classes.connection_recovery import ConnectionRecovery +from tests.unit.src.sqlbuild.executor.scheduling.classes._test_types import ( + UnreplaceableConnectionTestCase, +) +from tests.unit.src.sqlbuild.executor.scheduling.classes.helpers import ( + LostConnectionAdapter, + failure_completion, + later_execute, + success_completion, + timed_out_execute, +) + + +@pytest.mark.parametrize( + "test_case", + [ + UnreplaceableConnectionTestCase( + description="reconnect failure after a timeout keeps the slot and fails later items", + later_items=2, + expected_completions=( + ("orders", "timed out"), + ("customers", "failed: server closed the connection"), + ("customers", "failed: server closed the connection"), + ), + expected_reconnect_attempts=3, + expected_maximum_seconds=2.0, + ) + ], + ids=lambda case: case.description, +) +def test_given_unreplaceable_connection_when_returned_then_pool_keeps_it_and_later_items_fail( + test_case: UnreplaceableConnectionTestCase, +) -> None: + adapter: LostConnectionAdapter = LostConnectionAdapter() + pool: RecoveringConnectionPool = RecoveringConnectionPool( + recovery=ConnectionRecovery(adapter=adapter, connect=adapter.connect) + ) + connection: object = object() + pool.put(connection) + completions: queue.Queue[tuple[str, str]] = queue.Queue() + started: float = time.monotonic() + + run_worker_with_completion( + key="orders", + connection_pool=pool, + completion_queue=completions, + execute=timed_out_execute, + build_success=success_completion, + build_failure=failure_completion, + ) + for _ in range(test_case.later_items): + run_worker_with_completion( + key="customers", + connection_pool=pool, + completion_queue=completions, + execute=later_execute(adapter), + build_success=success_completion, + build_failure=failure_completion, + ) + published: tuple[tuple[str, str], ...] = tuple( + completions.get(timeout=1) for _ in test_case.expected_completions + ) + returned: Any = pool.get(timeout=1) + + assert published == test_case.expected_completions + assert returned is connection + assert adapter.reconnect_attempts == test_case.expected_reconnect_attempts + assert time.monotonic() - started < test_case.expected_maximum_seconds + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py index 308f29b005..d9bcba5fe7 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py @@ -69,3 +69,13 @@ class HeldClockTestCase: description: str held_seconds: float expected_maximum_running_seconds: float + + +@dataclass(frozen=True) +class SharedQueuedClockTestCase: + description: str + running_statements: int + queued_statements: int + waited_seconds: float + expected_minimum_running_seconds: float + expected_maximum_running_seconds: float diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py b/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py index 7f6aad8eec..1a7f7ae193 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/test_item_deadline_scope.py @@ -31,6 +31,7 @@ IdleExpiryTestCase, NestedDeadlineTestCase, NoPolicyTestCase, + SharedQueuedClockTestCase, SuspendedCleanupTestCase, ) from tests.unit.src.sqlbuild.runtime.execution_limits.helpers import ( @@ -311,5 +312,52 @@ def test_given_no_policy_when_entering_scope_then_no_deadline_applies( assert current is test_case.expected_deadline +@pytest.mark.parametrize( + "test_case", + [ + SharedQueuedClockTestCase( + description="one queued worker does not pause a model whose other worker runs", + running_statements=2, + queued_statements=1, + waited_seconds=0.3, + expected_minimum_running_seconds=0.25, + expected_maximum_running_seconds=10.0, + ), + SharedQueuedClockTestCase( + description="every running statement queued pauses the model", + running_statements=2, + queued_statements=2, + waited_seconds=0.3, + expected_minimum_running_seconds=0.0, + expected_maximum_running_seconds=0.2, + ), + ], + ids=lambda case: case.description, +) +def test_given_workers_sharing_a_model_deadline_when_some_are_queued_then_clock_follows_running( + test_case: SharedQueuedClockTestCase, +) -> None: + adapter: FakeCancellingAdapter = FakeCancellingAdapter(result=True, asynchronous=False) + + with item_deadline_scope( + kind=ItemKind.MODEL, name="orders", policy=_ONE_SECOND, adapter=adapter, connection=None + ) as deadline: + assert deadline is not None + for _ in range(test_case.running_statements): + deadline.statement_started() + for _ in range(test_case.queued_statements): + deadline.statement_queued() + before: float = deadline.running_seconds() + time.sleep(test_case.waited_seconds) + counted: float = deadline.running_seconds() - before + for _ in range(test_case.queued_statements): + deadline.statement_dequeued() + for _ in range(test_case.running_statements): + deadline.statement_finished(query_id=None) + + assert counted >= test_case.expected_minimum_running_seconds + assert counted < test_case.expected_maximum_running_seconds + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) From 6a42e243200eae30c06321ca641c9327551f6d7d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 11:26:16 +0100 Subject: [PATCH 11/11] test: list the new execution limit keys in config error goldens --- tests/goldens/compiler/failure/config-unknown-default-key.json | 2 +- tests/goldens/compiler/failure/local-config-unknown-key.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/goldens/compiler/failure/config-unknown-default-key.json b/tests/goldens/compiler/failure/config-unknown-default-key.json index 171f6252b1..9ad1c9ca74 100644 --- a/tests/goldens/compiler/failure/config-unknown-default-key.json +++ b/tests/goldens/compiler/failure/config-unknown-default-key.json @@ -8,7 +8,7 @@ "phase": "compile", "severity": "error", "code": "D001", - "message": "/sqlbuild_project.toml [defaults] contains unknown key(s): description. Allowed keys: append_cursor_inclusive, batch_concurrency, batch_size, contract, cursor_end, cursor_future_action, cursor_future_max_distance, cursor_start, cursor_start_max_action, cursor_start_max_ahead, cursor_watermark_mode, database, full_refresh, function_database, function_schema, incremental_mode, incremental_strategy, lookback, materialized, max_microbatches, merge_exclude_columns, microbatch_strategy, replay_on_change, row_diff_exclude_columns, row_diff_sample_rows, row_diff_sample_seed, row_diff_tolerances, schema, seed_database, seed_schema, tags, unaccounted_partition_policy", + "message": "/sqlbuild_project.toml [defaults] contains unknown key(s): description. Allowed keys: append_cursor_inclusive, batch_concurrency, batch_size, contract, cursor_end, cursor_future_action, cursor_future_max_distance, cursor_start, cursor_start_max_action, cursor_start_max_ahead, cursor_watermark_mode, database, full_refresh, function_database, function_schema, incremental_mode, incremental_strategy, lookback, materialized, max_duration, max_microbatches, merge_exclude_columns, microbatch_strategy, replay_on_change, row_diff_exclude_columns, row_diff_sample_rows, row_diff_sample_seed, row_diff_tolerances, schema, seed_database, seed_schema, tags, unaccounted_partition_policy", "help": "remove the key; see the reference for supported keys" } ] diff --git a/tests/goldens/compiler/failure/local-config-unknown-key.json b/tests/goldens/compiler/failure/local-config-unknown-key.json index 82ca337180..5aea9e667d 100644 --- a/tests/goldens/compiler/failure/local-config-unknown-key.json +++ b/tests/goldens/compiler/failure/local-config-unknown-key.json @@ -8,7 +8,7 @@ "phase": "compile", "severity": "error", "code": "D001", - "message": "/sqlbuild_local.toml settings contains unknown key(s): sql_analysiss. Allowed keys: auto_load_sources, column_contract_mode, concurrency, default_audit_run_scope, default_audit_severity, max_concurrency, microbatch_concurrency, microbatch_unaccounted_partition_policy, query_change_tracking, require_sql_analysis, sql_analysis, sql_validation, table_promotion_mode", + "message": "/sqlbuild_local.toml settings contains unknown key(s): sql_analysiss. Allowed keys: auto_load_sources, column_contract_mode, concurrency, default_audit_run_scope, default_audit_severity, max_concurrency, microbatch_concurrency, microbatch_unaccounted_partition_policy, query_change_tracking, require_sql_analysis, scenario_timeout, sql_analysis, sql_validation, table_promotion_mode, test_timeout, unit_test_max_real_models", "help": "did you mean 'sql_analysis'?" } ]