diff --git a/AGENTS.md b/AGENTS.md index c74814585b..6bb5192c5e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -120,6 +120,7 @@ - Treat a genuine duplicate as evidence of possible wider drift, not only a cleanup task: a missing shared owner, parallel implementations of one concept, logic on the wrong side of a boundary, or copies that have already diverged (use `--diff`; a fix present in only one copy may be a bug in the others). Diagnose and record that wider smell before consolidating, because removing the copies erases the only deterministic signal of it. - Consolidate genuine duplication the change introduces or touches. Report unrelated findings as follow-ups instead of refactoring them opportunistically. - Members marked `[forced]` are adapter contract overrides required by `test_strict_adapter.py`. Record other intentional mirrors in `[dupes]` in `fensu.toml` with a reason. +- Public adapter contract methods are duplicated per adapter on purpose: `test_strict_adapter.py` requires every first-class adapter to define them, and `[dupes]` exempts them. The private helpers behind them must not be copied. Keep them in shared builders, such as the snapshot SQL builders in `src/sqlbuild/adapter/contract/classes/`, and let each contract method delegate. Express real dialect differences as explicit parameters, such as a `SnapshotSqlDialect` field, not as a forked copy of the helper. ## Bounded Review Process diff --git a/Makefile b/Makefile index 701ed84de0..1043a3ee74 100644 --- a/Makefile +++ b/Makefile @@ -73,6 +73,7 @@ E2E_DUCKDB_BUILD_CORE_PATHS := \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_no_tests_no_audits_flags.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_plan_command_surface.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_python_hooks.py \ + tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_query_change_tracking.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_query_propagation.py \ tests/e2e/src/sqlbuild/cli/commands/main/build/test_remove_column_semantics.py \ diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs index c7b0a8fd77..ac05233798 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_references/extraction.rs @@ -1,6 +1,9 @@ //! Conservative native fast path for logical SQL reference extraction. use crate::constants::{DBT_REFERENCE_KIND, TABLE_FUNCTION_REFERENCE_KIND}; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::QuotePolicy; const PREFIXES: [(&str, &str); 6] = [ ("__dbt_ref(", "dbt_ref"), @@ -18,20 +21,13 @@ pub(crate) fn extract(sql: &str) -> Option> { let mut references: Vec = Vec::new(); let mut index = 0; while index < bytes.len() { - if bytes[index..].starts_with(b"--") { - index = sql[index + 2..] - .find('\n') - .map_or(bytes.len(), |offset| index + offset + 3); - continue; - } - if bytes[index..].starts_with(b"/*") { - let offset = sql[index + 2..].find("*/")?; - index += offset + 4; - continue; - } - if matches!(bytes[index], b'\'' | b'"' | b'`') { - index = quoted_end(bytes, index)?; - continue; + match non_code_end(bytes, index, QuotePolicy::COMPILER) { + Ok(Some(end)) => { + index = end; + continue; + } + Ok(None) => {} + Err(_) => return None, } let Some((prefix, kind)) = PREFIXES .iter() @@ -70,7 +66,9 @@ fn reference_name(sql: &str, start: usize) -> Option<(String, usize)> { let start = ascii_whitespace_end(bytes, start); let first = bytes.get(start).copied()?; if matches!(first, b'\'' | b'"') { - let end = quoted_end(bytes, start)?; + let Ok(end) = quote_end(bytes, start, QuotePolicy::COMPILER) else { + return None; + }; return Some((sql[start + 1..end - 1].to_owned(), end)); } if !is_identifier_start(first) { @@ -86,23 +84,6 @@ fn reference_name(sql: &str, start: usize) -> Option<(String, usize)> { Some((sql[start..end].to_owned(), end)) } -fn quoted_end(bytes: &[u8], start: usize) -> Option { - let quote = bytes[start]; - let mut index = start + 1; - while index < bytes.len() { - if bytes[index] != quote { - index += 1; - continue; - } - if matches!(quote, b'\'' | b'"') && bytes.get(index + 1).copied() == Some(quote) { - index += 2; - continue; - } - return Some(index + 1); - } - None -} - fn ascii_whitespace_end(bytes: &[u8], start: usize) -> usize { let mut end = start; while bytes diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs index 83e0e26683..03f0f43005 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/cte_sql.rs @@ -1,6 +1,7 @@ //! CTE definition and leading-WITH assembly shared by SQL-test planning and rendering. -use crate::compiler::_helpers::sql_tests::sql_scan::{comment_end, skip_whitespace}; +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::skip_whitespace::skip_whitespace; /// Render one `name AS (body)` definition, closing a trailing line comment first. pub(crate) fn cte_definition_sql(name: &str, sql: &str) -> String { @@ -53,7 +54,7 @@ pub(crate) fn leading_with_prefix_end(sql: &str) -> Option { fn skip_leading_ignorable(sql: &str, mut index: usize) -> usize { loop { index = skip_whitespace(sql, index); - match comment_end(sql, index) { + match comment_end(sql.as_bytes(), index) { Ok(Some(end)) => index = end, Ok(None) | Err(_) => return index, } diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs index 625ab2d292..ebcdd035c2 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/extraction.rs @@ -4,12 +4,16 @@ use std::collections::{HashMap, HashSet, VecDeque}; use serde::{Deserialize, Serialize}; -use crate::compiler::_helpers::sql_tests::sql_scan::{ - self, Unclosed, comment_end, quoted_end, skip_whitespace, -}; use crate::constants::{ DIRECT_DEPENDENCY_PATH_LENGTH, MACRO_TEST_MODE, TABLE_FUNCTION_TEST_MODE, UDF_TEST_MODE, }; +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::matching_paren::matching_paren as scan_matching_paren; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::main::skip_whitespace::skip_whitespace; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; const MODEL_PREFIXES: [&str; 8] = [ "__macro__", @@ -430,13 +434,13 @@ fn validate_expected( let branches = split_unions(&cte.1)?; let mut names: Vec> = Vec::new(); for branch in branches { - names.push(projection_names(branch, file)?); + names.push(projection_names(branch, file, label)?); } if let Some(first) = names.first() { for (index, branch) in names.iter().enumerate().skip(1) { if branch != first { return Err(format!( - "SQL test '{file}' must use the same __expected__ projection names and order in every set-operation branch; branch {} does not match branch 1", + "SQL test '{file}' must use the same {label} projection names and order in every set-operation branch; branch {} does not match branch 1", index + 1 )); } @@ -477,9 +481,11 @@ pub(crate) fn empty_fixture_marker_matches(sql: &str) -> Result { Ok(index == sql.len()) } -fn projection_names(branch: &str, file: &str) -> Result, String> { +fn projection_names(branch: &str, file: &str, label: &str) -> Result, String> { let start = skip_ignorable(branch, 0)?; - let select_end = consume_keyword(branch, start, "SELECT").ok_or_else(|| format!("SQL test '{file}' must define each __expected__ set-operation branch as a SELECT query"))?; + let select_end = consume_keyword(branch, start, "SELECT").ok_or_else(|| { + format!("SQL test '{file}' must define each {label} set-operation branch as a SELECT query") + })?; let mut end = branch.len(); for keyword in [ "FROM", "WHERE", "GROUP", "HAVING", "QUALIFY", "WINDOW", "ORDER", "LIMIT", "OFFSET", @@ -492,16 +498,16 @@ fn projection_names(branch: &str, file: &str) -> Result, String> { let expressions = split_top_level(&branch[select_end..end], b',')?; if expressions.is_empty() { return Err(format!( - "SQL test '{file}' must project at least one column in __expected__" + "SQL test '{file}' must project at least one column in {label}" )); } expressions .into_iter() - .map(|expression| projection_name(expression, file)) + .map(|expression| projection_name(expression, file, label)) .collect() } -fn projection_name(expression: &str, file: &str) -> Result { +fn projection_name(expression: &str, file: &str, label: &str) -> Result { if let Some(position) = find_last_top_level_keyword(expression, "AS")? { let alias_start = skip_ignorable(expression, position + 2)?; if let Some((alias, end)) = read_identifier(expression, alias_start) @@ -521,7 +527,7 @@ fn projection_name(expression: &str, file: &str) -> Result { return Ok(alias); } Err(format!( - "SQL test '{file}' must alias every non-trivial __expected__ projection" + "SQL test '{file}' must alias every non-trivial {label} projection" )) } @@ -862,8 +868,10 @@ fn split_unions(sql: &str) -> Result, String> { values.push(value); } index = skip_ignorable(sql, end)?; - if let Some(all_end) = consume_keyword(sql, index, "ALL") { - index = skip_ignorable(sql, all_end)?; + if let Some(quantifier_end) = consume_keyword(sql, index, "ALL") + .or_else(|| consume_keyword(sql, index, "DISTINCT")) + { + index = skip_ignorable(sql, quantifier_end)?; } start = index; continue; @@ -1060,7 +1068,9 @@ fn char_len(sql: &str, index: usize) -> usize { fn skip_ignorable(sql: &str, mut index: usize) -> Result { loop { index = skip_whitespace(sql, index); - match comment_end(sql, index).map_err(|error| scan_error_message(error, "SQL test"))? { + match comment_end(sql.as_bytes(), index) + .map_err(|error| scan_error_message(error, "SQL test"))? + { Some(end) => index = end, None => return Ok(index), } @@ -1068,26 +1078,21 @@ fn skip_ignorable(sql: &str, mut index: usize) -> Result { } fn skip_non_code(sql: &str, index: usize) -> Result { - if let Some(end) = - comment_end(sql, index).map_err(|error| scan_error_message(error, "SQL test"))? - { - return Ok(end); - } - if matches!(byte_at(sql, index), Some(b'\'') | Some(b'"') | Some(b'`')) { - return skip_quote(sql, index, "SQL test"); - } - Ok(index) + Ok(non_code_end(sql.as_bytes(), index, QuotePolicy::COMPILER) + .map_err(|error| scan_error_message(error, "SQL test"))? + .unwrap_or(index)) } fn skip_quote(sql: &str, start: usize, context: &str) -> Result { if byte_at(sql, start).is_none() { return Err(format!("{context} expected a quote")); } - quoted_end(sql, start).map_err(|error| scan_error_message(error, context)) + quote_end(sql.as_bytes(), start, QuotePolicy::COMPILER) + .map_err(|error| scan_error_message(error, context)) } fn matching_paren(sql: &str, open: usize, context: &str) -> Result { - sql_scan::matching_paren(sql, open).map_err(|error| match error { + scan_matching_paren(sql.as_bytes(), open, QuotePolicy::COMPILER).map_err(|error| match error { Unclosed::Parenthesis => scan_error_message(error, context), _ => scan_error_message(error, "SQL test"), }) diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs index da1babc312..ba1cda0325 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs @@ -5,7 +5,10 @@ use std::collections::HashSet; use regex::{Captures, Regex}; use crate::compiler::_helpers::sql_tests::planning::compile_error; -use crate::compiler::_helpers::sql_tests::sql_scan::{Unclosed, matching_paren, skip_whitespace}; +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::main::skip_whitespace::skip_whitespace; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; pub(crate) fn replace_callable_markers( sql: &str, @@ -51,7 +54,7 @@ where } fn matching_paren_end(sql: &str, open: usize) -> Result { - matching_paren(sql, open) + matching_paren(sql.as_bytes(), open, QuotePolicy::COMPILER) .map(|close| close + 1) .map_err(|error| { compile_error(match error { diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs index 7969a2dd71..ca1796ea9f 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs @@ -5,4 +5,3 @@ pub(crate) mod helper_scope; pub(crate) mod markers; pub(crate) mod planning; pub(crate) mod rendering; -pub(crate) mod sql_scan; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs deleted file mode 100644 index c1ba225599..0000000000 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/sql_scan.rs +++ /dev/null @@ -1,83 +0,0 @@ -//! Quote- and comment-aware SQL scanning shared by SQL-test extraction and planning. - -/// The SQL construct a scan reached the end of input inside. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum Unclosed { - BlockComment, - Quote, - Parenthesis, -} - -/// Return the first index at or after `index` that is not whitespace. -pub(crate) fn skip_whitespace(sql: &str, mut index: usize) -> usize { - while let Some(character) = sql.get(index..).and_then(|rest| rest.chars().next()) { - if !character.is_whitespace() { - break; - } - index += character.len_utf8(); - } - index -} - -/// Return the end of a line or block comment starting at `index`, if one starts there. -pub(crate) fn comment_end(sql: &str, index: usize) -> Result, Unclosed> { - let bytes = sql.as_bytes(); - match (bytes.get(index), bytes.get(index + 1)) { - (Some(b'-'), Some(b'-')) => Ok(Some( - sql[index..] - .find('\n') - .map_or(sql.len(), |offset| index + offset + 1), - )), - (Some(b'/'), Some(b'*')) => sql[index + 2..] - .find("*/") - .map(|offset| Some(index + 2 + offset + 2)) - .ok_or(Unclosed::BlockComment), - _ => Ok(None), - } -} - -/// Return the index after the quoted text starting at `start`, treating doubled quotes as escapes. -pub(crate) fn quoted_end(sql: &str, start: usize) -> Result { - let bytes = sql.as_bytes(); - let quote = bytes[start]; - let mut index = start + 1; - while index < bytes.len() { - if bytes[index] != quote { - index += 1; - } else if bytes.get(index + 1) == Some("e) { - index += 2; - } else { - return Ok(index + 1); - } - } - Err(Unclosed::Quote) -} - -/// Return the index of the parenthesis closing the one at `open`, skipping quotes and comments. -pub(crate) fn matching_paren(sql: &str, open: usize) -> Result { - let bytes = sql.as_bytes(); - let mut depth = 0usize; - let mut index = open; - while index < bytes.len() { - if let Some(end) = comment_end(sql, index)? { - index = end; - continue; - } - match bytes[index] { - b'\'' | b'"' | b'`' => { - index = quoted_end(sql, index)?; - continue; - } - b'(' => depth += 1, - b')' => { - depth -= 1; - if depth == 0 { - return Ok(index); - } - } - _ => {} - } - index += 1; - } - Err(Unclosed::Parenthesis) -} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs index ab82c272bd..422ac14668 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs @@ -200,6 +200,28 @@ pub(crate) fn empty_model_fixture_marker_preserves_direct_mode_validation() -> b true } +pub(crate) fn expected_projection_errors_name_the_expected_cte() -> bool { + let cases = [ + ( + r#"{"tests":[{"sql":"WITH __macro_actual__ AS (SELECT @status() AS status), __macro_expected__ AS (SELECT 1 + 1) SELECT 1","fileLabel":"tests/status.sql","mode":"macro"}]}"#, + "SQL test 'tests/status.sql' must alias every non-trivial __macro_expected__ projection", + ), + ( + r#"{"tests":[{"sql":"WITH __udf_actual__ AS (SELECT 1 AS value), __udf_expected__ AS (SELECT 1 AS value UNION ALL SELECT 2 AS other) SELECT 1","fileLabel":"tests/function.sql","mode":"udf"}]}"#, + "SQL test 'tests/function.sql' must use the same __udf_expected__ projection names and order in every set-operation branch; branch 2 does not match branch 1", + ), + ( + r#"{"tests":[{"sql":"WITH __source__raw_orders AS (SELECT 1 AS id), __expected__orders AS (SELECT 1 + 1) SELECT 1","fileLabel":"tests/orders.sql","mode":"model"}]}"#, + "SQL test 'tests/orders.sql' must alias every non-trivial __expected__ projection", + ), + ]; + for (request, expected_error) in cases { + let error = extract_batch_json(request).expect_err("expected projection is rejected"); + assert_eq!(error, expected_error); + } + true +} + pub(crate) fn concurrent_requests_initialize_shared_template_once() -> bool { let templates = Arc::new(Mutex::new(HashMap::new())); let starts = Arc::new(Barrier::new(4)); diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs index 6c3e4ea10d..671caec8be 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_extraction.rs @@ -1,6 +1,7 @@ use crate::compiler::tests::helpers::{ dependent_assertion_returns_authoritative_error, empty_model_fixture_marker_preserves_direct_mode_validation, + expected_projection_errors_name_the_expected_cte, mixed_expanded_tests_preserve_order_and_payloads, quoted_ctes_and_implicit_alias_preserve_payload, }; @@ -29,6 +30,11 @@ fn given_sql_test_cases_when_extracting_native_payloads_then_expected_behavior_h run: empty_model_fixture_marker_preserves_direct_mode_validation, expected_success: true, }, + SqlTestExtractionTestCase { + description: "expected projection errors name the expected CTE", + run: expected_projection_errors_name_the_expected_cte, + expected_success: true, + }, ]; for test_case in test_cases { diff --git a/crates/sqlbuild-rules-native/src/engine/tests.rs b/crates/sqlbuild-rules-native/src/engine/tests.rs index 1907da7600..08d7e5b848 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests.rs @@ -8,6 +8,8 @@ mod evaluation; mod helpers; #[path = "tests/test_scope_facts.rs"] mod scope_facts; +#[path = "tests/test_sql_scanning.rs"] +mod sql_scanning; #[path = "tests/test_sql_test_rules.rs"] mod sql_test_rules; #[path = "tests/test_types.rs"] diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs index e6af01313f..03d28d174c 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_evaluation.rs @@ -428,6 +428,57 @@ fn given_snowflake_expression_when_evaluating_rules_then_uses_project_dialect() Ok(()) } +#[test] +fn given_dialect_quoting_around_table_functions_when_evaluating_rules_then_model_parses() +-> Result<(), String> { + let project_dir = TempDir::new().map_err(|error| error.to_string())?; + let test_cases = [ + test_types::DialectEvaluationTestCase { + description: "Snowflake backslash-escaped quote before a table function", + dialect: "snowflake", + query_sql: concat!( + "SELECT 'it\\'s (' AS label, item_id ", + "FROM __table_fn(\"inventory_items\")(42)" + ), + expected_code: "SQBRCONTRACT101", + }, + test_types::DialectEvaluationTestCase { + description: "BigQuery backslash-escaped quote inside table function arguments", + dialect: "bigquery", + query_sql: concat!( + "SELECT item_id ", + "FROM __table_fn(\"inventory_items\", 'a\\')')(42)" + ), + expected_code: "SQBRCONTRACT101", + }, + test_types::DialectEvaluationTestCase { + description: "BigQuery backtick identifier containing a parenthesis in arguments", + dialect: "bigquery", + query_sql: concat!( + "SELECT item_id ", + "FROM __table_fn(\"inventory_items\", `batch)size`)(42)" + ), + expected_code: "SQBRCONTRACT101", + }, + ]; + + for test_case in test_cases { + let config = json!({"select": ["SQBRCONTRACT101"], "cache": {"enabled": false}}); + let mut request: Value = serde_json::from_str(&helpers::request(&project_dir, &config)) + .map_err(|error| error.to_string())?; + request["dialect"] = json!(test_case.dialect); + request["models"][0]["query_sql"] = json!(test_case.query_sql); + let result: Value = serde_json::from_str(&evaluate_json(&request.to_string())?) + .map_err(|error| error.to_string())?; + assert_eq!( + result["faults"][0]["code"], test_case.expected_code, + "{}: {result}", + test_case.description + ); + } + Ok(()) +} + #[test] fn given_enforced_contract_outputs_when_evaluating_explicit_type_rule_then_returns_expected_faults() -> Result<(), String> { diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs new file mode 100644 index 0000000000..baea4d051b --- /dev/null +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_sql_scanning.rs @@ -0,0 +1,306 @@ +use crate::compiler::_helpers::sql_references::extraction::extract; +use crate::engine::tests::test_types::SqlScannerTestCase; +use crate::rules::_helpers::evaluation::{normalize_rules_sql, rules_quote_policy}; +use crate::rules::_helpers::numeric_decisions::compact_sql; +use crate::sql_lint::_helpers::preparation::prepare; +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; + +#[test] +fn given_quoted_commented_and_malformed_fragments_when_scanning_then_every_scanner_is_characterised() +-> Result<(), String> { + let test_cases = [ + SqlScannerTestCase { + description: "plain code", + fragment: "a", + expected_compiler_paren: Ok(4), + expected_table_function: "SELECT * FROM __table_fn(\"o\", a x, 1)", + expected_table_function_token: "SELECT a __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT a a FROM t", + expected_compact: "a=aandb", + expected_snowflake_compact: "a=aandb", + expected_lint_site: Some("@m(a x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "single quote containing open paren", + fragment: "'a(b'", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a(b' x, 1)", + expected_table_function_token: "SELECT 'a(b' __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a(b' a FROM t", + expected_compact: "a='a(b'andb", + expected_snowflake_compact: "a='a(b'andb", + expected_lint_site: Some("@m('a(b' x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "single quote containing close paren", + fragment: "'a)b'", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a)b' x, 1)", + expected_table_function_token: "SELECT 'a)b' __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a)b' a FROM t", + expected_compact: "a='a)b'andb", + expected_snowflake_compact: "a='a)b'andb", + expected_lint_site: Some("@m('a)b' x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "doubled single quote", + fragment: "'it''s )'", + expected_compiler_paren: Ok(12), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'it''s )' x, 1)", + expected_table_function_token: "SELECT 'it''s )' __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'it''s )' a FROM t", + expected_compact: "a='it''s )'andb", + expected_snowflake_compact: "a='it''s )'andb", + expected_lint_site: Some("@m('it''s )' x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "backslash escaped quote before close paren", + fragment: "'a\\')'", + expected_compiler_paren: Ok(5), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a\\')' x)(1)", + expected_table_function_token: "SELECT 'a\\')' __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a\\')' a FROM t", + expected_compact: "a='a\\')' and b", + expected_snowflake_compact: "a='a\\')'andb", + expected_lint_site: Some("@m('a\\')' x)"), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "backslash escaped quote before open paren", + fragment: "'a\\'('", + expected_compiler_paren: Err(Unclosed::Quote), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'a\\'(' x)(1)", + expected_table_function_token: "SELECT 'a\\'(' __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'a\\'(' a FROM t", + expected_compact: "a='a\\'(' and b", + expected_snowflake_compact: "a='a\\'('andb", + expected_lint_site: Some("@m('a\\'(' x)"), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "double quote containing close paren", + fragment: "\"a)b\"", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", \"a)b\" x, 1)", + expected_table_function_token: "SELECT \"a)b\" __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT \"a)b\" a FROM t", + expected_compact: "a=\"a)b\"andb", + expected_snowflake_compact: "a=\"a)b\"andb", + expected_lint_site: Some("@m(\"a)b\" x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "backtick containing close paren", + fragment: "`a)b`", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", `a)b` x, 1)", + expected_table_function_token: "SELECT `a)b` __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT `a)b` a FROM t", + expected_compact: "a=`a)b`andb", + expected_snowflake_compact: "a=`a)b`andb", + expected_lint_site: Some("@m(`a)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "doubled backtick", + fragment: "`a``)b`", + expected_compiler_paren: Ok(10), + expected_table_function: "SELECT * FROM __table_fn(\"o\", `a``)b` x, 1)", + expected_table_function_token: "SELECT `a``)b` __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT `a``)b` a FROM t", + expected_compact: "a=`a``)b`andb", + expected_snowflake_compact: "a=`a``)b`andb", + expected_lint_site: Some("@m(`a``)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "line comment", + fragment: "-- )\n", + expected_compiler_paren: Ok(8), + expected_table_function: "SELECT * FROM __table_fn(\"o\", -- )\n x, 1)", + expected_table_function_token: "SELECT -- )\n __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * \n , a FROM t", + expected_snowflake_comma: "SELECT -- )\n a FROM t", + expected_compact: "a=andb", + expected_snowflake_compact: "a=andb", + expected_lint_site: Some("@m(-- )\n x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "block comment", + fragment: "/* ) */", + expected_compiler_paren: Ok(10), + expected_table_function: "SELECT * FROM __table_fn(\"o\", /* ) */ x, 1)", + expected_table_function_token: "SELECT /* ) */ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT /* ) */ a FROM t", + expected_compact: "a=andb", + expected_snowflake_compact: "a=andb", + expected_lint_site: Some("@m(/* ) */ x)"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "nested block comment", + fragment: "/* /* ) */ ) */", + expected_compiler_paren: Ok(12), + expected_table_function: "SELECT * FROM __table_fn(\"o\", /* /* ) */ ) */ x)(1)", + expected_table_function_token: "SELECT /* /* ) */ ) */ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * */ x) , a FROM t", + expected_snowflake_comma: "SELECT /* /* ) */ ) */ a FROM t", + expected_compact: "a=)*/andb", + expected_snowflake_compact: "a=)*/andb", + expected_lint_site: Some("@m(/* /* ) */ )"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "dollar quote", + fragment: "$$ ) $$", + expected_compiler_paren: Ok(4), + expected_table_function: "SELECT * FROM __table_fn(\"o\", $$ ) $$ x)(1)", + expected_table_function_token: "SELECT $$ ) $$ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * $$ x) , a FROM t", + expected_snowflake_comma: "SELECT $$ ) $$ a FROM t", + expected_compact: "a=$$)$$andb", + expected_snowflake_compact: "a=$$)$$andb", + expected_lint_site: Some("@m($$ )"), + expected_reference_fast_path: true, + }, + SqlScannerTestCase { + description: "unterminated quote", + fragment: "'abc", + expected_compiler_paren: Err(Unclosed::Quote), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'abc x)(1)", + expected_table_function_token: "SELECT 'abc __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * EXCLUDE ('abc x) , a FROM t", + expected_snowflake_comma: "SELECT 'abc a, FROM t", + expected_compact: "a='abc and b", + expected_snowflake_compact: "a='abc and b", + expected_lint_site: Some(""), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "unterminated block comment", + fragment: "/* abc", + expected_compiler_paren: Err(Unclosed::BlockComment), + expected_table_function: "SELECT * FROM __table_fn(\"o\", /* abc x)(1)", + expected_table_function_token: "SELECT /* abc __table_fn(\"o\")(1)", + expected_snowflake_exclude: "SELECT * EXCLUDE (/* abc x) , a FROM t", + expected_snowflake_comma: "SELECT /* abc a, FROM t", + expected_compact: "a=", + expected_snowflake_compact: "a=", + expected_lint_site: Some(""), + expected_reference_fast_path: false, + }, + SqlScannerTestCase { + description: "multibyte utf8", + fragment: "'é)ü' ñ", + expected_compiler_paren: Ok(13), + expected_table_function: "SELECT * FROM __table_fn(\"o\", 'é)ü' ñ x, 1)", + expected_table_function_token: "SELECT 'é)ü' ñ __table_fn(\"o\", 1)", + expected_snowflake_exclude: "SELECT * , a FROM t", + expected_snowflake_comma: "SELECT 'é)ü' ñ a FROM t", + expected_compact: "a='é)ü'ñandb", + expected_snowflake_compact: "a='é)ü'ñandb", + expected_lint_site: None, + expected_reference_fast_path: true, + }, + ]; + + for test_case in test_cases { + let fragment = test_case.fragment; + assert_eq!( + matching_paren( + format!("({fragment} x) y").as_bytes(), + 0, + QuotePolicy::COMPILER + ), + test_case.expected_compiler_paren, + "compiler matching paren: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql( + "duckdb", + &format!("SELECT * FROM __table_fn(\"o\", {fragment} x)(1)") + ), + test_case.expected_table_function, + "rules table-function parenthesis: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql("duckdb", &format!("SELECT {fragment} __table_fn(\"o\")(1)")), + test_case.expected_table_function_token, + "rules table-function token search: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql( + "snowflake", + &format!("SELECT * EXCLUDE ({fragment} x) , a FROM t") + ), + test_case.expected_snowflake_exclude, + "rules Snowflake EXCLUDE parenthesis: {}", + test_case.description + ); + assert_eq!( + normalize_rules_sql("snowflake", &format!("SELECT {fragment} a, FROM t")), + test_case.expected_snowflake_comma, + "rules Snowflake trailing comma: {}", + test_case.description + ); + assert_eq!( + compact_sql( + &format!("A = {fragment} AND B"), + rules_quote_policy("duckdb") + ), + test_case.expected_compact, + "rules numeric decision compaction: {}", + test_case.description + ); + assert_eq!( + compact_sql( + &format!("A = {fragment} AND B"), + rules_quote_policy("snowflake") + ), + test_case.expected_snowflake_compact, + "rules Snowflake numeric decision compaction: {}", + test_case.description + ); + let lint_sql = format!("SELECT @m({fragment} x) AS y"); + let lint_site = prepare(&lint_sql, &lint_sql, &[])?.map(|prepared| { + prepared + .1 + .first() + .map(|site| site.5.clone()) + .unwrap_or_default() + }); + assert_eq!( + lint_site.as_deref(), + test_case.expected_lint_site, + "SQL lint macro site: {}", + test_case.description + ); + assert_eq!( + extract(&format!("SELECT {fragment} __ref('b')")).is_some(), + test_case.expected_reference_fast_path, + "SQL reference fast path: {}", + test_case.description + ); + } + Ok(()) +} diff --git a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs index c3457bd571..9c319a4f50 100644 --- a/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/engine/tests/test_types.rs @@ -1,5 +1,7 @@ use serde_json::Value; +use crate::sql_scan::models::Unclosed; + pub(crate) struct NativeEvaluationTestCase { pub(crate) description: &'static str, pub(crate) config: Value, @@ -165,3 +167,19 @@ pub(crate) struct EmptyInputMinimumTestsTestCase { pub(crate) allowed_tests: Value, pub(crate) expected_messages: &'static [&'static str], } + +/// One SQL fragment scanned by every native quote- and comment-aware scanner. +pub(crate) struct SqlScannerTestCase { + pub(crate) description: &'static str, + pub(crate) fragment: &'static str, + pub(crate) expected_compiler_paren: Result, + pub(crate) expected_table_function: &'static str, + pub(crate) expected_table_function_token: &'static str, + pub(crate) expected_snowflake_exclude: &'static str, + pub(crate) expected_snowflake_comma: &'static str, + pub(crate) expected_compact: &'static str, + pub(crate) expected_snowflake_compact: &'static str, + /// `None` when SQL lint declines non-ASCII input; `Some("")` when no macro site closes. + pub(crate) expected_lint_site: Option<&'static str>, + pub(crate) expected_reference_fast_path: bool, +} diff --git a/crates/sqlbuild-rules-native/src/lib.rs b/crates/sqlbuild-rules-native/src/lib.rs index 773c00df7d..b37d85783b 100644 --- a/crates/sqlbuild-rules-native/src/lib.rs +++ b/crates/sqlbuild-rules-native/src/lib.rs @@ -12,3 +12,4 @@ mod scope_metadata; mod semantic_usage; mod semantic_validation; mod sql_lint; +mod sql_scan; diff --git a/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs b/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs index e997f1fdd4..6732ebc948 100644 --- a/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs +++ b/crates/sqlbuild-rules-native/src/rules/_helpers/evaluation.rs @@ -11,6 +11,10 @@ use crate::rules::_helpers::{ use crate::rules::models::{ FaultCollector, ModelEvaluationRequest, ProjectEvaluationRequest, ResolvedThresholdOverride, }; +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::QuotePolicy; use globset::{Glob, GlobSetBuilder}; use sqlparser::ast::{ BinaryOperator, Expr, GroupByExpr, JoinConstraint, JoinOperator, Query, Select, SelectItem, @@ -70,6 +74,7 @@ struct ParsedModel<'a> { model: &'a Model, classification: ModelClassification, authored_compact_sql: String, + quote_policy: QuotePolicy, } struct ModelClassification { @@ -145,12 +150,14 @@ fn evaluate_model_inner(request: ModelEvaluationRequest<'_>) -> Result String { + let policy = rules_quote_policy("generic"); let mut bytes = sql.as_bytes().to_vec(); - let mut quote: Option = None; let mut index = 0; while index < bytes.len() { let byte = bytes[index]; let next = bytes.get(index + 1).copied(); - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; - } - index += 1; + if policy.is_quote(byte) { + let Ok(end) = quote_end(&bytes, index, policy) else { + break; + }; + index = end; continue; } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - } else if byte == b'-' && next == Some(b'>') { + if byte == b'-' && next == Some(b'>') { bytes[index] = b','; bytes[index + 1] = b' '; index += 2; @@ -303,6 +304,17 @@ fn parse_with_dialect(sql: &str, dialect: &dyn Dialect) -> Result parser.parse_statements() } +pub(crate) fn rules_quote_policy(dialect_name: &str) -> QuotePolicy { + let dialect = rules_dialect(dialect_name); + let backslash_escapes = dialect.supports_string_literal_backslash_escape(); + QuotePolicy { + backtick_identifiers: dialect.is_delimited_identifier_start('`'), + single_quote_backslash_escapes: backslash_escapes, + double_quote_backslash_escapes: backslash_escapes + && !dialect.is_delimited_identifier_start('"'), + } +} + fn rules_dialect(name: &str) -> Box { match name.to_ascii_lowercase().as_str() { "bigquery" => Box::new(BigQueryDialect {}), @@ -317,57 +329,22 @@ fn rules_dialect(name: &str) -> Box { } pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { - let sql = normalize_table_function_calls(sql); + let policy = rules_quote_policy(dialect); + let sql = normalize_table_function_calls(sql, policy); if !dialect.eq_ignore_ascii_case("snowflake") { return sql; } let mut bytes = sql.as_bytes().to_vec(); let mut index = 0; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; while index < bytes.len() { let byte = bytes[index]; - let next = bytes.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; - } - continue; - } - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; + match non_code_end(&bytes, index, policy) { + Ok(Some(end)) => { + index = end; + continue; } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - index += 1; - continue; + Ok(None) => {} + Err(_) => break, } if byte == b',' && clause_follows(&bytes, index + 1) { bytes[index] = b' '; @@ -389,7 +366,7 @@ pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { if word.eq_ignore_ascii_case(b"EXCLUDE") && is_star_modifier_before(sql.as_bytes(), start) && bytes.get(following) == Some(&b'(') - && let Some(close) = matching_parenthesis(&bytes, following) + && let Ok(close) = matching_paren(&bytes, following, policy) { blank_preserving_newlines(&mut bytes[start..=close]); index = close + 1; @@ -398,7 +375,7 @@ pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { if word.eq_ignore_ascii_case(b"ARRAY") && is_cast_type_before(sql.as_bytes(), start) && bytes.get(following) == Some(&b'(') - && let Some(close) = matching_parenthesis(&bytes, following) + && let Ok(close) = matching_paren(&bytes, following, policy) { bytes[start..index].copy_from_slice(b"TEXT "); blank_preserving_newlines(&mut bytes[following..=close]); @@ -419,16 +396,17 @@ pub(crate) fn normalize_rules_sql(dialect: &str, sql: &str) -> String { String::from_utf8(bytes).unwrap_or(sql) } -fn normalize_table_function_calls(sql: &str) -> String { +fn normalize_table_function_calls(sql: &str, policy: QuotePolicy) -> String { const TABLE_FUNCTION: &[u8] = b"__table_fn("; let mut bytes = sql.as_bytes().to_vec(); let mut cursor = 0; while cursor + TABLE_FUNCTION.len() <= bytes.len() { - let Some(token_start) = find_token_outside_syntax(&bytes, TABLE_FUNCTION, cursor) else { + let Some(token_start) = find_token_outside_syntax(&bytes, TABLE_FUNCTION, cursor, policy) + else { break; }; let open = token_start + TABLE_FUNCTION.len() - 1; - let Some(close) = matching_parenthesis(&bytes, open) else { + let Ok(close) = matching_paren(&bytes, open, policy) else { break; }; let mut next = close + 1; @@ -452,53 +430,21 @@ fn blank_preserving_newlines(sql: &mut [u8]) { } } -fn find_token_outside_syntax(sql: &[u8], token: &[u8], start: usize) -> Option { +fn find_token_outside_syntax( + sql: &[u8], + token: &[u8], + start: usize, + policy: QuotePolicy, +) -> Option { let mut index = start; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; while index + token.len() <= sql.len() { - let byte = sql[index]; - let next = sql.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; + match non_code_end(sql, index, policy) { + Ok(Some(end)) => { + index = end; + continue; } - continue; - } - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; - } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - index += 1; - continue; + Ok(None) => {} + Err(_) => return None, } if sql[index..index + token.len()].eq_ignore_ascii_case(token) { return Some(index); @@ -548,70 +494,6 @@ fn is_cast_type_before(sql: &[u8], type_start: usize) -> bool { || !sql[cursor - BYTE_BEFORE_AS_OFFSET].is_ascii_alphanumeric()) } -fn matching_parenthesis(sql: &[u8], open: usize) -> Option { - let mut depth = 0; - let mut index = open; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; - while index < sql.len() { - let byte = sql[index]; - let next = sql.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; - } - continue; - } - if let Some(active_quote) = quote { - if byte == active_quote { - if next == Some(active_quote) { - index += 2; - continue; - } - quote = None; - } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - index += 1; - continue; - } - match byte { - b'(' => depth += 1, - b')' => { - depth -= 1; - if depth == 0 { - return Some(index); - } - } - _ => {} - } - index += 1; - } - None -} - fn is_lambda_parameter_before(sql: &[u8], type_start: usize) -> bool { let mut cursor = type_start; while cursor > 0 && sql[cursor - 1].is_ascii_whitespace() { @@ -1374,6 +1256,7 @@ fn evaluate_literal_rules( equality: comparison.equality, output_name: comparison.output_name.as_deref(), literal, + quote_policy: parsed.quote_policy, }) }); if magic { diff --git a/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs b/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs index fddcbcb86a..3d861d3fb3 100644 --- a/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs +++ b/crates/sqlbuild-rules-native/src/rules/_helpers/numeric_decisions.rs @@ -1,3 +1,7 @@ +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::QuotePolicy; + const CANONICAL_NUMERIC_DECISIONS: [&str; 3] = ["-1", "0", "1"]; pub(crate) struct AuthoredDecision<'a> { @@ -6,62 +10,38 @@ pub(crate) struct AuthoredDecision<'a> { pub(crate) equality: bool, pub(crate) output_name: Option<&'a str>, pub(crate) literal: &'a str, + pub(crate) quote_policy: QuotePolicy, } -pub(crate) fn compact_sql(sql: &str) -> String { +pub(crate) fn compact_sql(sql: &str, policy: QuotePolicy) -> String { let bytes = sql.as_bytes(); let mut compact = Vec::with_capacity(bytes.len()); let mut index = 0; - let mut quote: Option = None; - let mut line_comment = false; - let mut block_comment = false; while index < bytes.len() { let byte = bytes[index]; - let next = bytes.get(index + 1).copied(); - if line_comment { - line_comment = byte != b'\n'; - index += 1; - continue; - } - if block_comment { - if byte == b'*' && next == Some(b'/') { - block_comment = false; - index += 2; - } else { - index += 1; + let unclosed_end = match comment_end(bytes, index) { + Ok(Some(end)) => { + index = end; + continue; } - continue; - } - if let Some(active_quote) = quote { - compact.push(byte.to_ascii_lowercase()); - if byte == active_quote { - if next == Some(active_quote) { - compact.push(active_quote); - index += 2; - continue; + Ok(None) if policy.is_quote(byte) => { + quote_end(bytes, index, policy).unwrap_or(bytes.len()) + } + Ok(None) => { + if !byte.is_ascii_whitespace() { + compact.push(byte.to_ascii_lowercase()); } - quote = None; + index += 1; + continue; } - index += 1; - continue; - } - if byte == b'-' && next == Some(b'-') { - line_comment = true; - index += 2; - continue; - } - if byte == b'/' && next == Some(b'*') { - block_comment = true; - index += 2; - continue; - } - if matches!(byte, b'\'' | b'"') { - quote = Some(byte); - } - if !byte.is_ascii_whitespace() { - compact.push(byte.to_ascii_lowercase()); - } - index += 1; + Err(_) => break, + }; + compact.extend( + bytes[index..unclosed_end] + .iter() + .map(u8::to_ascii_lowercase), + ); + index = unclosed_end; } String::from_utf8(compact).unwrap_or_default() } @@ -70,7 +50,7 @@ pub(crate) fn is_authored_decision(decision: AuthoredDecision<'_>) -> bool { !CANONICAL_NUMERIC_DECISIONS.contains(&decision.literal) && decision .authored_compact_sql - .contains(&compact_sql(decision.comparison_sql)) + .contains(&compact_sql(decision.comparison_sql, decision.quote_policy)) && !decision.output_name.is_some_and(|name| { decision.equality && name.ends_with(&format!("_{}", decision.literal)) }) diff --git a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs index 8a9ddc97c2..872d157226 100644 --- a/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs +++ b/crates/sqlbuild-rules-native/src/sql_lint/_helpers/preparation.rs @@ -6,6 +6,8 @@ use std::sync::LazyLock; use regex::Regex; use crate::sql_lint::types::{InterpolationSite, PreparedSql}; +use crate::sql_scan::main::matching_paren::matching_paren as scan_matching_paren; +use crate::sql_scan::models::QuotePolicy; static SITES: LazyLock> = LazyLock::new(|| { Regex::new( @@ -135,49 +137,9 @@ fn scan_name(bytes: &[u8], mut index: usize, interpolation: bool) -> usize { index } -fn matching_paren(bytes: &[u8], mut index: usize) -> Option { - let mut depth = 0; - let mut quote = None; - while index < bytes.len() { - let byte = bytes[index]; - if let Some(quoted) = quote { - if byte == b'\\' && index + 1 < bytes.len() { - index += 2; - continue; - } - if byte == quoted { - if bytes.get(index + 1) == Some("ed) { - index += 2; - continue; - } - quote = None; - } - } else if bytes[index..].starts_with(b"--") { - index = bytes[index + 2..] - .iter() - .position(|byte| *byte == b'\n') - .map_or(bytes.len(), |offset| index + offset + 3); - continue; - } else if bytes[index..].starts_with(b"/*") { - index = bytes[index + 2..] - .windows(2) - .position(|pair| pair == b"*/") - .map_or(bytes.len(), |offset| index + offset + 4); - continue; - } else { - match byte { - b'\'' | b'"' => quote = Some(byte), - b'(' => depth += 1, - b')' => { - depth -= 1; - if depth == 0 { - return Some(index + 1); - } - } - _ => {} - } - } - index += 1; +fn matching_paren(bytes: &[u8], open: usize) -> Option { + match scan_matching_paren(bytes, open, QuotePolicy::SQL_LINT) { + Ok(close) => Some(close + 1), + Err(_) => None, } - None } diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs new file mode 100644 index 0000000000..b90d31b65d --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/comment_end.rs @@ -0,0 +1,21 @@ +//! SQL line and block comment boundaries. + +use crate::sql_scan::models::Unclosed; + +/// Return the end of a line or block comment starting at `index`, if one starts there. +pub(crate) fn comment_end(sql: &[u8], index: usize) -> Result, Unclosed> { + match (sql.get(index), sql.get(index + 1)) { + (Some(b'-'), Some(b'-')) => Ok(Some( + sql[index + 2..] + .iter() + .position(|byte| *byte == b'\n') + .map_or(sql.len(), |offset| index + 2 + offset + 1), + )), + (Some(b'/'), Some(b'*')) => sql[index + 2..] + .windows(2) + .position(|pair| pair == b"*/") + .map(|offset| Some(index + 2 + offset + 2)) + .ok_or(Unclosed::BlockComment), + _ => Ok(None), + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs new file mode 100644 index 0000000000..89ee9b98e3 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/matching_paren.rs @@ -0,0 +1,32 @@ +//! Quote- and comment-aware parenthesis matching. + +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::models::{QuotePolicy, Unclosed}; + +/// Return the index of the parenthesis closing the one at `open`, skipping quotes and comments. +pub(crate) fn matching_paren( + sql: &[u8], + open: usize, + policy: QuotePolicy, +) -> Result { + let mut depth = 0isize; + let mut index = open; + while index < sql.len() { + if let Some(end) = non_code_end(sql, index, policy)? { + index = end; + continue; + } + match sql[index] { + b'(' => depth += 1, + b')' => { + depth -= 1; + if depth == 0 { + return Ok(index); + } + } + _ => {} + } + index += 1; + } + Err(Unclosed::Parenthesis) +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs new file mode 100644 index 0000000000..4bf480ac79 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod comment_end; +pub(crate) mod matching_paren; +pub(crate) mod non_code_end; +pub(crate) mod quote_end; +pub(crate) mod skip_whitespace; diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs new file mode 100644 index 0000000000..cb99e49124 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/non_code_end.rs @@ -0,0 +1,20 @@ +//! Boundaries of the comment or quoted text starting at one offset. + +use crate::sql_scan::main::comment_end::comment_end; +use crate::sql_scan::main::quote_end::quote_end; +use crate::sql_scan::models::{QuotePolicy, Unclosed}; + +/// Return the end of a comment or quoted text starting at `index`, if one starts there. +pub(crate) fn non_code_end( + sql: &[u8], + index: usize, + policy: QuotePolicy, +) -> Result, Unclosed> { + if let Some(end) = comment_end(sql, index)? { + return Ok(Some(end)); + } + match sql.get(index) { + Some(byte) if policy.is_quote(*byte) => quote_end(sql, index, policy).map(Some), + _ => Ok(None), + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs new file mode 100644 index 0000000000..fdf256cea6 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/quote_end.rs @@ -0,0 +1,23 @@ +//! Quoted SQL text boundaries under a quoting policy. + +use crate::sql_scan::models::{QuotePolicy, Unclosed}; + +/// Return the index after the quoted text whose opening quote is at `start`. +pub(crate) fn quote_end(sql: &[u8], start: usize, policy: QuotePolicy) -> Result { + let quote = sql[start]; + let backslash_escapes = policy.backslash_escapes(quote); + let mut index = start + 1; + while index < sql.len() { + let byte = sql[index]; + if backslash_escapes && byte == b'\\' && index + 1 < sql.len() { + index += 2; + } else if byte != quote { + index += 1; + } else if sql.get(index + 1) == Some("e) { + index += 2; + } else { + return Ok(index + 1); + } + } + Err(Unclosed::Quote) +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs b/crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs new file mode 100644 index 0000000000..a60ea5e9c4 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/main/skip_whitespace.rs @@ -0,0 +1,12 @@ +//! Unicode whitespace skipping for SQL scanners. + +/// Return the first index at or after `index` that is not Unicode whitespace. +pub(crate) fn skip_whitespace(sql: &str, mut index: usize) -> usize { + while let Some(character) = sql.get(index..).and_then(|rest| rest.chars().next()) { + if !character.is_whitespace() { + break; + } + index += character.len_utf8(); + } + index +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/mod.rs b/crates/sqlbuild-rules-native/src/sql_scan/mod.rs new file mode 100644 index 0000000000..650595395f --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod main; +pub(crate) mod models; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-rules-native/src/sql_scan/models.rs b/crates/sqlbuild-rules-native/src/sql_scan/models.rs new file mode 100644 index 0000000000..74e0f5a908 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/models.rs @@ -0,0 +1,54 @@ +//! Lexical scanning outcomes and dialect-dependent quoting policy. + +/// The SQL construct a scan reached the end of input inside. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Unclosed { + BlockComment, + Quote, + Parenthesis, +} + +/// Dialect-dependent quoting rules; comments and doubled-quote escapes are fixed for every policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct QuotePolicy { + /// Backtick-delimited identifiers are quoted text. + pub(crate) backtick_identifiers: bool, + /// A backslash escapes the next byte inside single-quoted strings. + pub(crate) single_quote_backslash_escapes: bool, + /// A backslash escapes the next byte inside double-quoted text. + pub(crate) double_quote_backslash_escapes: bool, +} + +impl QuotePolicy { + /// Compiler policy, mirroring Python `sql_analysis`: backticks quote, backslashes are literal. + pub(crate) const COMPILER: Self = Self { + backtick_identifiers: true, + single_quote_backslash_escapes: false, + double_quote_backslash_escapes: false, + }; + + /// SQL lint policy, mirroring its site lexer: backslashes escape, backticks are code. + pub(crate) const SQL_LINT: Self = Self { + backtick_identifiers: false, + single_quote_backslash_escapes: true, + double_quote_backslash_escapes: true, + }; + + /// Return whether `byte` opens quoted text under this policy. + pub(crate) fn is_quote(self, byte: u8) -> bool { + match byte { + b'\'' | b'"' => true, + b'`' => self.backtick_identifiers, + _ => false, + } + } + + /// Return whether a backslash escapes the next byte inside text opened by `quote`. + pub(crate) fn backslash_escapes(self, quote: u8) -> bool { + match quote { + b'\'' => self.single_quote_backslash_escapes, + b'"' => self.double_quote_backslash_escapes, + _ => false, + } + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests.rs new file mode 100644 index 0000000000..0b2156d36c --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests.rs @@ -0,0 +1,4 @@ +#[path = "tests/test_sql_scan.rs"] +mod sql_scan; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs new file mode 100644 index 0000000000..99da540962 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_sql_scan.rs @@ -0,0 +1,155 @@ +use crate::sql_scan::main::matching_paren::matching_paren; +use crate::sql_scan::main::non_code_end::non_code_end; +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; +use crate::sql_scan::tests::test_types::{MatchingParenPolicyTestCase, NonCodeEndTestCase}; + +const QUOTES_ONLY: QuotePolicy = QuotePolicy { + backtick_identifiers: false, + single_quote_backslash_escapes: false, + double_quote_backslash_escapes: false, +}; + +const BACKSLASH_ONLY: QuotePolicy = QuotePolicy { + backtick_identifiers: false, + single_quote_backslash_escapes: true, + double_quote_backslash_escapes: false, +}; + +#[test] +fn given_quote_policies_when_matching_parentheses_then_policy_decides_quote_boundaries() { + let test_cases = [ + MatchingParenPolicyTestCase { + description: "compiler policy skips backtick identifiers", + sql: "(`a)b` x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(8), + }, + MatchingParenPolicyTestCase { + description: "quotes-only policy treats backticks as code", + sql: "(`a)b` x) y", + policy: QUOTES_ONLY, + expected_close: Ok(3), + }, + MatchingParenPolicyTestCase { + description: "compiler policy ends a quote at a backslash-escaped quote", + sql: r"('a\')' x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(5), + }, + MatchingParenPolicyTestCase { + description: "sql lint policy honours backslash escapes", + sql: r"('a\')' x) y", + policy: QuotePolicy::SQL_LINT, + expected_close: Ok(9), + }, + MatchingParenPolicyTestCase { + description: "single-quote backslash escapes do not apply to double quotes", + sql: r#"("a\")" x) y"#, + policy: BACKSLASH_ONLY, + expected_close: Ok(5), + }, + MatchingParenPolicyTestCase { + description: "doubled quotes are escapes under every policy", + sql: "('it''s )' x) y", + policy: QUOTES_ONLY, + expected_close: Ok(12), + }, + MatchingParenPolicyTestCase { + description: "block comments do not nest", + sql: "(/* /* ) */ ) */ x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(12), + }, + MatchingParenPolicyTestCase { + description: "dollar quotes are scanned as code", + sql: "($$ ) $$ x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(4), + }, + MatchingParenPolicyTestCase { + description: "multibyte text keeps byte offsets on character boundaries", + sql: "('é)ü' ñ x) y", + policy: QuotePolicy::COMPILER, + expected_close: Ok(13), + }, + MatchingParenPolicyTestCase { + description: "unterminated quote reports the quote", + sql: "('abc x) y", + policy: QuotePolicy::COMPILER, + expected_close: Err(Unclosed::Quote), + }, + MatchingParenPolicyTestCase { + description: "unterminated block comment reports the comment", + sql: "(/* abc x) y", + policy: QUOTES_ONLY, + expected_close: Err(Unclosed::BlockComment), + }, + MatchingParenPolicyTestCase { + description: "unclosed parenthesis reports the parenthesis", + sql: "(a (b)", + policy: QuotePolicy::SQL_LINT, + expected_close: Err(Unclosed::Parenthesis), + }, + ]; + + for test_case in test_cases { + assert_eq!( + matching_paren(test_case.sql.as_bytes(), 0, test_case.policy), + test_case.expected_close, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_comment_or_quote_starts_when_finding_non_code_end_then_returns_end_offset() { + let test_cases = [ + NonCodeEndTestCase { + description: "line comment ends after its newline", + sql: "-- note\nSELECT", + policy: QuotePolicy::COMPILER, + expected_end: Ok(Some(8)), + }, + NonCodeEndTestCase { + description: "line comment without newline runs to the end", + sql: "-- note", + policy: QuotePolicy::COMPILER, + expected_end: Ok(Some(7)), + }, + NonCodeEndTestCase { + description: "block comment ends after its terminator", + sql: "/* a */ b", + policy: QUOTES_ONLY, + expected_end: Ok(Some(7)), + }, + NonCodeEndTestCase { + description: "backtick is code when the policy disables it", + sql: "`a` b", + policy: QuotePolicy::SQL_LINT, + expected_end: Ok(None), + }, + NonCodeEndTestCase { + description: "trailing backslash leaves the quote unterminated", + sql: r"'a\", + policy: QuotePolicy::SQL_LINT, + expected_end: Err(Unclosed::Quote), + }, + NonCodeEndTestCase { + description: "plain code has no non-code end", + sql: "SELECT 1", + policy: QuotePolicy::COMPILER, + expected_end: Ok(None), + }, + ]; + + for test_case in test_cases { + assert_eq!( + non_code_end(test_case.sql.as_bytes(), 0, test_case.policy), + test_case.expected_end, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs new file mode 100644 index 0000000000..0951ec3e0e --- /dev/null +++ b/crates/sqlbuild-rules-native/src/sql_scan/tests/test_types.rs @@ -0,0 +1,16 @@ +use crate::sql_scan::models::QuotePolicy; +use crate::sql_scan::models::Unclosed; + +pub(crate) struct MatchingParenPolicyTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) policy: QuotePolicy, + pub(crate) expected_close: Result, +} + +pub(crate) struct NonCodeEndTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) policy: QuotePolicy, + pub(crate) expected_end: Result, Unclosed>, +} diff --git a/fensu.toml b/fensu.toml index 6d449018d8..a5409ee41b 100644 --- a/fensu.toml +++ b/fensu.toml @@ -93,11 +93,6 @@ rule = "FFR601" path = "src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py" reason = "The SQL Server adapter intentionally keeps overrideable backend behavior on one class." -[[targets.python.rule_exceptions]] -rule = "FFR601" -path = "src/sqlbuild/virtual/state/classes/postgres.py" -reason = "The Postgres state backend intentionally keeps its transactional storage contract on one class." - [[targets.python.rule_exceptions]] rule = "FFS120" path = "src/sqlbuild/integrations/dbt/_helpers/profile/render.py" diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index d2bf04d48a..db9a9d663b 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -9,7 +9,17 @@ from pathlib import Path from typing import Any, ClassVar, cast +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + render_is_distinct_from, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.retention_adapter import RetentionAdapterMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter from sqlbuild.adapter.contract.constants import ( @@ -37,14 +47,19 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) from sqlbuild.adapter.contract.types import ( CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -54,7 +69,6 @@ render_insert_source_freshness_records_sql, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS from sqlbuild.spec.contracts.models import SeedCsvSettings @@ -86,6 +100,14 @@ class BaseAdapter(RetentionAdapterMixin, StrictAdapter): state_tables_transient: ClassVar[bool] = False allows_implicit_managed_write_schema: ClassVar[bool] = False execution_duration_limit_seconds: ClassVar[int | None] = None + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.CORRELATED, + historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, + ) def supports_zero_copy_clone(self) -> bool: return False @@ -858,19 +880,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=self.render_current_timestamp(), - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -888,58 +908,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - _snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_apply_check_snapshot_changes( self, @@ -951,69 +934,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - _snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_current_timestamp(self) -> str: return "CURRENT_TIMESTAMP" @@ -1032,7 +961,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1056,7 +987,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1079,56 +1010,17 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = _historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, observed_at_column=observed_at_column, - valid_to_column=valid_to_column, valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1141,43 +1033,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = _historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_create_initial_historical_check_snapshot_destination( self, @@ -1193,7 +1057,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1218,7 +1084,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = _historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1226,48 +1092,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def create_table_as( self, @@ -2574,429 +2401,3 @@ def dimensions(node: SqlValue) -> tuple[int, ...] | None: def _quote_sql_string(value: str) -> str: return "'" + value.replace("'", "''") + "'" - - -def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, -) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - -def _snapshot_hard_delete_close_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, -) -> str: - missing_key_condition: str = _snapshot_key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - - -def _historical_check_snapshot_select_sql( - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, -) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - "), __versions AS (" - f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - -def _historical_timestamp_snapshot_select_sql( - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, -) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __versions AS (" - f"SELECT __changes.*, LEAD({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - -def _historical_timestamp_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, -) -> str: - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - if invalidate_hard_deletes: - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - latest_join_condition = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - -def _historical_timestamp_changes_select_sql( - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], -) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - -def _historical_timestamp_changes_new_records_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, -) -> str: - latest_join_condition: str = _snapshot_key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - -def _historical_check_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, -) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({delta_change_condition})" - ) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - -def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] -) -> str: - return " AND ".join(f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key) - - -def _historical_hard_deleted_at_sql( - *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str -) -> str: - present_condition: str = _snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - -def _historical_snapshot_combined_close_sql( - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, -) -> str: - close_candidate_condition: str = _snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - return ( - f"WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes WHERE __close_at IS NOT NULL" - ") " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"AND EXISTS (SELECT 1 FROM __close_candidates WHERE {close_candidate_condition})" - ) diff --git a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py index c87ab8f663..889dd0ff1c 100644 --- a/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/duckdb_backed_adapter.py @@ -19,7 +19,14 @@ _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -67,7 +74,6 @@ from sqlbuild.adapter.type_system.main.normalize_numeric_family import normalize_numeric_family from sqlbuild.adapter.type_system.main.types_equal import types_equal from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.diagnostics.main.log_sql import log_sql from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS @@ -401,20 +407,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -432,58 +435,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -501,7 +467,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=self._historical_timestamp_snapshot_select_sql( + sql=HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -527,7 +495,7 @@ def render_create_initial_historical_timestamp_changes_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=self._historical_timestamp_changes_select_sql( + sql=HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -550,7 +518,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -558,48 +526,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -612,43 +541,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -660,69 +561,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_check_snapshot_destination( self, @@ -740,7 +587,7 @@ def render_create_initial_historical_check_snapshot_destination( ) -> tuple[str, ...]: return self.render_create_table_as( destination=destination, - sql=self._historical_check_snapshot_select_sql( + sql=HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -765,7 +612,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -773,48 +620,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def recommended_max_sql_length(self) -> int | None: """DuckDB uses the framework default recommendation for lightweight unit-test SQL.""" @@ -2581,452 +2389,6 @@ def _render_seed_csv_options(self, csv_settings: SeedCsvSettings) -> str: def _duckdb_string_literal(self, value: str) -> str: return value.replace("'", "''") - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - "), __versions AS (" - f"SELECT __changes.*, LEAD({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __versions AS (" - f"SELECT __changes.*, LEAD({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column}" - f") AS __next_change_at, {hard_deleted_at_sql} AS __hard_deleted_at " - "FROM __changes" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deleted_at_sql: str = cls._historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - return ( - f"WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ") " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"AND EXISTS (SELECT 1 FROM __close_candidates WHERE {close_candidate_condition})" - ) - def _duckdb_string_literal(value: str) -> str: return "'" + value.replace("'", "''") + "'" diff --git a/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py new file mode 100644 index 0000000000..468015894f --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_check_snapshot_sql.py @@ -0,0 +1,179 @@ +"""Shared SQL for historical check snapshots.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import HistoricalSnapshotSql +from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( + HistoricalSnapshotStatementSql, +) +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql +from sqlbuild.adapter.contract.models import SnapshotSqlDialect + + +class HistoricalCheckSnapshotSql: + """Render SQL for historical check snapshots.""" + + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect + + def initial_select_sql( + self, + *, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> str: + """Render the initial-build query for a historical check snapshot.""" + + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + distinct_condition=self.dialect.distinct_condition, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=column, right=f"__prev_{column}") + for column in check_columns + ) + output_select_sql: str = ", ".join(output_columns) + ordered_sql: str = self._ordered_select_sql( + origin=origin, + unique_key=unique_key, + check_columns=check_columns, + observed_at_column=observed_at_column, + ) + return ( + f"WITH __ordered AS ({ordered_sql}), __changes AS (" + f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" + ") " + f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " + f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " + f"ORDER BY {observed_at_column}) AS {valid_to_column} " + "FROM __changes" + ) + + def apply_sql( + self, + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> tuple[str, ...]: + """Render the incremental statements for a historical check snapshot.""" + + statements: HistoricalSnapshotStatementSql = HistoricalSnapshotStatementSql( + dialect=self.dialect + ) + new_changes_sql: str + if invalidate_hard_deletes: + new_changes_sql = HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + distinct_condition=self.dialect.distinct_condition, + ).new_changes_ctes_sql(destination=destination) + else: + new_changes_sql = self._new_changes_ctes_sql( + statements=statements, + destination=destination, + origin=origin, + unique_key=unique_key, + check_columns=check_columns, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + ) + return statements.apply_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + change_time_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=invalidate_hard_deletes, + ) + + def _new_changes_ctes_sql( + self, + *, + statements: HistoricalSnapshotStatementSql, + destination: str, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + ) -> str: + delta_change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=column, right=f"__prev_{column}") + for column in check_columns + ) + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + latest_change_condition: str = " OR ".join( + self.dialect.distinct_condition( + left=f"__delta_changes.{column}", right=f"__latest.{column}" + ) + for column in check_columns + ) + latest_sql: str = statements.latest_version_cte_sql( + destination=destination, unique_key=unique_key, order_column=valid_from_column + ) + ordered_sql: str = self._ordered_select_sql( + origin=origin, + unique_key=unique_key, + check_columns=check_columns, + observed_at_column=observed_at_column, + ) + return ( + f"__ordered AS ({ordered_sql}), __delta_changes AS (" + "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " + f"OR ({delta_change_condition})" + f"), {latest_sql}, __new_changes AS (" + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{unique_key[0]} IS NULL OR (" + f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " + f"AND ({latest_change_condition})))" + ) + + @staticmethod + def _ordered_select_sql( + *, + origin: str, + unique_key: tuple[str, ...], + check_columns: tuple[str, ...], + observed_at_column: str, + ) -> str: + partition_sql: str = ", ".join(unique_key) + previous_columns_sql: str = "".join( + f", LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " + f"AS __prev_{column}" + for column in check_columns + ) + return ( + f"SELECT *, LAG({observed_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" + ) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py new file mode 100644 index 0000000000..a20f2933db --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_sql.py @@ -0,0 +1,300 @@ +"""Portable SQL for historical snapshots that invalidate hard deletes.""" + +from __future__ import annotations + +from collections.abc import Callable + +from sqlbuild.adapter.contract.constants import SNAPSHOT_VERSION_START_COLUMN + + +def render_is_distinct_from(*, left: str, right: str) -> str: + """Render a null-safe inequality with the SQL standard predicate.""" + + return f"{left} IS DISTINCT FROM {right}" + + +def historical_insert_validity_sql() -> str: + """Render ``valid_from, valid_to`` expressions selected from ``__new_changes``.""" + + version_end_sql: str = _version_end_sql(alias="__new_changes") + return f"__new_changes.{SNAPSHOT_VERSION_START_COLUMN}, {version_end_sql}" + + +def _version_end_sql(*, alias: str | None) -> str: + prefix: str = f"{alias}." if alias is not None else "" + next_start: str = f"{prefix}__next_version_start" + next_absence: str = f"{prefix}__next_absence_at" + return ( + f"CASE WHEN {next_start} IS NULL THEN {next_absence} " + f"WHEN {next_absence} IS NULL THEN {next_start} " + f"WHEN {next_absence} < {next_start} THEN {next_absence} " + f"ELSE {next_start} END" + ) + + +class HistoricalSnapshotSql: + """Render hard-delete historical snapshot SQL without QUALIFY or correlated subqueries.""" + + def __init__( + self, + *, + origin: str, + unique_key: tuple[str, ...], + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + check_columns: tuple[str, ...] = (), + updated_at_column: str | None = None, + distinct_condition: Callable[..., str] = render_is_distinct_from, + ) -> None: + self.origin: str = origin + self.unique_key: tuple[str, ...] = unique_key + self.observed_at_column: str = observed_at_column + self.valid_from_column: str = valid_from_column + self.valid_to_column: str = valid_to_column + self.check_columns: tuple[str, ...] = check_columns + self.updated_at_column: str | None = updated_at_column + self.distinct_condition: Callable[..., str] = distinct_condition + + def initial_select_sql(self, *, output_columns: tuple[str, ...]) -> str: + """Render the complete initial-build query.""" + + return ( + f"WITH {self.initial_ctes_sql()} " + f"SELECT {self.initial_select_list_sql(output_columns=output_columns)} " + "FROM __versions" + ) + + def initial_ctes_sql(self) -> str: + """Render the initial-build CTE list ending in ``__versions``.""" + + partition_sql: str = self._key_sql(alias=None) + return ( + f"{self._tracked_source_ctes_sql()}, __changes AS (" + f"SELECT *, {self._initial_version_start_sql()} AS {SNAPSHOT_VERSION_START_COLUMN} " + "FROM __tracked " + f"WHERE __is_change = 1 OR {self._initial_reappearance_sql()}" + "), __versions AS (" + f"SELECT *, LEAD({SNAPSHOT_VERSION_START_COLUMN}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {SNAPSHOT_VERSION_START_COLUMN}" + ") AS __next_version_start FROM __changes" + ")" + ) + + def initial_select_list_sql(self, *, output_columns: tuple[str, ...]) -> str: + """Render output columns and validity columns selected from ``__versions``.""" + + output_sql: str = ", ".join(output_columns) + return ( + f"{output_sql}, {SNAPSHOT_VERSION_START_COLUMN} AS {self.valid_from_column}, " + f"{_version_end_sql(alias=None)} AS {self.valid_to_column}" + ) + + def new_changes_ctes_sql(self, *, destination: str) -> str: + """Render the incremental CTE list defining ``__new_changes`` and ``__hard_deletes``.""" + + partition_sql: str = self._key_sql(alias=None) + first_key: str = self.unique_key[0] + latest_join_condition: str = self._key_condition(left="__tracked", right="__latest") + return ( + f"{self._tracked_source_ctes_sql()}, __latest AS (" + "SELECT * FROM (" + f"SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_sql} " + f"ORDER BY {self.valid_from_column} DESC) AS __latest_rn FROM {destination}" + ") AS __latest_ranked WHERE __latest_rn = 1" + "), __classified AS (" + "SELECT __tracked.*, " + f"CASE WHEN __latest.{first_key} IS NULL THEN 0 ELSE 1 END AS __has_latest, " + f"__latest.{self.valid_from_column} AS __latest_valid_from, " + f"{self._differs_from_latest_sql()} AS __differs_from_latest, " + f"{self._incremental_reappearance_sql()} AS __is_reappearance, " + f"{self._unchanged_from_previous_sql()} AS __is_unchanged " + f"FROM __tracked LEFT JOIN __latest ON {latest_join_condition}" + "), __changed_or_new AS (" + f"SELECT __classified.*, {self._incremental_version_start_sql()} " + f"AS {SNAPSHOT_VERSION_START_COLUMN} FROM __classified " + f"WHERE {self._incremental_start_condition_sql()}" + f"), {self._new_change_starts_sql()}, __new_changes AS (" + f"SELECT *, LEAD({SNAPSHOT_VERSION_START_COLUMN}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {SNAPSHOT_VERSION_START_COLUMN}" + ") AS __next_version_start FROM __new_change_starts" + f"), __hard_deletes AS ({self._open_version_hard_deletes_sql(destination=destination)})" + ) + + def _tracked_source_ctes_sql(self) -> str: + observed_at: str = self.observed_at_column + partition_sql: str = self._key_sql(alias=None) + window_sql: str = f"OVER (PARTITION BY {partition_sql} ORDER BY {observed_at})" + return ( + "__observed_group_sequence AS (" + "SELECT __observed_at, " + "LAG(__observed_at) OVER (ORDER BY __observed_at) AS __prev_group_observed_at, " + "LEAD(__observed_at) OVER (ORDER BY __observed_at) AS __next_group_observed_at " + f"FROM (SELECT DISTINCT {observed_at} AS __observed_at FROM {self.origin}) " + "AS __observed_group_values" + "), __ordered AS (" + "SELECT __source.*, __observed_group_sequence.__prev_group_observed_at, " + "__observed_group_sequence.__next_group_observed_at, " + f"LAG({observed_at}) {window_sql} AS __prev_observed_at, " + f"LEAD({observed_at}) {window_sql} AS __next_observed_at" + f"{self._previous_values_sql(window_sql=window_sql)} " + f"FROM {self.origin} AS __source LEFT JOIN __observed_group_sequence " + f"ON __observed_group_sequence.__observed_at = __source.{observed_at}" + "), __flagged AS (" + "SELECT *, " + "CASE WHEN __next_group_observed_at IS NOT NULL AND (__next_observed_at IS NULL " + "OR __next_observed_at <> __next_group_observed_at) " + "THEN __next_group_observed_at END AS __absent_after, " + f"CASE WHEN {self._change_condition_sql()} THEN 1 ELSE 0 END AS __is_change " + "FROM __ordered" + "), __tracked AS (" + "SELECT *, MIN(__absent_after) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at} " + "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING" + ") AS __next_absence_at FROM __flagged" + ")" + ) + + def _previous_values_sql(self, *, window_sql: str) -> str: + if self.updated_at_column is not None: + return f", LAG({self.updated_at_column}) {window_sql} AS __prev_updated_at" + return "".join( + f", LAG({column}) {window_sql} AS __prev_{column}" for column in self.check_columns + ) + + def _change_condition_sql(self) -> str: + if self.updated_at_column is not None: + return "__prev_updated_at IS NULL OR " + self.distinct_condition( + left=self.updated_at_column, right="__prev_updated_at" + ) + changed_sql: str = " OR ".join( + self.distinct_condition(left=column, right=f"__prev_{column}") + for column in self.check_columns + ) + return f"__prev_observed_at IS NULL OR ({changed_sql})" + + def _initial_reappearance_sql(self) -> str: + return "(__prev_observed_at IS NOT NULL AND __prev_observed_at <> __prev_group_observed_at)" + + def _initial_version_start_sql(self) -> str: + return self._version_start_sql( + unchanged_reappearance_sql=f"{self._initial_reappearance_sql()} AND __is_change = 0" + ) + + def _incremental_reappearance_sql(self) -> str: + return ( + "CASE WHEN __tracked.__prev_group_observed_at IS NOT NULL AND (" + f"(__tracked.__prev_observed_at IS NULL AND __latest.{self.unique_key[0]} IS NOT NULL) " + "OR __tracked.__prev_observed_at <> __tracked.__prev_group_observed_at" + ") THEN 1 ELSE 0 END" + ) + + def _incremental_version_start_sql(self) -> str: + return self._version_start_sql( + unchanged_reappearance_sql="__is_reappearance = 1 AND __is_unchanged = 1" + ) + + def _version_start_sql(self, *, unchanged_reappearance_sql: str) -> str: + if self.updated_at_column is None: + return self.observed_at_column + return ( + f"CASE WHEN {unchanged_reappearance_sql} THEN {self.observed_at_column} " + f"ELSE {self.updated_at_column} END" + ) + + def _unchanged_from_previous_sql(self) -> str: + if self.updated_at_column is None: + return "0" + changed_from_latest_sql: str = self.distinct_condition( + left=f"__tracked.{self.updated_at_column}", + right=f"__latest.{self.updated_at_column}", + ) + return ( + "CASE WHEN __tracked.__prev_observed_at IS NOT NULL AND __tracked.__is_change = 0 " + "THEN 1 " + "WHEN __tracked.__prev_observed_at IS NULL AND __latest." + f"{self.unique_key[0]} IS NOT NULL AND NOT ({changed_from_latest_sql}) THEN 1 " + "ELSE 0 END" + ) + + def _differs_from_latest_sql(self) -> str: + if not self.check_columns: + return "0" + differs_sql: str = " OR ".join( + self.distinct_condition(left=f"__tracked.{column}", right=f"__latest.{column}") + for column in self.check_columns + ) + return f"CASE WHEN {differs_sql} THEN 1 ELSE 0 END" + + def _incremental_start_condition_sql(self) -> str: + return ( + "(__has_latest = 0 AND (__is_change = 1 OR __is_reappearance = 1)) " + "OR (__has_latest = 1 AND __is_reappearance = 1 " + f"AND {self._incremental_version_start_sql()} > __latest_valid_from) " + "OR (__has_latest = 1 AND __is_reappearance = 0 AND __is_change = 1 " + f"AND {self._latest_change_condition_sql()})" + ) + + def _latest_change_condition_sql(self) -> str: + if self.updated_at_column is not None: + return f"{self.updated_at_column} > __latest_valid_from" + return ( + f"{self.observed_at_column} > __latest_valid_from " + "AND (__prev_observed_at >= __latest_valid_from OR __differs_from_latest = 1)" + ) + + def _new_change_starts_sql(self) -> str: + observed_at: str = self.observed_at_column + classified_key_sql: str = self._key_sql(alias="__classified") + return ( + "__first_new_starts AS (" + f"SELECT {self._key_sql(alias=None)}, MIN({SNAPSHOT_VERSION_START_COLUMN}) " + f"AS __first_version_start FROM __changed_or_new GROUP BY {self._key_sql(alias=None)}" + "), __reappearances AS (" + f"SELECT {classified_key_sql}, MIN(__classified.{observed_at}) AS __reappeared_at " + "FROM __classified " + f"JOIN __latest ON {self._key_condition(left='__classified', right='__latest')} " + "LEFT JOIN __observed_group_sequence AS __closing_group " + f"ON __closing_group.__observed_at = __latest.{self.valid_to_column} " + f"WHERE __latest.{self.valid_to_column} IS NOT NULL " + "AND __closing_group.__observed_at IS NULL " + f"AND __classified.{observed_at} > __latest.{self.valid_to_column} " + f"GROUP BY {classified_key_sql}" + "), __new_change_starts AS (" + "SELECT * FROM __changed_or_new " + "UNION ALL " + f"SELECT __classified.*, __classified.{observed_at} AS {SNAPSHOT_VERSION_START_COLUMN} " + "FROM __classified " + "JOIN __reappearances ON " + f"{self._key_condition(left='__classified', right='__reappearances')} " + f"AND __classified.{observed_at} = __reappearances.__reappeared_at " + "LEFT JOIN __first_new_starts ON " + f"{self._key_condition(left='__classified', right='__first_new_starts')} " + f"WHERE __first_new_starts.{self.unique_key[0]} IS NULL " + f"OR __classified.{observed_at} < __first_new_starts.__first_version_start" + ")" + ) + + def _open_version_hard_deletes_sql(self, *, destination: str) -> str: + target_key_sql: str = self._key_sql(alias="__target") + present_condition: str = self._key_condition(left="__present", right="__target") + observed_at: str = self.observed_at_column + return ( + f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " + f"FROM {destination} AS __target " + "JOIN __observed_group_sequence AS __observed_groups " + f"ON __observed_groups.__observed_at > __target.{observed_at} " + f"LEFT JOIN {self.origin} AS __present " + f"ON __present.{observed_at} = __observed_groups.__observed_at " + f"AND {present_condition} " + f"WHERE __target.{self.valid_to_column} IS NULL " + f"AND __present.{self.unique_key[0]} IS NULL " + f"GROUP BY {target_key_sql}" + ) + + def _key_sql(self, *, alias: str | None) -> str: + prefix: str = f"{alias}." if alias is not None else "" + return ", ".join(f"{prefix}{column}" for column in self.unique_key) + + def _key_condition(self, *, left: str, right: str) -> str: + return " AND ".join(f"{left}.{column} = {right}.{column}" for column in self.unique_key) diff --git a/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py new file mode 100644 index 0000000000..e0424356e1 --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_snapshot_statement_sql.py @@ -0,0 +1,249 @@ +"""Shared close and insert statements for historical snapshots.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + historical_insert_validity_sql, +) +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql +from sqlbuild.adapter.contract.constants import SNAPSHOT_VERSION_START_COLUMN +from sqlbuild.adapter.contract.models import SnapshotSqlDialect +from sqlbuild.adapter.contract.types import ( + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, + SnapshotLatestVersionStyle, +) + + +class HistoricalSnapshotStatementSql: + """Render latest-version, close and insert statements shared by historical snapshots.""" + + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect + + def latest_version_cte_sql( + self, *, destination: str, unique_key: tuple[str, ...], order_column: str + ) -> str: + """Render CTEs ending in ``__latest``, the newest stored version of each key.""" + + partition_sql: str = ", ".join(unique_key) + window_sql: str = ( + f"ROW_NUMBER() OVER (PARTITION BY {partition_sql} ORDER BY {order_column} DESC)" + ) + if self.dialect.latest_version == SnapshotLatestVersionStyle.QUALIFY: + return f"__latest AS (SELECT * FROM {destination} QUALIFY {window_sql} = 1)" + if self.dialect.latest_version == SnapshotLatestVersionStyle.DERIVED_TABLE: + return ( + f"__latest AS (SELECT * FROM (SELECT *, {window_sql} AS __rn FROM {destination}) " + "AS __q WHERE __rn = 1)" + ) + return ( + f"__latest_ordered AS (SELECT *, {window_sql} AS __rn FROM {destination}), " + "__latest AS (SELECT * FROM __latest_ordered WHERE __rn = 1)" + ) + + def apply_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + change_time_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> tuple[str, ...]: + """Render the close and insert statements that apply ``__new_changes``.""" + + close_sql: str + if invalidate_hard_deletes: + close_sql = self._combined_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + change_time_column=SNAPSHOT_VERSION_START_COLUMN, + ) + else: + close_sql = self._change_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + change_time_column=change_time_column, + ) + insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) + output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) + partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) + version_columns_sql: str = ( + f"__new_changes.{change_time_column}, LEAD(__new_changes.{change_time_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY __new_changes.{change_time_column})" + ) + if invalidate_hard_deletes: + version_columns_sql = historical_insert_validity_sql() + insert_sql: str = self._insert_sql( + destination=destination, + insert_column_sql=insert_column_sql, + new_changes_sql=new_changes_sql, + select_sql=f"SELECT {output_select_sql}, {version_columns_sql} FROM __new_changes", + ) + return (close_sql, insert_sql) + + def _change_close_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + change_time_column: str, + ) -> str: + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.UPDATE_FROM: + return self._grouped_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_candidates_sql=( + f"SELECT {', '.join(unique_key)}, {change_time_column} AS __close_at " + "FROM __new_changes" + ), + ) + key_condition: str = SnapshotSql.key_condition( + left_alias="__target", right_alias="__new_changes", unique_key=unique_key + ) + return self._correlated_close_sql( + ctes_sql=new_changes_sql, + destination=destination, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_at_sql=( + f"SELECT MIN(__new_changes.{change_time_column}) " + f"FROM __new_changes WHERE {key_condition}" + ), + exists_sql=f"SELECT 1 FROM __new_changes WHERE {key_condition}", + ) + + def _combined_close_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + change_time_column: str, + ) -> str: + candidate_key_sql: str = ", ".join(unique_key) + close_candidates_sql: str = ( + f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " + "UNION ALL " + f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " + "WHERE __close_at IS NOT NULL" + ) + if self.dialect.historical_close in { + HistoricalSnapshotCloseStyle.UPDATE_FROM, + HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES, + }: + return self._grouped_close_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_candidates_sql=close_candidates_sql, + ) + close_candidate_condition: str = SnapshotSql.key_condition( + left_alias="__close_candidates", right_alias="__target", unique_key=unique_key + ) + return self._correlated_close_sql( + ctes_sql=f"{new_changes_sql}, __close_candidates AS ({close_candidates_sql})", + destination=destination, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + close_at_sql=( + "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " + f"WHERE {close_candidate_condition}" + ), + exists_sql=f"SELECT 1 FROM __close_candidates WHERE {close_candidate_condition}", + ) + + def _correlated_close_sql( + self, + *, + ctes_sql: str, + destination: str, + valid_from_column: str, + valid_to_column: str, + close_at_sql: str, + exists_sql: str, + ) -> str: + update_sql: str = ( + f"WITH {ctes_sql} UPDATE {destination} AS __target " + f"SET {valid_to_column} = ({close_at_sql}) " + ) + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.TSQL: + update_sql = ( + f";WITH {ctes_sql} UPDATE __target SET {valid_to_column} = ({close_at_sql}) " + f"FROM {destination} AS __target " + ) + return ( + f"{update_sql}WHERE __target.{valid_to_column} IS NULL " + f"AND __target.{valid_from_column} < ({close_at_sql}) " + f"AND EXISTS ({exists_sql})" + ) + + def _grouped_close_sql( + self, + *, + destination: str, + new_changes_sql: str, + unique_key: tuple[str, ...], + valid_from_column: str, + valid_to_column: str, + close_candidates_sql: str, + ) -> str: + candidate_key_sql: str = ", ".join(unique_key) + close_candidates_query: str = ( + f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " + f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " + f"FROM __close_candidates GROUP BY {candidate_key_sql}" + ) + close_condition: str = ( + f"__target.{valid_to_column} IS NULL " + f"AND __target.{valid_from_column} < __close_candidates.__close_at" + ) + if self.dialect.historical_close == HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES: + merge_key_condition: str = SnapshotSql.key_condition( + left_alias="__target", right_alias="__close_candidates", unique_key=unique_key + ) + return ( + f"MERGE INTO {destination} AS __target " + f"USING ({close_candidates_query}) AS __close_candidates " + f"ON {merge_key_condition} AND {close_condition} " + f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __close_candidates.__close_at" + ) + update_key_condition: str = SnapshotSql.key_condition( + left_alias="__close_candidates", right_alias="__target", unique_key=unique_key + ) + return ( + f"UPDATE {destination} AS __target " + f"SET {valid_to_column} = __close_candidates.__close_at " + f"FROM ({close_candidates_query}) AS __close_candidates " + f"WHERE {close_condition} AND {update_key_condition}" + ) + + def _insert_sql( + self, *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str + ) -> str: + insert_sql: str = f"INSERT INTO {destination} ({insert_column_sql})" + if self.dialect.historical_insert == HistoricalSnapshotInsertStyle.INSERT_WITH: + return f"{insert_sql} WITH {new_changes_sql} {select_sql}" + if self.dialect.historical_insert == HistoricalSnapshotInsertStyle.TSQL: + return f";WITH {new_changes_sql} {insert_sql} {select_sql}" + return f"WITH {new_changes_sql} {insert_sql} {select_sql}" diff --git a/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py new file mode 100644 index 0000000000..4b9f0157bd --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/historical_timestamp_snapshot_sql.py @@ -0,0 +1,203 @@ +"""Shared SQL for historical timestamp snapshots.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import HistoricalSnapshotSql +from sqlbuild.adapter.contract.classes.historical_snapshot_statement_sql import ( + HistoricalSnapshotStatementSql, +) +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql +from sqlbuild.adapter.contract.models import SnapshotSqlDialect + + +class HistoricalTimestampSnapshotSql: + """Render SQL for historical timestamp and historical changes snapshots.""" + + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect + + def initial_select_sql( + self, + *, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> str: + """Render the initial-build query for a historical timestamp snapshot.""" + + if invalidate_hard_deletes: + return HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + distinct_condition=self.dialect.distinct_condition, + ).initial_select_sql(output_columns=output_columns) + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(output_columns) + ordered_ctes_sql: str = self._ordered_ctes_sql( + origin=origin, + unique_key=unique_key, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + changes_cte="__changes", + ) + return ( + f"WITH {ordered_ctes_sql} " + f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " + f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " + f"ORDER BY {updated_at_column}) AS {valid_to_column} " + "FROM __changes" + ) + + def apply_sql( + self, + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + observed_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + invalidate_hard_deletes: bool, + ) -> tuple[str, ...]: + """Render the incremental statements for a historical timestamp snapshot.""" + + statements: HistoricalSnapshotStatementSql = HistoricalSnapshotStatementSql( + dialect=self.dialect + ) + new_changes_sql: str + if invalidate_hard_deletes: + new_changes_sql = HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + distinct_condition=self.dialect.distinct_condition, + ).new_changes_ctes_sql(destination=destination) + else: + ordered_ctes_sql: str = self._ordered_ctes_sql( + origin=origin, + unique_key=unique_key, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + changes_cte="__delta_changes", + ) + latest_sql: str = statements.latest_version_cte_sql( + destination=destination, unique_key=unique_key, order_column=valid_from_column + ) + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key + ) + new_changes_sql = ( + f"{ordered_ctes_sql}, {latest_sql}, __new_changes AS (" + "SELECT __delta_changes.* FROM __delta_changes " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{unique_key[0]} IS NULL " + f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column})" + ) + return statements.apply_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + change_time_column=updated_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=invalidate_hard_deletes, + ) + + @staticmethod + def changes_initial_select_sql( + *, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + ) -> str: + """Render the initial-build query for a historical changes snapshot.""" + + partition_sql: str = ", ".join(unique_key) + output_select_sql: str = ", ".join(output_columns) + return ( + f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " + f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " + f"ORDER BY {updated_at_column}) AS {valid_to_column} " + f"FROM {origin}" + ) + + def changes_apply_sql( + self, + *, + destination: str, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + valid_from_column: str, + valid_to_column: str, + output_columns: tuple[str, ...], + ) -> tuple[str, ...]: + """Render the incremental statements for a historical changes snapshot.""" + + statements: HistoricalSnapshotStatementSql = HistoricalSnapshotStatementSql( + dialect=self.dialect + ) + latest_sql: str = statements.latest_version_cte_sql( + destination=destination, unique_key=unique_key, order_column=updated_at_column + ) + latest_join_condition: str = SnapshotSql.key_condition( + left_alias="__source", right_alias="__latest", unique_key=unique_key + ) + new_changes_sql: str = ( + f"{latest_sql}, __new_changes AS (" + f"SELECT __source.* FROM {origin} AS __source " + f"LEFT JOIN __latest ON {latest_join_condition} " + f"WHERE __latest.{unique_key[0]} IS NULL " + f"OR __source.{updated_at_column} > __latest.{updated_at_column})" + ) + return statements.apply_sql( + destination=destination, + new_changes_sql=new_changes_sql, + unique_key=unique_key, + change_time_column=updated_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + invalidate_hard_deletes=False, + ) + + def _ordered_ctes_sql( + self, + *, + origin: str, + unique_key: tuple[str, ...], + updated_at_column: str, + observed_at_column: str, + changes_cte: str, + ) -> str: + partition_sql: str = ", ".join(unique_key) + changed_sql: str = self.dialect.distinct_condition( + left=updated_at_column, right="__prev_updated_at" + ) + return ( + "__ordered AS (" + f"SELECT *, LAG({updated_at_column}) OVER (" + f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" + f") AS __prev_updated_at FROM {origin}" + f"), {changes_cte} AS (" + f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL OR {changed_sql})" + ) diff --git a/src/sqlbuild/adapter/contract/classes/snapshot_sql.py b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py new file mode 100644 index 0000000000..7187fe778e --- /dev/null +++ b/src/sqlbuild/adapter/contract/classes/snapshot_sql.py @@ -0,0 +1,287 @@ +"""Shared SQL for current-state snapshots and fragments common to every snapshot.""" + +from __future__ import annotations + +from sqlbuild.adapter.contract.models import SnapshotChangeTarget, SnapshotSqlDialect +from sqlbuild.adapter.contract.types import SnapshotUpdateStyle +from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy + + +class SnapshotSql: + """Render current-state snapshot statements and shared snapshot fragments.""" + + def __init__(self, *, dialect: SnapshotSqlDialect) -> None: + self.dialect: SnapshotSqlDialect = dialect + + @staticmethod + def key_condition(*, left_alias: str, right_alias: str, unique_key: tuple[str, ...]) -> str: + """Render the equality join between two aliases over the snapshot unique key.""" + + return " AND ".join( + f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key + ) + + @staticmethod + def initial_valid_from_expr( + *, + snapshot_strategy: str | None, + updated_at_column: str | None, + observed_at_column: str | None, + initial_valid_from: str | None, + source_alias: str | None, + current_timestamp: str, + ) -> str: + """Render the valid_from expression for a key's first snapshot version.""" + + prefix: str = f"{source_alias}." if source_alias is not None else "" + if initial_valid_from == InitialValidFrom.EXECUTION_TIME: + return current_timestamp + if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: + return f"{prefix}{observed_at_column}" + if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: + return f"{prefix}{updated_at_column}" + if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: + return f"{prefix}{updated_at_column}" + return current_timestamp + + def initial_select_sql( + self, + *, + origin: str, + snapshot_strategy: str | None, + updated_at_column: str | None, + observed_at_column: str | None, + valid_from_column: str, + valid_to_column: str, + initial_valid_from: str | None, + current_timestamp: str, + ) -> str: + """Render the initial-build query for a current-state snapshot.""" + + valid_from_expr: str = self.initial_valid_from_expr( + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + initial_valid_from=initial_valid_from, + source_alias=None, + current_timestamp=current_timestamp, + ) + return ( + f"SELECT *, {valid_from_expr} AS {valid_from_column}, " + f"CAST(NULL AS {self.dialect.timestamp_type}) AS {valid_to_column} FROM {origin}" + ) + + def timestamp_changes_sql( + self, + *, + target: SnapshotChangeTarget, + updated_at_column: str, + observed_at_column: str | None, + initial_valid_from: str | None, + invalidate_hard_deletes: bool, + current_timestamp: str, + ) -> tuple[str, ...]: + """Render changes, inserting before closing when hard deletes need active-row evidence.""" + + initial_valid_from_expr: str = self.initial_valid_from_expr( + snapshot_strategy="timestamp", + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + initial_valid_from=initial_valid_from, + source_alias="__source", + current_timestamp=current_timestamp, + ) + first_key: str = target.unique_key[0] + close_sql: str = self._keyed_close_sql( + target=target, + valid_to_sql=f"__source.{updated_at_column}", + changed_sql=f"__source.{updated_at_column} > __target.{updated_at_column}", + ) + version_valid_from_expr: str = ( + f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql: str = "" + insert_guard_sql: str | None = None + if invalidate_hard_deletes: + key_sql: str = ", ".join(target.unique_key) + history_condition: str = self.key_condition( + left_alias="__history", right_alias="__source", unique_key=target.unique_key + ) + version_valid_from_expr = ( + f"CASE WHEN __active.{first_key} IS NULL AND __history.__closed_at IS NOT NULL " + f"THEN {current_timestamp} " + f"WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " + f"ELSE __source.{updated_at_column} END" + ) + history_join_sql = ( + f"LEFT JOIN (SELECT {key_sql}, MAX({target.valid_to_column}) AS __closed_at " + f"FROM {target.destination} GROUP BY {key_sql}) AS __history " + f"ON {history_condition} " + ) + inserted_key_condition: str = self.key_condition( + left_alias="__inserted", right_alias="__source", unique_key=target.unique_key + ) + insert_guard_sql = ( + f"NOT EXISTS (SELECT 1 FROM {target.destination} AS __inserted " + f"WHERE {inserted_key_condition} " + f"AND __inserted.{target.valid_to_column} IS NULL " + f"AND __inserted.{updated_at_column} = __source.{updated_at_column})" + ) + insert_sql: str = self._insert_sql( + target=target, + valid_from_sql=version_valid_from_expr, + joins_sql=history_join_sql, + changed_sql=f"__source.{updated_at_column} > __active.{updated_at_column}", + insert_guard_sql=insert_guard_sql, + ) + return self._with_hard_delete_close( + statements=(insert_sql, close_sql) + if invalidate_hard_deletes + else (close_sql, insert_sql), + target=target, + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=current_timestamp, + ) + + def check_changes_sql( + self, + *, + target: SnapshotChangeTarget, + check_columns: tuple[str, ...], + updated_at_column: str | None, + observed_at_column: str | None, + initial_valid_from: str | None, + invalidate_hard_deletes: bool, + current_timestamp: str, + ) -> tuple[str, ...]: + """Render the incremental statements for a current-state check snapshot.""" + + initial_valid_from_expr: str = self.initial_valid_from_expr( + snapshot_strategy="check", + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + initial_valid_from=initial_valid_from, + source_alias="__source", + current_timestamp=current_timestamp, + ) + change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=f"__source.{column}", right=f"__target.{column}") + for column in check_columns + ) + close_sql: str = self._keyed_close_sql( + target=target, valid_to_sql=current_timestamp, changed_sql=f"({change_condition})" + ) + active_change_condition: str = " OR ".join( + self.dialect.distinct_condition(left=f"__source.{column}", right=f"__active.{column}") + for column in check_columns + ) + insert_sql: str = self._insert_sql( + target=target, + valid_from_sql=( + f"CASE WHEN __active.{target.unique_key[0]} IS NULL " + f"THEN {initial_valid_from_expr} ELSE {current_timestamp} END" + ), + joins_sql="", + changed_sql=f"({active_change_condition})", + ) + return self._with_hard_delete_close( + statements=(close_sql, insert_sql), + target=target, + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=current_timestamp, + ) + + def _keyed_close_sql( + self, *, target: SnapshotChangeTarget, valid_to_sql: str, changed_sql: str + ) -> str: + key_condition: str = self.key_condition( + left_alias="__target", right_alias="__source", unique_key=target.unique_key + ) + valid_to_column: str = target.valid_to_column + if self.dialect.update_style == SnapshotUpdateStyle.MERGE: + return ( + f"MERGE INTO {target.destination} AS __target " + f"USING {target.origin} AS __source " + f"ON {key_condition} " + f"AND __target.{valid_to_column} IS NULL " + f"AND {changed_sql} " + f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = {valid_to_sql}" + ) + if self.dialect.update_style == SnapshotUpdateStyle.TSQL: + return ( + f"UPDATE __target SET {valid_to_column} = {valid_to_sql} " + f"FROM {target.destination} AS __target " + f"JOIN {target.origin} AS __source ON {key_condition} " + f"WHERE __target.{valid_to_column} IS NULL " + f"AND {changed_sql}" + ) + return ( + f"UPDATE {target.destination} AS __target " + f"SET {valid_to_column} = {valid_to_sql} " + f"FROM {target.origin} AS __source " + f"WHERE {key_condition} " + f"AND __target.{valid_to_column} IS NULL " + f"AND {changed_sql}" + ) + + def _insert_sql( + self, + *, + target: SnapshotChangeTarget, + valid_from_sql: str, + joins_sql: str, + changed_sql: str, + insert_guard_sql: str | None = None, + ) -> str: + insert_column_sql: str = ", ".join( + (*target.output_columns, target.valid_from_column, target.valid_to_column) + ) + output_select_sql: str = ", ".join(f"__source.{column}" for column in target.output_columns) + active_join_condition: str = self.key_condition( + left_alias="__active", right_alias="__source", unique_key=target.unique_key + ) + insert_condition: str = f"__active.{target.unique_key[0]} IS NULL OR {changed_sql}" + if insert_guard_sql is not None: + insert_condition = f"({insert_condition}) AND {insert_guard_sql}" + return ( + f"INSERT INTO {target.destination} ({insert_column_sql}) " + f"SELECT {output_select_sql}, {valid_from_sql}, " + f"CAST(NULL AS {self.dialect.timestamp_type}) " + f"FROM {target.origin} AS __source " + f"LEFT JOIN {target.destination} AS __active " + f"ON {active_join_condition} AND __active.{target.valid_to_column} IS NULL " + f"{joins_sql}" + f"WHERE {insert_condition}" + ) + + def _with_hard_delete_close( + self, + *, + statements: tuple[str, ...], + target: SnapshotChangeTarget, + invalidate_hard_deletes: bool, + current_timestamp: str, + ) -> tuple[str, ...]: + if not invalidate_hard_deletes: + return statements + missing_key_condition: str = self.key_condition( + left_alias="__source", right_alias="__target", unique_key=target.unique_key + ) + update_sql: str = ( + f"UPDATE {target.destination} AS __target " + f"SET {target.valid_to_column} = {current_timestamp} " + ) + if self.dialect.update_style == SnapshotUpdateStyle.TSQL: + update_sql = ( + f"UPDATE __target SET {target.valid_to_column} = {current_timestamp} " + f"FROM {target.destination} AS __target " + ) + hard_delete_sql: str = ( + f"{update_sql}WHERE __target.{target.valid_to_column} IS NULL " + "AND NOT EXISTS (" + f"SELECT 1 FROM {target.origin} AS __source " + f"WHERE {missing_key_condition} AND __source.{target.unique_key[0]} IS NOT NULL" + ")" + ) + return (*statements, hard_delete_sql) diff --git a/src/sqlbuild/adapter/contract/constants.py b/src/sqlbuild/adapter/contract/constants.py index 1eb7331809..f4f7f1e53a 100644 --- a/src/sqlbuild/adapter/contract/constants.py +++ b/src/sqlbuild/adapter/contract/constants.py @@ -27,3 +27,5 @@ QUALIFIED_NAME_SEPARATOR: str = "." PYTHON_INIT_FILE_NAME: str = "__init__.py" PYTHON_IDENTIFIER_REPLACEMENT: str = "_" + +SNAPSHOT_VERSION_START_COLUMN: str = "__version_start" diff --git a/src/sqlbuild/adapter/contract/models.py b/src/sqlbuild/adapter/contract/models.py index cc9aead38d..d376cbebc0 100644 --- a/src/sqlbuild/adapter/contract/models.py +++ b/src/sqlbuild/adapter/contract/models.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from datetime import datetime from decimal import Decimal @@ -12,9 +12,13 @@ from sqlbuild.adapter.contract.types import ( CursorKind, FunctionNullabilityRule, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LifeCycleEventKind, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TypeFamily, ) from sqlbuild.compiler.compile.types import FunctionLanguage @@ -202,6 +206,18 @@ class SnapshotChangeTarget: output_columns: tuple[str, ...] +@dataclass(frozen=True) +class SnapshotSqlDialect: + """Dialect differences in adapter-rendered snapshot SQL.""" + + timestamp_type: str + distinct_condition: Callable[..., str] + update_style: SnapshotUpdateStyle + latest_version: SnapshotLatestVersionStyle + historical_close: HistoricalSnapshotCloseStyle + historical_insert: HistoricalSnapshotInsertStyle + + @dataclass(frozen=True) class FunctionDefinition: """Warehouse function definition inputs for create/render adapter methods.""" diff --git a/src/sqlbuild/adapter/contract/types.py b/src/sqlbuild/adapter/contract/types.py index 4fea9e8db9..56b0f4de32 100644 --- a/src/sqlbuild/adapter/contract/types.py +++ b/src/sqlbuild/adapter/contract/types.py @@ -100,3 +100,36 @@ class TypeDialect(StrEnum): DATABRICKS = "databricks" POSTGRES = "postgres" TSQL = "tsql" + + +class SnapshotLatestVersionStyle(StrEnum): + """How incremental snapshot SQL selects each key's latest stored version.""" + + QUALIFY = "qualify" + DERIVED_TABLE = "derived_table" + ORDERED_CTE = "ordered_cte" + + +class SnapshotUpdateStyle(StrEnum): + """How current-state snapshot SQL closes active versions.""" + + UPDATE_FROM = "update_from" + MERGE = "merge" + TSQL = "tsql" + + +class HistoricalSnapshotCloseStyle(StrEnum): + """How historical snapshot SQL closes versions superseded by new observations.""" + + CORRELATED = "correlated" + UPDATE_FROM = "update_from" + MERGE_HARD_DELETES = "merge_hard_deletes" + TSQL = "tsql" + + +class HistoricalSnapshotInsertStyle(StrEnum): + """Where historical snapshot inserts place their common table expressions.""" + + WITH_INSERT = "with_insert" + INSERT_WITH = "insert_with" + TSQL = "tsql" diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index d5081c4a1d..7715a4fbde 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -19,10 +19,23 @@ _typed_collection_items, _typed_scalar_payload, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + render_is_distinct_from, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, +) from sqlbuild.adapter.contract.exceptions import ( AdapterUserError, UnsupportedTypedSqlRenderingError, @@ -49,6 +62,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -56,10 +70,14 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -95,7 +113,6 @@ TIMESTAMP_TYPE_TOKEN, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.diagnostics.main.log_sql import log_sql from sqlbuild.runtime.observability.classes.statement_lifecycle import StatementLifecycle @@ -121,6 +138,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.BIGQUERY.value sql_analysis_dialect_name: ClassVar[str | None] = "bigquery" max_identifier_length: ClassVar[int] = 1024 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.UPDATE_FROM, + historical_insert=HistoricalSnapshotInsertStyle.INSERT_WITH, + ) def inspect_retention( self, *, connection: _BigQueryConnection, request: RetentionRequest @@ -1005,20 +1030,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -1036,58 +1058,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -1103,7 +1088,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1127,7 +1114,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1150,7 +1137,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1158,44 +1145,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = self._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1208,39 +1160,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, - valid_to_column=valid_to_column, - ) - close_sql: str = self._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), + output_columns=output_columns, ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -1252,69 +1180,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" - ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), + ) def render_create_initial_historical_check_snapshot_destination( self, @@ -1330,7 +1204,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1355,7 +1231,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1363,44 +1239,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = self._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {observed_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") FROM __new_changes" - ), - ) - return (close_sql, insert_sql) def __init__(self) -> None: self._location: str | None = None @@ -2966,581 +2807,3 @@ def _render_cursor_bound_string(value: str) -> str: except ValueError: return f"TIMESTAMP '{value}'" return value - - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - "SELECT __source.*, " - f"(SELECT MAX(__groups.__observed_at) FROM (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - ") AS __groups " - f"WHERE __groups.__observed_at < __source.{observed_at_column}" - ") AS __prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin} AS __source" - "), __changes AS (" - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{observed_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{observed_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{updated_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{updated_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_ordered_join_condition: str = cls._snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - reappearing_partition_sql: str = ", ".join( - f"__ordered.{column}" for column in unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - " UNION DISTINCT " - "SELECT __ordered.* FROM __ordered " - f"JOIN __latest ON {latest_ordered_join_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {reappearing_partition_sql} " - f"ORDER BY __ordered.{observed_at_column}" - ") = 1" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_hard_deletes_select_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, - ) -> str: - target_key_sql: str = ", ".join(f"__target.{column}" for column in unique_key) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " - f"FROM {destination} AS __target " - f"JOIN (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"ON __observed_groups.__observed_at > __target.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __present.{first_key} IS NULL " - f"GROUP BY {target_key_sql}" - ) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - candidate_key_sql: str = ", ".join(unique_key) - return cls._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at " - "FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ), - ) - - @classmethod - def _historical_snapshot_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - close_candidates_sql: str, - ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" - ) - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __close_candidates.__close_at " - f"FROM ({close_candidates_query}) AS __close_candidates " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"AND {close_candidate_condition}" - ) - - @staticmethod - def _historical_snapshot_insert_sql( - *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str - ) -> str: - return ( - f"INSERT INTO {destination} ({insert_column_sql}) WITH {new_changes_sql} {select_sql}" - ) diff --git a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py index e7fb0bc6ef..1b67dc4f65 100644 --- a/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py +++ b/src/sqlbuild/adapters/databricks/classes/databricks_adapter.py @@ -19,10 +19,23 @@ _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + render_is_distinct_from, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, +) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value from sqlbuild.adapter.contract.models import ( @@ -47,6 +60,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -54,11 +68,15 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, RelationType, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -83,7 +101,6 @@ TABLE_RELATION_METADATA_TYPES, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.diagnostics.main.log_sql import log_sql from sqlbuild.spec.contracts.constants import DEFAULT_SEED_CSV_SETTINGS @@ -107,6 +124,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.DATABRICKS.value sql_analysis_dialect_name: ClassVar[str | None] = "databricks" max_identifier_length: ClassVar[int] = 255 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.MERGE, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.MERGE_HARD_DELETES, + historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, + ) def inspect_retention(self, *, connection: Any, request: RetentionRequest) -> RetentionState: target: str = self._retention_relation_target(request=request) @@ -855,20 +880,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -886,58 +908,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"MERGE INTO {destination} AS __target " - f"USING {origin} AS __source " - f"ON {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column} " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __source.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -953,7 +938,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -977,7 +964,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1000,7 +987,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1008,48 +995,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1062,43 +1010,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -1110,69 +1030,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"MERGE INTO {destination} AS __target " - f"USING {origin} AS __source " - f"ON {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition}) " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = {current_timestamp}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_check_snapshot_destination( self, @@ -1188,7 +1054,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1213,7 +1081,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1221,48 +1089,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def supports_zero_copy_clone(self) -> bool: return True @@ -3095,553 +2924,3 @@ def _to_int(value: object) -> int: if isinstance(value, int): return value return int(str(value)) - - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - source_group_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __source_with_prev_group AS (" - "SELECT __source.*, MAX(__observed_groups.__observed_at) " - f"AS __prev_group_observed_at FROM {origin} AS __source " - "LEFT JOIN __observed_groups " - f"ON __observed_groups.__observed_at < __source.{observed_at_column} " - f"GROUP BY {source_group_sql}" - "), __ordered AS (" - "SELECT *, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM __source_with_prev_group" - "), __changes AS (" - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{observed_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{observed_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{updated_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{updated_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_ordered_join_condition: str = cls._snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - reappearing_partition_sql: str = ", ".join( - f"__ordered.{column}" for column in unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - " UNION DISTINCT " - "SELECT __ordered.* FROM __ordered " - f"JOIN __latest ON {latest_ordered_join_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {reappearing_partition_sql} " - f"ORDER BY __ordered.{observed_at_column}" - ") = 1" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_hard_deletes_select_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, - ) -> str: - target_key_sql: str = ", ".join(f"__target.{column}" for column in unique_key) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " - f"FROM {destination} AS __target " - f"JOIN (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"ON __observed_groups.__observed_at > __target.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __present.{first_key} IS NULL " - f"GROUP BY {target_key_sql}" - ) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( - left_alias="__target", right_alias="__close_candidates", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ") " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" - ) - return ( - f"MERGE INTO {destination} AS __target " - f"USING ({close_candidates_query}) AS __close_candidates " - f"ON {close_candidate_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"WHEN MATCHED THEN UPDATE SET {valid_to_column} = __close_candidates.__close_at" - ) diff --git a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py index 31a20b204c..0e2edbcdf2 100644 --- a/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py +++ b/src/sqlbuild/adapters/postgres/classes/postgres_adapter.py @@ -15,20 +15,22 @@ _build_names_filter, _build_schemas_filter, _encode_typed_json, - _historical_check_snapshot_select_sql, - _historical_hard_deleted_at_sql, - _historical_snapshot_combined_close_sql, - _historical_timestamp_changes_select_sql, - _historical_timestamp_snapshot_select_sql, _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, - _snapshot_hard_delete_close_sql, - _snapshot_initial_valid_from_expr, - _snapshot_key_condition, _validate_rectangular_typed_array, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + render_is_distinct_from, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin from sqlbuild.adapter.contract.constants import ( @@ -57,6 +59,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -64,8 +67,12 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -102,6 +109,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.POSTGRES.value sql_analysis_dialect_name: ClassVar[str | None] = "postgres" max_identifier_length: ClassVar[int] = 63 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.DERIVED_TABLE, + historical_close=HistoricalSnapshotCloseStyle.CORRELATED, + historical_insert=HistoricalSnapshotInsertStyle.WITH_INSERT, + ) def supports_zero_copy_clone(self) -> bool: return False @@ -1120,19 +1135,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=self.render_current_timestamp(), - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -1150,58 +1163,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - _snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_apply_check_snapshot_changes( self, @@ -1213,69 +1189,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - _snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -1291,7 +1213,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1315,7 +1239,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = _historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1339,7 +1263,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = _historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -2416,7 +2342,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._pg_historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2424,48 +2350,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -2478,43 +2365,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._pg_historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_check_snapshot_changes( self, @@ -2529,7 +2388,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._pg_historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -2537,225 +2396,6 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = _historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = ( - f"WITH {new_changes_sql} " - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f"WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) - - @staticmethod - def _pg_historical_timestamp_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - latest_sql: str = ( - "__latest AS (SELECT * FROM (" - f"SELECT *, ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" - ) - if invalidate_hard_deletes: - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{col}' for col in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - latest_join_condition = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @staticmethod - def _pg_historical_timestamp_changes_new_records_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - del valid_to_column - latest_join_condition: str = _snapshot_key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (SELECT * FROM (" - f"SELECT *, ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - f") AS __rn FROM {destination}) AS __q WHERE __rn = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - @staticmethod - def _pg_historical_check_new_changes_cte_sql( - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL" - f" OR ({delta_change_condition})" - ) - latest_sql: str = ( - "__latest AS (SELECT * FROM (" - f"SELECT *, ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - f") AS __rn FROM {destination}) AS __q WHERE __rn = 1)" - ) - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - f"), {latest_sql}, __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 58fc8b415e..21cce1ee90 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -17,10 +17,23 @@ _render_ansi_typed_scalar, _render_typed_value_list, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + render_is_distinct_from, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, +) from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.contract.main.normalize_seed_csv_value import normalize_seed_csv_value from sqlbuild.adapter.contract.models import ( @@ -44,6 +57,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -51,10 +65,14 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, RetentionChangePhase, RetentionScope, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.relation_age_timestamp import relation_age_timestamp_utc @@ -84,7 +102,6 @@ VIEW_RELATION_TYPE_TOKEN, ) from sqlbuild.compiler.compile.types import FunctionLanguage -from sqlbuild.compiler.planner.types import InitialValidFrom, SnapshotStrategy from sqlbuild.compiler.source_freshness.models import SourceFreshnessRecord from sqlbuild.cost.main.collection import collect_snowflake_cost from sqlbuild.cost.models import RunCostSummary @@ -111,6 +128,14 @@ class SnowflakeAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter): connection_routing_keys: ClassVar[frozenset[str]] = frozenset( {"source", "profile", "target", "project_dir", "profiles_dir"} ) + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="TIMESTAMP", + distinct_condition=render_is_distinct_from, + update_style=SnapshotUpdateStyle.UPDATE_FROM, + latest_version=SnapshotLatestVersionStyle.QUALIFY, + historical_close=HistoricalSnapshotCloseStyle.UPDATE_FROM, + historical_insert=HistoricalSnapshotInsertStyle.INSERT_WITH, + ) def inspect_retention(self, *, connection: Any, request: RetentionRequest) -> RetentionState: self._validate_relation_retention_request(request=request) @@ -1257,21 +1282,18 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=current_timestamp, - ) return self._render_create_table_as_type( - destination=destination, table_type=table_type, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS TIMESTAMP) AS {valid_to_column} FROM {origin}" + destination=destination, + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) @@ -1289,58 +1311,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_timestamp_snapshot_destination( self, @@ -1356,7 +1341,9 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_snapshot_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1382,7 +1369,7 @@ def render_create_initial_historical_timestamp_changes_destination( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - historical_sql: str = self._historical_timestamp_changes_select_sql( + historical_sql: str = HistoricalTimestampSnapshotSql.changes_initial_select_sql( origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, @@ -1407,7 +1394,7 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1415,44 +1402,9 @@ def render_apply_historical_timestamp_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = self._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1465,39 +1417,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, - valid_to_column=valid_to_column, - ) - close_sql: str = self._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, valid_from_column=valid_from_column, valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {updated_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") FROM __new_changes" - ), + output_columns=output_columns, ) - return (close_sql, insert_sql) def render_apply_check_snapshot_changes( self, @@ -1509,69 +1437,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = self._snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = self._snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __target.{column}" for column in check_columns - ) - close_sql: str = ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"FROM {origin} AS __source " - f"WHERE {key_condition} " - f"AND __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = self._snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - f"__source.{column} IS DISTINCT FROM __active.{column}" for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS TIMESTAMP) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_create_initial_historical_check_snapshot_destination( self, @@ -1587,7 +1461,9 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - historical_sql: str = self._historical_check_snapshot_select_sql( + historical_sql: str = HistoricalCheckSnapshotSql( + dialect=self._snapshot_sql_dialect + ).initial_select_sql( origin=origin, unique_key=unique_key, check_columns=check_columns, @@ -1614,7 +1490,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1622,44 +1498,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = self._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {', '.join(unique_key)}, {observed_at_column} AS __close_at " - "FROM __new_changes" - ), - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = self._historical_snapshot_insert_sql( - destination=destination, - insert_column_sql=insert_column_sql, - new_changes_sql=new_changes_sql, - select_sql=( - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") FROM __new_changes" - ), - ) - return (close_sql, insert_sql) def supports_python_functions(self) -> bool: return True @@ -3269,581 +3110,3 @@ def schema_exists( return cursor.fetchone() is not None finally: cursor.close() - - @staticmethod - def _snapshot_initial_valid_from_expr( - *, - snapshot_strategy: str | None, - updated_at_column: str | None, - observed_at_column: str | None, - initial_valid_from: str | None, - source_alias: str | None, - current_timestamp: str, - ) -> str: - prefix: str = f"{source_alias}." if source_alias is not None else "" - if initial_valid_from == InitialValidFrom.EXECUTION_TIME: - return current_timestamp - if initial_valid_from == InitialValidFrom.OBSERVED_AT and observed_at_column is not None: - return f"{prefix}{observed_at_column}" - if initial_valid_from == InitialValidFrom.UPDATED_AT and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - if snapshot_strategy == SnapshotStrategy.TIMESTAMP and updated_at_column is not None: - return f"{prefix}{updated_at_column}" - return current_timestamp - - @classmethod - def _snapshot_hard_delete_close_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = {current_timestamp} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - f")" - ) - - @classmethod - def _historical_check_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - "SELECT __source.*, " - f"(SELECT MAX(__groups.__observed_at) FROM (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - ") AS __groups " - f"WHERE __groups.__observed_at < __source.{observed_at_column}" - ") AS __prev_group_observed_at, " - f"LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin} AS __source" - "), __changes AS (" - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({change_condition}) " - "OR __prev_observed_at IS DISTINCT FROM __prev_group_observed_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{observed_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{observed_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_observed_at IS NULL OR ({change_condition})" - ") " - f"SELECT {output_select_sql}, {observed_at_column} AS {valid_from_column}, " - f"LEAD({observed_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {observed_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_snapshot_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - if invalidate_hard_deletes: - hard_delete_join_condition: str = cls._snapshot_key_condition( - left_alias="__hard_delete_candidates", - right_alias="__changes", - unique_key=unique_key, - ) - hard_delete_key_sql: str = ", ".join(f"__changes.{column}" for column in unique_key) - hard_delete_group_sql: str = ", ".join( - [ - *(f"__changes.{column}" for column in unique_key), - f"__changes.{observed_at_column}", - ] - ) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__changes", unique_key=unique_key - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __observed_groups AS (" - f"SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}" - "), __hard_delete_candidates AS (" - f"SELECT {hard_delete_key_sql}, __changes.{observed_at_column}, " - "MIN(__observed_groups.__observed_at) AS __hard_deleted_at " - "FROM __changes " - "JOIN __observed_groups " - f"ON __observed_groups.__observed_at > __changes.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __present.{unique_key[0]} IS NULL " - f"GROUP BY {hard_delete_group_sql}" - "), __versions AS (" - f"SELECT __changes.*, LEAD(__changes.{updated_at_column}) OVER (" - f"PARTITION BY {', '.join(f'__changes.{column}' for column in unique_key)} " - f"ORDER BY __changes.{updated_at_column}" - ") AS __next_change_at, __hard_delete_candidates.__hard_deleted_at " - "FROM __changes LEFT JOIN __hard_delete_candidates " - f"ON {hard_delete_join_condition} " - f"AND __hard_delete_candidates.{observed_at_column} = " - f"__changes.{observed_at_column}" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - "CASE " - "WHEN __next_change_at IS NULL THEN __hard_deleted_at " - "WHEN __hard_deleted_at IS NULL THEN __next_change_at " - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at " - f"ELSE __next_change_at END AS {valid_to_column} " - "FROM __versions" - ) - return ( - "WITH __ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - ") " - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - "FROM __changes" - ) - - @classmethod - def _historical_timestamp_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - first_key: str = unique_key[0] - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL " - f"OR {updated_at_column} IS DISTINCT FROM __prev_updated_at" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - ")" - ) - - @classmethod - def _historical_timestamp_changes_select_sql( - cls, - *, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_from_column: str, - valid_to_column: str, - output_columns: tuple[str, ...], - ) -> str: - partition_sql: str = ", ".join(unique_key) - output_select_sql: str = ", ".join(column for column in output_columns) - return ( - f"SELECT {output_select_sql}, {updated_at_column} AS {valid_from_column}, " - f"LEAD({updated_at_column}) OVER (PARTITION BY {partition_sql} " - f"ORDER BY {updated_at_column}) AS {valid_to_column} " - f"FROM {origin}" - ) - - @classmethod - def _historical_timestamp_changes_new_records_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - valid_to_column: str, - ) -> str: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__source", right_alias="__latest", unique_key=unique_key - ) - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - return ( - "__latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {updated_at_column} DESC" - ") = 1" - "), __new_changes AS (" - f"SELECT __source.* FROM {origin} AS __source " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __latest.{updated_at_column}" - ")" - ) - - @classmethod - def _historical_check_new_changes_cte_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - f"LAG({column}) OVER (PARTITION BY {partition_sql} ORDER BY {observed_at_column}) " - f"AS __prev_{column}" - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - f"{column} IS DISTINCT FROM __prev_{column}" for column in check_columns - ) - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - first_key: str = unique_key[0] - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - if invalidate_hard_deletes: - latest_join_condition: str = cls._snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_ordered_join_condition: str = cls._snapshot_key_condition( - left_alias="__ordered", right_alias="__latest", unique_key=unique_key - ) - reappearing_partition_sql: str = ", ".join( - f"__ordered.{column}" for column in unique_key - ) - latest_change_condition: str = " OR ".join( - f"__delta_changes.{column} IS DISTINCT FROM __latest.{column}" - for column in check_columns - ) - hard_deletes_sql: str = cls._historical_hard_deletes_select_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - valid_to_column=valid_to_column, - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR (__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - " UNION DISTINCT " - "SELECT __ordered.* FROM __ordered " - f"JOIN __latest ON {latest_ordered_join_condition} " - f"WHERE __latest.{valid_to_column} IS NOT NULL " - f"AND __ordered.{observed_at_column} > __latest.{valid_to_column} " - "QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {reappearing_partition_sql} " - f"ORDER BY __ordered.{observed_at_column}" - ") = 1" - "), __hard_deletes AS (" - f"{hard_deletes_sql}" - ")" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest AS (" - f"SELECT * FROM {destination} QUALIFY ROW_NUMBER() OVER (" - f"PARTITION BY {partition_sql} ORDER BY {valid_from_column} DESC" - ") = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition})" - ")" - ")" - ) - - @staticmethod - def _snapshot_key_condition( - *, left_alias: str, right_alias: str, unique_key: tuple[str, ...] - ) -> str: - return " AND ".join( - f"{left_alias}.{column} = {right_alias}.{column}" for column in unique_key - ) - - @classmethod - def _historical_hard_deleted_at_sql( - cls, *, origin: str, unique_key: tuple[str, ...], observed_at_column: str, row_alias: str - ) -> str: - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias=row_alias, unique_key=unique_key - ) - return ( - "(SELECT MIN(__observed_groups.__observed_at) " - f"FROM (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"WHERE __observed_groups.__observed_at > {row_alias}.{observed_at_column} " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __present " - f"WHERE __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition}" - "))" - ) - - @classmethod - def _historical_hard_deletes_select_sql( - cls, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - observed_at_column: str, - valid_to_column: str, - ) -> str: - target_key_sql: str = ", ".join(f"__target.{column}" for column in unique_key) - present_condition: str = cls._snapshot_key_condition( - left_alias="__present", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"SELECT {target_key_sql}, MIN(__observed_groups.__observed_at) AS __close_at " - f"FROM {destination} AS __target " - f"JOIN (SELECT DISTINCT {observed_at_column} AS __observed_at FROM {origin}) " - "AS __observed_groups " - f"ON __observed_groups.__observed_at > __target.{observed_at_column} " - f"LEFT JOIN {origin} AS __present " - f"ON __present.{observed_at_column} = __observed_groups.__observed_at " - f"AND {present_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __present.{first_key} IS NULL " - f"GROUP BY {target_key_sql}" - ) - - @classmethod - def _historical_snapshot_combined_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - candidate_key_sql: str = ", ".join(unique_key) - return cls._historical_snapshot_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - close_candidates_sql=( - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at " - "FROM __new_changes " - "UNION ALL " - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ), - ) - - @classmethod - def _historical_snapshot_close_sql( - cls, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - close_candidates_sql: str, - ) -> str: - close_candidate_condition: str = cls._snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - close_candidates_query: str = ( - f"WITH {new_changes_sql}, __close_candidates AS ({close_candidates_sql}) " - f"SELECT {candidate_key_sql}, MIN(__close_at) AS __close_at " - "FROM __close_candidates GROUP BY " - f"{candidate_key_sql}" - ) - return ( - f"UPDATE {destination} AS __target " - f"SET {valid_to_column} = __close_candidates.__close_at " - f"FROM ({close_candidates_query}) AS __close_candidates " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < __close_candidates.__close_at " - f"AND {close_candidate_condition}" - ) - - @staticmethod - def _historical_snapshot_insert_sql( - *, destination: str, insert_column_sql: str, new_changes_sql: str, select_sql: str - ) -> str: - return ( - f"INSERT INTO {destination} ({insert_column_sql}) WITH {new_changes_sql} {select_sql}" - ) diff --git a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py index cb9f464a30..62e09380e0 100644 --- a/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py +++ b/src/sqlbuild/adapters/sqlserver/classes/sqlserver_adapter.py @@ -15,19 +15,28 @@ _build_names_filter, _build_schemas_filter, _encode_typed_json, - _historical_hard_deleted_at_sql, - _historical_timestamp_changes_new_records_cte_sql, _quote_sql_string, _render_ansi_typed_scalar, _render_typed_value_list, - _snapshot_initial_valid_from_expr, - _snapshot_key_condition, _typed_scalar_payload, ) +from sqlbuild.adapter.contract.classes.historical_check_snapshot_sql import ( + HistoricalCheckSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_snapshot_sql import ( + HistoricalSnapshotSql, +) +from sqlbuild.adapter.contract.classes.historical_timestamp_snapshot_sql import ( + HistoricalTimestampSnapshotSql, +) from sqlbuild.adapter.contract.classes.microbatch import MicrobatchMixin +from sqlbuild.adapter.contract.classes.snapshot_sql import SnapshotSql from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.classes.unkeyed_diff import UnkeyedDiffMixin -from sqlbuild.adapter.contract.constants import DIFF_LEFT_SIDE, DIFF_RIGHT_SIDE +from sqlbuild.adapter.contract.constants import ( + DIFF_LEFT_SIDE, + DIFF_RIGHT_SIDE, +) from sqlbuild.adapter.contract.exceptions import ( AdapterUserError, UnsupportedTypedSqlRenderingError, @@ -52,6 +61,7 @@ RowDiffTolerances, SchemaDiffResult, SnapshotChangeTarget, + SnapshotSqlDialect, TableFreshnessMetadata, TableFreshnessRequest, ) @@ -59,8 +69,12 @@ BuiltinAdapter, CursorKind, FrameworkType, + HistoricalSnapshotCloseStyle, + HistoricalSnapshotInsertStyle, LoaderLogicalType, PromotionStrategy, + SnapshotLatestVersionStyle, + SnapshotUpdateStyle, TablePromotionMode, ) from sqlbuild.adapter.relations.main.get_columns_for_relations import ( @@ -88,6 +102,14 @@ from sqlbuild.sql_values.types import SqlValueKind +def _render_tsql_distinct_condition(*, left: str, right: str) -> str: + return ( + f"({left} <> {right} " + f"OR ({left} IS NULL AND {right} IS NOT NULL) " + f"OR ({left} IS NOT NULL AND {right} IS NULL))" + ) + + class SqlServerAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter): """Microsoft SQL Server adapter backed by pymssql.""" @@ -104,6 +126,14 @@ def get_columns_for_relations( adapter_name: ClassVar[str] = BuiltinAdapter.SQLSERVER.value sql_analysis_dialect_name: ClassVar[str | None] = "tsql" max_identifier_length: ClassVar[int] = 128 + _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( + timestamp_type="DATETIME2", + distinct_condition=_render_tsql_distinct_condition, + update_style=SnapshotUpdateStyle.TSQL, + latest_version=SnapshotLatestVersionStyle.ORDERED_CTE, + historical_close=HistoricalSnapshotCloseStyle.TSQL, + historical_insert=HistoricalSnapshotInsertStyle.TSQL, + ) def supports_table_freshness_metadata(self) -> bool: return False @@ -131,11 +161,7 @@ def get_tables_freshness_metadata( ) def _distinct_condition(self, *, left: str, right: str) -> str: - return ( - f"({left} <> {right} " - f"OR ({left} IS NULL AND {right} IS NOT NULL) " - f"OR ({left} IS NOT NULL AND {right} IS NULL))" - ) + return _render_tsql_distinct_condition(left=left, right=right) def _lag_expr(self, *, column: str, partition_sql: str, observed_at_column: str) -> str: return ( @@ -143,196 +169,6 @@ def _lag_expr(self, *, column: str, partition_sql: str, observed_at_column: str) f"ORDER BY {observed_at_column}) AS __prev_{column}" ) - def _historical_snapshot_combined_close_sql( - self, - *, - destination: str, - new_changes_sql: str, - unique_key: tuple[str, ...], - valid_from_column: str, - valid_to_column: str, - change_time_column: str, - ) -> str: - close_candidate_condition: str = _snapshot_key_condition( - left_alias="__close_candidates", right_alias="__target", unique_key=unique_key - ) - candidate_key_sql: str = ", ".join(unique_key) - hard_delete_select_sql: str = ( - f"SELECT {candidate_key_sql}, __close_at FROM __hard_deletes " - "WHERE __close_at IS NOT NULL" - ) - return ( - f";WITH {new_changes_sql}, __close_candidates AS (" - f"SELECT {candidate_key_sql}, {change_time_column} AS __close_at FROM __new_changes " - "UNION ALL " - f"{hard_delete_select_sql}" - ") " - f"UPDATE __target SET {valid_to_column} = (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - "SELECT MIN(__close_candidates.__close_at) FROM __close_candidates " - f"WHERE {close_candidate_condition}" - ") " - f"AND EXISTS (SELECT 1 FROM __close_candidates WHERE {close_candidate_condition})" - ) - - def _snapshot_hard_delete_close_sql( - self, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - valid_to_column: str, - current_timestamp: str, - ) -> str: - missing_key_condition: str = _snapshot_key_condition( - left_alias="__source", right_alias="__target", unique_key=unique_key - ) - first_key: str = unique_key[0] - return ( - f"UPDATE __target SET {valid_to_column} = {current_timestamp} " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - "AND NOT EXISTS (" - f"SELECT 1 FROM {origin} AS __source " - f"WHERE {missing_key_condition} AND __source.{first_key} IS NOT NULL" - ")" - ) - - def _historical_timestamp_new_changes_cte_sql( - self, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - updated_at_column: str, - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - first_key: str = unique_key[0] - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - updated_changed: str = self._distinct_condition( - left=f"{updated_at_column}", right="__prev_updated_at" - ) - hard_deletes_sql: str = "" - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - hard_deletes_sql = ( - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_updated_at FROM {origin}" - "), __delta_changes AS (" - f"SELECT * FROM __ordered WHERE __prev_updated_at IS NULL OR {updated_changed}" - "), __latest_ordered AS (" - f"SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_sql} " - f"ORDER BY {valid_from_column} DESC) AS __rn FROM {destination}" - "), __latest AS (" - "SELECT * FROM __latest_ordered WHERE __rn = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL " - f"OR __delta_changes.{updated_at_column} > __latest.{valid_from_column}" - f"{hard_deletes_sql}" - ")" - ) - - def _historical_check_new_changes_cte_sql( - self, - *, - destination: str, - origin: str, - unique_key: tuple[str, ...], - check_columns: tuple[str, ...], - observed_at_column: str, - valid_from_column: str, - valid_to_column: str, - invalidate_hard_deletes: bool, - ) -> str: - partition_sql: str = ", ".join(unique_key) - previous_columns_sql: str = ", ".join( - self._lag_expr( - column=column, - partition_sql=partition_sql, - observed_at_column=observed_at_column, - ) - for column in check_columns - ) - if previous_columns_sql: - previous_columns_sql = f", {previous_columns_sql}" - delta_change_condition: str = " OR ".join( - self._distinct_condition(left=column, right=f"__prev_{column}") - for column in check_columns - ) - latest_join_condition: str = _snapshot_key_condition( - left_alias="__delta_changes", right_alias="__latest", unique_key=unique_key - ) - latest_change_condition: str = " OR ".join( - self._distinct_condition(left=f"__delta_changes.{column}", right=f"__latest.{column}") - for column in check_columns - ) - first_key: str = unique_key[0] - hard_deletes_sql: str = "" - if invalidate_hard_deletes: - hard_deleted_at_sql: str = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__target", - ) - hard_deletes_sql = ( - "), __hard_deletes AS (" - f"SELECT {', '.join(f'__target.{column}' for column in unique_key)}, " - f"{hard_deleted_at_sql} AS __close_at FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL" - ) - changed_or_first_sql: str = ( - "SELECT * FROM __ordered WHERE __prev_observed_at IS NULL " - f"OR ({delta_change_condition})" - ) - return ( - "__ordered AS (" - f"SELECT *, LAG({observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY {observed_at_column}" - f") AS __prev_observed_at{previous_columns_sql} FROM {origin}" - "), __delta_changes AS (" - f"{changed_or_first_sql}" - "), __latest_ordered AS (" - f"SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_sql} " - f"ORDER BY {valid_from_column} DESC) AS __rn FROM {destination}" - "), __latest AS (" - "SELECT * FROM __latest_ordered WHERE __rn = 1" - "), __new_changes AS (" - "SELECT __delta_changes.* FROM __delta_changes " - f"LEFT JOIN __latest ON {latest_join_condition} " - f"WHERE __latest.{first_key} IS NULL OR (" - f"__delta_changes.{observed_at_column} > __latest.{valid_from_column} " - f"AND ({latest_change_condition}))" - f"{hard_deletes_sql}" - ")" - ) - def connect(self, config: dict[str, Any]) -> _SqlServerConnection: try: import pymssql @@ -1808,70 +1644,15 @@ def render_apply_check_snapshot_changes( initial_valid_from: str | None, invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - destination: str = target.destination - origin: str = target.origin - unique_key: tuple[str, ...] = target.unique_key - valid_from_column: str = target.valid_from_column - valid_to_column: str = target.valid_to_column - output_columns: tuple[str, ...] = target.output_columns - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy="check", + return SnapshotSql(dialect=self._snapshot_sql_dialect).check_changes_sql( + target=target, + check_columns=check_columns, updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - change_condition: str = " OR ".join( - self._distinct_condition(left=f"__source.{column}", right=f"__target.{column}") - for column in check_columns - ) - close_sql: str = ( - f"UPDATE __target SET {valid_to_column} = {current_timestamp} " - f"FROM {destination} AS __target " - f"JOIN {origin} AS __source ON {key_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND ({change_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - active_change_condition: str = " OR ".join( - self._distinct_condition(left=f"__source.{column}", right=f"__active.{column}") - for column in check_columns - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE {current_timestamp} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS DATETIME2) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL OR ({active_change_condition})" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_apply_historical_check_snapshot_changes( self, @@ -1886,7 +1667,7 @@ def render_apply_historical_check_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_check_new_changes_cte_sql( + return HistoricalCheckSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, @@ -1894,49 +1675,9 @@ def render_apply_historical_check_snapshot_changes( observed_at_column=observed_at_column, valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=observed_at_column, - ) - else: - close_sql = ( - f";WITH {new_changes_sql} " - f"UPDATE __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{observed_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f";WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{observed_at_column}, " - f"LEAD(__new_changes.{observed_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{observed_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_changes( self, @@ -1949,44 +1690,15 @@ def render_apply_historical_timestamp_changes( valid_to_column: str, output_columns: tuple[str, ...], ) -> tuple[str, ...]: - new_changes_sql: str = _historical_timestamp_changes_new_records_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).changes_apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, + valid_from_column=valid_from_column, valid_to_column=valid_to_column, + output_columns=output_columns, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - close_sql: str = ( - f";WITH {new_changes_sql} " - f"UPDATE __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f";WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_historical_timestamp_snapshot_changes( self, @@ -2001,57 +1713,17 @@ def render_apply_historical_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - new_changes_sql: str = self._historical_timestamp_new_changes_cte_sql( + return HistoricalTimestampSnapshotSql(dialect=self._snapshot_sql_dialect).apply_sql( destination=destination, origin=origin, unique_key=unique_key, updated_at_column=updated_at_column, observed_at_column=observed_at_column, - valid_to_column=valid_to_column, valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, invalidate_hard_deletes=invalidate_hard_deletes, ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__new_changes", unique_key=unique_key - ) - if invalidate_hard_deletes: - close_sql: str = self._historical_snapshot_combined_close_sql( - destination=destination, - new_changes_sql=new_changes_sql, - unique_key=unique_key, - valid_from_column=valid_from_column, - valid_to_column=valid_to_column, - change_time_column=updated_at_column, - ) - else: - close_sql = ( - f";WITH {new_changes_sql} " - f"UPDATE __target " - f"SET {valid_to_column} = (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"FROM {destination} AS __target " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __target.{valid_from_column} < (" - f"SELECT MIN(__new_changes.{updated_at_column}) " - f"FROM __new_changes WHERE {key_condition}" - f") " - f"AND EXISTS (SELECT 1 FROM __new_changes WHERE {key_condition})" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__new_changes.{column}" for column in output_columns) - partition_sql: str = ", ".join(f"__new_changes.{column}" for column in unique_key) - insert_sql: str = ( - f";WITH {new_changes_sql} " - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, __new_changes.{updated_at_column}, " - f"LEAD(__new_changes.{updated_at_column}) OVER (" - f"PARTITION BY {partition_sql} ORDER BY __new_changes.{updated_at_column}" - f") " - f"FROM __new_changes" - ) - return (close_sql, insert_sql) def render_apply_timestamp_snapshot_changes( self, @@ -2067,57 +1739,21 @@ def render_apply_timestamp_snapshot_changes( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - current_timestamp: str = self.render_current_timestamp() - initial_valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy="timestamp", + return SnapshotSql(dialect=self._snapshot_sql_dialect).timestamp_changes_sql( + target=SnapshotChangeTarget( + destination=destination, + origin=origin, + unique_key=unique_key, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + output_columns=output_columns, + ), updated_at_column=updated_at_column, observed_at_column=observed_at_column, initial_valid_from=initial_valid_from, - source_alias="__source", - current_timestamp=current_timestamp, - ) - key_condition: str = _snapshot_key_condition( - left_alias="__target", right_alias="__source", unique_key=unique_key - ) - close_sql: str = ( - f"UPDATE __target SET {valid_to_column} = __source.{updated_at_column} " - f"FROM {destination} AS __target " - f"JOIN {origin} AS __source ON {key_condition} " - f"WHERE __target.{valid_to_column} IS NULL " - f"AND __source.{updated_at_column} > __target.{updated_at_column}" - ) - insert_column_sql: str = ", ".join((*output_columns, valid_from_column, valid_to_column)) - output_select_sql: str = ", ".join(f"__source.{column}" for column in output_columns) - active_join_condition: str = _snapshot_key_condition( - left_alias="__active", right_alias="__source", unique_key=unique_key - ) - first_key: str = unique_key[0] - version_valid_from_expr: str = ( - f"CASE WHEN __active.{first_key} IS NULL THEN {initial_valid_from_expr} " - f"ELSE __source.{updated_at_column} END" + invalidate_hard_deletes=invalidate_hard_deletes, + current_timestamp=self.render_current_timestamp(), ) - insert_sql: str = ( - f"INSERT INTO {destination} ({insert_column_sql}) " - f"SELECT {output_select_sql}, {version_valid_from_expr}, CAST(NULL AS DATETIME2) " - f"FROM {origin} AS __source " - f"LEFT JOIN {destination} AS __active " - f"ON {active_join_condition} AND __active.{valid_to_column} IS NULL " - f"WHERE __active.{first_key} IS NULL " - f"OR __source.{updated_at_column} > __active.{updated_at_column}" - ) - statements: tuple[str, ...] = (close_sql, insert_sql) - if invalidate_hard_deletes: - statements = ( - *statements, - self._snapshot_hard_delete_close_sql( - destination=destination, - origin=origin, - unique_key=unique_key, - valid_to_column=valid_to_column, - current_timestamp=current_timestamp, - ), - ) - return statements def render_clone( self, @@ -2200,6 +1836,22 @@ def render_create_initial_historical_check_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: + if invalidate_hard_deletes: + historical: HistoricalSnapshotSql = HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + check_columns=check_columns, + distinct_condition=self._distinct_condition, + ) + return ( + f"DROP TABLE IF EXISTS {destination}", + f";WITH {historical.initial_ctes_sql()} " + f"SELECT {historical.initial_select_list_sql(output_columns=output_columns)} " + f"INTO {destination} FROM __versions", + ) partition_sql: str = ", ".join(unique_key) previous_columns_sql: str = ", ".join( self._lag_expr( @@ -2217,13 +1869,6 @@ def render_create_initial_historical_check_snapshot_destination( ) output_select_sql: str = ", ".join(output_columns) hard_deleted_at_sql: str = "CAST(NULL AS DATETIME2)" - if invalidate_hard_deletes: - hard_deleted_at_sql = _historical_hard_deleted_at_sql( - origin=origin, - unique_key=unique_key, - observed_at_column=observed_at_column, - row_alias="__changes", - ) valid_to_expr: str = ( "CASE " "WHEN __next_change_at IS NULL THEN __hard_deleted_at " @@ -2284,7 +1929,22 @@ def render_create_initial_historical_timestamp_snapshot_destination( output_columns: tuple[str, ...], invalidate_hard_deletes: bool, ) -> tuple[str, ...]: - del invalidate_hard_deletes + if invalidate_hard_deletes: + historical: HistoricalSnapshotSql = HistoricalSnapshotSql( + origin=origin, + unique_key=unique_key, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + updated_at_column=updated_at_column, + distinct_condition=self._distinct_condition, + ) + return ( + f"DROP TABLE IF EXISTS {destination}", + f";WITH {historical.initial_ctes_sql()} " + f"SELECT {historical.initial_select_list_sql(output_columns=output_columns)} " + f"INTO {destination} FROM __versions", + ) partition_sql: str = ", ".join(unique_key) changed_condition: str = self._distinct_condition( left=updated_at_column, right="__prev_updated_at" @@ -2316,19 +1976,17 @@ def render_create_initial_snapshot_destination( valid_to_column: str, initial_valid_from: str | None, ) -> tuple[str, ...]: - valid_from_expr: str = _snapshot_initial_valid_from_expr( - snapshot_strategy=snapshot_strategy, - updated_at_column=updated_at_column, - observed_at_column=observed_at_column, - initial_valid_from=initial_valid_from, - source_alias=None, - current_timestamp=self.render_current_timestamp(), - ) return self.render_create_table_as( destination=destination, - sql=( - f"SELECT *, {valid_from_expr} AS {valid_from_column}, " - f"CAST(NULL AS DATETIME2) AS {valid_to_column} FROM {origin}" + sql=SnapshotSql(dialect=self._snapshot_sql_dialect).initial_select_sql( + origin=origin, + snapshot_strategy=snapshot_strategy, + updated_at_column=updated_at_column, + observed_at_column=observed_at_column, + valid_from_column=valid_from_column, + valid_to_column=valid_to_column, + initial_valid_from=initial_valid_from, + current_timestamp=self.render_current_timestamp(), ), ) diff --git a/src/sqlbuild/cli/commands/_helpers/audit/invocation.py b/src/sqlbuild/cli/commands/_helpers/audit/invocation.py index 85d92da554..a5ad81f7d5 100644 --- a/src/sqlbuild/cli/commands/_helpers/audit/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/audit/invocation.py @@ -2,58 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - AuditCommandRequest, - AuditInvocation, -) -from sqlbuild.cli.progress.main._build_command_progress_reporters import ( - build_command_progress_reporters, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_reported_command_invocation, ) -from sqlbuild.cli.progress.models import CommandProgressReporters -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import AuditCommandRequest, AuditInvocation def resolve_audit_invocation(*, request: AuditCommandRequest) -> AuditInvocation: """Resolve discovery, adapter, connection, and reporters for audit.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - reporters: CommandProgressReporters = build_command_progress_reporters( - adapter_name=adapter_context.adapter_name, - stream=progress_stream, - use_color=use_color, - ) - return AuditInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - connection_progress=reporters.connection, - planning_progress=reporters.planning, - ) + return resolve_reported_command_invocation(request=request, invocation_type=AuditInvocation) diff --git a/src/sqlbuild/cli/commands/_helpers/check/core.py b/src/sqlbuild/cli/commands/_helpers/check/core.py index 177b887708..5b09357aab 100644 --- a/src/sqlbuild/cli/commands/_helpers/check/core.py +++ b/src/sqlbuild/cli/commands/_helpers/check/core.py @@ -23,6 +23,7 @@ ) from sqlbuild.compiler.compile.models import CompiledRelationLocation from sqlbuild.compiler.discovery.models import DiscoveredCheckFunction, DiscoveredProjectInputs +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.pipeline.main.relation_targets import build_python_relation_targets from sqlbuild.compiler.pipeline.models import CompilePipelineResult from sqlbuild.compiler.planner.exceptions import PlannerInputError @@ -78,16 +79,10 @@ def check_dependency_closure( ) -> frozenset[str]: """Return transitive executable Python dependencies of selected checks.""" - dependencies: set[str] = set() - pending: list[str] = list(check_names) - while pending: - name: str = pending.pop() - for upstream_name in graph.upstream_deps.get(name, ()): - if upstream_name in dependencies: - continue - dependencies.add(upstream_name) - pending.append(upstream_name) - return frozenset(dependencies - check_names) + dependencies: frozenset[str] = transitive_closure_many( + starts=check_names, edges=graph.upstream_deps, include_starts=False + ) + return dependencies - check_names def _validate_check_selectors(*, select: tuple[str, ...]) -> None: diff --git a/src/sqlbuild/cli/commands/_helpers/check/invocation.py b/src/sqlbuild/cli/commands/_helpers/check/invocation.py index dc144d6c55..e0d0b8556d 100644 --- a/src/sqlbuild/cli/commands/_helpers/check/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/check/invocation.py @@ -2,58 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - CheckCommandRequest, - CheckInvocation, -) -from sqlbuild.cli.progress.main._build_command_progress_reporters import ( - build_command_progress_reporters, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_reported_command_invocation, ) -from sqlbuild.cli.progress.models import CommandProgressReporters -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import CheckCommandRequest, CheckInvocation def resolve_check_invocation(*, request: CheckCommandRequest) -> CheckInvocation: """Resolve discovery, adapter, connection, and reporters for check.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - reporters: CommandProgressReporters = build_command_progress_reporters( - adapter_name=adapter_context.adapter_name, - stream=progress_stream, - use_color=use_color, - ) - return CheckInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - connection_progress=reporters.connection, - planning_progress=reporters.planning, - ) + return resolve_reported_command_invocation(request=request, invocation_type=CheckInvocation) diff --git a/src/sqlbuild/cli/commands/_helpers/compile/output.py b/src/sqlbuild/cli/commands/_helpers/compile/output.py index 7edd8c5456..75944bcc6d 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/output.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/output.py @@ -27,7 +27,8 @@ RelatedLocation, ) from sqlbuild.compiler.compile.types import CompiledResourceType, DiagnosticSeverity -from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.lineage.models import ProjectColumnLineage from sqlbuild.compiler.pipeline.models import ProjectGraph from sqlbuild.compiler.python_nodes.main.hook_identities import build_hook_identities @@ -280,27 +281,7 @@ def _model_resource( def _hook_resources(*, hooks: object) -> list[dict[str, object]]: - if not isinstance(hooks, (list, tuple)): - return [] - resources: list[dict[str, object]] = [] - hook: object - for hook in hooks: - if isinstance(hook, SqlHookEntry): - resource: dict[str, object] = {"type": "sql", "statement": hook.statement} - if hook.name is not None: - resource["name"] = hook.name - if hook.relative_path is not None: - resource["relative_path"] = hook.relative_path.as_posix() - if hook.definition_sql is not None: - resource["definition_sql"] = hook.definition_sql - if hook.kwargs is not None: - resource["kwargs"] = hook.kwargs - if hook.description is not None: - resource["description"] = hook.description - resources.append(resource) - elif isinstance(hook, PythonHookEntry): - resources.append({"type": "python", "name": hook.name, "kwargs": hook.kwargs}) - return resources + return serialize_hook_entries(value=hooks, sql_fields=SQL_HOOK_OUTPUT_FIELDS) def _hook_definitions(project: CompiledProject) -> list[dict[str, object]]: diff --git a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py index 5f93cd9f6b..3a75fbb85b 100644 --- a/src/sqlbuild/cli/commands/_helpers/lineage/selection.py +++ b/src/sqlbuild/cli/commands/_helpers/lineage/selection.py @@ -34,6 +34,9 @@ CompiledProject, ) from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.path_nodes import path_nodes +from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.lineage.main.columns import build_project_column_lineage from sqlbuild.compiler.lineage.models import ( ColumnLineageEdge, @@ -124,9 +127,17 @@ def _expand( anchors: tuple[CompiledObjectKey, ...] = tuple(keys) expanded: set[CompiledObjectKey] = set() if direction in {UPSTREAM_DIRECTION, BOTH_DIRECTIONS}: - expanded.update(_walk_bounded(anchors=anchors, deps=graph.upstream_deps, max_depth=depth)) + expanded.update( + transitive_closure_many( + starts=anchors, edges=graph.upstream_deps, include_starts=False, max_depth=depth + ) + ) if direction in {DOWNSTREAM_DIRECTION, BOTH_DIRECTIONS}: - expanded.update(_walk_bounded(anchors=anchors, deps=graph.downstream_deps, max_depth=depth)) + expanded.update( + transitive_closure_many( + starts=anchors, edges=graph.downstream_deps, include_starts=False, max_depth=depth + ) + ) return frozenset(expanded) @@ -233,7 +244,7 @@ def _column_lineage_candidate_selection( deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = ( graph.downstream_deps if direction == DOWNSTREAM_DIRECTION else graph.upstream_deps ) - selected.update(_walk_bounded(anchors=(key,), deps=deps, max_depth=depth)) + selected.update(transitive_closure(start=key, edges=deps, max_depth=depth)) model_names: frozenset[str] = frozenset( selected_key.name for selected_key in selected @@ -242,7 +253,7 @@ def _column_lineage_candidate_selection( if depth is None: return _ColumnLineageCandidateSelection(model_names=model_names, truncated=False) extended: set[CompiledObjectKey] = {key} - extended.update(_walk_bounded(anchors=(key,), deps=deps, max_depth=depth + 1)) + extended.update(transitive_closure(start=key, edges=deps, max_depth=depth + 1)) return _ColumnLineageCandidateSelection( model_names=model_names, truncated=extended != selected, @@ -466,58 +477,21 @@ def _trim_selected_keys( retained.update(anchors.upstream) retained.update(anchors.downstream) retained.update( - _walk_bounded(anchors=anchors.upstream, deps=upstream_deps, max_depth=max_depth) + transitive_closure_many( + starts=anchors.upstream, edges=upstream_deps, include_starts=False, max_depth=max_depth + ) ) retained.update( - _walk_bounded(anchors=anchors.downstream, deps=downstream_deps, max_depth=max_depth) + transitive_closure_many( + starts=anchors.downstream, + edges=downstream_deps, + include_starts=False, + max_depth=max_depth, + ) ) return frozenset(selected_keys & retained) -def _walk_bounded( - *, - anchors: Iterable[CompiledObjectKey], - deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], - max_depth: int | None, -) -> frozenset[CompiledObjectKey]: - if max_depth is None: - result: set[CompiledObjectKey] = set() - for anchor in anchors: - result.update(_walk_all(key=anchor, deps=deps)) - return frozenset(result) - if max_depth == 0: - return frozenset() - visited: set[CompiledObjectKey] = set() - queue: list[tuple[CompiledObjectKey, int]] = [(anchor, 0) for anchor in anchors] - while queue: - current, current_depth = queue.pop(0) - if current_depth >= max_depth: - continue - for neighbor in deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - queue.append((neighbor, current_depth + 1)) - return frozenset(visited) - - -def _walk_all( - *, - key: CompiledObjectKey, - deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [key] - while stack: - current: CompiledObjectKey = stack.pop() - for neighbor in deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) - - def _lookup_name(*, name: str, all_keys: dict[str, CompiledObjectKey]) -> CompiledObjectKey: key: CompiledObjectKey | None = all_keys.get(name) if key is None: @@ -562,13 +536,14 @@ def _resolve_selectors( ) ) scoped: set[CompiledObjectKey] = selected - excluded - for key in tuple(scoped): - for upstream_key in _walk_all(key=key, deps=upstream): - if upstream_key.resource_type in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - scoped.add(upstream_key) + for upstream_key in transitive_closure_many( + starts=tuple(scoped), edges=upstream, include_starts=False + ): + if upstream_key.resource_type in { + CompiledResourceType.UDF, + CompiledResourceType.TABLE_FN, + }: + scoped.add(upstream_key) return frozenset(scoped) @@ -616,9 +591,9 @@ def _resolve_single( _find_path_keys(start=start_key, end=end_key, downstream=downstream) ) if parsed.upstream: - result.update(_walk_all(key=start_key, deps=upstream)) + result.update(transitive_closure(start=start_key, edges=upstream)) if parsed.downstream: - result.update(_walk_all(key=end_key, deps=downstream)) + result.update(transitive_closure(start=end_key, edges=downstream)) return frozenset(result) if parsed.kind == SelectorKind.TAG: matched_keys: frozenset[CompiledObjectKey] = tag_index.get(parsed.value, frozenset()) @@ -648,11 +623,14 @@ def _apply_selector_expansion( downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: result: set[CompiledObjectKey] = set(matched_keys) - for key in matched_keys: - if parsed.upstream: - result.update(_walk_all(key=key, deps=upstream)) - if parsed.downstream: - result.update(_walk_all(key=key, deps=downstream)) + if parsed.upstream: + result.update( + transitive_closure_many(starts=matched_keys, edges=upstream, include_starts=False) + ) + if parsed.downstream: + result.update( + transitive_closure_many(starts=matched_keys, edges=downstream, include_starts=False) + ) return frozenset(result) @@ -729,28 +707,16 @@ def _find_path_keys( end: CompiledObjectKey, downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: - reachable_from_start: frozenset[CompiledObjectKey] = _walk_all(key=start, deps=downstream) - if end not in reachable_from_start: + path_keys: frozenset[CompiledObjectKey] | None = path_nodes( + start=start, end=end, downstream=downstream + ) + if path_keys is None: raise CliUserError( f"'{end.resource_type}:{end.name}' is not downstream of " f"'{start.resource_type}:{start.name}'", code="C319", ) - upstream: dict[CompiledObjectKey, list[CompiledObjectKey]] = {} - for key, dep_keys in downstream.items(): - for dep_key in dep_keys: - upstream.setdefault(dep_key, []).append(key) - upstream_from_end: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [end] - while stack: - current: CompiledObjectKey = stack.pop() - if current in upstream_from_end: - continue - upstream_from_end.add(current) - for parent in upstream.get(current, ()): - if parent in reachable_from_start or parent == start: - stack.append(parent) - return frozenset(reachable_from_start & upstream_from_end | {start, end}) + return path_keys def _match_path( diff --git a/src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py b/src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py new file mode 100644 index 0000000000..62e105d459 --- /dev/null +++ b/src/sqlbuild/cli/commands/_helpers/runtime/command_invocation.py @@ -0,0 +1,116 @@ +"""Shared invocation resolution for project commands that connect to one target.""" + +from __future__ import annotations + +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol, TextIO + +from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( + resolve_adapter_connection_context, +) +from sqlbuild.cli.commands.models import AdapterConnectionContext +from sqlbuild.cli.progress.main._build_command_progress_reporters import ( + build_command_progress_reporters, +) +from sqlbuild.cli.progress.models import CommandProgressReporters +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs +from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs +from sqlbuild.presentation.main.supports_color import supports_color + + +class _CommandInvocationRequest(Protocol): + """CLI request fields needed to resolve a project command invocation.""" + + @property + def project_dir(self) -> Path | None: ... + + @property + def selected_target(self) -> str | None: ... + + @property + def cli_vars(self) -> dict[str, object] | None: ... + + @property + def json_output(self) -> bool: ... + + @property + def no_color(self) -> bool: ... + + +@dataclass(frozen=True) +class _CommandInvocationContext: + effective_project_dir: Path + discovered_inputs: DiscoveredProjectInputs + adapter_context: AdapterConnectionContext + use_color: bool + progress_stream: TextIO + + +def resolve_command_invocation[InvocationT]( + *, request: _CommandInvocationRequest, invocation_type: Callable[..., InvocationT] +) -> InvocationT: + """Resolve discovery, adapter, connection, and output context for one command.""" + + context: _CommandInvocationContext = _resolve_command_invocation_context(request=request) + return invocation_type( + effective_project_dir=context.effective_project_dir, + discovered_inputs=context.discovered_inputs, + adapter_name=context.adapter_context.adapter_name, + adapter=context.adapter_context.adapter, + connection_config=context.adapter_context.connection_config, + use_color=context.use_color, + progress_stream=context.progress_stream, + ) + + +def resolve_reported_command_invocation[InvocationT]( + *, request: _CommandInvocationRequest, invocation_type: Callable[..., InvocationT] +) -> InvocationT: + """Resolve command context plus shared connection and planning progress reporters.""" + + context: _CommandInvocationContext = _resolve_command_invocation_context(request=request) + reporters: CommandProgressReporters = build_command_progress_reporters( + adapter_name=context.adapter_context.adapter_name, + stream=context.progress_stream, + use_color=context.use_color, + ) + return invocation_type( + effective_project_dir=context.effective_project_dir, + discovered_inputs=context.discovered_inputs, + adapter_name=context.adapter_context.adapter_name, + adapter=context.adapter_context.adapter, + connection_config=context.adapter_context.connection_config, + use_color=context.use_color, + progress_stream=context.progress_stream, + connection_progress=reporters.connection, + planning_progress=reporters.planning, + ) + + +def _resolve_command_invocation_context( + *, request: _CommandInvocationRequest +) -> _CommandInvocationContext: + effective_project_dir: Path = ( + request.project_dir if request.project_dir is not None else Path.cwd() + ) + discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( + project_dir=effective_project_dir + ) + adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( + discovered_inputs=discovered_inputs, + effective_project_dir=effective_project_dir, + selected_target=request.selected_target, + cli_vars=request.cli_vars, + ) + machine_output: bool = request.json_output + use_color: bool = not request.no_color and not machine_output and supports_color() + return _CommandInvocationContext( + effective_project_dir=effective_project_dir, + discovered_inputs=discovered_inputs, + adapter_context=adapter_context, + use_color=use_color, + progress_stream=sys.stderr if machine_output else sys.stdout, + ) diff --git a/src/sqlbuild/cli/commands/_helpers/seed/invocation.py b/src/sqlbuild/cli/commands/_helpers/seed/invocation.py index 9e4a75332b..3fdb7028db 100644 --- a/src/sqlbuild/cli/commands/_helpers/seed/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/seed/invocation.py @@ -2,47 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - SeedCommandRequest, - SeedInvocation, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_command_invocation, ) -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import SeedCommandRequest, SeedInvocation def resolve_seed_invocation(*, request: SeedCommandRequest) -> SeedInvocation: """Resolve discovery, adapter, connection, and output context for seed.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - return SeedInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - ) + return resolve_command_invocation(request=request, invocation_type=SeedInvocation) diff --git a/src/sqlbuild/cli/commands/_helpers/test/invocation.py b/src/sqlbuild/cli/commands/_helpers/test/invocation.py index f03cb38837..6572d84451 100644 --- a/src/sqlbuild/cli/commands/_helpers/test/invocation.py +++ b/src/sqlbuild/cli/commands/_helpers/test/invocation.py @@ -2,58 +2,13 @@ from __future__ import annotations -import sys -from pathlib import Path -from typing import TextIO - -from sqlbuild.cli.commands._helpers.runtime.adapter_context import ( - resolve_adapter_connection_context, -) -from sqlbuild.cli.commands.models import ( - AdapterConnectionContext, - TestCommandRequest, - TestInvocation, -) -from sqlbuild.cli.progress.main._build_command_progress_reporters import ( - build_command_progress_reporters, +from sqlbuild.cli.commands._helpers.runtime.command_invocation import ( + resolve_reported_command_invocation, ) -from sqlbuild.cli.progress.models import CommandProgressReporters -from sqlbuild.compiler.discovery.main.discover import discover_project_inputs -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.presentation.main.supports_color import supports_color +from sqlbuild.cli.commands.models import TestCommandRequest, TestInvocation def resolve_test_invocation(*, request: TestCommandRequest) -> TestInvocation: """Resolve discovery, adapter, connection, and reporters for test.""" - effective_project_dir: Path = ( - request.project_dir if request.project_dir is not None else Path.cwd() - ) - discovered_inputs: DiscoveredProjectInputs = discover_project_inputs( - project_dir=effective_project_dir - ) - adapter_context: AdapterConnectionContext = resolve_adapter_connection_context( - discovered_inputs=discovered_inputs, - effective_project_dir=effective_project_dir, - selected_target=request.selected_target, - cli_vars=request.cli_vars, - ) - machine_output: bool = request.json_output - use_color: bool = not request.no_color and not machine_output and supports_color() - progress_stream: TextIO = sys.stderr if machine_output else sys.stdout - reporters: CommandProgressReporters = build_command_progress_reporters( - adapter_name=adapter_context.adapter_name, - stream=progress_stream, - use_color=use_color, - ) - return TestInvocation( - effective_project_dir=effective_project_dir, - discovered_inputs=discovered_inputs, - adapter_name=adapter_context.adapter_name, - adapter=adapter_context.adapter, - connection_config=adapter_context.connection_config, - use_color=use_color, - progress_stream=progress_stream, - connection_progress=reporters.connection, - planning_progress=reporters.planning, - ) + return resolve_reported_command_invocation(request=request, invocation_type=TestInvocation) diff --git a/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py b/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py index 2d39fc3357..b838e2837f 100644 --- a/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py +++ b/src/sqlbuild/compiler/auditing/_helpers/schema_audits.py @@ -14,12 +14,39 @@ from sqlbuild.compiler.auditing.exceptions import MeasurementAuditError from sqlbuild.compiler.auditing.models import MeasurementThresholdBound, MeasurementThresholds from sqlbuild.compiler.auditing.types import AuditSeverity, ThresholdOperator +from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool +from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.resource_names.main._validate_resource_identity import ( validate_resource_identity, ) from sqlbuild.spec.contracts.models import SchemaAuditInstance +def parse_audit_instances_impl( + *, + raw_audits: object | None, + file_path: Path, + label: str, + error_class: type[Exception], + null_as_empty: bool, +) -> tuple[SchemaAuditInstance, ...]: + """Parse an authored audit list, treating null as empty only when requested.""" + + if raw_audits is None and null_as_empty: + return () + if not isinstance(raw_audits, list): + raise error_class(f"{file_path} {label} audits must be a list") + return tuple( + parse_audit_instance_impl( + raw_audit=raw_audit, + file_path=file_path, + label=label, + error_class=error_class, + ) + for raw_audit in raw_audits + ) + + def parse_audit_instance_impl( *, raw_audit: object, @@ -62,7 +89,7 @@ def parse_audit_instance_impl( argument_mapping: dict[str, object] = cast(dict[str, object], raw_arguments) option_label: str = f"{label} audit '{definition_name}'" - name: str | None = _optional_named_string( + name: str | None = optional_named_string( raw_value=argument_mapping.get("name"), file_path=file_path, label=option_label, @@ -75,14 +102,14 @@ def parse_audit_instance_impl( kind=f"{label} audit instance", path=file_path, ) - description: str | None = _optional_named_string( + description: str | None = optional_named_string( raw_value=argument_mapping.get("description"), file_path=file_path, label=option_label, key="description", error_class=error_class, ) - raw_severity: str | None = _optional_named_string( + raw_severity: str | None = optional_named_string( raw_value=argument_mapping.get("severity"), file_path=file_path, label=option_label, @@ -98,19 +125,20 @@ def parse_audit_instance_impl( raise error_class( f"{file_path} {option_label} 'severity' must be one of: {allowed}" ) from error - run_scope: str | None = _optional_named_string( + run_scope: str | None = optional_named_string( raw_value=argument_mapping.get("run_scope"), file_path=file_path, label=option_label, key="run_scope", error_class=error_class, ) - always_run: bool = _optional_named_bool( + always_run: bool = optional_named_bool( raw_value=argument_mapping.get("always_run"), file_path=file_path, label=option_label, key="always_run", error_class=error_class, + default=False, ) thresholds: MeasurementThresholds | None = parse_measurement_thresholds( raw_value=argument_mapping.get("thresholds"), @@ -244,33 +272,3 @@ def _parse_threshold_bound(raw_value: object | None) -> MeasurementThresholdBoun if isinstance(raw_limit, bool) or not isinstance(raw_limit, (int, float)): raise MeasurementAuditError(f"{operator.value} threshold requires one numeric value") return MeasurementThresholdBound(operator=operator, limit=float(raw_limit)) - - -def _optional_named_string( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: type[Exception], -) -> str | None: - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise error_class(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - -def _optional_named_bool( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: type[Exception], -) -> bool: - if raw_value is None: - return False - if not isinstance(raw_value, bool): - raise error_class(f"{file_path} {label} '{key}' must be a boolean") - return raw_value diff --git a/src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py b/src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py new file mode 100644 index 0000000000..4b451e5a09 --- /dev/null +++ b/src/sqlbuild/compiler/auditing/main/_parse_audit_instances.py @@ -0,0 +1,27 @@ +"""Schema-attached audit list parsing entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.auditing._helpers.schema_audits import parse_audit_instances_impl +from sqlbuild.spec.contracts.models import SchemaAuditInstance + + +def parse_audit_instances( + *, + raw_audits: object | None, + file_path: Path, + label: str, + error_class: type[Exception], + null_as_empty: bool, +) -> tuple[SchemaAuditInstance, ...]: + """Parse an authored audit list, treating null as empty only when requested.""" + + return parse_audit_instances_impl( + raw_audits=raw_audits, + file_path=file_path, + label=label, + error_class=error_class, + null_as_empty=null_as_empty, + ) diff --git a/src/sqlbuild/compiler/authored_values/__init__.py b/src/sqlbuild/compiler/authored_values/__init__.py new file mode 100644 index 0000000000..4c2eb9cf6a --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/__init__.py @@ -0,0 +1 @@ +"""Validation of optional and required authored configuration values.""" diff --git a/src/sqlbuild/compiler/authored_values/_helpers/__init__.py b/src/sqlbuild/compiler/authored_values/_helpers/__init__.py new file mode 100644 index 0000000000..9c587e407e --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/_helpers/__init__.py @@ -0,0 +1 @@ +"""Authored configuration value validation implementations.""" diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py b/src/sqlbuild/compiler/authored_values/_helpers/fields.py similarity index 62% rename from src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py rename to src/sqlbuild/compiler/authored_values/_helpers/fields.py index a41b03f36b..a67aad1518 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/primitives.py +++ b/src/sqlbuild/compiler/authored_values/_helpers/fields.py @@ -1,42 +1,12 @@ -"""Shared YAML field-extraction primitives for discovery parsers.""" +"""Authored configuration value validation implementations.""" from __future__ import annotations from pathlib import Path from typing import cast -from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance -from sqlbuild.spec.contracts.models import SchemaAuditInstance - -def parse_audit_instances( - *, - entry: dict[str, object], - file_path: Path, - label: str, - error_class: type[Exception], -) -> tuple[SchemaAuditInstance, ...]: - """Parse a list of schema-attached audit instances from a YAML entry.""" - - raw_audits: object = entry.get("audits", []) - if not isinstance(raw_audits, list): - raise error_class(f"{file_path} {label} audits must be a list") - - parsed_audits: list[SchemaAuditInstance] = [] - raw_audit: object - for raw_audit in raw_audits: - parsed_audits.append( - parse_audit_instance( - raw_audit=raw_audit, - file_path=file_path, - label=label, - error_class=error_class, - ) - ) - return tuple(parsed_audits) - - -def require_non_empty_string( +def require_non_empty_string_impl( *, entry: dict[str, object], key: str, @@ -44,7 +14,7 @@ def require_non_empty_string( label: str, error_class: type[Exception], ) -> str: - """Extract a required non-empty string from a YAML mapping.""" + """Extract a required non-empty string from a mapping.""" raw_value: object | None = entry.get(key) if not isinstance(raw_value, str) or not raw_value.strip(): @@ -52,7 +22,7 @@ def require_non_empty_string( return raw_value -def optional_non_empty_string( +def optional_non_empty_string_impl( *, entry: dict[str, object], key: str, @@ -60,9 +30,9 @@ def optional_non_empty_string( label: str, error_class: type[Exception], ) -> str | None: - """Extract an optional non-empty string from a YAML mapping.""" + """Extract an optional non-empty string from a mapping.""" - return optional_named_string( + return optional_named_string_impl( raw_value=entry.get(key), file_path=file_path, label=label, @@ -71,7 +41,7 @@ def optional_non_empty_string( ) -def optional_named_string( +def optional_named_string_impl( *, raw_value: object | None, file_path: Path, @@ -88,25 +58,45 @@ def optional_named_string( return raw_value -def optional_bool( +def optional_named_bool_impl[D: (bool, None)]( *, - entry: dict[str, object], - key: str, + raw_value: object | None, file_path: Path, label: str, + key: str, error_class: type[Exception], -) -> bool | None: - """Extract an optional boolean from a YAML mapping.""" + default: D, +) -> bool | D: + """Validate and return an optional named boolean value.""" - raw_value: object | None = entry.get(key) if raw_value is None: - return None + return default if not isinstance(raw_value, bool): raise error_class(f"{file_path} {label} '{key}' must be a boolean") return raw_value -def optional_mapping( +def optional_bool_impl( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> bool | None: + """Extract an optional boolean from a mapping.""" + + return optional_named_bool_impl( + raw_value=entry.get(key), + file_path=file_path, + label=label, + key=key, + error_class=error_class, + default=None, + ) + + +def optional_mapping_impl( *, entry: dict[str, object], key: str, @@ -114,7 +104,7 @@ def optional_mapping( label: str, error_class: type[Exception], ) -> dict[str, object]: - """Extract an optional mapping from a YAML mapping.""" + """Extract an optional mapping from a mapping, defaulting to empty.""" raw_value: object | None = entry.get(key) if raw_value is None: @@ -124,7 +114,7 @@ def optional_mapping( return cast(dict[str, object], raw_value) -def optional_string_tuple( +def optional_string_tuple_impl( *, entry: dict[str, object], key: str, @@ -132,7 +122,7 @@ def optional_string_tuple( label: str, error_class: type[Exception], ) -> tuple[str, ...]: - """Extract an optional list of strings from a YAML mapping.""" + """Extract an optional list of strings from a mapping.""" raw_value: object | None = entry.get(key) if raw_value is None: diff --git a/src/sqlbuild/compiler/authored_values/main/__init__.py b/src/sqlbuild/compiler/authored_values/main/__init__.py new file mode 100644 index 0000000000..1f01d5721c --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/__init__.py @@ -0,0 +1 @@ +"""Authored configuration value public entries.""" diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_bool.py b/src/sqlbuild/compiler/authored_values/main/_optional_bool.py new file mode 100644 index 0000000000..2c2f3f1372 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_bool.py @@ -0,0 +1,22 @@ +"""Public optional bool entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_bool_impl + + +def optional_bool( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> bool | None: + """Extract an optional boolean from a mapping.""" + + return optional_bool_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_mapping.py b/src/sqlbuild/compiler/authored_values/main/_optional_mapping.py new file mode 100644 index 0000000000..c44bd6e5be --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_mapping.py @@ -0,0 +1,22 @@ +"""Public optional mapping entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_mapping_impl + + +def optional_mapping( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> dict[str, object]: + """Extract an optional mapping from a mapping, defaulting to empty.""" + + return optional_mapping_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py b/src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py new file mode 100644 index 0000000000..8e57a6da90 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_named_bool.py @@ -0,0 +1,28 @@ +"""Public optional named bool entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_named_bool_impl + + +def optional_named_bool[D: (bool, None)]( + *, + raw_value: object | None, + file_path: Path, + label: str, + key: str, + error_class: type[Exception], + default: D, +) -> bool | D: + """Validate an optional named boolean, returning default when absent.""" + + return optional_named_bool_impl( + raw_value=raw_value, + file_path=file_path, + label=label, + key=key, + error_class=error_class, + default=default, + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_named_string.py b/src/sqlbuild/compiler/authored_values/main/_optional_named_string.py new file mode 100644 index 0000000000..2756d0745a --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_named_string.py @@ -0,0 +1,22 @@ +"""Public optional named string entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_named_string_impl + + +def optional_named_string( + *, + raw_value: object | None, + file_path: Path, + label: str, + key: str, + error_class: type[Exception], +) -> str | None: + """Validate and return an optional named string value.""" + + return optional_named_string_impl( + raw_value=raw_value, file_path=file_path, label=label, key=key, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py b/src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py new file mode 100644 index 0000000000..28f1af346f --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_non_empty_string.py @@ -0,0 +1,22 @@ +"""Public optional non empty string entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_non_empty_string_impl + + +def optional_non_empty_string( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> str | None: + """Extract an optional non-empty string from a mapping.""" + + return optional_non_empty_string_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py b/src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py new file mode 100644 index 0000000000..716bf965c4 --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_optional_string_tuple.py @@ -0,0 +1,22 @@ +"""Public optional string tuple entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import optional_string_tuple_impl + + +def optional_string_tuple( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> tuple[str, ...]: + """Extract an optional list of strings from a mapping.""" + + return optional_string_tuple_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py b/src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py new file mode 100644 index 0000000000..bd7f087b9a --- /dev/null +++ b/src/sqlbuild/compiler/authored_values/main/_require_non_empty_string.py @@ -0,0 +1,22 @@ +"""Public require non empty string entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values._helpers.fields import require_non_empty_string_impl + + +def require_non_empty_string( + *, + entry: dict[str, object], + key: str, + file_path: Path, + label: str, + error_class: type[Exception], +) -> str: + """Extract a required non-empty string from a mapping.""" + + return require_non_empty_string_impl( + entry=entry, key=key, file_path=file_path, label=label, error_class=error_class + ) diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py b/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py index a0e46e240b..43294b01a3 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/tests.py @@ -25,7 +25,7 @@ def extract_expected_branch_column_names_with_sql_analysis( - *, sql: str, file_label: str + *, sql: str, file_label: str, label: str = "__expected__" ) -> tuple[tuple[str, ...], ...] | None: """Return expected SELECT branch names using required Polyglot analysis.""" @@ -40,32 +40,34 @@ def extract_expected_branch_column_names_with_sql_analysis( _extract_branch_names( expression=polyglot_module.parse_one(branch_sql, dialect="generic"), file_label=file_label, + label=label, )[0] for branch_sql in branches ) return None - return _extract_branch_names(expression=parsed_expression, file_label=file_label) + return _extract_branch_names(expression=parsed_expression, file_label=file_label, label=label) -def _extract_branch_names(*, expression: Any, file_label: str) -> tuple[tuple[str, ...], ...]: +def _extract_branch_names( + *, expression: Any, file_label: str, label: str +) -> tuple[tuple[str, ...], ...]: expression = _unwrap_expression(expression=expression) if expression.__class__.__name__ == POLYGLOT_UNION_EXPRESSION_NAME: left_expression: Any = expression.args["left"] right_expression: Any = expression.args["right"] return ( - *_extract_branch_names(expression=left_expression, file_label=file_label), - *_extract_branch_names(expression=right_expression, file_label=file_label), + *_extract_branch_names(expression=left_expression, file_label=file_label, label=label), + *_extract_branch_names(expression=right_expression, file_label=file_label, label=label), ) if expression.__class__.__name__ == POLYGLOT_SELECT_EXPRESSION_NAME: if _is_synthetic_values_set_branch(expression): raise CompileInputError( - f"SQL test '{file_label}' must define each __expected__ set-operation " + f"SQL test '{file_label}' must define each {label} set-operation " "branch as a SELECT query" ) - return (_extract_select_names(expression=expression, file_label=file_label),) + return (_extract_select_names(expression=expression, file_label=file_label, label=label),) raise CompileInputError( - f"SQL test '{file_label}' must define each __expected__ set-operation " - "branch as a SELECT query" + f"SQL test '{file_label}' must define each {label} set-operation branch as a SELECT query" ) @@ -88,13 +90,13 @@ def _unwrap_expression(*, expression: Any) -> Any: return expression -def _extract_select_names(*, expression: Any, file_label: str) -> tuple[str, ...]: +def _extract_select_names(*, expression: Any, file_label: str, label: str) -> tuple[str, ...]: names: list[str] = [] projection: Any for projection in expression.expressions: if projection.is_star: raise CompileInputError( - f"SQL test '{file_label}' must not use SELECT * in __expected__ CTEs" + f"SQL test '{file_label}' must not use SELECT * in {label} CTEs" ) alias_name: str = str(projection.alias or "") if alias_name: @@ -104,11 +106,11 @@ def _extract_select_names(*, expression: Any, file_label: str) -> tuple[str, ... names.append(str(projection.name)) continue raise CompileInputError( - f"SQL test '{file_label}' must alias every non-trivial __expected__ projection" + f"SQL test '{file_label}' must alias every non-trivial {label} projection" ) if not names: raise CompileInputError( - f"SQL test '{file_label}' must project at least one column in __expected__" + f"SQL test '{file_label}' must project at least one column in {label}" ) return tuple(names) diff --git a/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py b/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py index 67e92c85ed..ea7c56e8c0 100644 --- a/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py +++ b/src/sqlbuild/compiler/compile/_helpers/attachment/model_config.py @@ -6,7 +6,8 @@ from pathlib import Path from typing import cast -from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances +from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.compile._helpers.audit_factories.core import ( merge_validated_model_audits, parse_model_header_audit_factories, @@ -265,11 +266,12 @@ def build_model_header_schema_entry( ): return None - model_description: str | None = _optional_model_header_string( + model_description: str | None = optional_named_string( raw_value=raw_description, file_path=file_path, label="model", key="description", + error_class=CompileInputError, ) description: str | None = model_description or model_schema_description local_columns: tuple[SchemaColumn, ...] = _parse_model_header_columns( @@ -289,10 +291,12 @@ def build_model_header_schema_entry( model_name=model_name, file_path=file_path, ) - audits: tuple[SchemaAuditInstance, ...] = _parse_model_header_audits( + audits: tuple[SchemaAuditInstance, ...] = parse_audit_instances( raw_audits=raw_audits, file_path=file_path, label="model", + error_class=CompileInputError, + null_as_empty=True, ) generated_audits: tuple[SchemaAuditInstance, ...] = parse_model_header_audit_factories( raw_audit_factories=raw_audit_factories, @@ -496,34 +500,6 @@ def _validate_model_schema_audit_augmentation( ) -def _parse_model_header_audits( - *, raw_audits: object | None, file_path: Path, label: str -) -> tuple[SchemaAuditInstance, ...]: - if raw_audits is None: - return () - if not isinstance(raw_audits, list): - raise CompileInputError(f"{file_path} {label} audits must be a list") - return tuple( - parse_audit_instance( - raw_audit=raw_audit, - file_path=file_path, - label=label, - error_class=CompileInputError, - ) - for raw_audit in raw_audits - ) - - -def _optional_model_header_string( - *, raw_value: object | None, file_path: Path, label: str, key: str -) -> str | None: - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise CompileInputError(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - def _resolve_model_schema( *, values: dict[str, object], diff --git a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py index 471e1b93f0..111313f063 100644 --- a/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/scenarios/core.py @@ -2,18 +2,14 @@ from __future__ import annotations -from functools import lru_cache - from sqlbuild.compiler.compile._helpers.analysis.ctes import ( extract_top_level_ctes_with_sql_analysis, ) from sqlbuild.compiler.compile._helpers.sql_tests.core import ( - _consume_keyword, - _read_identifier, _require_prefixed_name, _skip_ignorable, _try_consume_keyword, - _validate_ceremonial_select, + extract_top_level_ctes_with_scanner, validate_independent_expected_and_assertion_ctes, ) from sqlbuild.compiler.compile.constants import ( @@ -24,8 +20,6 @@ REF_TEST_CTE_PREFIX, SEED_TEST_CTE_PREFIX, SOURCE_TEST_CTE_PREFIX, - SQL_ARGUMENT_SEPARATOR_TOKEN, - SQL_OPEN_PAREN_TOKEN, SQL_WITH_KEYWORD, ) from sqlbuild.compiler.compile.exceptions import CompileInputError @@ -33,18 +27,23 @@ CompileSqlScenarioCte, CompileSqlScenarioCtes, ) -from sqlbuild.compiler.sql_analysis.main._find_matching_paren import find_matching_paren _CONTEXT: str = "SQL scenario" +_WITH_REQUIREMENT: str = ( + "fixture CTEs and at least one __expected__ or __assert__ CTE" +) def extract_sql_scenario_ctes(*, sql: str, file_label: str) -> CompileSqlScenarioCtes: """Extract top-level SQL-native scenario fixture, expected, and assertion CTEs.""" try: - ctes: tuple[CompileSqlScenarioCte, ...] = _extract_sql_scenario_ctes_with_scanner( + ctes: tuple[CompileSqlScenarioCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + with_requirement=_WITH_REQUIREMENT, + cte_type=CompileSqlScenarioCte, ) except CompileInputError as scanner_error: cte_values: tuple[tuple[str, str], ...] | None = extract_top_level_ctes_with_sql_analysis( @@ -61,12 +60,15 @@ def extract_sql_scenario_ctes(*, sql: str, file_label: str) -> CompileSqlScenari def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> tuple[str, ...]: """Extract explicit expected-model relationships without inspecting CTE bodies.""" - start: int = _skip_ignorable(sql=sql, start=0) + start: int = _skip_ignorable(sql=sql, start=0, context_label=_CONTEXT) if _try_consume_keyword(sql=sql, start=start, keyword=SQL_WITH_KEYWORD) is None: return () - ctes: tuple[CompileSqlScenarioCte, ...] = _extract_sql_scenario_ctes_with_scanner( + ctes: tuple[CompileSqlScenarioCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + with_requirement=_WITH_REQUIREMENT, + cte_type=CompileSqlScenarioCte, ) return tuple( _require_prefixed_name( @@ -74,65 +76,13 @@ def extract_sql_scenario_expected_model_names(*, sql: str, file_label: str) -> t prefix=EXPECTED_TEST_CTE_PREFIX, label="__expected__", file_label=file_label, + context_label=_CONTEXT, ) for cte in ctes if cte.name.startswith(EXPECTED_TEST_CTE_PREFIX) ) -@lru_cache(maxsize=4096) -def _extract_sql_scenario_ctes_with_scanner( - *, sql: str, file_label: str -) -> tuple[CompileSqlScenarioCte, ...]: - index: int = _skip_ignorable(sql=sql, start=0) - index = _consume_keyword(sql=sql, start=index, keyword="WITH", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) - recursive_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="RECURSIVE") - if recursive_end is not None: - index = _skip_ignorable(sql=sql, start=recursive_end) - - ctes: list[CompileSqlScenarioCte] = [] - seen_cte_names: set[str] = set() - while True: - cte_name, index = _read_identifier(sql=sql, start=index, file_label=file_label) - if cte_name in seen_cte_names: - raise CompileInputError( - f"SQL scenario '{file_label}' defines duplicate CTE '{cte_name}'" - ) - seen_cte_names.add(cte_name) - - index = _skip_ignorable(sql=sql, start=index) - if index < len(sql) and sql[index] == SQL_OPEN_PAREN_TOKEN: - index = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) + 1 - index = _skip_ignorable(sql=sql, start=index) - index = _consume_keyword(sql=sql, start=index, keyword="AS", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) - if index >= len(sql) or sql[index] != SQL_OPEN_PAREN_TOKEN: - raise CompileInputError( - f"SQL scenario '{file_label}' CTE '{cte_name}' must use AS (...)" - ) - cte_body_start: int = index + 1 - cte_body_end: int = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) - ctes.append( - CompileSqlScenarioCte( - name=cte_name, - sql_body=sql[cte_body_start:cte_body_end].strip(), - ) - ) - index = _skip_ignorable(sql=sql, start=cte_body_end + 1) - if index < len(sql) and sql[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: - index = _skip_ignorable(sql=sql, start=index + 1) - continue - break - - _validate_ceremonial_select( - sql=sql, - start=index, - file_label=file_label, - ) - return tuple(ctes) - - def _classify_sql_scenario_ctes( *, ctes: tuple[CompileSqlScenarioCte, ...], file_label: str ) -> CompileSqlScenarioCtes: @@ -162,6 +112,7 @@ def _classify_sql_scenario_ctes( prefix=SOURCE_TEST_CTE_PREFIX, label="__source__", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -173,6 +124,7 @@ def _classify_sql_scenario_ctes( prefix=REF_TEST_CTE_PREFIX, label="__ref__", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -184,6 +136,7 @@ def _classify_sql_scenario_ctes( prefix=SEED_TEST_CTE_PREFIX, label="__seed__", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -195,6 +148,7 @@ def _classify_sql_scenario_ctes( prefix=DBT_REF_TEST_CTE_PREFIX, label="__dbt_ref__ or __dbt_ref____", file_label=file_label, + context_label=_CONTEXT, ) ) authored_ctes.append(cte) @@ -206,6 +160,7 @@ def _classify_sql_scenario_ctes( prefix=EXPECTED_TEST_CTE_PREFIX, label="__expected__", file_label=file_label, + context_label=_CONTEXT, ) ) expected_ctes.append(cte) @@ -217,6 +172,7 @@ def _classify_sql_scenario_ctes( prefix=ASSERT_SCENARIO_CTE_PREFIX, label="__assert__", file_label=file_label, + context_label=_CONTEXT, ) ) assertion_ctes.append(cte) diff --git a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py index 6586f608d7..9887de4642 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py +++ b/src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py @@ -2,6 +2,8 @@ from __future__ import annotations +from collections.abc import Callable +from dataclasses import dataclass from functools import lru_cache from typing import Any @@ -28,9 +30,7 @@ SOURCE_TEST_CTE_PREFIX, SQL_ARGUMENT_SEPARATOR_TOKEN, SQL_CEREMONIAL_SELECT_VALUE, - SQL_CLOSE_PAREN_TOKEN, SQL_OPEN_PAREN_TOKEN, - SQL_QUOTE_TOKENS, SQL_SINGLE_QUOTE_TOKEN, SQL_STATEMENT_TERMINATOR_TOKEN, SQL_WILDCARD_TOKEN, @@ -57,6 +57,7 @@ is_identifier_character, ) from sqlbuild.compiler.sql_analysis.main._is_identifier_start import is_identifier_start +from sqlbuild.compiler.sql_analysis.main._iter_code_positions import iter_code_positions from sqlbuild.compiler.sql_analysis.main._skip_block_comment import skip_block_comment from sqlbuild.compiler.sql_analysis.main._skip_line_comment import skip_line_comment from sqlbuild.compiler.sql_analysis.main._skip_quoted_text import ( @@ -65,10 +66,56 @@ from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql _CONTEXT: str = "SQL test" +_SQL_TEST_WITH_REQUIREMENT: str = "mock CTEs and one __expected__ CTE" _DIRECT_DEPENDENCY_PATH_LENGTH: int = 2 _SQL_IDENTIFIER_QUOTE_TOKENS: frozenset[str] = frozenset({'"', "`"}) +@dataclass(frozen=True) +class _DirectLogicModeSpec: + mode: SqlTestMode + actual_cte_name: str + expected_cte_name: str + foreign_ctes: tuple[tuple[frozenset[str], str], ...] + + +_MACRO_TEST_CTE_NAMES: frozenset[str] = frozenset( + {MACRO_ACTUAL_TEST_CTE_NAME, MACRO_EXPECTED_TEST_CTE_NAME} +) +_UDF_TEST_CTE_NAMES: frozenset[str] = frozenset( + {UDF_ACTUAL_TEST_CTE_NAME, UDF_EXPECTED_TEST_CTE_NAME} +) +_TABLE_FN_TEST_CTE_NAMES: frozenset[str] = frozenset( + {TABLE_FN_ACTUAL_TEST_CTE_NAME, TABLE_FN_EXPECTED_TEST_CTE_NAME} +) +_DIRECT_LOGIC_MODE_SPECS: dict[SqlTestMode, _DirectLogicModeSpec] = { + SqlTestMode.MACRO: _DirectLogicModeSpec( + mode=SqlTestMode.MACRO, + actual_cte_name=MACRO_ACTUAL_TEST_CTE_NAME, + expected_cte_name=MACRO_EXPECTED_TEST_CTE_NAME, + foreign_ctes=( + (_UDF_TEST_CTE_NAMES, "UDF-test"), + (_TABLE_FN_TEST_CTE_NAMES, "table_fn-test"), + ), + ), + SqlTestMode.UDF: _DirectLogicModeSpec( + mode=SqlTestMode.UDF, + actual_cte_name=UDF_ACTUAL_TEST_CTE_NAME, + expected_cte_name=UDF_EXPECTED_TEST_CTE_NAME, + foreign_ctes=( + (_MACRO_TEST_CTE_NAMES, "macro-test"), + (_TABLE_FN_TEST_CTE_NAMES, "table_fn-test"), + ), + ), + SqlTestMode.TABLE_FN: _DirectLogicModeSpec( + mode=SqlTestMode.TABLE_FN, + actual_cte_name=TABLE_FN_ACTUAL_TEST_CTE_NAME, + expected_cte_name=TABLE_FN_EXPECTED_TEST_CTE_NAME, + foreign_ctes=((_MACRO_TEST_CTE_NAMES | _UDF_TEST_CTE_NAMES, "another direct-logic"),), + ), +} + + def extract_sql_test_ctes( *, sql: str, file_label: str, mode: SqlTestMode = DEFAULT_SQL_TEST_MODE ) -> CompileSqlTestCtes: @@ -87,9 +134,12 @@ def extract_unclassified_sql_test_ctes( """Extract raw top-level CTEs before mode-specific classification.""" try: - ctes: tuple[CompileSqlTestCte, ...] = _extract_sql_test_ctes_with_scanner( + ctes: tuple[CompileSqlTestCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + with_requirement=_SQL_TEST_WITH_REQUIREMENT, + cte_type=CompileSqlTestCte, ) except CompileInputError as scanner_error: cte_values: tuple[tuple[str, str], ...] | None = extract_top_level_ctes_with_sql_analysis( @@ -121,9 +171,12 @@ def extract_sql_test_expected_model_names( start: int = _skip_ignorable(sql=sql, start=0) if _try_consume_keyword(sql=sql, start=start, keyword=SQL_WITH_KEYWORD) is None: return () - ctes: tuple[CompileSqlTestCte, ...] = _extract_sql_test_ctes_with_scanner( + ctes: tuple[CompileSqlTestCte, ...] = extract_top_level_ctes_with_scanner( sql=sql, file_label=file_label, + context_label=_CONTEXT, + with_requirement=_SQL_TEST_WITH_REQUIREMENT, + cte_type=CompileSqlTestCte, ) return tuple( _require_prefixed_name( @@ -151,43 +204,66 @@ def extract_assertion_target_model_names(*, assertion_sql: tuple[str, ...]) -> t @lru_cache(maxsize=4096) -def _extract_sql_test_ctes_with_scanner( - *, sql: str, file_label: str -) -> tuple[CompileSqlTestCte, ...]: - index: int = _skip_ignorable(sql=sql, start=0) - index = _consume_keyword(sql=sql, start=index, keyword="WITH", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) +def extract_top_level_ctes_with_scanner[CteT]( + *, + sql: str, + file_label: str, + context_label: str, + with_requirement: str, + cte_type: Callable[..., CteT], +) -> tuple[CteT, ...]: + """Scan top-level `WITH` CTEs followed by the ceremonial `SELECT 1` of a test file.""" + + with_end: int | None = _try_consume_keyword( + sql=sql, + start=_skip_ignorable(sql=sql, start=0, context_label=context_label), + keyword=SQL_WITH_KEYWORD, + ) + if with_end is None: + raise CompileInputError( + f"{context_label} '{file_label}' must declare {with_requirement} before `SELECT 1`" + ) + index: int = _skip_ignorable(sql=sql, start=with_end, context_label=context_label) recursive_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="RECURSIVE") if recursive_end is not None: - index = _skip_ignorable(sql=sql, start=recursive_end) + index = _skip_ignorable(sql=sql, start=recursive_end, context_label=context_label) - ctes: list[CompileSqlTestCte] = [] + ctes: list[CteT] = [] seen_cte_names: set[str] = set() while True: - cte_name, index = _read_identifier(sql=sql, start=index, file_label=file_label) + cte_name, index = _read_identifier( + sql=sql, start=index, file_label=file_label, context_label=context_label + ) if cte_name in seen_cte_names: - raise CompileInputError(f"SQL test '{file_label}' defines duplicate CTE '{cte_name}'") + raise CompileInputError( + f"{context_label} '{file_label}' defines duplicate CTE '{cte_name}'" + ) seen_cte_names.add(cte_name) - index = _skip_ignorable(sql=sql, start=index) + index = _skip_ignorable(sql=sql, start=index, context_label=context_label) if index < len(sql) and sql[index] == SQL_OPEN_PAREN_TOKEN: - index = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) + 1 - index = _skip_ignorable(sql=sql, start=index) - index = _consume_keyword(sql=sql, start=index, keyword="AS", file_label=file_label) - index = _skip_ignorable(sql=sql, start=index) + index = find_matching_paren(sql=sql, open_paren_index=index, context=context_label) + 1 + index = _skip_ignorable(sql=sql, start=index, context_label=context_label) + index = _consume_keyword( + sql=sql, + start=index, + keyword="AS", + file_label=file_label, + context_label=context_label, + ) + index = _skip_ignorable(sql=sql, start=index, context_label=context_label) if index >= len(sql) or sql[index] != SQL_OPEN_PAREN_TOKEN: - raise CompileInputError(f"SQL test '{file_label}' CTE '{cte_name}' must use AS (...)") - cte_body_start: int = index + 1 - cte_body_end: int = find_matching_paren(sql=sql, open_paren_index=index, context=_CONTEXT) - ctes.append( - CompileSqlTestCte( - name=cte_name, - sql_body=sql[cte_body_start:cte_body_end].strip(), + raise CompileInputError( + f"{context_label} '{file_label}' CTE '{cte_name}' must use AS (...)" ) + cte_body_start: int = index + 1 + cte_body_end: int = find_matching_paren( + sql=sql, open_paren_index=index, context=context_label ) - index = _skip_ignorable(sql=sql, start=cte_body_end + 1) + ctes.append(cte_type(name=cte_name, sql_body=sql[cte_body_start:cte_body_end].strip())) + index = _skip_ignorable(sql=sql, start=cte_body_end + 1, context_label=context_label) if index < len(sql) and sql[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: - index = _skip_ignorable(sql=sql, start=index + 1) + index = _skip_ignorable(sql=sql, start=index + 1, context_label=context_label) continue break @@ -195,6 +271,7 @@ def _extract_sql_test_ctes_with_scanner( sql=sql, start=index, file_label=file_label, + context_label=context_label, ) return tuple(ctes) @@ -205,177 +282,86 @@ def _classify_sql_test_ctes( match mode: case SqlTestMode.MODEL: return _classify_model_sql_test_ctes(ctes=ctes, file_label=file_label) - case SqlTestMode.MACRO: - return _classify_macro_sql_test_ctes(ctes=ctes, file_label=file_label) - case SqlTestMode.UDF: - return _classify_udf_sql_test_ctes(ctes=ctes, file_label=file_label) - case SqlTestMode.TABLE_FN: - return _classify_table_fn_sql_test_ctes(ctes=ctes, file_label=file_label) + case SqlTestMode.MACRO | SqlTestMode.UDF | SqlTestMode.TABLE_FN: + return _classify_direct_logic_sql_test_ctes( + ctes=ctes, file_label=file_label, spec=_DIRECT_LOGIC_MODE_SPECS[mode] + ) case _: raise CompileInputError(f"SQL test '{file_label}' has unsupported mode '{mode}'") -def _classify_macro_sql_test_ctes( - *, ctes: tuple[CompileSqlTestCte, ...], file_label: str +def _classify_direct_logic_sql_test_ctes( + *, ctes: tuple[CompileSqlTestCte, ...], file_label: str, spec: _DirectLogicModeSpec ) -> CompileSqlTestCtes: + mode: str = spec.mode.value authored_ctes: list[CompileSqlTestCte] = [] - macro_actual_cte: CompileSqlTestCte | None = None - macro_expected_cte: CompileSqlTestCte | None = None + actual_cte: CompileSqlTestCte | None = None + expected_cte: CompileSqlTestCte | None = None cte: CompileSqlTestCte for cte in ctes: - if cte.name == MACRO_ACTUAL_TEST_CTE_NAME: - if macro_actual_cte is not None: + if cte.name == spec.actual_cte_name: + if actual_cte is not None: raise CompileInputError( - f"SQL test '{file_label}' mode 'macro' must define exactly one " - f"{MACRO_ACTUAL_TEST_CTE_NAME} CTE" + f"SQL test '{file_label}' mode '{mode}' must define exactly one " + f"{spec.actual_cte_name} CTE" ) - macro_actual_cte = cte + actual_cte = cte continue - if cte.name == MACRO_EXPECTED_TEST_CTE_NAME: - if macro_expected_cte is not None: + if cte.name == spec.expected_cte_name: + if expected_cte is not None: raise CompileInputError( - f"SQL test '{file_label}' mode 'macro' must define exactly one " - f"{MACRO_EXPECTED_TEST_CTE_NAME} CTE" + f"SQL test '{file_label}' mode '{mode}' must define exactly one " + f"{spec.expected_cte_name} CTE" ) _validate_expected_cte_query(cte=cte, file_label=file_label, label=cte.name) - macro_expected_cte = cte + expected_cte = cte continue if _is_model_mode_cte(cte.name): raise CompileInputError( - f"SQL test '{file_label}' is mode 'macro' but defines model-test CTE '{cte.name}'" - ) - if cte.name in {UDF_ACTUAL_TEST_CTE_NAME, UDF_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'macro' but defines UDF-test CTE '{cte.name}'" + f"SQL test '{file_label}' is mode '{mode}' but defines model-test CTE '{cte.name}'" ) - if cte.name in {TABLE_FN_ACTUAL_TEST_CTE_NAME, TABLE_FN_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'macro' but defines table_fn-test CTE " - f"'{cte.name}'" - ) - if cte.name in RESERVED_SQL_TEST_CTE_NAMES: - raise CompileInputError( - f"SQL test '{file_label}' uses reserved helper CTE name '{cte.name}'" - ) - authored_ctes.append(cte) - - direct_logic_payload: CompileDirectLogicSqlTestCtes = _validate_macro_test_ctes( - authored_ctes=tuple(authored_ctes), - macro_actual_cte=macro_actual_cte, - macro_expected_cte=macro_expected_cte, - file_label=file_label, - ) - return CompileSqlTestCtes(mode=SqlTestMode.MACRO, payload=direct_logic_payload) - - -def _classify_udf_sql_test_ctes( - *, ctes: tuple[CompileSqlTestCte, ...], file_label: str -) -> CompileSqlTestCtes: - authored_ctes: list[CompileSqlTestCte] = [] - udf_actual_cte: CompileSqlTestCte | None = None - udf_expected_cte: CompileSqlTestCte | None = None - - cte: CompileSqlTestCte - for cte in ctes: - if cte.name == UDF_ACTUAL_TEST_CTE_NAME: - if udf_actual_cte is not None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'udf' must define exactly one " - f"{UDF_ACTUAL_TEST_CTE_NAME} CTE" - ) - udf_actual_cte = cte - continue - if cte.name == UDF_EXPECTED_TEST_CTE_NAME: - if udf_expected_cte is not None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'udf' must define exactly one " - f"{UDF_EXPECTED_TEST_CTE_NAME} CTE" - ) - _validate_expected_cte_query(cte=cte, file_label=file_label, label=cte.name) - udf_expected_cte = cte - continue - if _is_model_mode_cte(cte.name): - raise CompileInputError( - f"SQL test '{file_label}' is mode 'udf' but defines model-test CTE '{cte.name}'" - ) - if cte.name in {MACRO_ACTUAL_TEST_CTE_NAME, MACRO_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'udf' but defines macro-test CTE '{cte.name}'" - ) - if cte.name in {TABLE_FN_ACTUAL_TEST_CTE_NAME, TABLE_FN_EXPECTED_TEST_CTE_NAME}: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'udf' but defines table_fn-test CTE '{cte.name}'" - ) - if cte.name in RESERVED_SQL_TEST_CTE_NAMES: - raise CompileInputError( - f"SQL test '{file_label}' uses reserved helper CTE name '{cte.name}'" - ) - authored_ctes.append(cte) - - direct_logic_payload: CompileDirectLogicSqlTestCtes = _validate_udf_test_ctes( - authored_ctes=tuple(authored_ctes), - udf_actual_cte=udf_actual_cte, - udf_expected_cte=udf_expected_cte, - file_label=file_label, - ) - return CompileSqlTestCtes(mode=SqlTestMode.UDF, payload=direct_logic_payload) - - -def _classify_table_fn_sql_test_ctes( - *, ctes: tuple[CompileSqlTestCte, ...], file_label: str -) -> CompileSqlTestCtes: - authored_ctes: list[CompileSqlTestCte] = [] - table_fn_actual_cte: CompileSqlTestCte | None = None - table_fn_expected_cte: CompileSqlTestCte | None = None - - cte: CompileSqlTestCte - for cte in ctes: - if cte.name == TABLE_FN_ACTUAL_TEST_CTE_NAME: - if table_fn_actual_cte is not None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'table_fn' must define exactly one " - f"{TABLE_FN_ACTUAL_TEST_CTE_NAME} CTE" - ) - table_fn_actual_cte = cte - continue - if cte.name == TABLE_FN_EXPECTED_TEST_CTE_NAME: - if table_fn_expected_cte is not None: + foreign_names: frozenset[str] + foreign_label: str + for foreign_names, foreign_label in spec.foreign_ctes: + if cte.name in foreign_names: raise CompileInputError( - f"SQL test '{file_label}' mode 'table_fn' must define exactly one " - f"{TABLE_FN_EXPECTED_TEST_CTE_NAME} CTE" + f"SQL test '{file_label}' is mode '{mode}' but defines {foreign_label} CTE " + f"'{cte.name}'" ) - _validate_expected_cte_query(cte=cte, file_label=file_label, label=cte.name) - table_fn_expected_cte = cte - continue - if _is_model_mode_cte(cte.name): - raise CompileInputError( - f"SQL test '{file_label}' is mode 'table_fn' but defines model-test CTE " - f"'{cte.name}'" - ) - if cte.name in { - MACRO_ACTUAL_TEST_CTE_NAME, - MACRO_EXPECTED_TEST_CTE_NAME, - UDF_ACTUAL_TEST_CTE_NAME, - UDF_EXPECTED_TEST_CTE_NAME, - }: - raise CompileInputError( - f"SQL test '{file_label}' is mode 'table_fn' but defines another " - f"direct-logic CTE '{cte.name}'" - ) if cte.name in RESERVED_SQL_TEST_CTE_NAMES: raise CompileInputError( f"SQL test '{file_label}' uses reserved helper CTE name '{cte.name}'" ) authored_ctes.append(cte) - direct_logic_payload: CompileDirectLogicSqlTestCtes = _validate_table_fn_test_ctes( - authored_ctes=tuple(authored_ctes), - table_fn_actual_cte=table_fn_actual_cte, - table_fn_expected_cte=table_fn_expected_cte, - file_label=file_label, + if actual_cte is None or expected_cte is None: + raise CompileInputError( + f"SQL test '{file_label}' mode '{mode}' must define exactly one " + f"{spec.actual_cte_name} CTE and exactly one " + f"{spec.expected_cte_name} CTE" + ) + if spec.mode is SqlTestMode.MACRO: + _validate_macro_test_bodies( + helper_ctes=tuple(authored_ctes), expected_cte=expected_cte, file_label=file_label + ) + else: + _validate_call_free_direct_logic_bodies( + helper_ctes=tuple(authored_ctes), + expected_cte=expected_cte, + file_label=file_label, + mode=spec.mode, + actual_cte_name=spec.actual_cte_name, + ) + return CompileSqlTestCtes( + mode=spec.mode, + payload=CompileDirectLogicSqlTestCtes( + mode=spec.mode, + helper_ctes=tuple(authored_ctes), + actual_cte=actual_cte, + expected_cte=expected_cte, + ), ) - return CompileSqlTestCtes(mode=SqlTestMode.TABLE_FN, payload=direct_logic_payload) def _classify_model_sql_test_ctes( @@ -563,112 +549,51 @@ def _is_model_mode_cte(cte_name: str) -> bool: ) -def _validate_macro_test_ctes( +def _validate_macro_test_bodies( *, - authored_ctes: tuple[CompileSqlTestCte, ...], - macro_actual_cte: CompileSqlTestCte | None, - macro_expected_cte: CompileSqlTestCte | None, + helper_ctes: tuple[CompileSqlTestCte, ...], + expected_cte: CompileSqlTestCte, file_label: str, -) -> CompileDirectLogicSqlTestCtes: - if macro_actual_cte is None or macro_expected_cte is None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'macro' must define exactly one " - f"{MACRO_ACTUAL_TEST_CTE_NAME} CTE and exactly one " - f"{MACRO_EXPECTED_TEST_CTE_NAME} CTE" - ) +) -> None: helper_cte: CompileSqlTestCte - for helper_cte in authored_ctes: + for helper_cte in helper_ctes: macro_names: tuple[str, ...] = find_macro_call_names(helper_cte.sql_body) if macro_names: raise CompileInputError( f"SQL test '{file_label}' mode 'macro' helper CTE '{helper_cte.name}' " "must not call macros; call macros only in __macro_actual__" ) - expected_macro_names: tuple[str, ...] = find_macro_call_names(macro_expected_cte.sql_body) + expected_macro_names: tuple[str, ...] = find_macro_call_names(expected_cte.sql_body) if expected_macro_names: raise CompileInputError( f"SQL test '{file_label}' mode 'macro' CTE {MACRO_EXPECTED_TEST_CTE_NAME} " "must not call macros" ) - return CompileDirectLogicSqlTestCtes( - mode=SqlTestMode.MACRO, - helper_ctes=authored_ctes, - actual_cte=macro_actual_cte, - expected_cte=macro_expected_cte, - ) -def _validate_udf_test_ctes( +def _validate_call_free_direct_logic_bodies( *, - authored_ctes: tuple[CompileSqlTestCte, ...], - udf_actual_cte: CompileSqlTestCte | None, - udf_expected_cte: CompileSqlTestCte | None, + helper_ctes: tuple[CompileSqlTestCte, ...], + expected_cte: CompileSqlTestCte, file_label: str, -) -> CompileDirectLogicSqlTestCtes: - if udf_actual_cte is None or udf_expected_cte is None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'udf' must define exactly one " - f"{UDF_ACTUAL_TEST_CTE_NAME} CTE and exactly one " - f"{UDF_EXPECTED_TEST_CTE_NAME} CTE" - ) - helper_cte: CompileSqlTestCte - for helper_cte in authored_ctes: - _validate_no_direct_logic_calls( - sql=helper_cte.sql_body, - file_label=file_label, - mode=SqlTestMode.UDF, - cte_label=f"helper CTE '{helper_cte.name}'", - allowed_location=UDF_ACTUAL_TEST_CTE_NAME, - ) - _validate_no_direct_logic_calls( - sql=udf_expected_cte.sql_body, - file_label=file_label, - mode=SqlTestMode.UDF, - cte_label=f"CTE {UDF_EXPECTED_TEST_CTE_NAME}", - allowed_location=UDF_ACTUAL_TEST_CTE_NAME, - ) - return CompileDirectLogicSqlTestCtes( - mode=SqlTestMode.UDF, - helper_ctes=authored_ctes, - actual_cte=udf_actual_cte, - expected_cte=udf_expected_cte, - ) - - -def _validate_table_fn_test_ctes( - *, - authored_ctes: tuple[CompileSqlTestCte, ...], - table_fn_actual_cte: CompileSqlTestCte | None, - table_fn_expected_cte: CompileSqlTestCte | None, - file_label: str, -) -> CompileDirectLogicSqlTestCtes: - if table_fn_actual_cte is None or table_fn_expected_cte is None: - raise CompileInputError( - f"SQL test '{file_label}' mode 'table_fn' must define exactly one " - f"{TABLE_FN_ACTUAL_TEST_CTE_NAME} CTE and exactly one " - f"{TABLE_FN_EXPECTED_TEST_CTE_NAME} CTE" - ) + mode: SqlTestMode, + actual_cte_name: str, +) -> None: helper_cte: CompileSqlTestCte - for helper_cte in authored_ctes: + for helper_cte in helper_ctes: _validate_no_direct_logic_calls( sql=helper_cte.sql_body, file_label=file_label, - mode=SqlTestMode.TABLE_FN, + mode=mode, cte_label=f"helper CTE '{helper_cte.name}'", - allowed_location=TABLE_FN_ACTUAL_TEST_CTE_NAME, + allowed_location=actual_cte_name, ) _validate_no_direct_logic_calls( - sql=table_fn_expected_cte.sql_body, + sql=expected_cte.sql_body, file_label=file_label, - mode=SqlTestMode.TABLE_FN, - cte_label=f"CTE {TABLE_FN_EXPECTED_TEST_CTE_NAME}", - allowed_location=TABLE_FN_ACTUAL_TEST_CTE_NAME, - ) - return CompileDirectLogicSqlTestCtes( - mode=SqlTestMode.TABLE_FN, - helper_ctes=authored_ctes, - actual_cte=table_fn_actual_cte, - expected_cte=table_fn_expected_cte, + mode=mode, + cte_label=f"CTE {expected_cte.name}", + allowed_location=actual_cte_name, ) @@ -754,29 +679,32 @@ def _validate_expected_cte_query( branch_column_names: tuple[tuple[str, ...], ...] = _extract_expected_branch_column_names( sql=cte.sql_body, file_label=file_label, + label=label, ) first_branch_column_names: tuple[str, ...] = branch_column_names[0] branch_index: int for branch_index, column_names in enumerate(branch_column_names[1:], start=2): if column_names != first_branch_column_names: raise CompileInputError( - f"SQL test '{file_label}' must use the same __expected__ " + f"SQL test '{file_label}' must use the same {label} " f"projection names and order in every set-operation branch; branch {branch_index} " "does not match branch 1" ) def _extract_expected_branch_column_names( - *, sql: str, file_label: str + *, sql: str, file_label: str, label: str ) -> tuple[tuple[str, ...], ...]: sql_analysis_column_names: tuple[tuple[str, ...], ...] | None = ( - extract_expected_branch_column_names_with_sql_analysis(sql=sql, file_label=file_label) + extract_expected_branch_column_names_with_sql_analysis( + sql=sql, file_label=file_label, label=label + ) ) if sql_analysis_column_names is not None: return sql_analysis_column_names branches: tuple[str, ...] = _split_set_operation_branches(sql) return tuple( - _extract_expected_select_column_names(branch_sql=branch, file_label=file_label) + _extract_expected_select_column_names(branch_sql=branch, file_label=file_label, label=label) for branch in branches ) @@ -784,38 +712,25 @@ def _extract_expected_branch_column_names( def _split_set_operation_branches(sql: str) -> tuple[str, ...]: branches: list[str] = [] branch_start: int = 0 - index: int = 0 - depth: int = 0 - while index < len(sql): - if sql.startswith("--", index): - index = skip_line_comment(sql=sql, start=index) - continue - if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - index += 1 - continue - if sql[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - index += 1 + resume: int = 0 + index: int + depth: int + for index, depth in iter_code_positions(sql=sql, context=_CONTEXT): + if index < resume or depth != 0: continue union_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="UNION") - if depth == 0 and union_end is not None: - branch_sql: str = sql[branch_start:index].strip() - if branch_sql: - branches.append(branch_sql) - index = _skip_ignorable(sql=sql, start=union_end) - all_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="ALL") - if all_end is not None: - index = _skip_ignorable(sql=sql, start=all_end) - branch_start = index + if union_end is None: continue - index += 1 + branch_sql: str = sql[branch_start:index].strip() + if branch_sql: + branches.append(branch_sql) + resume = _skip_ignorable(sql=sql, start=union_end) + quantifier_end: int | None = _try_consume_keyword( + sql=sql, start=resume, keyword="ALL" + ) or _try_consume_keyword(sql=sql, start=resume, keyword="DISTINCT") + if quantifier_end is not None: + resume = _skip_ignorable(sql=sql, start=quantifier_end) + branch_start = resume final_branch_sql: str = sql[branch_start:].strip() if final_branch_sql: @@ -823,12 +738,14 @@ def _split_set_operation_branches(sql: str) -> tuple[str, ...]: return tuple(branches) -def _extract_expected_select_column_names(*, branch_sql: str, file_label: str) -> tuple[str, ...]: +def _extract_expected_select_column_names( + *, branch_sql: str, file_label: str, label: str +) -> tuple[str, ...]: index: int = _skip_ignorable(sql=branch_sql, start=0) select_end: int | None = _try_consume_keyword(sql=branch_sql, start=index, keyword="SELECT") if select_end is None: raise CompileInputError( - f"SQL test '{file_label}' must define each __expected__ set-operation " + f"SQL test '{file_label}' must define each {label} set-operation " "branch as a SELECT query" ) select_list_end: int = _find_select_list_end(sql=branch_sql, start=select_end) @@ -836,66 +753,37 @@ def _extract_expected_select_column_names(*, branch_sql: str, file_label: str) - expressions: tuple[str, ...] = _split_top_level_commas(raw_select_list) if not expressions: raise CompileInputError( - f"SQL test '{file_label}' must project at least one column in __expected__" + f"SQL test '{file_label}' must project at least one column in {label}" ) return tuple( - _extract_expected_projection_name(expression=expression, file_label=file_label) + _extract_expected_projection_name(expression=expression, file_label=file_label, label=label) for expression in expressions ) def _find_select_list_end(*, sql: str, start: int) -> int: - index: int = start - depth: int = 0 - while index < len(sql): - if sql.startswith("--", index): - index = skip_line_comment(sql=sql, start=index) - continue - if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - index += 1 - continue - if sql[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - index += 1 - continue - if depth == 0 and _try_consume_keyword(sql=sql, start=index, keyword="FROM") is not None: - return index - index += 1 + index: int + depth: int + for index, depth in iter_code_positions(sql=sql[start:], context=_CONTEXT): + if ( + depth == 0 + and _try_consume_keyword(sql=sql, start=start + index, keyword="FROM") is not None + ): + return start + index return len(sql) def _split_top_level_commas(raw_value: str) -> tuple[str, ...]: values: list[str] = [] value_start: int = 0 - index: int = 0 - depth: int = 0 - while index < len(raw_value): - if raw_value.startswith("--", index): - index = skip_line_comment(sql=raw_value, start=index) - continue - if raw_value.startswith("/*", index): - index = skip_block_comment(sql=raw_value, start=index, context=_CONTEXT) - continue - if raw_value[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=raw_value, start=index, context=_CONTEXT) - continue - if raw_value[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - elif raw_value[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - elif raw_value[index] == SQL_ARGUMENT_SEPARATOR_TOKEN and depth == 0: + index: int + depth: int + for index, depth in iter_code_positions(sql=raw_value, context=_CONTEXT): + if depth == 0 and raw_value[index] == SQL_ARGUMENT_SEPARATOR_TOKEN: item: str = raw_value[value_start:index].strip() if item: values.append(item) value_start = index + 1 - index += 1 final_item: str = raw_value[value_start:].strip() if final_item: @@ -903,7 +791,7 @@ def _split_top_level_commas(raw_value: str) -> tuple[str, ...]: return tuple(values) -def _extract_expected_projection_name(*, expression: str, file_label: str) -> str: +def _extract_expected_projection_name(*, expression: str, file_label: str, label: str) -> str: alias_name: str | None = _extract_as_alias(expression) if alias_name is not None: return alias_name @@ -911,46 +799,29 @@ def _extract_expected_projection_name(*, expression: str, file_label: str) -> st if _is_simple_identifier(stripped_expression): return stripped_expression raise CompileInputError( - f"SQL test '{file_label}' must alias every non-trivial __expected__ projection" + f"SQL test '{file_label}' must alias every non-trivial {label} projection" ) def _extract_as_alias(expression: str) -> str | None: - index: int = 0 - depth: int = 0 last_alias_name: str | None = None - while index < len(expression): - if expression.startswith("--", index): - index = skip_line_comment(sql=expression, start=index) - continue - if expression.startswith("/*", index): - index = skip_block_comment(sql=expression, start=index, context=_CONTEXT) - continue - if expression[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=expression, start=index, context=_CONTEXT) - continue - if expression[index] == SQL_OPEN_PAREN_TOKEN: - depth += 1 - index += 1 - continue - if expression[index] == SQL_CLOSE_PAREN_TOKEN: - depth -= 1 - index += 1 + index: int + depth: int + for index, depth in iter_code_positions(sql=expression, context=_CONTEXT): + if depth != 0: continue as_end: int | None = _try_consume_keyword(sql=expression, start=index, keyword="AS") - if depth == 0 and as_end is not None: - alias_index: int = _skip_ignorable(sql=expression, start=as_end) - if alias_index < len(expression) and is_identifier_start(expression[alias_index]): - alias_name, alias_end = _read_identifier( - sql=expression, - start=alias_index, - file_label="projection", - ) - if not expression[alias_end:].strip(): - last_alias_name = alias_name - index = as_end + if as_end is None: continue - index += 1 + alias_index: int = _skip_ignorable(sql=expression, start=as_end) + if alias_index < len(expression) and is_identifier_start(expression[alias_index]): + alias_name, alias_end = _read_identifier( + sql=expression, + start=alias_index, + file_label="projection", + ) + if not expression[alias_end:].strip(): + last_alias_name = alias_name return last_alias_name @@ -961,33 +832,24 @@ def _is_simple_identifier(value: str) -> bool: def _contains_select_star(sql: str) -> bool: - index: int = 0 - while index < len(sql): - if sql.startswith("--", index): - index = skip_line_comment(sql=sql, start=index) - continue - if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) - continue - if sql[index] in SQL_QUOTE_TOKENS: - index = skip_quoted_text(sql=sql, start=index, context=_CONTEXT) - continue + index: int + for index, _depth in iter_code_positions(sql=sql, context=_CONTEXT): select_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="SELECT") - if select_end is not None: - value_index: int = _skip_ignorable(sql=sql, start=select_end) - if value_index < len(sql) and sql[value_index] == SQL_WILDCARD_TOKEN: - return True - index = select_end + if select_end is None: continue - index += 1 + value_index: int = _skip_ignorable(sql=sql, start=select_end) + if value_index < len(sql) and sql[value_index] == SQL_WILDCARD_TOKEN: + return True return False -def _require_prefixed_name(*, cte_name: str, prefix: str, label: str, file_label: str) -> str: +def _require_prefixed_name( + *, cte_name: str, prefix: str, label: str, file_label: str, context_label: str = _CONTEXT +) -> str: extracted_name: str = cte_name.removeprefix(prefix) if extracted_name: return extracted_name - raise CompileInputError(f"SQL test '{file_label}' must use {label} to identify a target") + raise CompileInputError(f"{context_label} '{file_label}' must use {label} to identify a target") def _validate_ceremonial_select( @@ -995,35 +857,35 @@ def _validate_ceremonial_select( sql: str, start: int, file_label: str, + context_label: str = _CONTEXT, ) -> None: - if _is_ceremonial_select_statement(sql=sql, start=start): + if _is_ceremonial_select_statement(sql=sql, start=start, context_label=context_label): return - raise CompileInputError(_ceremonial_select_error(file_label)) + raise CompileInputError( + f"{context_label} '{file_label}' must end with a ceremonial top-level `SELECT 1` " + "after its CTEs" + ) -def _is_ceremonial_select_statement(*, sql: str, start: int) -> bool: - index: int = _skip_ignorable(sql=sql, start=start) +def _is_ceremonial_select_statement(*, sql: str, start: int, context_label: str = _CONTEXT) -> bool: + index: int = _skip_ignorable(sql=sql, start=start, context_label=context_label) select_end: int | None = _try_consume_keyword(sql=sql, start=index, keyword="SELECT") if select_end is None: return False - index = _skip_ignorable(sql=sql, start=select_end) + index = _skip_ignorable(sql=sql, start=select_end, context_label=context_label) if index >= len(sql) or sql[index] != SQL_CEREMONIAL_SELECT_VALUE: return False - index = _skip_ignorable(sql=sql, start=index + 1) - return _is_statement_end(sql=sql, start=index) + index = _skip_ignorable(sql=sql, start=index + 1, context_label=context_label) + return _is_statement_end(sql=sql, start=index, context_label=context_label) -def _is_statement_end(*, sql: str, start: int) -> bool: +def _is_statement_end(*, sql: str, start: int, context_label: str = _CONTEXT) -> bool: index: int = start if index < len(sql) and sql[index] == SQL_STATEMENT_TERMINATOR_TOKEN: - index = _skip_ignorable(sql=sql, start=index + 1) + index = _skip_ignorable(sql=sql, start=index + 1, context_label=context_label) return index == len(sql) -def _ceremonial_select_error(file_label: str) -> str: - return f"SQL test '{file_label}' must end with a ceremonial top-level `SELECT 1` after its CTEs" - - def validate_independent_expected_and_assertion_ctes( *, ctes: tuple[CompileSqlTestCte | CompileSqlScenarioCte, ...], @@ -1206,16 +1068,13 @@ def _dependency_path( return None -def _consume_keyword(*, sql: str, start: int, keyword: str, file_label: str) -> int: +def _consume_keyword( + *, sql: str, start: int, keyword: str, file_label: str, context_label: str = _CONTEXT +) -> int: keyword_end: int | None = _try_consume_keyword(sql=sql, start=start, keyword=keyword) if keyword_end is not None: return keyword_end - if keyword == SQL_WITH_KEYWORD: - raise CompileInputError( - f"SQL test '{file_label}' must declare mock CTEs and one __expected__ " - "CTE before `SELECT 1`" - ) - raise CompileInputError(f"SQL test '{file_label}' expected keyword {keyword}") + raise CompileInputError(f"{context_label} '{file_label}' expected keyword {keyword}") def _try_consume_keyword(*, sql: str, start: int, keyword: str) -> int | None: @@ -1229,16 +1088,18 @@ def _try_consume_keyword(*, sql: str, start: int, keyword: str) -> int | None: return keyword_end -def _read_identifier(*, sql: str, start: int, file_label: str) -> tuple[str, int]: +def _read_identifier( + *, sql: str, start: int, file_label: str, context_label: str = _CONTEXT +) -> tuple[str, int]: if start >= len(sql) or not is_identifier_start(sql[start]): - raise CompileInputError(f"SQL test '{file_label}' expected a CTE name") + raise CompileInputError(f"{context_label} '{file_label}' expected a CTE name") index: int = start + 1 while index < len(sql) and is_identifier_character(sql[index]): index += 1 return sql[start:index], index -def _skip_ignorable(*, sql: str, start: int) -> int: +def _skip_ignorable(*, sql: str, start: int, context_label: str = _CONTEXT) -> int: index: int = start while index < len(sql): if sql[index].isspace(): @@ -1248,7 +1109,7 @@ def _skip_ignorable(*, sql: str, start: int) -> int: index = skip_line_comment(sql=sql, start=index) continue if sql.startswith("/*", index): - index = skip_block_comment(sql=sql, start=index, context=_CONTEXT) + index = skip_block_comment(sql=sql, start=index, context=context_label) continue return index return index diff --git a/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py b/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py index 7a954a3968..ebc6b82114 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py +++ b/src/sqlbuild/compiler/discovery/_helpers/filesystem/core.py @@ -5,7 +5,7 @@ import importlib.util import inspect import sys -from collections.abc import Callable +from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from importlib.machinery import ModuleSpec from pathlib import Path @@ -57,7 +57,6 @@ SchemaParseError, ) from sqlbuild.compiler.discovery.models import ( - ConstantDeclaration, DiscoveredAdapterFile, DiscoveredAssetFunction, DiscoveredAuditFactory, @@ -87,7 +86,6 @@ DiscoveredSqlTestFile, DiscoveredTaskFunction, DiscoveryFileFault, - EnumDeclaration, ) from sqlbuild.compiler.resource_names.main._validate_resource_identity import ( validate_resource_identity, @@ -127,7 +125,7 @@ ) from sqlbuild.runtime.event_exporting.models import LifecycleEventSinkDefinition from sqlbuild.runtime.observability.models import LifecycleEvent -from sqlbuild.spec.contracts.models import SchemaModelEntry, SchemaSeedEntry, SourceEntry +from sqlbuild.spec.contracts.models import SchemaModelEntry, SchemaSeedEntry _PYTHON_NODE_KIND_FOLDERS: tuple[str, ...] = ("loaders", "tasks", "assets", "checks") _PYTHON_NODE_FACTORY_FOLDERS: tuple[str, ...] = (*_PYTHON_NODE_KIND_FOLDERS, "factories") @@ -465,39 +463,14 @@ def discover_enum_files( ) -> tuple[DiscoveredEnumFile, ...]: """Discover global and scoped enum declaration files.""" - discovered_files: list[DiscoveredEnumFile] = [] - facts: _DeclarationFileFacts - for facts in _discover_declaration_file_facts( + return _discover_declaration_kind_files( project_dir=project_dir, - declaration_kind=DeclarationKind.ENUM if isolate_declaration_kind else None, - ): - if facts.declaration_kind is not DeclarationKind.ENUM: - continue - try: - contents: str = facts.file_path.read_text(encoding="utf-8") - declarations: tuple[EnumDeclaration, ...] = parse_enum_declaration_file( - contents=contents, - file_path=facts.file_path, - relative_path=facts.relative_path, - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=facts.file_path, error=error)) - continue - discovered_files.append( - DiscoveredEnumFile( - file_path=facts.file_path, - relative_path=facts.relative_path, - contents=contents, - declarations=declarations, - scope_kind=facts.scope_kind, - ownership_root=facts.ownership_root, - owning_path=facts.owning_path, - declaration_root=facts.declaration_root, - ) - ) - return tuple(discovered_files) + declaration_kind=DeclarationKind.ENUM, + isolate_declaration_kind=isolate_declaration_kind, + parse_declarations=parse_enum_declaration_file, + discovered_file_type=DiscoveredEnumFile, + on_fault=on_fault, + ) def discover_constant_files( @@ -508,39 +481,102 @@ def discover_constant_files( ) -> tuple[DiscoveredConstantFile, ...]: """Discover global and scoped constant declaration files.""" - discovered_files: list[DiscoveredConstantFile] = [] - facts: _DeclarationFileFacts - for facts in _discover_declaration_file_facts( + return _discover_declaration_kind_files( project_dir=project_dir, - declaration_kind=DeclarationKind.CONSTANT if isolate_declaration_kind else None, - ): - if facts.declaration_kind is not DeclarationKind.CONSTANT: - continue - try: - contents: str = facts.file_path.read_text(encoding="utf-8") - declarations: tuple[ConstantDeclaration, ...] = parse_constant_declaration_file( - contents=contents, - file_path=facts.file_path, - relative_path=facts.relative_path, + declaration_kind=DeclarationKind.CONSTANT, + isolate_declaration_kind=isolate_declaration_kind, + parse_declarations=parse_constant_declaration_file, + discovered_file_type=DiscoveredConstantFile, + on_fault=on_fault, + ) + + +def _discover_declaration_kind_files[DeclarationT, DiscoveredT]( + *, + project_dir: Path, + declaration_kind: DeclarationKind, + isolate_declaration_kind: bool, + parse_declarations: Callable[..., tuple[DeclarationT, ...]], + discovered_file_type: Callable[..., DiscoveredT], + on_fault: Callable[[DiscoveryFileFault], None] | None, +) -> tuple[DiscoveredT, ...]: + def parse(facts: _DeclarationFileFacts) -> DiscoveredT: + contents: str = facts.file_path.read_text(encoding="utf-8") + declarations: tuple[DeclarationT, ...] = parse_declarations( + contents=contents, + file_path=facts.file_path, + relative_path=facts.relative_path, + ) + return discovered_file_type( + file_path=facts.file_path, + relative_path=facts.relative_path, + contents=contents, + declarations=declarations, + scope_kind=facts.scope_kind, + ownership_root=facts.ownership_root, + owning_path=facts.owning_path, + declaration_root=facts.declaration_root, + ) + + return _parse_discovered_files( + project_dir=project_dir, + items=( + facts + for facts in _discover_declaration_file_facts( + project_dir=project_dir, + declaration_kind=declaration_kind if isolate_declaration_kind else None, ) + if facts.declaration_kind is declaration_kind + ), + item_path=lambda facts: facts.file_path, + parse=parse, + on_fault=on_fault, + ) + + +def _parse_discovered_files[ItemT, DiscoveredT]( + *, + project_dir: Path, + items: Iterable[ItemT], + item_path: Callable[[ItemT], Path], + parse: Callable[[ItemT], DiscoveredT], + on_fault: Callable[[DiscoveryFileFault], None] | None, +) -> tuple[DiscoveredT, ...]: + """Parse each discovered file, isolating per-file read and parse faults when requested.""" + + discovered: list[DiscoveredT] = [] + item: ItemT + for item in items: + try: + discovered.append(parse(item)) except (OSError, UnicodeError, ValueError, SyntaxError) as error: if on_fault is None: raise - on_fault(_discovery_fault(project_dir=project_dir, path=facts.file_path, error=error)) - continue - discovered_files.append( - DiscoveredConstantFile( - file_path=facts.file_path, - relative_path=facts.relative_path, - contents=contents, - declarations=declarations, - scope_kind=facts.scope_kind, - ownership_root=facts.ownership_root, - owning_path=facts.owning_path, - declaration_root=facts.declaration_root, - ) - ) - return tuple(discovered_files) + on_fault(_discovery_fault(project_dir=project_dir, path=item_path(item), error=error)) + return tuple(discovered) + + +def _parse_discovered_paths[DiscoveredT]( + *, + project_dir: Path, + file_paths: Iterable[Path], + parse: Callable[[Path], DiscoveredT], + on_fault: Callable[[DiscoveryFileFault], None] | None, +) -> tuple[DiscoveredT, ...]: + return _parse_discovered_files( + project_dir=project_dir, + items=file_paths, + item_path=lambda file_path: file_path, + parse=parse, + on_fault=on_fault, + ) + + +def _unscoped_files(*, root: Path, pattern: str, project_dir: Path) -> Iterator[Path]: + file_path: Path + for file_path in sorted(root.rglob(pattern)): + if not _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): + yield file_path def discover_model_schema_files(*, project_dir: Path) -> tuple[DiscoveredModelSchemaFile, ...]: @@ -578,31 +614,25 @@ def discover_sql_function_files( if not function_root.is_dir(): return () - discovered_function_files: list[DiscoveredSqlFunctionFile] = [] - file_path: Path - for file_path in sorted(function_root.rglob("*.sql")): - if _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): - continue - try: - contents: str = file_path.read_text(encoding="utf-8") - header_values: dict[str, object] - body_sql: str - header_values, body_sql = parse_function_sql(contents=contents, file_path=file_path) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - discovered_function_files.append( - DiscoveredSqlFunctionFile( - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - contents=contents, - header_values=header_values, - body_sql=body_sql, - ) + def parse(file_path: Path) -> DiscoveredSqlFunctionFile: + contents: str = file_path.read_text(encoding="utf-8") + header_values: dict[str, object] + body_sql: str + header_values, body_sql = parse_function_sql(contents=contents, file_path=file_path) + return DiscoveredSqlFunctionFile( + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + contents=contents, + header_values=header_values, + body_sql=body_sql, ) - return tuple(discovered_function_files) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=_unscoped_files(root=function_root, pattern="*.sql", project_dir=project_dir), + parse=parse, + on_fault=on_fault, + ) def discover_python_function_files( @@ -691,28 +721,19 @@ def discover_source_files( yaml_paths: tuple[Path, ...] = tuple( sorted(path for path in sources_root.iterdir() if path.suffix in YAML_FILE_SUFFIXES) ) - discovered_source_files: list[DiscoveredSourceFile] = [] - file_path: Path - for file_path in yaml_paths: - try: - contents: str = file_path.read_text(encoding="utf-8") - source_entries: tuple[SourceEntry, ...] = parse_sources_yml( - contents=contents, file_path=file_path - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - discovered_source_files.append( - DiscoveredSourceFile( - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - contents=contents, - source_entries=source_entries, - ) + + def parse(file_path: Path) -> DiscoveredSourceFile: + contents: str = file_path.read_text(encoding="utf-8") + return DiscoveredSourceFile( + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + contents=contents, + source_entries=parse_sources_yml(contents=contents, file_path=file_path), ) - return tuple(discovered_source_files) + + return _parse_discovered_paths( + project_dir=project_dir, file_paths=yaml_paths, parse=parse, on_fault=on_fault + ) def discover_seed_files(*, project_dir: Path) -> tuple[DiscoveredSeedFile, ...]: @@ -766,25 +787,19 @@ def discover_scenario_files( if not scenarios_root.is_dir(): return () - discovered_scenario_files: list[DiscoveredSqlScenarioFile] = [] - file_path: Path - for file_path in sorted(scenarios_root.rglob("*.sql")): - if _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): - continue - try: - contents: str = file_path.read_text(encoding="utf-8") - scenario: DiscoveredSqlScenarioFile = parse_sql_scenario_file( - contents=contents, - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - discovered_scenario_files.append(scenario) - return tuple(discovered_scenario_files) + def parse(file_path: Path) -> DiscoveredSqlScenarioFile: + return parse_sql_scenario_file( + contents=file_path.read_text(encoding="utf-8"), + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + ) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=_unscoped_files(root=scenarios_root, pattern="*.sql", project_dir=project_dir), + parse=parse, + on_fault=on_fault, + ) def discover_audit_files( @@ -796,27 +811,21 @@ def discover_audit_files( if not audits_root.is_dir(): return () - discovered_audit_files: list[DiscoveredAuditFile] = [] - file_path: Path - for file_path in sorted(audits_root.rglob("*.sql")): - if _is_in_scoped_declaration_tree(file_path=file_path, project_dir=project_dir): - continue - try: - contents: str = file_path.read_text(encoding="utf-8") - discovered_audit_files.append( - DiscoveredAuditFile( - file_path=file_path, - relative_path=_project_relative_path(path=file_path, project_dir=project_dir), - contents=contents, - blocks=parse_sql_audit_file(contents=contents, file_path=file_path), - ) - ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - return tuple(discovered_audit_files) + def parse(file_path: Path) -> DiscoveredAuditFile: + contents: str = file_path.read_text(encoding="utf-8") + return DiscoveredAuditFile( + file_path=file_path, + relative_path=_project_relative_path(path=file_path, project_dir=project_dir), + contents=contents, + blocks=parse_sql_audit_file(contents=contents, file_path=file_path), + ) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=_unscoped_files(root=audits_root, pattern="*.sql", project_dir=project_dir), + parse=parse, + on_fault=on_fault, + ) def discover_macro_files( @@ -1003,29 +1012,26 @@ def discover_sql_hook_files( if not hooks_root.is_dir(): return () - discovered_hooks: list[DiscoveredSqlHookFile] = [] - file_path: Path - for file_path in sorted(hooks_root.rglob("*.sql")): - if file_path.name.startswith("_") or _is_in_scoped_declaration_tree( - file_path=file_path, project_dir=project_dir - ): - continue - try: - relative_path: Path = _project_relative_path(path=file_path, project_dir=project_dir) - contents: str = file_path.read_text(encoding="utf-8") - discovered_hooks.append( - parse_sql_hook_file( - contents=contents, - file_path=file_path, - relative_path=relative_path, - ) + def parse(file_path: Path) -> DiscoveredSqlHookFile: + relative_path: Path = _project_relative_path(path=file_path, project_dir=project_dir) + return parse_sql_hook_file( + contents=file_path.read_text(encoding="utf-8"), + file_path=file_path, + relative_path=relative_path, + ) + + return _parse_discovered_paths( + project_dir=project_dir, + file_paths=( + file_path + for file_path in _unscoped_files( + root=hooks_root, pattern="*.sql", project_dir=project_dir ) - except (OSError, UnicodeError, ValueError, SyntaxError) as error: - if on_fault is None: - raise - on_fault(_discovery_fault(project_dir=project_dir, path=file_path, error=error)) - continue - return tuple(discovered_hooks) + if not file_path.name.startswith("_") + ), + parse=parse, + on_fault=on_fault, + ) def discover_provider_classes(*, project_dir: Path) -> tuple[DiscoveredProvider, ...]: @@ -1804,82 +1810,56 @@ def _normalize_factory_result( def _load_loader_module(*, file_path: Path, project_dir: Path) -> ModuleType: - module_name: str = "sqlbuild_project_loader_" + "_".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + return _exec_project_module( + module_name="sqlbuild_project_loader_" + + "_".join( + _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + ), + file_path=file_path, + project_dir=project_dir, + file_label="source loader", + error_type=LoaderDiscoveryError, ) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise LoaderDiscoveryError(f"Could not load source loader file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise LoaderDiscoveryError( - f"Failed to import source loader file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def _load_python_node_module(*, file_path: Path, project_dir: Path, node_folder: str) -> ModuleType: - module_name: str = "sqlbuild_project_python_node_" + "_".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + return _exec_project_module( + module_name="sqlbuild_project_python_node_" + + "_".join( + _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + ), + file_path=file_path, + project_dir=project_dir, + file_label="Python node", + error_type=PythonNodeDiscoveryError, ) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise PythonNodeDiscoveryError(f"Could not load Python node file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise PythonNodeDiscoveryError( - f"Failed to import Python node file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def _load_provider_module(*, file_path: Path, project_dir: Path) -> ModuleType: - module_name: str = ".".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts - ) - _evict_stale_project_package_modules( - root_module=module_name.split(".", maxsplit=1)[0], + return _load_project_package_module( + file_path=file_path, project_dir=project_dir, + file_label="provider", + error_type=ProviderDiscoveryError, ) - existing_module: ModuleType | None = sys.modules.get(module_name) - if existing_module is not None: - existing_file: object = getattr(existing_module, "__file__", None) - if isinstance(existing_file, str) and Path(existing_file).resolve() == file_path.resolve(): - return existing_module - sys.modules.pop(module_name, None) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise ProviderDiscoveryError(f"Could not load provider file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise ProviderDiscoveryError( - f"Failed to import provider file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def _load_sink_module(*, file_path: Path, project_dir: Path) -> ModuleType: + return _load_project_package_module( + file_path=file_path, + project_dir=project_dir, + file_label="sink", + error_type=EventExporterDiscoveryError, + ) + + +def _load_project_package_module( + *, + file_path: Path, + project_dir: Path, + file_label: str, + error_type: type[Exception], +) -> ModuleType: module_name: str = ".".join( _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts ) @@ -1893,9 +1873,26 @@ def _load_sink_module(*, file_path: Path, project_dir: Path) -> ModuleType: if isinstance(existing_file, str) and Path(existing_file).resolve() == file_path.resolve(): return existing_module sys.modules.pop(module_name, None) + return _exec_project_module( + module_name=module_name, + file_path=file_path, + project_dir=project_dir, + file_label=file_label, + error_type=error_type, + ) + + +def _exec_project_module( + *, + module_name: str, + file_path: Path, + project_dir: Path, + file_label: str, + error_type: type[Exception], +) -> ModuleType: spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) if spec is None or spec.loader is None: - raise EventExporterDiscoveryError(f"Could not load sink file {file_path}") + raise error_type(f"Could not load {file_label} file {file_path}") module: ModuleType = importlib.util.module_from_spec(spec) old_path: list[str] = list(sys.path) sys.path.insert(0, str(project_dir)) @@ -1903,8 +1900,8 @@ def _load_sink_module(*, file_path: Path, project_dir: Path) -> ModuleType: sys.modules[module_name] = module spec.loader.exec_module(module) except Exception as error: - raise EventExporterDiscoveryError( - f"Failed to import sink file {file_path.relative_to(project_dir)}: {error}" + raise error_type( + f"Failed to import {file_label} file {file_path.relative_to(project_dir)}: {error}" ) from error finally: sys.path = old_path @@ -1936,25 +1933,15 @@ def _evict_stale_project_package_modules(*, root_module: str, project_dir: Path) def _load_materialization_module(*, file_path: Path, project_dir: Path) -> ModuleType: - module_name: str = ".".join( - _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + return _exec_project_module( + module_name=".".join( + _project_relative_path(path=file_path, project_dir=project_dir).with_suffix("").parts + ), + file_path=file_path, + project_dir=project_dir, + file_label="materialization", + error_type=PythonNodeDiscoveryError, ) - spec: ModuleSpec | None = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None or spec.loader is None: - raise PythonNodeDiscoveryError(f"Could not load materialization file {file_path}") - module: ModuleType = importlib.util.module_from_spec(spec) - old_path: list[str] = list(sys.path) - sys.path.insert(0, str(project_dir)) - try: - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as error: - raise PythonNodeDiscoveryError( - f"Failed to import materialization file {file_path.relative_to(project_dir)}: {error}" - ) from error - finally: - sys.path = old_path - return module def discover_adapter_file(*, project_dir: Path) -> DiscoveredAdapterFile | None: diff --git a/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py b/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py index 2175e6b6ee..81ced7e510 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py +++ b/src/sqlbuild/compiler/discovery/_helpers/integrations/loaders.py @@ -5,8 +5,11 @@ from pathlib import Path from typing import cast -from sqlbuild.compiler.discovery._helpers.yml.primitives import ( +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( optional_non_empty_string, +) +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( require_non_empty_string, ) from sqlbuild.compiler.discovery.constants import ( @@ -191,8 +194,12 @@ def _parse_dlt_source_group( if not isinstance(raw_config, dict): raise SourceParseError(f"{file_path} dlt source config must be a mapping") group_config: dict[str, object] = cast(dict[str, object], raw_config) - destination_config: dict[str, object] = _optional_mapping_value( - entry=group, key="destination", file_path=file_path, label="dlt source" + destination_config: dict[str, object] = optional_mapping( + entry=group, + key="destination", + file_path=file_path, + label="dlt source", + error_class=SourceParseError, ) _validate_dlt_destination_config(config=destination_config, file_path=file_path) group_schema: str | None = optional_non_empty_string( @@ -274,8 +281,12 @@ def _parse_dlt_resource_entry( write_disposition=write_disposition, primary_key=resource.get("primary_key"), merge_key=resource.get("merge_key"), - incremental=_optional_mapping_value( - entry=resource, key="incremental", file_path=file_path, label=f"dlt resource '{name}'" + incremental=optional_mapping( + entry=resource, + key="incremental", + file_path=file_path, + label=f"dlt resource '{name}'", + error_class=SourceParseError, ), ) return SourceEntry( @@ -372,14 +383,3 @@ def _raw_dlt_resource_config(*, source_type: str, resource: dict[str, object]) - if source_type == DLT_SOURCE_TYPE_REST_API: return {"endpoint": resource["endpoint"]} return {key: value for key, value in resource.items() if key not in excluded} - - -def _optional_mapping_value( - *, entry: dict[str, object], key: str, file_path: Path, label: str -) -> dict[str, object]: - value: object | None = entry.get(key) - if value is None: - return {} - if not isinstance(value, dict): - raise SourceParseError(f"{file_path} {label} '{key}' must be a mapping") - return cast(dict[str, object], value) diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py b/src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py new file mode 100644 index 0000000000..e92f5e5db8 --- /dev/null +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/hook_entries.py @@ -0,0 +1,53 @@ +"""Serialization of resolved model lifecycle hook entries.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry + + +def serialize_hook_entries_impl( + *, + value: object, + sql_fields: tuple[str, ...], + python_hook_fields: Mapping[str, Mapping[str, object]], +) -> list[dict[str, object]]: + """Serialize resolved hook entries, keeping only the requested SQL hook fields.""" + + if not isinstance(value, list | tuple): + return [] + hooks: list[dict[str, object]] = [] + entry: object + for entry in value: + if isinstance(entry, SqlHookEntry): + hooks.append(_serialize_sql_hook(entry=entry, sql_fields=sql_fields)) + elif isinstance(entry, PythonHookEntry): + python_hook: dict[str, object] = { + "type": "python", + "name": entry.name, + "kwargs": entry.kwargs, + } + python_hook.update(python_hook_fields.get(entry.name, {})) + hooks.append(python_hook) + return hooks + + +def _serialize_sql_hook(*, entry: SqlHookEntry, sql_fields: tuple[str, ...]) -> dict[str, object]: + field_values: dict[str, object | None] = { + "statement": entry.statement, + "name": entry.name, + "relative_path": ( + entry.relative_path.as_posix() if entry.relative_path is not None else None + ), + "definition_sql": entry.definition_sql, + "kwargs": entry.kwargs, + "description": entry.description, + } + hook: dict[str, object] = {"type": "sql"} + field_name: str + for field_name in sql_fields: + field_value: object | None = field_values[field_name] + if field_value is not None: + hook[field_name] = field_value + return hook diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py b/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py index 35661dcfcf..e42eb709bc 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/schema_columns.py @@ -6,7 +6,9 @@ from pathlib import Path from typing import cast -from sqlbuild.compiler.auditing.main._parse_audit_instance import parse_audit_instance +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances +from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool +from sqlbuild.compiler.authored_values.main._optional_named_string import optional_named_string from sqlbuild.compiler.compile.constants import NOT_NULL_AUDIT_NAME from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.discovery.exceptions import DeclarationParseError @@ -67,21 +69,23 @@ def parse_schema_columns( f"{file_path} {label} column '{raw_column_name}' has unknown metadata keys: " f"{', '.join(sorted(unknown_keys))}" ) - nullable: bool | None = _optional_bool( + nullable: bool | None = optional_named_bool( raw_value=column_metadata.get("nullable"), file_path=file_path, label=f"{label} column '{raw_column_name}'", key="nullable", error_class=error_class, + default=None, ) column_location: SourceLocation | None = locations.get(raw_column_name) audits: tuple[SchemaAuditInstance, ...] = tuple( replace(audit, location=column_location) - for audit in _parse_audits( + for audit in parse_audit_instances( raw_audits=column_metadata.get("audits"), file_path=file_path, label=f"{label} column '{raw_column_name}'", error_class=error_class, + null_as_empty=True, ) ) if nullable is True and any( @@ -101,7 +105,7 @@ def parse_schema_columns( parsed_columns.append( SchemaColumn( name=raw_column_name, - type=_optional_string( + type=optional_named_string( raw_value=column_metadata.get("type"), file_path=file_path, label=f"{label} column '{raw_column_name}'", @@ -109,7 +113,7 @@ def parse_schema_columns( error_class=error_class, ), nullable=nullable, - description=_optional_string( + description=optional_named_string( raw_value=column_metadata.get("description"), file_path=file_path, label=f"{label} column '{raw_column_name}'", @@ -121,51 +125,3 @@ def parse_schema_columns( ) ) return tuple(parsed_columns) - - -def _parse_audits( - *, raw_audits: object | None, file_path: Path, label: str, error_class: _SchemaColumnParseError -) -> tuple[SchemaAuditInstance, ...]: - if raw_audits is None: - return () - if not isinstance(raw_audits, list): - raise error_class(f"{file_path} {label} audits must be a list") - return tuple( - parse_audit_instance( - raw_audit=raw_audit, - file_path=file_path, - label=label, - error_class=error_class, - ) - for raw_audit in raw_audits - ) - - -def _optional_string( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: _SchemaColumnParseError, -) -> str | None: - if raw_value is None: - return None - if not isinstance(raw_value, str) or not raw_value.strip(): - raise error_class(f"{file_path} {label} '{key}' must be a non-empty string") - return raw_value - - -def _optional_bool( - *, - raw_value: object | None, - file_path: Path, - label: str, - key: str, - error_class: _SchemaColumnParseError, -) -> bool | None: - if raw_value is None: - return None - if not isinstance(raw_value, bool): - raise error_class(f"{file_path} {label} '{key}' must be a boolean") - return raw_value diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py b/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py index 45447c5501..ccbd79ac00 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/schema.py @@ -8,12 +8,14 @@ import yaml from yaml import YAMLError -from sqlbuild.compiler.discovery._helpers.yml.primitives import ( - optional_bool, - optional_mapping, +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances +from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( optional_non_empty_string, - optional_string_tuple, - parse_audit_instances, +) +from sqlbuild.compiler.authored_values.main._optional_string_tuple import optional_string_tuple +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( require_non_empty_string, ) from sqlbuild.compiler.discovery.constants import NOT_NULL_AUDIT_NAME, SEEDS_DIRECTORY_NAME @@ -158,7 +160,11 @@ def _parse_model_entry(*, entry: dict[str, object], file_path: Path) -> SchemaMo ), columns=_parse_columns(entry=entry, file_path=file_path, label="model"), audits=parse_audit_instances( - entry=entry, file_path=file_path, label="model", error_class=SchemaParseError + raw_audits=entry.get("audits", []), + file_path=file_path, + label="model", + error_class=SchemaParseError, + null_as_empty=False, ), tags=optional_string_tuple( entry=entry, @@ -328,10 +334,11 @@ def _parse_columns( error_class=SchemaParseError, ) audits: tuple[SchemaAuditInstance, ...] = parse_audit_instances( - entry=column, + raw_audits=column.get("audits", []), file_path=file_path, label=column_label, error_class=SchemaParseError, + null_as_empty=False, ) _validate_nullable_audits( file_path=file_path, diff --git a/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py b/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py index ebfd0fe8cb..35ac84ddbc 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py +++ b/src/sqlbuild/compiler/discovery/_helpers/yml/sources.py @@ -8,18 +8,20 @@ import yaml from yaml import YAMLError +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances +from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( + optional_non_empty_string, +) +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( + require_non_empty_string, +) from sqlbuild.compiler.discovery._helpers.integrations.loaders import ( integration_loader_name, parse_dlt_sources, parse_source_integration_loader, ) -from sqlbuild.compiler.discovery._helpers.yml.primitives import ( - optional_bool, - optional_mapping, - optional_non_empty_string, - parse_audit_instances, - require_non_empty_string, -) from sqlbuild.compiler.discovery.constants import ( NOT_NULL_AUDIT_NAME, SOURCE_AGE_POLICY_CONFIG_KEY, @@ -208,7 +210,11 @@ def _parse_source_entry(*, entry: dict[str, object], file_path: Path) -> SourceE ), columns=columns, audits=parse_audit_instances( - entry=entry, file_path=file_path, label="source", error_class=SourceParseError + raw_audits=entry.get("audits", []), + file_path=file_path, + label="source", + error_class=SourceParseError, + null_as_empty=False, ), ) _validate_source_entry(entry=source_entry, file_path=file_path) @@ -589,10 +595,11 @@ def _parse_columns(*, entry: dict[str, object], file_path: Path) -> tuple[Source error_class=SourceParseError, ) audits: tuple[SchemaAuditInstance, ...] = parse_audit_instances( - entry=column, + raw_audits=column.get("audits", []), file_path=file_path, label=column_label, error_class=SourceParseError, + null_as_empty=False, ) _validate_nullable_audits( file_path=file_path, diff --git a/src/sqlbuild/compiler/discovery/constants.py b/src/sqlbuild/compiler/discovery/constants.py index c91382938c..34d287fe32 100644 --- a/src/sqlbuild/compiler/discovery/constants.py +++ b/src/sqlbuild/compiler/discovery/constants.py @@ -62,3 +62,13 @@ ) NOT_NULL_AUDIT_NAME: str = "not_null" + +SQL_HOOK_OUTPUT_FIELDS: tuple[str, ...] = ( + "statement", + "name", + "relative_path", + "definition_sql", + "kwargs", + "description", +) +SQL_HOOK_IDENTITY_FIELDS: tuple[str, ...] = ("statement", "name", "definition_sql", "kwargs") diff --git a/src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py b/src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py new file mode 100644 index 0000000000..dcf654f9c1 --- /dev/null +++ b/src/sqlbuild/compiler/discovery/main/serialize_hook_entries.py @@ -0,0 +1,20 @@ +"""Model lifecycle hook serialization entrypoint.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from sqlbuild.compiler.discovery._helpers.sql.hook_entries import serialize_hook_entries_impl + + +def serialize_hook_entries( + *, + value: object, + sql_fields: tuple[str, ...], + python_hook_fields: Mapping[str, Mapping[str, object]] | None = None, +) -> list[dict[str, object]]: + """Serialize resolved hook entries with an explicit SQL hook field allowlist.""" + + return serialize_hook_entries_impl( + value=value, sql_fields=sql_fields, python_hook_fields=python_hook_fields or {} + ) diff --git a/src/sqlbuild/compiler/graph/_helpers/algorithms.py b/src/sqlbuild/compiler/graph/_helpers/algorithms.py index e6bb069987..5f9fb1d266 100644 --- a/src/sqlbuild/compiler/graph/_helpers/algorithms.py +++ b/src/sqlbuild/compiler/graph/_helpers/algorithms.py @@ -2,7 +2,8 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections import deque +from collections.abc import Callable, Iterable, Mapping from typing import Any @@ -30,12 +31,27 @@ def transitive_closure_impl[K]( ) -> frozenset[K]: """Return all graph keys reachable from a starting key.""" - visited: set[K] = set() - frontier: list[tuple[K, int]] = [(start, 0)] + return transitive_closure_many_impl( + starts=(start,), edges=edges, include_starts=False, max_depth=max_depth + ) + + +def transitive_closure_many_impl[K]( + *, + starts: Iterable[K], + edges: Mapping[K, Iterable[K]], + include_starts: bool, + max_depth: int | None = None, +) -> frozenset[K]: + """Return graph keys within an optional edge distance of any starting key.""" + + start_keys: tuple[K, ...] = tuple(starts) + visited: set[K] = set(start_keys) if include_starts else set() + frontier: deque[tuple[K, int]] = deque((start_key, 0) for start_key in start_keys) while frontier: current: K depth: int - current, depth = frontier.pop() + current, depth = frontier.popleft() if max_depth is not None and depth >= max_depth: continue neighbor: K diff --git a/src/sqlbuild/compiler/graph/main/transitive_closure_many.py b/src/sqlbuild/compiler/graph/main/transitive_closure_many.py new file mode 100644 index 0000000000..fc4025dc0d --- /dev/null +++ b/src/sqlbuild/compiler/graph/main/transitive_closure_many.py @@ -0,0 +1,21 @@ +"""Directed graph multi-start transitive closure entrypoint.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping + +from sqlbuild.compiler.graph._helpers.algorithms import transitive_closure_many_impl + + +def transitive_closure_many[K]( + *, + starts: Iterable[K], + edges: Mapping[K, Iterable[K]], + include_starts: bool, + max_depth: int | None = None, +) -> frozenset[K]: + """Return keys within max_depth edges of any start, plus the starts when requested.""" + + return transitive_closure_many_impl( + starts=starts, edges=edges, include_starts=include_starts, max_depth=max_depth + ) diff --git a/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py b/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py index 2ded851ef4..6acd78d8cd 100644 --- a/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py +++ b/src/sqlbuild/compiler/manifest/_helpers/model_nodes.py @@ -5,7 +5,8 @@ from pathlib import Path from sqlbuild.compiler.compile.models import CompiledModel -from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.fingerprints.main.compute_query_hash import compute_query_hash from sqlbuild.compiler.manifest._helpers.shared import ( build_columns_dict, @@ -151,33 +152,8 @@ def _extract_meta( def _serialize_hooks( *, value: object, python_hook_metadata: dict[str, dict[str, object]] ) -> list[dict[str, object]]: - if not isinstance(value, list | tuple): - return [] - hooks: list[dict[str, object]] = [] - entry: object - for entry in value: - if isinstance(entry, SqlHookEntry): - hook: dict[str, object] = { - "type": "sql", - "statement": entry.statement, - } - if entry.name is not None: - hook["name"] = entry.name - if entry.relative_path is not None: - hook["relative_path"] = entry.relative_path.as_posix() - if entry.definition_sql is not None: - hook["definition_sql"] = entry.definition_sql - if entry.kwargs is not None: - hook["kwargs"] = entry.kwargs - if entry.description is not None: - hook["description"] = entry.description - hooks.append(hook) - elif isinstance(entry, PythonHookEntry): - python_hook: dict[str, object] = { - "type": "python", - "name": entry.name, - "kwargs": entry.kwargs, - } - python_hook.update(python_hook_metadata.get(entry.name, {})) - hooks.append(python_hook) - return hooks + return serialize_hook_entries( + value=value, + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields=python_hook_metadata, + ) diff --git a/src/sqlbuild/compiler/pipeline/_helpers/diff.py b/src/sqlbuild/compiler/pipeline/_helpers/diff.py index 70f0e6a489..9edc690bd2 100644 --- a/src/sqlbuild/compiler/pipeline/_helpers/diff.py +++ b/src/sqlbuild/compiler/pipeline/_helpers/diff.py @@ -15,6 +15,7 @@ from sqlbuild.compiler.graph.main._build_lineage_upstream_deps import ( build_lineage_upstream_deps, ) +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.pipeline.constants import ( MODEL_PATH_ROOT, PATH_SELECTOR_PREFIX, @@ -222,28 +223,10 @@ def _expand_keys( downstream: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[CompiledObjectKey]: expanded: set[CompiledObjectKey] = set(keys) - key: CompiledObjectKey - for key in keys: - if upstream_requested: - expanded.update(_expand_graph(key=key, graph=upstream)) - if downstream_requested: - expanded.update(_expand_graph(key=key, graph=downstream)) + if upstream_requested: + expanded.update(transitive_closure_many(starts=keys, edges=upstream, include_starts=False)) + if downstream_requested: + expanded.update( + transitive_closure_many(starts=keys, edges=downstream, include_starts=False) + ) return frozenset(expanded) - - -def _expand_graph( - *, - key: CompiledObjectKey, - graph: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [key] - while stack: - current: CompiledObjectKey = stack.pop() - neighbor: CompiledObjectKey - for neighbor in graph.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) diff --git a/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py b/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py index 7a338b590a..39f0623cc6 100644 --- a/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py +++ b/src/sqlbuild/compiler/planner/_helpers/graph/model_closure.py @@ -4,6 +4,7 @@ from sqlbuild.compiler.compile.models import CompiledObjectKey from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many def build_downstream_model_name_closure( @@ -31,17 +32,8 @@ def _build_model_name_closure( start_keys: tuple[CompiledObjectKey, ...], deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], ) -> frozenset[str]: - model_names: set[str] = set() - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = list(start_keys) - while stack: - current: CompiledObjectKey = stack.pop() - if current in visited: - continue - visited.add(current) - if current.resource_type == CompiledResourceType.MODEL: - model_names.add(current.name) - neighbor: CompiledObjectKey - for neighbor in deps.get(current, ()): # pragma: no branch - stack.append(neighbor) - return frozenset(model_names) + return frozenset( + key.name + for key in transitive_closure_many(starts=start_keys, edges=deps, include_starts=True) + if key.resource_type == CompiledResourceType.MODEL + ) diff --git a/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py b/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py index 3cb2578ba0..2e3218ca50 100644 --- a/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py +++ b/src/sqlbuild/compiler/planner/_helpers/graph/source_load_nodes.py @@ -4,6 +4,7 @@ from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner._helpers.graph.loader_dag import ( build_intermediate_source_map, build_upstream_intermediate_source_map, @@ -22,7 +23,7 @@ def build_source_load_map( """Return source entries available to source-load and source-read planning.""" relevant_keys: frozenset[CompiledObjectKey] | None = ( - _upstream_closure(selected_keys=selected_keys, upstream_deps=upstream_deps) + transitive_closure_many(starts=selected_keys, edges=upstream_deps, include_starts=True) if upstream_deps is not None else None ) @@ -38,24 +39,6 @@ def build_source_load_map( return source_map -def _upstream_closure( - *, - selected_keys: frozenset[CompiledObjectKey], - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - """Return an adapter-neutral dependency closure for selected metadata reads.""" - - closure: set[CompiledObjectKey] = set(selected_keys) - pending: list[CompiledObjectKey] = list(selected_keys) - while pending: - key: CompiledObjectKey = pending.pop() - for upstream_key in upstream_deps.get(key, ()): - if upstream_key not in closure: - closure.add(upstream_key) - pending.append(upstream_key) - return frozenset(closure) - - def build_source_load_entries( *, execution_order: tuple[CompiledObjectKey, ...], diff --git a/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py b/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py index cfc3097ea0..9897711956 100644 --- a/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py +++ b/src/sqlbuild/compiler/planner/_helpers/output/audit_scheduling.py @@ -15,7 +15,7 @@ AttachedAuditTargetKind, CompiledResourceType, ) -from sqlbuild.compiler.planner._helpers.graph.core import expand_downstream +from sqlbuild.compiler.planner._helpers.graph.core import expand_downstream, expand_upstream from sqlbuild.compiler.planner.exceptions import PlannerInputError from sqlbuild.compiler.planner.types import MaterializationType from sqlbuild.compiler.references.types import SqlReferenceKind @@ -110,8 +110,8 @@ def _model_lifecycle_attachment_is_safe( attached_key: CompiledObjectKey = CompiledObjectKey( resource_type=CompiledResourceType.MODEL, name=audit.attached_target_name ) - attached_upstream: frozenset[CompiledObjectKey] = _expand_upstream( - key=attached_key, upstream_deps=upstream_deps + attached_upstream: frozenset[CompiledObjectKey] = expand_upstream( + key=attached_key, upstream=upstream_deps ) dep_key: CompiledObjectKey @@ -211,23 +211,3 @@ def _find_single_safe_latest_owner( return deepest.name return None - - -def _expand_upstream( - *, - key: CompiledObjectKey, - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - """Return all transitive upstream keys reachable from the given key.""" - - visited: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [key] - while stack: - current: CompiledObjectKey = stack.pop() - neighbor: CompiledObjectKey - for neighbor in upstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) diff --git a/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py b/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py index c4628c9172..333af898ce 100644 --- a/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py +++ b/src/sqlbuild/compiler/planner/_helpers/planning/scopes.py @@ -5,6 +5,7 @@ from dataclasses import replace from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner._helpers.graph.scope import build_planner_scope from sqlbuild.compiler.planner.models import ( PlannerPolicies, @@ -35,9 +36,10 @@ def resolve_planner_scopes( exclude=(), auto_load_sources=policies.auto_load_sources, ) - stale_warning_keys: frozenset[CompiledObjectKey] = _upstream_closure( - selected_keys=selected_scope.selected_keys, - upstream_deps=full_scope.upstream_deps, + stale_warning_keys: frozenset[CompiledObjectKey] = transitive_closure_many( + starts=selected_scope.selected_keys, + edges=full_scope.upstream_deps, + include_starts=True, ) stale_warning_scope: PlannerScope = replace( full_scope, @@ -60,21 +62,3 @@ def resolve_planner_scopes( stale_warning_scope=stale_warning_scope, inspection_scope=selected_scope, ) - - -def _upstream_closure( - *, - selected_keys: frozenset[CompiledObjectKey], - upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]], -) -> frozenset[CompiledObjectKey]: - closure: set[CompiledObjectKey] = set(selected_keys) - pending: list[CompiledObjectKey] = list(selected_keys) - while pending: - key: CompiledObjectKey = pending.pop() - upstream_key: CompiledObjectKey - for upstream_key in upstream_deps.get(key, ()): - if upstream_key in closure: - continue - closure.add(upstream_key) - pending.append(upstream_key) - return frozenset(closure) diff --git a/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py b/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py index 45c5b2af52..ebc3905442 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py @@ -32,6 +32,7 @@ ) from sqlbuild.compiler.fingerprints.main.read import read_latest_fingerprints from sqlbuild.compiler.fingerprints.models import Fingerprint, FingerprintSet +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner._helpers.graph.buildability import check_buildability from sqlbuild.compiler.planner._helpers.graph.core import build_execution_upstream_deps from sqlbuild.compiler.planner._helpers.graph.loader_dag import ( @@ -291,16 +292,7 @@ def _relevant_state_keys( upstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = ( build_execution_upstream_deps(project) ) - closure: set[CompiledObjectKey] = set(selected_keys) - pending: list[CompiledObjectKey] = list(selected_keys) - while pending: - key: CompiledObjectKey = pending.pop() - upstream_key: CompiledObjectKey - for upstream_key in upstream_deps.get(key, ()): - if upstream_key not in closure: - closure.add(upstream_key) - pending.append(upstream_key) - return frozenset(closure) + return transitive_closure_many(starts=selected_keys, edges=upstream_deps, include_starts=True) def _resolve_database( diff --git a/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py b/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py index 2f8de0cdd3..ad8445d44f 100644 --- a/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py +++ b/src/sqlbuild/compiler/planner/main/identity/version_identity_model_metadata.py @@ -6,7 +6,8 @@ from sqlbuild.compiler.compile.models import CompiledModel from sqlbuild.compiler.compile.types import CompiledResourceType -from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from sqlbuild.compiler.discovery.constants import SQL_HOOK_IDENTITY_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries from sqlbuild.compiler.planner._helpers.identity.model_metadata import contract_output_signature from sqlbuild.compiler.planner.constants import ( MODEL_CUSTOM_CONFIG_KEY, @@ -77,35 +78,11 @@ def _model_execution_signature( def _hook_execution_signature( *, value: object, hook_version_hashes: dict[str, str] ) -> list[dict[str, object]]: - if not isinstance(value, list | tuple): - return [] - hooks: list[dict[str, object]] = [] - entry: object - for entry in value: - if isinstance(entry, SqlHookEntry): - hook: dict[str, object] = { - "type": "sql", - "statement": entry.statement, - } - if entry.name is not None: - hook["name"] = entry.name - if entry.relative_path is not None: - hook["relative_path"] = entry.relative_path.as_posix() - if entry.definition_sql is not None: - hook["definition_sql"] = entry.definition_sql - if entry.kwargs is not None: - hook["kwargs"] = entry.kwargs - if entry.description is not None: - hook["description"] = entry.description - hooks.append(hook) - elif isinstance(entry, PythonHookEntry): - python_hook: dict[str, object] = { - "type": "python", - "name": entry.name, - "kwargs": entry.kwargs, - } - version_hash: str | None = hook_version_hashes.get(entry.name) - if version_hash is not None: - python_hook["version_hash"] = version_hash - hooks.append(python_hook) - return hooks + return serialize_hook_entries( + value=value, + sql_fields=SQL_HOOK_IDENTITY_FIELDS, + python_hook_fields={ + hook_name: {"version_hash": version_hash} + for hook_name, version_hash in hook_version_hashes.items() + }, + ) diff --git a/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py b/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py index f18448b7a1..fbe89b7bd4 100644 --- a/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py +++ b/src/sqlbuild/compiler/python_nodes/_helpers/inventory.py @@ -4,7 +4,9 @@ import inspect from collections.abc import Callable -from typing import get_type_hints +from dataclasses import replace +from pathlib import Path +from typing import Protocol, get_type_hints from sqlbuild.compiler.discovery.models import ( DiscoveredAssetFunction, @@ -174,34 +176,45 @@ def build_python_node_path_index(*, nodes: tuple[DiscoveredPythonNode, ...]) -> return {node.name: node.relative_path.parent.as_posix() for node in nodes} +class _DiscoveredPythonFunction(Protocol): + @property + def file_path(self) -> Path: ... + + @property + def relative_path(self) -> Path: ... + + @property + def name(self) -> str: ... + + @property + def function(self) -> Callable[..., object]: ... + + @property + def depends_on(self) -> tuple[Callable[..., object] | SqlResourceRef, ...]: ... + + +class _DocumentedPythonFunction(_DiscoveredPythonFunction, Protocol): + @property + def tags(self) -> tuple[str, ...]: ... + + @property + def group(self) -> str | None: ... + + @property + def description(self) -> str | None: ... + + @property + def meta(self) -> dict[str, object] | None: ... + + def _build_loader_node( *, loader: DiscoveredLoaderFunction, provider_by_name: dict[str, DiscoveredProvider] ) -> DiscoveredPythonNode: - return DiscoveredPythonNode( - kind=PythonNodeKind.LOADER, - file_path=loader.file_path, - relative_path=loader.relative_path, - name=loader.name, - function=loader.function, - depends_on=loader.depends_on, - sql_deps=_sql_deps(loader.depends_on), - provider_usages=_provider_usages( - function=loader.function, provider_by_name=provider_by_name - ), - loader=DiscoveredPythonLoaderMetadata( - destination=loader.destination, - write_strategy=loader.write_strategy, - cursor_column=loader.cursor_column, - unique_key=loader.unique_key, - columns=loader.columns, - contract=loader.contract, - connection_mode=loader.connection_mode, - ), - identity=build_python_identity( - node_type=PythonNodeKind.LOADER.value, - node_name=loader.name, - function=loader.function, - project_dir=loader.file_path.parent, + return replace( + _build_python_node( + kind=PythonNodeKind.LOADER, + discovered=loader, + provider_by_name=provider_by_name, decorator_config={ "contract": loader.contract, "cursor_column": loader.cursor_column, @@ -212,112 +225,117 @@ def _build_loader_node( else None, }, ), + loader=DiscoveredPythonLoaderMetadata( + destination=loader.destination, + write_strategy=loader.write_strategy, + cursor_column=loader.cursor_column, + unique_key=loader.unique_key, + columns=loader.columns, + contract=loader.contract, + connection_mode=loader.connection_mode, + ), ) def _build_task_node( *, task: DiscoveredTaskFunction, provider_by_name: dict[str, DiscoveredProvider] ) -> DiscoveredPythonNode: - return DiscoveredPythonNode( - kind=PythonNodeKind.TASK, - file_path=task.file_path, - relative_path=task.relative_path, - name=task.name, - function=task.function, - depends_on=task.depends_on, - sql_deps=_sql_deps(task.depends_on), - tags=task.tags, - group=task.group, - description=task.description, - meta=task.meta, - provider_usages=_provider_usages(function=task.function, provider_by_name=provider_by_name), - task=DiscoveredPythonTaskMetadata(retry=task.retry), - identity=build_python_identity( - node_type=PythonNodeKind.TASK.value, - node_name=task.name, - function=task.function, - project_dir=task.file_path.parent, - decorator_config={ - "description": task.description, - "group": task.group, - "meta": task.meta, - "tags": task.tags, - }, + return replace( + _build_documented_python_node( + kind=PythonNodeKind.TASK, + discovered=task, + provider_by_name=provider_by_name, + decorator_config={}, ), + task=DiscoveredPythonTaskMetadata(retry=task.retry), ) def _build_asset_node( *, asset: DiscoveredAssetFunction, provider_by_name: dict[str, DiscoveredProvider] ) -> DiscoveredPythonNode: - return DiscoveredPythonNode( - kind=PythonNodeKind.ASSET, - file_path=asset.file_path, - relative_path=asset.relative_path, - name=asset.name, - function=asset.function, - depends_on=asset.depends_on, - sql_deps=_sql_deps(asset.depends_on), - tags=asset.tags, - group=asset.group, - description=asset.description, - meta=asset.meta, - provider_usages=_provider_usages( - function=asset.function, provider_by_name=provider_by_name + return replace( + _build_documented_python_node( + kind=PythonNodeKind.ASSET, + discovered=asset, + provider_by_name=provider_by_name, + decorator_config={"columns": asset.columns, "column_lineage": asset.column_lineage}, ), asset=DiscoveredPythonAssetMetadata( columns=asset.columns, column_lineage=asset.column_lineage, retry=asset.retry, ), - identity=build_python_identity( - node_type=PythonNodeKind.ASSET.value, - node_name=asset.name, - function=asset.function, - project_dir=asset.file_path.parent, + ) + + +def _build_check_node( + *, check: DiscoveredCheckFunction, provider_by_name: dict[str, DiscoveredProvider] +) -> DiscoveredPythonNode: + return replace( + _build_documented_python_node( + kind=PythonNodeKind.CHECK, + discovered=check, + provider_by_name=provider_by_name, + decorator_config={"severity": check.severity.value}, + ), + check=DiscoveredPythonCheckMetadata(severity=check.severity), + ) + + +def _build_documented_python_node( + *, + kind: PythonNodeKind, + discovered: _DocumentedPythonFunction, + provider_by_name: dict[str, DiscoveredProvider], + decorator_config: dict[str, object], +) -> DiscoveredPythonNode: + return replace( + _build_python_node( + kind=kind, + discovered=discovered, + provider_by_name=provider_by_name, decorator_config={ - "columns": asset.columns, - "column_lineage": asset.column_lineage, - "description": asset.description, - "group": asset.group, - "meta": asset.meta, - "tags": asset.tags, + **decorator_config, + "description": discovered.description, + "group": discovered.group, + "meta": discovered.meta, + "tags": discovered.tags, }, ), + tags=discovered.tags, + group=discovered.group, + description=discovered.description, + meta=discovered.meta, ) -def _build_check_node( - *, check: DiscoveredCheckFunction, provider_by_name: dict[str, DiscoveredProvider] +def _build_python_node( + *, + kind: PythonNodeKind, + discovered: _DiscoveredPythonFunction, + provider_by_name: dict[str, DiscoveredProvider], + decorator_config: dict[str, object], ) -> DiscoveredPythonNode: + """Build the kind-independent node view; decorator_config is exactly the identity input.""" + return DiscoveredPythonNode( - kind=PythonNodeKind.CHECK, - file_path=check.file_path, - relative_path=check.relative_path, - name=check.name, - function=check.function, - depends_on=check.depends_on, - sql_deps=_sql_deps(check.depends_on), - tags=check.tags, - group=check.group, - description=check.description, - meta=check.meta, + kind=kind, + file_path=discovered.file_path, + relative_path=discovered.relative_path, + name=discovered.name, + function=discovered.function, + depends_on=discovered.depends_on, + sql_deps=_sql_deps(discovered.depends_on), provider_usages=_provider_usages( - function=check.function, provider_by_name=provider_by_name + function=discovered.function, provider_by_name=provider_by_name ), - check=DiscoveredPythonCheckMetadata(severity=check.severity), identity=build_python_identity( - node_type=PythonNodeKind.CHECK.value, - node_name=check.name, - function=check.function, - project_dir=check.file_path.parent, - decorator_config={ - "description": check.description, - "group": check.group, - "meta": check.meta, - "severity": check.severity.value, - "tags": check.tags, - }, + node_type=kind.value, + node_name=discovered.name, + function=discovered.function, + project_dir=discovered.file_path.parent, + decorator_config=decorator_config, ), ) diff --git a/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py b/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py index dfe07fc7eb..007fb60e72 100644 --- a/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py +++ b/src/sqlbuild/compiler/python_nodes/_helpers/selectors.py @@ -4,6 +4,7 @@ from fnmatch import fnmatchcase +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.constants import PATH_SELECTOR_EXPLICIT_ROOT_ERROR from sqlbuild.compiler.planner.exceptions import PlannerInputError from sqlbuild.compiler.planner.main.selection.selector_parse import parse_project_selector @@ -83,15 +84,7 @@ def _resolve_single(*, raw: str, graph: PythonNodeGraph) -> frozenset[str]: if not node_names: raise PlannerInputError(f"unknown Python node selector '{parsed.value}'", code="S007") - result: set[str] = set(node_names) - node_name: str - if parsed.upstream: - for node_name in node_names: - result.update(_expand_upstream(name=node_name, graph=graph)) - if parsed.downstream: - for node_name in node_names: - result.update(_expand_downstream(name=node_name, graph=graph)) - return frozenset(result) + return _expand_selected(names=node_names, parsed=parsed, graph=graph) def _resolve_path(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozenset[str]: @@ -105,15 +98,7 @@ def _resolve_path(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozense if not matched_names: raise PlannerInputError(f"no Python nodes found under path '{folder}'", code="S009") - result: set[str] = set(matched_names) - node_name: str - if parsed.upstream: - for node_name in matched_names: - result.update(_expand_upstream(name=node_name, graph=graph)) - if parsed.downstream: - for node_name in matched_names: - result.update(_expand_downstream(name=node_name, graph=graph)) - return frozenset(result) + return _expand_selected(names=matched_names, parsed=parsed, graph=graph) def _path_matches(*, indexed_folder: str, selector_folder: str) -> bool: @@ -137,15 +122,7 @@ def _resolve_tag(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozenset if not tagged_names: raise PlannerInputError(f"no Python nodes found with tag '{parsed.value}'", code="S008") - result: set[str] = set(tagged_names) - node_name: str - if parsed.upstream: - for node_name in tagged_names: - result.update(_expand_upstream(name=node_name, graph=graph)) - if parsed.downstream: - for node_name in tagged_names: - result.update(_expand_downstream(name=node_name, graph=graph)) - return frozenset(result) + return _expand_selected(names=tagged_names, parsed=parsed, graph=graph) def _lookup_node_names(*, parsed: ParsedSelector, graph: PythonNodeGraph) -> frozenset[str]: @@ -180,29 +157,16 @@ def _is_name_pattern(value: str) -> bool: return any(character in value for character in "*?[") -def _expand_upstream(*, name: str, graph: PythonNodeGraph) -> frozenset[str]: - visited: set[str] = set() - stack: list[str] = [name] - while stack: - current: str = stack.pop() - neighbor: str - for neighbor in graph.upstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) - - -def _expand_downstream(*, name: str, graph: PythonNodeGraph) -> frozenset[str]: - visited: set[str] = set() - stack: list[str] = [name] - while stack: - current: str = stack.pop() - neighbor: str - for neighbor in graph.downstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - stack.append(neighbor) - return frozenset(visited) +def _expand_selected( + *, names: frozenset[str], parsed: ParsedSelector, graph: PythonNodeGraph +) -> frozenset[str]: + result: set[str] = set(names) + if parsed.upstream: + result.update( + transitive_closure_many(starts=names, edges=graph.upstream_deps, include_starts=False) + ) + if parsed.downstream: + result.update( + transitive_closure_many(starts=names, edges=graph.downstream_deps, include_starts=False) + ) + return frozenset(result) diff --git a/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py b/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py index f6253409c0..bfd1320426 100644 --- a/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py +++ b/src/sqlbuild/compiler/sql_analysis/_helpers/scanning.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from collections.abc import Iterator from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.sql_analysis.constants import ( @@ -80,6 +81,31 @@ def find_matching_paren_impl(*, sql: str, open_paren_index: int, context: str = raise CompileInputError(f"{context} contains an unclosed parenthesis") +def iter_code_positions_impl(*, sql: str, context: str = "SQL") -> Iterator[tuple[int, int]]: + """Yield code offsets outside comments, quotes, and parentheses with their nesting depth.""" + + depth: int = 0 + index: int = 0 + while index < len(sql): + if sql.startswith("--", index): + index = skip_line_comment_impl(sql=sql, start=index) + continue + if sql.startswith("/*", index): + index = skip_block_comment_impl(sql=sql, start=index, context=context) + continue + character: str = sql[index] + if character in SQL_QUOTE_CHARACTERS: + index = skip_quoted_text_impl(sql=sql, start=index, context=context) + continue + if character == SQL_OPEN_PARENTHESIS: + depth += 1 + elif character == SQL_CLOSE_PARENTHESIS: + depth -= 1 + else: + yield index, depth + index += 1 + + def is_identifier_start_impl(character: str) -> bool: """Return whether a character can begin an SQL identifier.""" diff --git a/src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py b/src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py new file mode 100644 index 0000000000..b2a491333d --- /dev/null +++ b/src/sqlbuild/compiler/sql_analysis/main/_iter_code_positions.py @@ -0,0 +1,11 @@ +"""Top-level SQL code position scanning entrypoint.""" + +from collections.abc import Iterator + +from sqlbuild.compiler.sql_analysis._helpers.scanning import iter_code_positions_impl + + +def iter_code_positions(*, sql: str, context: str = "SQL") -> Iterator[tuple[int, int]]: + """Yield each code character offset outside comments and quotes with its parenthesis depth.""" + + return iter_code_positions_impl(sql=sql, context=context) diff --git a/src/sqlbuild/executor/build/_helpers/blocking.py b/src/sqlbuild/executor/build/_helpers/blocking.py index 56ecc74815..a26680457e 100644 --- a/src/sqlbuild/executor/build/_helpers/blocking.py +++ b/src/sqlbuild/executor/build/_helpers/blocking.py @@ -3,6 +3,7 @@ from __future__ import annotations from sqlbuild.compiler.compile.models import CompiledObjectKey +from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure def downstream_blocked_keys( @@ -13,17 +14,4 @@ def downstream_blocked_keys( ) -> frozenset[CompiledObjectKey]: """Return all selected transitive downstream keys to block.""" - blocked: set[CompiledObjectKey] = set() - stack: list[CompiledObjectKey] = [failed_key] - visited: set[CompiledObjectKey] = set() - while stack: - current: CompiledObjectKey = stack.pop() - neighbor: CompiledObjectKey - for neighbor in downstream_deps.get(current, ()): - if neighbor in visited: - continue - visited.add(neighbor) - if neighbor in selected_keys: - blocked.add(neighbor) - stack.append(neighbor) - return frozenset(blocked) + return transitive_closure(start=failed_key, edges=downstream_deps) & selected_keys diff --git a/src/sqlbuild/executor/pipeline/_helpers/scenario.py b/src/sqlbuild/executor/pipeline/_helpers/scenario.py index cb0eb50b0f..12bc3f66d3 100644 --- a/src/sqlbuild/executor/pipeline/_helpers/scenario.py +++ b/src/sqlbuild/executor/pipeline/_helpers/scenario.py @@ -4,9 +4,10 @@ import logging import time -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path -from typing import Any +from typing import Any, Protocol from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.cli.progress.classes.native_progress_projector import ( @@ -63,6 +64,18 @@ def _scenario_failure_help(exc: Exception) -> str | None: return _SCENARIO_INTERNAL_ERROR_HELP +class _ScenarioPipelineResult(Protocol): + @property + def status(self) -> ExecutionStatus: ... + + @property + def error_code(self) -> str | None: ... + + +class _ScenarioFailureResult[ResultT](Protocol): + def __call__(self, *, scenario_name: str, exc: Exception) -> ResultT: ... + + def run_scenario_test_pipeline( *, pipeline_result: CompilePipelineResult, @@ -80,73 +93,41 @@ def run_scenario_test_pipeline( ) -> tuple[ScenarioRunResult, ...]: """Execute selected scenarios from a compiled project.""" - hooks: ConnectionHooks = connection_hooks if connection_hooks is not None else ConnectionHooks() - if hooks.on_connection_start is not None: - hooks.on_connection_start(1) - start: float = time.monotonic() - try: - connection: Any = adapter.connect(connection_config) - except Exception: - if hooks.on_connection_error is not None: - hooks.on_connection_error(1, elapsed_seconds=time.monotonic() - start) - raise - if hooks.on_connection_complete is not None: - hooks.on_connection_complete(1, elapsed_seconds=time.monotonic() - start) - try: - results: list[ScenarioRunResult] = [] - scenario: CompiledSqlScenario - for scenario in scenarios: - scenario_plan: ScenarioExecutionPlan | None = None - resource_id: str = f"sql_scenario:{scenario.name}" - projector: NativeProgressProjector | None = _prepare_scenario_presentation( - scenario_name=scenario.name, - has_completion_callback=on_scenario_complete is not None, - ) - with ResourceAttemptLifecycle( - resource_id=resource_id, - resource_kind="scenario", - resource_name=scenario.name, + with _scenario_target_connection( + adapter=adapter, + connection_config=connection_config, + connection_hooks=connection_hooks, + ) as connection: + + def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: + return execute_scenario_run( + scenario_plan=scenario_plan, + adapter=adapter, + connection=connection, run_id=pipeline_result.project.run_id, - ) as lifecycle: - if on_scenario_start is not None: - on_scenario_start(scenario) - try: - scenario_plan = build_scenario_plan( - scenario=scenario, - pipeline_result=pipeline_result, - adapter=adapter, - project_name=project_name, - ) - result: ScenarioRunResult = execute_scenario_run( - scenario_plan=scenario_plan, - adapter=adapter, - connection=connection, - run_id=pipeline_result.project.run_id, - retain=retain, - ) - except Exception as exc: - result = ScenarioRunResult( - scenario_name=scenario.name, - status=ExecutionStatus.FAILED, - retained=retain, - error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), - error_help=_scenario_failure_help(exc), - error_message=error_message(exc), - ) - if result.status == ExecutionStatus.FAILED: - lifecycle.failed(error_code=result.error_code) - if projector is not None and on_scenario_complete is not None: - _ = projector.consume_resource_terminal( - resource_name=scenario.name, - resource_id=resource_id, - resource_attempt_id=lifecycle.resource_attempt_id, - ) - results.append(result) - if on_scenario_complete is not None: - on_scenario_complete(scenario, scenario_plan, result) - return tuple(results) - finally: - adapter.close(connection) + retain=retain, + ) + + def failed(*, scenario_name: str, exc: Exception) -> ScenarioRunResult: + return ScenarioRunResult( + scenario_name=scenario_name, + status=ExecutionStatus.FAILED, + retained=retain, + error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), + error_help=_scenario_failure_help(exc), + error_message=error_message(exc), + ) + + return _run_scenarios( + pipeline_result=pipeline_result, + scenarios=scenarios, + adapter=adapter, + project_name=project_name, + execute=execute, + failed=failed, + on_scenario_start=on_scenario_start, + on_scenario_complete=on_scenario_complete, + ) def run_scenario_local_test_pipeline( @@ -167,60 +148,37 @@ def run_scenario_local_test_pipeline( ) -> tuple[ScenarioRunResult, ...]: """Load selected local scenario snapshots into run-scoped DuckDB databases.""" - results: list[ScenarioRunResult] = [] - scenario: CompiledSqlScenario - for scenario in scenarios: - scenario_plan: ScenarioExecutionPlan | None = None - resource_id: str = f"sql_scenario:{scenario.name}" - projector: NativeProgressProjector | None = _prepare_scenario_presentation( - scenario_name=scenario.name, - has_completion_callback=on_scenario_complete is not None, + def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioRunResult: + return execute_local_scenario_load_only_run( + project_dir=project_dir, + scenario_plan=scenario_plan, + adapter=adapter, + strict=strict, + capture_adapter=capture_adapter, + capture_dialect=capture_dialect, ) - with ResourceAttemptLifecycle( - resource_id=resource_id, - resource_kind="scenario", - resource_name=scenario.name, - run_id=pipeline_result.project.run_id, - ) as lifecycle: - if on_scenario_start is not None: - on_scenario_start(scenario) - try: - scenario_plan = build_scenario_plan( - scenario=scenario, - pipeline_result=pipeline_result, - adapter=adapter, - project_name=project_name, - ) - result: ScenarioRunResult = execute_local_scenario_load_only_run( - project_dir=project_dir, - scenario_plan=scenario_plan, - adapter=adapter, - strict=strict, - capture_adapter=capture_adapter, - capture_dialect=capture_dialect, - ) - except Exception as exc: - result = ScenarioRunResult( - scenario_name=scenario.name, - status=ExecutionStatus.FAILED, - local_status=ScenarioLocalRunStatus.ERROR, - retained=False, - error_code=error_code(error=exc, fallback_code=SCENARIO_LOCAL_INTERNAL), - error_help=_scenario_failure_help(exc), - error_message=error_message(exc), - ) - if result.status == ExecutionStatus.FAILED: - lifecycle.failed(error_code=result.error_code) - if projector is not None and on_scenario_complete is not None: - _ = projector.consume_resource_terminal( - resource_name=scenario.name, - resource_id=resource_id, - resource_attempt_id=lifecycle.resource_attempt_id, - ) - results.append(result) - if on_scenario_complete is not None: - on_scenario_complete(scenario, scenario_plan, result) - return tuple(results) + + def failed(*, scenario_name: str, exc: Exception) -> ScenarioRunResult: + return ScenarioRunResult( + scenario_name=scenario_name, + status=ExecutionStatus.FAILED, + local_status=ScenarioLocalRunStatus.ERROR, + retained=False, + error_code=error_code(error=exc, fallback_code=SCENARIO_LOCAL_INTERNAL), + error_help=_scenario_failure_help(exc), + error_message=error_message(exc), + ) + + return _run_scenarios( + pipeline_result=pipeline_result, + scenarios=scenarios, + adapter=adapter, + project_name=project_name, + execute=execute, + failed=failed, + on_scenario_start=on_scenario_start, + on_scenario_complete=on_scenario_complete, + ) def run_scenario_capture_pipeline( @@ -241,6 +199,55 @@ def run_scenario_capture_pipeline( ) -> tuple[ScenarioSnapshotCaptureRunResult, ...]: """Capture selected scenario inputs into durable local snapshot files.""" + with _scenario_target_connection( + adapter=adapter, + connection_config=connection_config, + connection_hooks=connection_hooks, + ) as connection: + + def execute(scenario_plan: ScenarioExecutionPlan) -> ScenarioSnapshotCaptureRunResult: + return execute_scenario_snapshot_capture_run( + project_dir=project_dir, + scenario_plan=scenario_plan, + adapter=adapter, + connection=connection, + run_id=pipeline_result.project.run_id, + settings=settings, + local_type_overrides=scenario_local_type_overrides_for_dialect( + scenario_config=pipeline_result.project.scenario, + sql_analysis_dialect=adapter.sql_analysis_dialect(), + ), + ) + + def failed(*, scenario_name: str, exc: Exception) -> ScenarioSnapshotCaptureRunResult: + return ScenarioSnapshotCaptureRunResult( + scenario_name=scenario_name, + status=ExecutionStatus.FAILED, + retained=settings.retain, + error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), + error_help=_scenario_failure_help(exc), + error_message=error_message(exc), + ) + + return _run_scenarios( + pipeline_result=pipeline_result, + scenarios=scenarios, + adapter=adapter, + project_name=project_name, + execute=execute, + failed=failed, + on_scenario_start=on_scenario_start, + on_scenario_complete=on_scenario_complete, + ) + + +@contextmanager +def _scenario_target_connection( + *, + adapter: BaseAdapter, + connection_config: dict[str, object], + connection_hooks: ConnectionHooks | None, +) -> Iterator[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(1) @@ -257,69 +264,66 @@ def run_scenario_capture_pipeline( if hooks.on_connection_complete is not None: hooks.on_connection_complete(1, elapsed_seconds=time.monotonic() - start) try: - results: list[ScenarioSnapshotCaptureRunResult] = [] - scenario: CompiledSqlScenario - for scenario in scenarios: - scenario_plan: ScenarioExecutionPlan | None = None - resource_id: str = f"sql_scenario:{scenario.name}" - projector: NativeProgressProjector | None = _prepare_scenario_presentation( - scenario_name=scenario.name, - has_completion_callback=on_scenario_complete is not None, - ) - with ResourceAttemptLifecycle( - resource_id=resource_id, - resource_kind="scenario", - resource_name=scenario.name, - run_id=pipeline_result.project.run_id, - ) as lifecycle: - if on_scenario_start is not None: - on_scenario_start(scenario) - try: - scenario_plan = build_scenario_plan( - scenario=scenario, - pipeline_result=pipeline_result, - adapter=adapter, - project_name=project_name, - ) - result: ScenarioSnapshotCaptureRunResult = ( - execute_scenario_snapshot_capture_run( - project_dir=project_dir, - scenario_plan=scenario_plan, - adapter=adapter, - connection=connection, - run_id=pipeline_result.project.run_id, - settings=settings, - local_type_overrides=scenario_local_type_overrides_for_dialect( - scenario_config=pipeline_result.project.scenario, - sql_analysis_dialect=adapter.sql_analysis_dialect(), - ), - ) - ) - except Exception as exc: - result = ScenarioSnapshotCaptureRunResult( - scenario_name=scenario.name, - status=ExecutionStatus.FAILED, - retained=settings.retain, - error_code=error_code(error=exc, fallback_code=SCENARIO_EXEC_INTERNAL), - error_help=_scenario_failure_help(exc), - error_message=error_message(exc), - ) - if result.status == ExecutionStatus.FAILED: - lifecycle.failed(error_code=result.error_code) - if projector is not None and on_scenario_complete is not None: - _ = projector.consume_resource_terminal( - resource_name=scenario.name, - resource_id=resource_id, - resource_attempt_id=lifecycle.resource_attempt_id, - ) - results.append(result) - if on_scenario_complete is not None: - on_scenario_complete(scenario, scenario_plan, result) - return tuple(results) + yield connection finally: adapter.close(connection) +def _run_scenarios[ResultT: _ScenarioPipelineResult]( + *, + pipeline_result: CompilePipelineResult, + scenarios: tuple[CompiledSqlScenario, ...], + adapter: BaseAdapter, + project_name: str, + execute: Callable[[ScenarioExecutionPlan], ResultT], + failed: _ScenarioFailureResult[ResultT], + on_scenario_start: Callable[[CompiledSqlScenario], None] | None, + on_scenario_complete: Callable[ + [CompiledSqlScenario, ScenarioExecutionPlan | None, ResultT], None + ] + | None, +) -> tuple[ResultT, ...]: + results: list[ResultT] = [] + scenario: CompiledSqlScenario + for scenario in scenarios: + scenario_plan: ScenarioExecutionPlan | None = None + resource_id: str = f"sql_scenario:{scenario.name}" + projector: NativeProgressProjector | None = _prepare_scenario_presentation( + scenario_name=scenario.name, + has_completion_callback=on_scenario_complete is not None, + ) + with ResourceAttemptLifecycle( + resource_id=resource_id, + resource_kind="scenario", + resource_name=scenario.name, + run_id=pipeline_result.project.run_id, + ) as lifecycle: + if on_scenario_start is not None: + on_scenario_start(scenario) + try: + scenario_plan = build_scenario_plan( + scenario=scenario, + pipeline_result=pipeline_result, + adapter=adapter, + project_name=project_name, + ) + result: ResultT = execute(scenario_plan) + except Exception as exc: + result = failed(scenario_name=scenario.name, exc=exc) + if result.status == ExecutionStatus.FAILED: + lifecycle.failed(error_code=result.error_code) + if projector is not None and on_scenario_complete is not None: + _ = projector.consume_resource_terminal( + resource_name=scenario.name, + resource_id=resource_id, + resource_attempt_id=lifecycle.resource_attempt_id, + ) + results.append(result) + if on_scenario_complete is not None: + on_scenario_complete(scenario, scenario_plan, result) + return tuple(results) + + def _prepare_scenario_presentation( *, scenario_name: str, has_completion_callback: bool ) -> NativeProgressProjector | None: diff --git a/src/sqlbuild/executor/scenario/_helpers/capture/core.py b/src/sqlbuild/executor/scenario/_helpers/capture/core.py index 73c8897f8f..b210b1bd9c 100644 --- a/src/sqlbuild/executor/scenario/_helpers/capture/core.py +++ b/src/sqlbuild/executor/scenario/_helpers/capture/core.py @@ -8,6 +8,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.compiler.planner.models import ScenarioExecutionPlan from sqlbuild.executor.build.models import SeedExecutionResult +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details from sqlbuild.executor.scenario._helpers.lifecycle.fixtures import ( execute_scenario_fixtures, execute_scenario_seed_entries, @@ -33,6 +34,8 @@ from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle from sqlbuild.runtime.observability.main.run_scope import run_scope +_CAPTURE_FAILED_MESSAGE: str = "scenario snapshot capture failed" + def execute_scenario_snapshot_capture_steps( *, @@ -97,8 +100,10 @@ def _execute_scenario_snapshot_capture_steps( adapter=adapter, connection=connection, ) - fixture_error: str | None = _first_error(fixture_results) - if fixture_error is not None: + fixture_failure: ScenarioFailureDetails = first_failure_details( + results=fixture_results, fallback_message=_CAPTURE_FAILED_MESSAGE + ) + if fixture_failure.error_message is not None: return _finish_capture_run( scenario_plan=scenario_plan, adapter=adapter, @@ -106,11 +111,7 @@ def _execute_scenario_snapshot_capture_steps( retain=settings.retain, prepare_cleanup_result=prepare_result, fixture_results=fixture_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(fixture_results), - error_help=_first_error_help(fixture_results), - error_message=fixture_error, - ), + failure=fixture_failure, ) seed_results: tuple[SeedExecutionResult, ...] = execute_scenario_seed_entries( @@ -120,8 +121,10 @@ def _execute_scenario_snapshot_capture_steps( connection=connection, run_id=run_id, ) - seed_error: str | None = _first_error(seed_results) - if seed_error is not None: + seed_failure: ScenarioFailureDetails = first_failure_details( + results=seed_results, fallback_message=_CAPTURE_FAILED_MESSAGE + ) + if seed_failure.error_message is not None: return _finish_capture_run( scenario_plan=scenario_plan, adapter=adapter, @@ -130,11 +133,7 @@ def _execute_scenario_snapshot_capture_steps( prepare_cleanup_result=prepare_result, fixture_results=fixture_results, seed_results=seed_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(seed_results), - error_help=_first_error_help(seed_results), - error_message=seed_error, - ), + failure=seed_failure, ) capture_plan: ScenarioSnapshotCapturePlan = build_scenario_snapshot_capture_plan( @@ -228,31 +227,3 @@ def _finish_capture_run( error_help=error_help, error_message=error_message, ) - - -def _first_error(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - return str(getattr(result, "error_message", "scenario snapshot capture failed")) - return None - - -def _first_error_code(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_code: object | None = getattr(result, "error_code", None) - if isinstance(error_code, str) and error_code: - return error_code - return None - - -def _first_error_help(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_help: object | None = getattr(result, "error_help", None) - if isinstance(error_help, str) and error_help: - return error_help - return None diff --git a/src/sqlbuild/executor/scenario/_helpers/execution/run.py b/src/sqlbuild/executor/scenario/_helpers/execution/run.py index 5d7fab2665..666b58b377 100644 --- a/src/sqlbuild/executor/scenario/_helpers/execution/run.py +++ b/src/sqlbuild/executor/scenario/_helpers/execution/run.py @@ -13,6 +13,7 @@ execute_scenario_assertion_expectations, execute_scenario_expected_expectations, ) +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details from sqlbuild.executor.scenario._helpers.lifecycle.fixtures import ( execute_scenario_fixtures, execute_scenario_seed_entries, @@ -31,6 +32,8 @@ from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle +_STEP_FAILED_MESSAGE: str = "scenario step failed" + def execute_scenario_run_steps( *, @@ -91,10 +94,8 @@ def _execute_scenario_run_steps( retain=retain, prepare_cleanup_result=prepare_result, results=ScenarioStepResults(fixture_results=fixture_results), - failure=ScenarioFailureDetails( - error_code=_first_error_code(fixture_results), - error_help=_first_error_help(fixture_results), - error_message=_first_error(fixture_results), + failure=first_failure_details( + results=fixture_results, fallback_message=_STEP_FAILED_MESSAGE ), ) @@ -116,10 +117,8 @@ def _execute_scenario_run_steps( fixture_results=fixture_results, seed_results=seed_results, ), - failure=ScenarioFailureDetails( - error_code=_first_error_code(seed_results), - error_help=_first_error_help(seed_results), - error_message=_first_error(seed_results), + failure=first_failure_details( + results=seed_results, fallback_message=_STEP_FAILED_MESSAGE ), ) @@ -141,10 +140,8 @@ def _execute_scenario_run_steps( seed_results=seed_results, model_results=model_results, ), - failure=ScenarioFailureDetails( - error_code=_first_error_code(model_results), - error_help=_first_error_help(model_results), - error_message=_first_error(model_results), + failure=first_failure_details( + results=model_results, fallback_message=_STEP_FAILED_MESSAGE ), ) @@ -160,9 +157,6 @@ def _execute_scenario_run_steps( adapter=adapter, connection=connection, ) - failed_check_message: str | None = _first_error((*expected_results, *assertion_results)) - failed_check_code: str | None = _first_error_code((*expected_results, *assertion_results)) - failed_check_help: str | None = _first_error_help((*expected_results, *assertion_results)) return _finish_scenario( scenario_plan=scenario_plan, adapter=adapter, @@ -176,10 +170,9 @@ def _execute_scenario_run_steps( expected_results=expected_results, assertion_results=assertion_results, ), - failure=ScenarioFailureDetails( - error_code=failed_check_code, - error_help=failed_check_help, - error_message=failed_check_message, + failure=first_failure_details( + results=(*expected_results, *assertion_results), + fallback_message=_STEP_FAILED_MESSAGE, ), ) @@ -265,34 +258,3 @@ def _scenario_failure( def _has_failed(results: tuple[object, ...]) -> bool: return any(getattr(result, "status", None) == ExecutionStatus.FAILED for result in results) - - -def _first_error(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_message: object | None = getattr(result, "error_message", None) - if isinstance(error_message, str) and error_message: - return error_message - return "scenario step failed" - return None - - -def _first_error_code(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_code: object | None = getattr(result, "error_code", None) - if isinstance(error_code, str) and error_code: - return error_code - return None - - -def _first_error_help(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - error_help: object | None = getattr(result, "error_help", None) - if isinstance(error_help, str) and error_help: - return error_help - return None diff --git a/src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py b/src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py new file mode 100644 index 0000000000..575dc78a2f --- /dev/null +++ b/src/sqlbuild/executor/scenario/_helpers/lifecycle/failures.py @@ -0,0 +1,41 @@ +"""Failure-detail projection over scenario step results.""" + +from __future__ import annotations + +from sqlbuild.executor.scenario.models import ScenarioFailureDetails +from sqlbuild.executor.scheduling.types import ExecutionStatus + + +def first_failure_details( + *, + results: tuple[object, ...], + fallback_message: str, + fallback_code: str | None = None, +) -> ScenarioFailureDetails: + """Return the first failed step's message plus the first failed code and help present.""" + + failed: tuple[object, ...] = tuple( + result for result in results if getattr(result, "status", None) == ExecutionStatus.FAILED + ) + error_message: str | None = None + if failed: + first_message: str | None = _non_empty_text(getattr(failed[0], "error_message", None)) + error_message = first_message if first_message is not None else fallback_message + return ScenarioFailureDetails( + error_code=_first_text(results=failed, field_name="error_code") or fallback_code, + error_help=_first_text(results=failed, field_name="error_help"), + error_message=error_message, + ) + + +def _first_text(*, results: tuple[object, ...], field_name: str) -> str | None: + result: object + for result in results: + text: str | None = _non_empty_text(getattr(result, field_name, None)) + if text is not None: + return text + return None + + +def _non_empty_text(value: object) -> str | None: + return value if isinstance(value, str) and value else None diff --git a/src/sqlbuild/executor/scenario/_helpers/local/execution.py b/src/sqlbuild/executor/scenario/_helpers/local/execution.py index 21d774797d..52b73bfd98 100644 --- a/src/sqlbuild/executor/scenario/_helpers/local/execution.py +++ b/src/sqlbuild/executor/scenario/_helpers/local/execution.py @@ -30,6 +30,7 @@ execute_scenario_assertion_expectations, execute_scenario_expected_expectations, ) +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details from sqlbuild.executor.scenario._helpers.local.snapshots import ( load_scenario_snapshot_into_duckdb, local_snapshot_table_name, @@ -71,6 +72,8 @@ ) from sqlbuild.spec.contracts.models import SourceEntry +_LOCAL_STEP_FAILED_MESSAGE: str = "local scenario step failed" + def execute_local_scenario_load_only_run( *, @@ -302,10 +305,10 @@ def _execute_local_plan( duckdb_path=duckdb_path, loaded_relations=loaded_relations, function_results=function_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(function_results) or SCENARIO_LOCAL_FUNCTION_FAILED, - error_help=_first_error_help(function_results), - error_message=_first_error(function_results), + failure=first_failure_details( + results=function_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, + fallback_code=SCENARIO_LOCAL_FUNCTION_FAILED, ), ) @@ -327,10 +330,10 @@ def _execute_local_plan( loaded_relations=loaded_relations, function_results=function_results, model_results=local_model_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(local_model_results) or SCENARIO_LOCAL_MODEL_FAILED, - error_help=_first_error_help(local_model_results), - error_message=_first_error(local_model_results), + failure=first_failure_details( + results=local_model_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, + fallback_code=SCENARIO_LOCAL_MODEL_FAILED, ), ) @@ -362,10 +365,10 @@ def _execute_local_plan( model_results=model_results, expected_results=expected_results, assertion_results=assertion_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(check_results) or SCENARIO_LOCAL_MODEL_FAILED, - error_help=_first_error_help(check_results), - error_message=_first_error(check_results), + failure=first_failure_details( + results=check_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, + fallback_code=SCENARIO_LOCAL_MODEL_FAILED, ), ) if _has_failed(check_results): @@ -379,10 +382,9 @@ def _execute_local_plan( model_results=model_results, expected_results=expected_results, assertion_results=assertion_results, - failure=ScenarioFailureDetails( - error_code=_first_error_code(check_results), - error_help=_first_error_help(check_results), - error_message=_first_error(check_results), + failure=first_failure_details( + results=check_results, + fallback_message=_LOCAL_STEP_FAILED_MESSAGE, ), ) return _local_result( @@ -719,35 +721,6 @@ def _has_local_check_error(results: tuple[object, ...]) -> bool: ) -def _first_error(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - message: object | None = getattr(result, "error_message", None) - return message if isinstance(message, str) and message else "local scenario step failed" - return None - - -def _first_error_code(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - code: object | None = getattr(result, "error_code", None) - if isinstance(code, str) and code: - return code - return None - - -def _first_error_help(results: tuple[object, ...]) -> str | None: - result: object - for result in results: - if getattr(result, "status", None) == ExecutionStatus.FAILED: - help_text: object | None = getattr(result, "error_help", None) - if isinstance(help_text, str) and help_text: - return help_text - return None - - def _remove_local_duckdb_files(duckdb_path: Path) -> None: duckdb_path.unlink(missing_ok=True) duckdb_path.with_name(f"{duckdb_path.name}.wal").unlink(missing_ok=True) diff --git a/src/sqlbuild/integrations/dagster/_helpers/selection.py b/src/sqlbuild/integrations/dagster/_helpers/selection.py index 3d61e529dd..05ae1aa30f 100644 --- a/src/sqlbuild/integrations/dagster/_helpers/selection.py +++ b/src/sqlbuild/integrations/dagster/_helpers/selection.py @@ -2,10 +2,11 @@ from __future__ import annotations -from collections import deque from collections.abc import Mapping, Sequence from typing import Any +from sqlbuild.compiler.graph.main.path_nodes import path_nodes +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.main.selection.selector_parse import parse_project_selector from sqlbuild.compiler.planner.models import ParsedSelector, PathSelector from sqlbuild.compiler.planner.types import SelectorKind @@ -116,18 +117,25 @@ def _resolve_atomic( if isinstance(parsed, PathSelector): start_id: str = _node_id_for_name(nodes_by_id=nodes_by_id, name=parsed.start_name) end_id: str = _node_id_for_name(nodes_by_id=nodes_by_id, name=parsed.end_name) - matched: set[str] = _shortest_path(start_id=start_id, end_id=end_id, downstream=downstream) + matched: set[str] = _path_node_ids(start_id=start_id, end_id=end_id, downstream=downstream) if parsed.upstream: - matched.update(_expand(start_ids={start_id}, adjacency=upstream)) + matched.update( + transitive_closure_many(starts=(start_id,), edges=upstream, include_starts=True) + ) if parsed.downstream: - matched.update(_expand(start_ids={end_id}, adjacency=downstream)) + matched.update( + transitive_closure_many(starts=(end_id,), edges=downstream, include_starts=True) + ) return matched matched = _match_parsed_selector(parsed=parsed, dag=dag, nodes_by_id=nodes_by_id) + anchors: frozenset[str] = frozenset(matched) if parsed.upstream: - matched.update(_expand(start_ids=matched, adjacency=upstream)) + matched.update(transitive_closure_many(starts=anchors, edges=upstream, include_starts=True)) if parsed.downstream: - matched.update(_expand(start_ids=matched, adjacency=downstream)) + matched.update( + transitive_closure_many(starts=anchors, edges=downstream, include_starts=True) + ) return matched @@ -199,19 +207,6 @@ def _dependency_indexes( return upstream, downstream -def _expand(*, start_ids: set[str], adjacency: Mapping[str, set[str]]) -> set[str]: - expanded: set[str] = set(start_ids) - pending: list[str] = list(start_ids) - while pending: - node_id: str = pending.pop() - for adjacent_id in adjacency.get(node_id, set()): - if adjacent_id in expanded: - continue - expanded.add(adjacent_id) - pending.append(adjacent_id) - return expanded - - def _node_id_for_name(*, nodes_by_id: Mapping[str, Mapping[str, Any]], name: str) -> str: matching_ids: list[str] = [ node_id for node_id, node in nodes_by_id.items() if str(node.get("name")) == name @@ -223,16 +218,14 @@ def _node_id_for_name(*, nodes_by_id: Mapping[str, Mapping[str, Any]], name: str return matching_ids[0] -def _shortest_path(*, start_id: str, end_id: str, downstream: Mapping[str, set[str]]) -> set[str]: - pending: deque[tuple[str, tuple[str, ...]]] = deque([(start_id, (start_id,))]) - visited: set[str] = {start_id} - while pending: - node_id, path = pending.popleft() - if node_id == end_id: - return set(path) - for adjacent_id in downstream.get(node_id, set()): - if adjacent_id in visited: - continue - visited.add(adjacent_id) - pending.append((adjacent_id, (*path, adjacent_id))) - raise DagsterDagInputError(f"no SQLBuild DAG path exists between {start_id!r} and {end_id!r}") +def _path_node_ids(*, start_id: str, end_id: str, downstream: Mapping[str, set[str]]) -> set[str]: + on_path: frozenset[str] | None = path_nodes( + start=start_id, + end=end_id, + downstream={node_id: tuple(sorted(targets)) for node_id, targets in downstream.items()}, + ) + if on_path is None: + raise DagsterDagInputError( + f"no SQLBuild DAG path exists between {start_id!r} and {end_id!r}" + ) + return set(on_path) diff --git a/src/sqlbuild/virtual/planner/_helpers/planning.py b/src/sqlbuild/virtual/planner/_helpers/planning.py index 4ea8f8589f..27f96153ea 100644 --- a/src/sqlbuild/virtual/planner/_helpers/planning.py +++ b/src/sqlbuild/virtual/planner/_helpers/planning.py @@ -7,6 +7,7 @@ from sqlbuild.compiler.compile.models import CompiledObjectKey from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.pipeline.models import ProjectGraph from sqlbuild.compiler.planner.exceptions import PlannerInputError from sqlbuild.compiler.planner.main.identity.seed_identity import build_seed_identity @@ -168,21 +169,13 @@ def build_source_freshness_incomplete_model_names( incomplete_source_keys: tuple[Any, ...] = tuple( source.key for source in graph.project.sources if source.name not in source_version_hashes ) - incomplete_model_names: set[str] = set() - source_key: Any - for source_key in incomplete_source_keys: - stack: list[Any] = list(graph.downstream_deps.get(source_key, ())) - visited: set[Any] = set() - while stack: - current: Any = stack.pop() - if current in visited: - continue - visited.add(current) - if current.resource_type == CompiledResourceType.MODEL: - incomplete_model_names.add(current.name) - downstream_key: Any - for downstream_key in graph.downstream_deps.get(current, ()): # pragma: no branch - stack.append(downstream_key) + incomplete_model_names: frozenset[str] = frozenset( + key.name + for key in transitive_closure_many( + starts=incomplete_source_keys, edges=graph.downstream_deps, include_starts=False + ) + if key.resource_type == CompiledResourceType.MODEL + ) return tuple(sorted(incomplete_model_names)) diff --git a/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py b/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py index 60ef0b0b18..8de231ee47 100644 --- a/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py +++ b/src/sqlbuild/virtual/planner/_helpers/python_node_closure.py @@ -2,7 +2,7 @@ from __future__ import annotations -from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from sqlbuild.compiler.planner.models import PlanOutput from sqlbuild.compiler.python_nodes.models import DiscoveredPythonNode, PythonNodeGraph from sqlbuild.compiler.python_nodes.types import PythonNodeKind @@ -49,10 +49,9 @@ def python_upstream_closure( ) -> frozenset[str]: """Return all Python upstreams for selected Python nodes.""" - names: set[str] = set() - for node_name in selected_python_names: - names.update(transitive_closure(start=node_name, edges=python_graph.upstream_deps)) - return frozenset(names) + return transitive_closure_many( + starts=selected_python_names, edges=python_graph.upstream_deps, include_starts=False + ) def python_downstream_closure( diff --git a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py b/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py deleted file mode 100644 index 565604db6d..0000000000 --- a/src/sqlbuild/virtual/state/classes/_duckdb_conditional_publish.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Conditional DuckDB virtual-environment publication.""" - -from __future__ import annotations - -from typing import Any, cast - -from sqlbuild.virtual.state._helpers.state_storage.validation import ( - validate_conditional_virtual_environment_publication, -) -from sqlbuild.virtual.state.constants import ( - LOCK_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, -) -from sqlbuild.virtual.state.models import ( - StateLockLease, - VirtualEnvironmentCheckpointFunctionRefRecord, - VirtualEnvironmentCheckpointModelRefRecord, - VirtualEnvironmentCheckpointRecord, - VirtualEnvironmentCheckpointSeedRefRecord, - VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentRecord, -) - - -class DuckDbConditionalPublishMixin: - """Publish virtual refs in the same transaction that validates lease ownership.""" - - def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - leases: tuple[StateLockLease, ...], - checkpoint: VirtualEnvironmentCheckpointRecord | None = None, - checkpoint_refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...] = (), - checkpoint_function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - checkpoint_seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> bool: - backend: Any = cast(Any, self) - validate_conditional_virtual_environment_publication( - record=record, - refs_by_node_type=refs_by_node_type, - checkpoint=checkpoint, - checkpoint_refs=checkpoint_refs, - checkpoint_function_refs=checkpoint_function_refs, - checkpoint_seed_refs=checkpoint_seed_refs, - ) - connection.execute("BEGIN") - try: - lock_table: str = backend._qualified_name(schema=schema, table=LOCK_TABLE) - lease: StateLockLease - for lease in leases: - owned: tuple[Any, ...] | None = connection.execute( - f"UPDATE {lock_table} SET updated_at = CURRENT_TIMESTAMP " - "WHERE lock_key = ? AND owner_id = ? AND expires_at > CURRENT_TIMESTAMP " - "RETURNING lock_key", - [lease.lock_key, lease.owner_id], - ).fetchone() - if owned is None: - connection.execute("ROLLBACK") - return False - backend._upsert_virtual_environment_record( - connection=connection, schema=schema, record=record - ) - backend._replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - if checkpoint is not None: - self._insert_virtual_environment_checkpoint_rows( - connection=connection, - schema=schema, - checkpoint=checkpoint, - refs=checkpoint_refs, - function_refs=checkpoint_function_refs, - seed_refs=checkpoint_seed_refs, - ) - connection.execute("COMMIT") - return True - except BaseException: - connection.execute("ROLLBACK") - raise - - def _insert_virtual_environment_checkpoint_rows( - self, - *, - connection: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...], - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...], - ) -> None: - backend: Any = cast(Any, self) - checkpoint_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE - ) - model_ref_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE - ) - function_ref_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE - ) - seed_ref_table: str = backend._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE - ) - connection.execute( - f"INSERT INTO {checkpoint_table} " - "(checkpoint_id, virtual_environment_name, created_at) " - "VALUES (?, ?, CURRENT_TIMESTAMP)", - [checkpoint.checkpoint_id, checkpoint.virtual_environment_name], - ) - ref: VirtualEnvironmentCheckpointModelRefRecord - for ref in refs: - connection.execute( - f"INSERT INTO {model_ref_table} " - "(checkpoint_id, model_name, version_hash) VALUES (?, ?, ?)", - [ref.checkpoint_id, ref.model_name, ref.version_hash], - ) - function_ref: VirtualEnvironmentCheckpointFunctionRefRecord - for function_ref in function_refs: - connection.execute( - f"INSERT INTO {function_ref_table} " - "(checkpoint_id, function_name, version_hash) VALUES (?, ?, ?)", - [function_ref.checkpoint_id, function_ref.function_name, function_ref.version_hash], - ) - seed_ref: VirtualEnvironmentCheckpointSeedRefRecord - for seed_ref in seed_refs: - connection.execute( - f"INSERT INTO {seed_ref_table} " - "(checkpoint_id, seed_name, version_hash) VALUES (?, ?, ?)", - [seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], - ) diff --git a/src/sqlbuild/virtual/state/classes/_sql_state_backend.py b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py new file mode 100644 index 0000000000..0606be0831 --- /dev/null +++ b/src/sqlbuild/virtual/state/classes/_sql_state_backend.py @@ -0,0 +1,1715 @@ +"""Shared SQL for relational virtual-state backends.""" + +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Sequence +from contextlib import AbstractContextManager +from datetime import datetime +from typing import Any, ClassVar, cast + +from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json +from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json +from sqlbuild.executor.node_results.models import ( + NodeResultEnvelope, + NodeResultQuery, + NodeResultRecord, +) +from sqlbuild.virtual.state._helpers.state_storage.datetime import ( + from_naive_utc_wall_clock, +) +from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id +from sqlbuild.virtual.state._helpers.state_storage.validation import ( + validate_conditional_virtual_environment_publication, +) +from sqlbuild.virtual.state.classes.state_backend import StateBackend +from sqlbuild.virtual.state.constants import ( + CURRENT_STATE_SCHEMA_VERSION, + FUNCTION_VERSION_TABLE, + LOCK_TABLE, + MODEL_VERSION_TABLE, + NODE_RESULTS_TABLE, + NON_UNIQUE_STATE_INDEXES, + PHYSICAL_RELATION_ANCESTRY_TABLE, + PHYSICAL_RELATION_TABLE, + PYTHON_NODE_VERSION_TABLE, + RECONCILE_EVENT_TABLE, + SEED_VERSION_TABLE, + SOURCE_FRESHNESS_OBSERVATION_TABLE, + STATE_BOOLEAN_TRUE, + STATE_MIGRATION_EVENTS_TABLE, + STATE_OPERATION_EVENT_TABLE, + STATE_OPERATION_TABLE, + STATE_TABLE_COLUMNS, + STATE_TABLE_INDEXES, + STATE_TABLES, + STATE_VERSION_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, + VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, + VIRTUAL_ENVIRONMENT_TABLE, +) +from sqlbuild.virtual.state.exceptions import ( + StateBackendConfigError, + StateBackupNotFoundError, + StateSchemaInvalidError, +) +from sqlbuild.virtual.state.models import ( + FunctionVersionRecord, + ModelVersionRecord, + PhysicalRelationAncestryRecord, + PhysicalRelationRecord, + PythonNodeVersionRecord, + ReconcileEventRecord, + SeedVersionRecord, + SourceFreshnessRecord, + StateBackupRecord, + StateLockLease, + StateLockRecord, + StateOperationEventRecord, + StateOperationRecord, + StateSchemaValidationResult, + VirtualEnvironmentCheckpointFunctionRefRecord, + VirtualEnvironmentCheckpointModelRefRecord, + VirtualEnvironmentCheckpointRecord, + VirtualEnvironmentCheckpointSeedRefRecord, + VirtualEnvironmentFunctionRefRecord, + VirtualEnvironmentModelRefRecord, + VirtualEnvironmentNodeRefRecord, + VirtualEnvironmentPythonNodeRefRecord, + VirtualEnvironmentRecord, + VirtualEnvironmentRetentionRecord, + VirtualEnvironmentSeedRefRecord, +) +from sqlbuild.virtual.state.types import ( + ModelVersionStatus, + PhysicalArtifactType, + StateColumnType, + StateMigrationAction, + StateMigrationStatus, + StateOperationStatus, + StateOperationType, + VirtualEnvironmentStatus, +) + + +class SqlStateBackend(StateBackend): + """Relational state backend owning SQL shared by DuckDB and Postgres.""" + + _placeholder: ClassVar[str] + + @abstractmethod + def _fetch_one( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> tuple[Any, ...] | None: + """Run one read statement and return its first row.""" + ... + + @abstractmethod + def _fetch_all( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> list[tuple[Any, ...]]: + """Run one read statement and return all rows.""" + ... + + @abstractmethod + def _write_transaction(self, *, connection: Any) -> AbstractContextManager[Any]: + """Open a write transaction and yield the executor that runs its statements.""" + ... + + @abstractmethod + def _statement_executor(self, *, connection: Any) -> AbstractContextManager[Any]: + """Open an executor for explicitly delimited statements on a connection.""" + ... + + @abstractmethod + def _execute_in( + self, *, executor: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + """Run one statement on an executor.""" + ... + + @abstractmethod + def _lease_is_owned(self, *, executor: Any, schema: str, lease: StateLockLease) -> bool: + """Fence one lock lease inside an open transaction and report whether it is held.""" + ... + + @abstractmethod + def _fetch_one_in( + self, *, executor: Any, sql: str, params: Sequence[object] + ) -> tuple[Any, ...] | None: + """Run one read statement on a transaction executor and return its first row.""" + ... + + def upsert_model_version( + self, *, connection: Any, schema: str, record: ModelVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=MODEL_VERSION_TABLE, + key_values={"model_name": record.model_name, "version_hash": record.version_hash}, + row_values={ + "definition_identity_hash": record.definition_identity_hash, + "identity_metadata_hash": record.identity_metadata_hash, + "definition_text_b64": record.definition_text_b64, + "identity_metadata_json_b64": record.identity_metadata_json_b64, + "compiled_sql_b64": record.compiled_sql_b64, + "status": record.status.value, + }, + ) + + def upsert_function_version( + self, *, connection: Any, schema: str, record: FunctionVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=FUNCTION_VERSION_TABLE, + key_values={ + "function_name": record.function_name, + "version_hash": record.version_hash, + }, + row_values={ + "language": record.language, + "returns": record.returns, + "arguments_json_b64": record.arguments_json_b64, + "return_columns_json_b64": record.return_columns_json_b64, + "packages_json_b64": record.packages_json_b64, + "runtime_version": record.runtime_version, + "entry_point": record.entry_point, + "body_sql_b64": record.body_sql_b64, + "definition_text_b64": record.definition_text_b64, + "status": record.status.value, + }, + ) + + def upsert_seed_version( + self, *, connection: Any, schema: str, record: SeedVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=SEED_VERSION_TABLE, + key_values={"seed_name": record.seed_name, "version_hash": record.version_hash}, + row_values={ + "identity_metadata_hash": record.identity_metadata_hash, + "identity_metadata_json_b64": record.identity_metadata_json_b64, + "status": record.status.value, + }, + ) + + def upsert_python_node_version( + self, *, connection: Any, schema: str, record: PythonNodeVersionRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=PYTHON_NODE_VERSION_TABLE, + key_values={ + "node_type": record.node_type, + "node_name": record.node_name, + "version_hash": record.version_hash, + }, + row_values={ + "definition_hash": record.definition_hash, + "identity_metadata_hash": record.identity_metadata_hash, + "definition_json_b64": record.definition_json_b64, + "identity_metadata_json_b64": record.identity_metadata_json_b64, + "status": record.status.value, + }, + ) + + def upsert_physical_relation( + self, *, connection: Any, schema: str, record: PhysicalRelationRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=PHYSICAL_RELATION_TABLE, + key_values={ + "artifact_type": record.artifact_type.value, + "artifact_name": record.artifact_name, + "version_hash": record.version_hash, + }, + row_values={ + "database_name": record.database_name, + "schema_name": record.schema_name, + "relation_name": record.relation_name, + "relation_type": record.relation_type, + }, + ) + + def upsert_physical_relation_ancestry( + self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=PHYSICAL_RELATION_ANCESTRY_TABLE, + key_values={"model_name": record.model_name, "version_hash": record.version_hash}, + row_values={ + "parent_model_name": record.parent_model_name, + "parent_version_hash": record.parent_version_hash, + "seed_strategy": record.seed_strategy, + }, + ) + + def upsert_virtual_environment( + self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._upsert_virtual_environment_record(executor=executor, schema=schema, record=record) + + def upsert_state_operation( + self, *, connection: Any, schema: str, record: StateOperationRecord + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=STATE_OPERATION_TABLE, + key_values={"operation_id": record.operation_id}, + row_values={ + "operation_type": record.operation_type.value, + "status": record.status.value, + "virtual_environment_name": record.virtual_environment_name, + }, + ) + + def _upsert_virtual_environment_record( + self, *, executor: Any, schema: str, record: VirtualEnvironmentRecord + ) -> None: + self._replace_row_preserving_created_at( + executor=executor, + schema=schema, + table_name=VIRTUAL_ENVIRONMENT_TABLE, + key_values={"virtual_environment_name": record.virtual_environment_name}, + row_values={ + "status": record.status.value, + "baseline_virtual_environment_name": record.baseline_virtual_environment_name, + "finalized_at": record.finalized_at, + }, + ) + + def list_virtual_environments( + self, *, connection: Any, schema: str + ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, status, updated_at " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " + "ORDER BY updated_at DESC, virtual_environment_name DESC", + ) + return tuple( + VirtualEnvironmentRetentionRecord( + virtual_environment_name=row[0], + status=VirtualEnvironmentStatus(row[1]), + updated_at=row[2], + ) + for row in rows + ) + + def renew_lock( + self, + *, + connection: Any, + schema: str, + lock_key: str, + owner_id: str, + expires_at: datetime, + ) -> bool: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql=f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " + f"SET expires_at = {p}, updated_at = CURRENT_TIMESTAMP " + f"WHERE lock_key = {p} AND owner_id = {p} AND expires_at > CURRENT_TIMESTAMP " + "RETURNING lock_key", + params=[expires_at, lock_key, owner_id], + ) + return row is not None + + def list_active_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT lock_key, owner_id, expires_at FROM " + f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " + "WHERE expires_at > CURRENT_TIMESTAMP ORDER BY lock_key", + ) + return tuple( + StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows + ) + + def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT lock_key, owner_id, expires_at FROM " + f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " + "WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY lock_key", + ) + return tuple( + StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows + ) + + @abstractmethod + def _replace_virtual_environment_node_ref_groups( + self, + *, + executor: Any, + schema: str, + virtual_environment_name: str, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + ) -> None: + """Replace node refs per node type inside an open write transaction.""" + ... + + def replace_virtual_environment_node_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + node_type: str, + refs: tuple[VirtualEnvironmentNodeRefRecord, ...], + ) -> None: + self.replace_virtual_environment_node_ref_groups( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + refs_by_node_type={node_type: refs}, + ) + + def replace_virtual_environment_node_ref_groups( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._replace_virtual_environment_node_ref_groups( + executor=executor, + schema=schema, + virtual_environment_name=virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + + def upsert_virtual_environment_and_replace_node_ref_groups( + self, + *, + connection: Any, + schema: str, + record: VirtualEnvironmentRecord, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._upsert_virtual_environment_record(executor=executor, schema=schema, record=record) + self._replace_virtual_environment_node_ref_groups( + executor=executor, + schema=schema, + virtual_environment_name=record.virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + + def replace_virtual_environment_model_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs: tuple[VirtualEnvironmentModelRefRecord, ...], + ) -> None: + self.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="model", + refs=tuple( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type="model", + node_name=ref.model_name, + version_hash=ref.version_hash, + ) + for ref in refs + ), + ) + + def replace_virtual_environment_function_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs: tuple[VirtualEnvironmentFunctionRefRecord, ...], + ) -> None: + ref: VirtualEnvironmentFunctionRefRecord + for ref in refs: + if ref.node_type not in { + CompiledResourceType.UDF, + CompiledResourceType.TABLE_FN, + }: + raise StateBackendConfigError("Function ref node_type must be 'udf' or 'table_fn'") + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]] = {} + for node_type in ("udf", "table_fn"): + node_refs: list[VirtualEnvironmentNodeRefRecord] = [] + for ref in refs: + if ref.node_type == node_type: + node_refs.append( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type=ref.node_type, + node_name=ref.function_name, + version_hash=ref.version_hash, + ) + ) + refs_by_node_type[node_type] = tuple(node_refs) + self.replace_virtual_environment_node_ref_groups( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + + def replace_virtual_environment_seed_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + refs: tuple[VirtualEnvironmentSeedRefRecord, ...], + ) -> None: + self.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="seed", + refs=tuple( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type="seed", + node_name=ref.seed_name, + version_hash=ref.version_hash, + ) + for ref in refs + ), + ) + + def upsert_virtual_environment_python_node_ref( + self, + *, + connection: Any, + schema: str, + ref: VirtualEnvironmentPythonNodeRefRecord, + ) -> None: + self.upsert_virtual_environment_node_ref( + connection=connection, + schema=schema, + ref=VirtualEnvironmentNodeRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type=ref.node_type, + node_name=ref.node_name, + version_hash=ref.version_hash, + ), + ) + + def _validate_node_ref_replacement( + self, + *, + virtual_environment_name: str, + node_type: str, + refs: tuple[VirtualEnvironmentNodeRefRecord, ...], + ) -> None: + seen_node_names: set[str] = set() + ref: VirtualEnvironmentNodeRefRecord + for ref in refs: + if ref.virtual_environment_name != virtual_environment_name: + raise StateBackendConfigError( + "Node ref virtual_environment_name must match replacement " + "virtual_environment_name" + ) + if ref.node_type != node_type: + raise StateBackendConfigError("Node ref node_type must match replacement node_type") + if ref.node_name in seen_node_names: + raise StateBackendConfigError( + f"Duplicate node ref for node type '{node_type}' and name '{ref.node_name}'" + ) + seen_node_names.add(ref.node_name) + + def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( + self, + *, + connection: Any, + schema: str, + record: VirtualEnvironmentRecord, + refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], + leases: tuple[StateLockLease, ...], + checkpoint: VirtualEnvironmentCheckpointRecord | None = None, + checkpoint_refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...] = (), + checkpoint_function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), + checkpoint_seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), + ) -> bool: + validate_conditional_virtual_environment_publication( + record=record, + refs_by_node_type=refs_by_node_type, + checkpoint=checkpoint, + checkpoint_refs=checkpoint_refs, + checkpoint_function_refs=checkpoint_function_refs, + checkpoint_seed_refs=checkpoint_seed_refs, + ) + with self._statement_executor(connection=connection) as executor: + self._execute_in(executor=executor, sql="BEGIN") + try: + lease: StateLockLease + for lease in leases: + if not self._lease_is_owned(executor=executor, schema=schema, lease=lease): + self._execute_in(executor=executor, sql="ROLLBACK") + return False + self._upsert_virtual_environment_record( + executor=executor, schema=schema, record=record + ) + self._replace_virtual_environment_node_ref_groups( + executor=executor, + schema=schema, + virtual_environment_name=record.virtual_environment_name, + refs_by_node_type=refs_by_node_type, + ) + if checkpoint is not None: + self._insert_virtual_environment_checkpoint_rows( + executor=executor, + schema=schema, + checkpoint=checkpoint, + refs=checkpoint_refs, + function_refs=checkpoint_function_refs, + seed_refs=checkpoint_seed_refs, + ) + self._execute_in(executor=executor, sql="COMMIT") + return True + except BaseException: + self._execute_in(executor=executor, sql="ROLLBACK") + raise + + def create_virtual_environment_checkpoint( + self, + *, + connection: Any, + schema: str, + checkpoint: VirtualEnvironmentCheckpointRecord, + refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], + function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), + seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), + ) -> None: + with self._write_transaction(connection=connection) as executor: + self._insert_virtual_environment_checkpoint_rows( + executor=executor, + schema=schema, + checkpoint=checkpoint, + refs=refs, + function_refs=function_refs, + seed_refs=seed_refs, + ) + + def delete_virtual_environment_checkpoint( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> None: + p: str = self._placeholder + table_names: tuple[str, ...] = ( + VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, + VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, + ) + with self._write_transaction(connection=connection) as executor: + table_name: str + for table_name in table_names: + self._execute_in( + executor=executor, + sql=( + f"DELETE FROM {self._qualified_name(schema=schema, table=table_name)} " + f"WHERE checkpoint_id = {p}" + ), + params=[checkpoint_id], + ) + + def _insert_virtual_environment_checkpoint_rows( + self, + *, + executor: Any, + schema: str, + checkpoint: VirtualEnvironmentCheckpointRecord, + refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], + function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...], + seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...], + ) -> None: + p: str = self._placeholder + checkpoint_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE + ) + model_ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE + ) + function_ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE + ) + seed_ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE + ) + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {checkpoint_table} " + "(checkpoint_id, virtual_environment_name, created_at) " + f"VALUES ({p}, {p}, CURRENT_TIMESTAMP)" + ), + params=[checkpoint.checkpoint_id, checkpoint.virtual_environment_name], + ) + ref: VirtualEnvironmentCheckpointModelRefRecord + for ref in refs: + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {model_ref_table} " + f"(checkpoint_id, model_name, version_hash) VALUES ({p}, {p}, {p})" + ), + params=[ref.checkpoint_id, ref.model_name, ref.version_hash], + ) + function_ref: VirtualEnvironmentCheckpointFunctionRefRecord + for function_ref in function_refs: + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {function_ref_table} " + f"(checkpoint_id, function_name, version_hash) VALUES ({p}, {p}, {p})" + ), + params=[ + function_ref.checkpoint_id, + function_ref.function_name, + function_ref.version_hash, + ], + ) + seed_ref: VirtualEnvironmentCheckpointSeedRefRecord + for seed_ref in seed_refs: + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {seed_ref_table} " + f"(checkpoint_id, seed_name, version_hash) VALUES ({p}, {p}, {p})" + ), + params=[seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], + ) + + def insert_node_result( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + record: NodeResultRecord, + ) -> None: + p: str = self._placeholder + self._execute( + connection=connection, + sql=f"INSERT INTO {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " + "(virtual_environment_name, node_type, node_name, target_database, target_schema, " + "target_name, run_id, status, payload_json_b64, metadata_json_b64, error_message, " + "materialized, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p}, {p})", + params=[ + virtual_environment_name, + record.node_type, + record.node_name, + record.target_database, + record.target_schema, + record.target_name, + record.run_id, + record.status, + encode_node_result_json( + value=record.payload, label="payload", node_name=record.node_name + ), + encode_node_result_json( + value=record.metadata, label="metadata", node_name=record.node_name + ), + record.error_message, + self._materialized_storage(record.materialized), + record.ts, + ], + ) + + def read_node_results( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + query: NodeResultQuery, + ) -> tuple[NodeResultEnvelope, ...]: + p: str = self._placeholder + if query.limit < 1: + return () + predicates: list[str] = [ + f"virtual_environment_name = {p}", + f"node_type = {p}", + f"node_name = {p}", + self._optional_equality_sql( + column="target_database", value=query.target_database, placeholder=f"{p}" + ), + self._optional_equality_sql( + column="target_schema", value=query.target_schema, placeholder=f"{p}" + ), + self._optional_equality_sql( + column="target_name", value=query.target_name, placeholder=f"{p}" + ), + ] + params: list[object] = [virtual_environment_name, query.node_type, query.node_name] + for value in (query.target_database, query.target_schema, query.target_name): + if value is not None: + params.append(value) + if query.statuses is not None: + placeholders: str = ", ".join(f"{p}" for _ in query.statuses) + predicates.append(f"status IN ({placeholders})") + params.extend(query.statuses) + if query.run_id is not None: + predicates.append(f"run_id = {p}") + params.append(query.run_id) + params.append(query.limit) + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT node_type, node_name, run_id, status, payload_json_b64, metadata_json_b64, " + "error_message, materialized, created_at " + f"FROM {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " + f"WHERE {' AND '.join(predicates)} " + f"ORDER BY created_at DESC, run_id DESC LIMIT {p}", + params=params, + ) + return tuple(self._node_result_row_to_envelope(row) for row in rows) + + def create_state_operation_event( + self, *, connection: Any, schema: str, record: StateOperationEventRecord + ) -> None: + p: str = self._placeholder + self._execute( + connection=connection, + sql="INSERT INTO " + f"{self._qualified_name(schema=schema, table=STATE_OPERATION_EVENT_TABLE)} " + "(event_id, operation_id, action, status, message, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, {p}, CURRENT_TIMESTAMP)", + params=[ + record.event_id, + record.operation_id, + record.action, + record.status.value, + record.message, + ], + ) + + def create_reconcile_event( + self, *, connection: Any, schema: str, record: ReconcileEventRecord + ) -> None: + p: str = self._placeholder + self._execute( + connection=connection, + sql=f"INSERT INTO {self._qualified_name(schema=schema, table=RECONCILE_EVENT_TABLE)} " + "(event_id, action, status, message, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, CURRENT_TIMESTAMP)", + params=[record.event_id, record.action.value, record.status.value, record.message], + ) + + def _node_result_row_to_envelope(self, row: tuple[Any, ...]) -> NodeResultEnvelope: + node_name: str = str(row[1]) + metadata: object = decode_node_result_json( + value=str(row[5]), label="metadata", node_name=node_name + ) + normalized_metadata: dict[str, object] = ( + {str(key): value for key, value in metadata.items()} + if isinstance(metadata, dict) + else {} + ) + return NodeResultEnvelope( + node_type=str(row[0]), + node_name=node_name, + run_id=str(row[2]), + status=str(row[3]), + payload=decode_node_result_json( + value=str(row[4]), label="payload", node_name=node_name + ), + metadata=normalized_metadata, + error_message=str(row[6]) if row[6] is not None else None, + materialized=self._parse_materialized(row[7]), + ts=row[8], + ) + + def _optional_equality_sql(self, *, column: str, value: object | None, placeholder: str) -> str: + if value is None: + return f"{column} IS NULL" + return f"{column} = {placeholder}" + + def _materialized_storage(self, value: bool | None) -> str | None: + if value is None: + return None + return "true" if value else "false" + + def _parse_materialized(self, value: object) -> bool | None: + if value is None: + return None + return str(value).lower() == STATE_BOOLEAN_TRUE + + def _execute( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + with self._statement_executor(connection=connection) as executor: + self._execute_in(executor=executor, sql=sql, params=params) + + def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: + p: str = self._placeholder + with self._write_transaction(connection=connection) as executor: + self._execute_in( + executor=executor, + sql=f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}", + ) + self._execute_in( + executor=executor, + sql="CREATE TABLE IF NOT EXISTS " + f"{self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} (" + "schema_version INTEGER NOT NULL, " + "sqlbuild_version TEXT NOT NULL, " + "updated_at TIMESTAMP NOT NULL" + ")", + ) + self._execute_in( + executor=executor, + sql="CREATE TABLE IF NOT EXISTS " + f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} (" + "event_id TEXT NOT NULL, " + "action TEXT NOT NULL, " + "backup_id TEXT, " + "status TEXT NOT NULL, " + "message TEXT, " + "created_at TIMESTAMP NOT NULL" + ")", + ) + self._create_additional_state_tables(executor=executor, schema=schema) + self._execute_in( + executor=executor, + sql=f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)}", + ) + self._execute_in( + executor=executor, + sql=f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} " + "(schema_version, sqlbuild_version, updated_at) " + f"VALUES ({p}, {p}, CURRENT_TIMESTAMP)", + params=[CURRENT_STATE_SCHEMA_VERSION, sqlbuild_version], + ) + self._record_event( + executor=executor, + schema=schema, + action=StateMigrationAction.INIT, + backup_id_value=None, + status=StateMigrationStatus.SUCCESS, + message=None, + ) + + def create_backup(self, *, connection: Any, schema: str) -> str: + validation: StateSchemaValidationResult = self.inspect_schema( + connection=connection, schema=schema + ) + if not validation.valid: + raise StateSchemaInvalidError("Cannot backup invalid state schema") + backup_id_value: str = backup_id() + backup_schema: str = self._backup_schema_name( + schema=schema, + backup_id_value=backup_id_value, + ) + with self._write_transaction(connection=connection) as executor: + self._execute_in( + executor=executor, sql=f"CREATE SCHEMA {self._quote_identifier(backup_schema)}" + ) + table_name: str + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql="CREATE TABLE " + f"{self._qualified_name(schema=backup_schema, table=table_name)} " + "AS " + f"SELECT * FROM {self._qualified_name(schema=schema, table=table_name)}", + ) + self._record_event( + executor=executor, + schema=schema, + action=StateMigrationAction.BACKUP, + backup_id_value=backup_id_value, + status=StateMigrationStatus.SUCCESS, + message=None, + ) + return backup_id_value + + def rollback(self, *, connection: Any, schema: str, backup_id: str | None = None) -> str: + backup_id_value: str = backup_id or self._latest_backup_id( + connection=connection, schema=schema + ) + backup_schema: str = self._backup_schema_name( + schema=schema, + backup_id_value=backup_id_value, + ) + if not self._schema_exists(connection=connection, schema=backup_schema): + raise StateBackupNotFoundError(f"State backup schema '{backup_schema}' does not exist") + with self._write_transaction(connection=connection) as executor: + table_name: str + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql="DROP TABLE IF EXISTS " + f"{self._qualified_name(schema=schema, table=table_name)}", + ) + self._execute_in( + executor=executor, + sql=f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}", + ) + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql=f"CREATE TABLE {self._qualified_name(schema=schema, table=table_name)} AS " + f"SELECT * FROM {self._qualified_name(schema=backup_schema, table=table_name)}", + ) + self._create_state_indexes(executor=executor, schema=schema) + self._record_event( + executor=executor, + schema=schema, + action=StateMigrationAction.ROLLBACK, + backup_id_value=backup_id_value, + status=StateMigrationStatus.SUCCESS, + message=None, + ) + return backup_id_value + + def reset(self, *, connection: Any, schema: str) -> None: + with self._write_transaction(connection=connection) as executor: + for table_name in STATE_TABLES: + self._execute_in( + executor=executor, + sql="DROP TABLE IF EXISTS " + f"{self._qualified_name(schema=schema, table=table_name)}", + ) + + def _create_additional_state_tables(self, *, executor: Any, schema: str) -> None: + table_name: str + columns: dict[str, StateColumnType] + for table_name, columns in STATE_TABLE_COLUMNS.items(): + if table_name in {STATE_VERSION_TABLE, STATE_MIGRATION_EVENTS_TABLE}: + continue + column_sql: str = ", ".join( + f"{self._quote_identifier(column_name)} {self._state_column_sql_type(column_type)}" + for column_name, column_type in columns.items() + ) + self._execute_in( + executor=executor, + sql="CREATE TABLE IF NOT EXISTS " + f"{self._qualified_name(schema=schema, table=table_name)} " + f"({column_sql})", + ) + column_name: str + column_type: StateColumnType + for column_name, column_type in columns.items(): + self._execute_in( + executor=executor, + sql=f"ALTER TABLE {self._qualified_name(schema=schema, table=table_name)} " + f"ADD COLUMN IF NOT EXISTS {self._quote_identifier(column_name)} " + f"{self._state_column_sql_type(column_type)}", + ) + self._create_state_indexes(executor=executor, schema=schema) + + def _create_state_indexes(self, *, executor: Any, schema: str) -> None: + table_name: str + indexes: dict[str, tuple[str, ...]] + for table_name, indexes in STATE_TABLE_INDEXES.items(): + index_name: str + columns: tuple[str, ...] + for index_name, columns in indexes.items(): + column_sql: str = ", ".join(self._quote_identifier(column) for column in columns) + unique_sql: str = "" if index_name in NON_UNIQUE_STATE_INDEXES else "UNIQUE " + self._execute_in( + executor=executor, + sql=f"CREATE {unique_sql}INDEX IF NOT EXISTS " + f"{self._quote_identifier(index_name)} " + f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})", + ) + + def _record_event( + self, + *, + executor: Any, + schema: str, + action: StateMigrationAction, + backup_id_value: str | None, + status: StateMigrationStatus, + message: str | None, + ) -> None: + p: str = self._placeholder + self._execute_in( + executor=executor, + sql="INSERT INTO " + f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} " + "(event_id, action, backup_id, status, message, created_at) " + f"VALUES ({p}, {p}, {p}, {p}, {p}, CURRENT_TIMESTAMP)", + params=[event_id(), action.value, backup_id_value, status.value, message], + ) + + def list_state_backups(self, *, connection: Any, schema: str) -> tuple[StateBackupRecord, ...]: + p: str = self._placeholder + prefix: str = f"{schema}__backup_%" + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT s.schema_name, e.backup_id, MAX(e.created_at) " + "FROM information_schema.schemata s " + "LEFT JOIN " + f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} e " + f"ON s.schema_name = {p} || e.backup_id " + f"WHERE s.schema_name LIKE {p} " + "GROUP BY s.schema_name, e.backup_id ORDER BY s.schema_name DESC", + params=[f"{schema}__backup_", prefix], + ) + return tuple( + StateBackupRecord( + backup_id=row[1] or str(row[0]).removeprefix(f"{schema}__backup_"), + schema_name=row[0], + created_at=row[2], + ) + for row in rows + ) + + def _latest_backup_id(self, *, connection: Any, schema: str) -> str: + p: str = self._placeholder + prefix: str = f"{schema}__backup_%" + rows: list[tuple[str]] = self._fetch_all( + connection=connection, + sql=f"SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE {p} " + "ORDER BY schema_name DESC LIMIT 1", + params=[prefix], + ) + if not rows: + raise StateBackupNotFoundError("No state backup is available for rollback") + return rows[0][0].removeprefix(f"{schema}__backup_") + + def _schema_exists(self, *, connection: Any, schema: str) -> bool: + p: str = self._placeholder + rows: list[tuple[str]] = self._fetch_all( + connection=connection, + sql=f"SELECT schema_name FROM information_schema.schemata WHERE schema_name = {p}", + params=[schema], + ) + return bool(rows) + + def _state_column_sql_type(self, column_type: StateColumnType) -> str: + match column_type: + case StateColumnType.INTEGER: + return "INTEGER" + case StateColumnType.TEXT: + return "TEXT" + case StateColumnType.TIMESTAMP: + return "TIMESTAMP" + raise StateBackendConfigError(f"Unsupported state column type: {column_type}") + + def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: + return f"{schema}__backup_{backup_id_value}" + + def close(self, connection: Any) -> None: + connection.close() + + def delete_virtual_environment( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> None: + p: str = self._placeholder + table_names: tuple[str, ...] = ( + VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, + SOURCE_FRESHNESS_OBSERVATION_TABLE, + VIRTUAL_ENVIRONMENT_TABLE, + ) + with self._write_transaction(connection=connection) as executor: + table_name: str + for table_name in table_names: + self._execute_in( + executor=executor, + sql=( + f"DELETE FROM {self._qualified_name(schema=schema, table=table_name)} " + f"WHERE virtual_environment_name = {p}" + ), + params=[virtual_environment_name], + ) + + def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: + version_table: str = self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE) + ref_table: str = self._qualified_name( + schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE + ) + row: tuple[Any, ...] = cast( + tuple[Any, ...], + self._fetch_one( + connection=connection, + sql=( + "SELECT COUNT(*) " + f"FROM {version_table} versions " + "WHERE NOT EXISTS (" + "SELECT 1 " + f"FROM {ref_table} refs " + "WHERE refs.node_type = versions.node_type " + "AND refs.node_name = versions.node_name " + "AND refs.version_hash = versions.version_hash)" + ), + ), + ) + return int(row[0]) + + def _validate_source_freshness_records( + self, + *, + virtual_environment_name: str, + records: tuple[SourceFreshnessRecord, ...], + ) -> None: + seen_source_names: set[str] = set() + record: SourceFreshnessRecord + for record in records: + if record.virtual_environment_name != virtual_environment_name: + raise StateBackendConfigError( + "Source freshness record virtual_environment_name must match replacement " + "virtual_environment_name" + ) + if record.source_name in seen_source_names: + raise StateBackendConfigError( + f"Duplicate source freshness record for source '{record.source_name}'" + ) + seen_source_names.add(record.source_name) + + def _replace_row_preserving_created_at( + self, + *, + executor: Any, + schema: str, + table_name: str, + key_values: dict[str, object], + row_values: dict[str, object], + ) -> None: + p: str = self._placeholder + table: str = self._qualified_name(schema=schema, table=table_name) + where_sql: str = " AND ".join(f"{column} = {p}" for column in key_values) + key_params: list[object] = list(key_values.values()) + existing_created_at: datetime | None = self._created_at_for_key( + executor=executor, + schema=schema, + table_name=table_name, + where_sql=where_sql, + params=key_params, + ) + self._execute_in( + executor=executor, sql=f"DELETE FROM {table} WHERE {where_sql}", params=key_params + ) + columns: tuple[str, ...] = (*key_values, *row_values) + value_placeholders: str = ", ".join(p for _ in columns) + self._execute_in( + executor=executor, + sql=( + f"INSERT INTO {table} ({', '.join(columns)}, created_at, updated_at) " + f"VALUES ({value_placeholders}, COALESCE({p}, CURRENT_TIMESTAMP), " + "CURRENT_TIMESTAMP)" + ), + params=[*key_params, *row_values.values(), existing_created_at], + ) + + def _created_at_for_key( + self, + *, + executor: Any, + schema: str, + table_name: str, + where_sql: str, + params: list[object], + ) -> datetime | None: + row: tuple[Any, ...] | None = self._fetch_one_in( + executor=executor, + sql=( + f"SELECT created_at FROM {self._qualified_name(schema=schema, table=table_name)} " + f"WHERE {where_sql}" + ), + params=params, + ) + if row is None: + return None + return row[0] + + def _quote_identifier(self, identifier: str) -> str: + return '"' + identifier.replace('"', '""') + '"' + + def _qualified_name(self, *, schema: str, table: str) -> str: + return f"{self._quote_identifier(schema)}.{self._quote_identifier(table)}" + + def get_model_version( + self, *, connection: Any, schema: str, model_name: str, version_hash: str + ) -> ModelVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT model_name, version_hash, definition_identity_hash, " + "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " + "compiled_sql_b64, status " + f"FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " + f"WHERE model_name = {p} AND version_hash = {p}", + params=[model_name, version_hash], + ) + if row is None: + return None + return ModelVersionRecord( + model_name=row[0], + version_hash=row[1], + definition_identity_hash=row[2], + identity_metadata_hash=row[3], + definition_text_b64=row[4], + identity_metadata_json_b64=row[5], + compiled_sql_b64=row[6], + status=ModelVersionStatus(row[7]), + ) + + def get_function_version( + self, *, connection: Any, schema: str, function_name: str, version_hash: str + ) -> FunctionVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT function_name, version_hash, language, returns, arguments_json_b64, " + "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " + "body_sql_b64, definition_text_b64, status " + f"FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " + f"WHERE function_name = {p} AND version_hash = {p}", + params=[function_name, version_hash], + ) + if row is None: + return None + return FunctionVersionRecord( + function_name=row[0], + version_hash=row[1], + language=row[2], + returns=row[3], + arguments_json_b64=row[4], + return_columns_json_b64=row[5], + packages_json_b64=row[6], + runtime_version=row[7], + entry_point=row[8], + body_sql_b64=row[9], + definition_text_b64=row[10], + status=ModelVersionStatus(row[11]), + ) + + def get_seed_version( + self, *, connection: Any, schema: str, seed_name: str, version_hash: str + ) -> SeedVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT seed_name, version_hash, identity_metadata_hash, " + "identity_metadata_json_b64, status " + f"FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " + f"WHERE seed_name = {p} AND version_hash = {p}", + params=[seed_name, version_hash], + ) + if row is None: + return None + return SeedVersionRecord( + seed_name=row[0], + version_hash=row[1], + identity_metadata_hash=row[2], + identity_metadata_json_b64=row[3], + status=ModelVersionStatus(row[4]), + ) + + def get_python_node_version( + self, + *, + connection: Any, + schema: str, + node_type: str, + node_name: str, + version_hash: str, + ) -> PythonNodeVersionRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT node_type, node_name, version_hash, definition_hash, " + "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, status " + f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " + f"WHERE node_type = {p} AND node_name = {p} AND version_hash = {p}", + params=[node_type, node_name, version_hash], + ) + if row is None: + return None + return PythonNodeVersionRecord( + node_type=row[0], + node_name=row[1], + version_hash=row[2], + definition_hash=row[3], + identity_metadata_hash=row[4], + definition_json_b64=row[5], + identity_metadata_json_b64=row[6], + status=ModelVersionStatus(row[7]), + ) + + def get_physical_relation_for_artifact( + self, + *, + connection: Any, + schema: str, + artifact_type: PhysicalArtifactType, + artifact_name: str, + version_hash: str, + ) -> PhysicalRelationRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " + "relation_name, relation_type " + f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " + f"WHERE artifact_type = {p} AND artifact_name = {p} AND version_hash = {p}", + params=[artifact_type.value, artifact_name, version_hash], + ) + if row is None: + return None + return PhysicalRelationRecord( + artifact_type=PhysicalArtifactType(row[0]), + artifact_name=row[1], + version_hash=row[2], + database_name=row[3], + schema_name=row[4], + relation_name=row[5], + relation_type=row[6], + ) + + def list_physical_relations_for_artifact( + self, + *, + connection: Any, + schema: str, + artifact_type: PhysicalArtifactType, + artifact_name: str, + ) -> tuple[PhysicalRelationRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " + "relation_name, relation_type " + f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " + f"WHERE artifact_type = {p} AND artifact_name = {p} " + "ORDER BY updated_at DESC, version_hash DESC", + params=[artifact_type.value, artifact_name], + ) + return tuple( + PhysicalRelationRecord( + artifact_type=PhysicalArtifactType(row[0]), + artifact_name=row[1], + version_hash=row[2], + database_name=row[3], + schema_name=row[4], + relation_name=row[5], + relation_type=row[6], + ) + for row in rows + ) + + def get_physical_relation_ancestry( + self, *, connection: Any, schema: str, model_name: str, version_hash: str + ) -> PhysicalRelationAncestryRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT model_name, version_hash, parent_model_name, parent_version_hash, " + "seed_strategy " + f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " + f"WHERE model_name = {p} AND version_hash = {p}", + params=[model_name, version_hash], + ) + if row is None: + return None + return PhysicalRelationAncestryRecord( + model_name=row[0], + version_hash=row[1], + parent_model_name=row[2], + parent_version_hash=row[3], + seed_strategy=row[4], + ) + + def get_virtual_environment( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> VirtualEnvironmentRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT virtual_environment_name, status, baseline_virtual_environment_name, " + "finalized_at " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " + f"WHERE virtual_environment_name = {p}", + params=[virtual_environment_name], + ) + if row is None: + return None + return VirtualEnvironmentRecord( + virtual_environment_name=row[0], + status=VirtualEnvironmentStatus(row[1]), + baseline_virtual_environment_name=row[2], + finalized_at=row[3], + ) + + def get_virtual_environment_node_refs( + self, + *, + connection: Any, + schema: str, + virtual_environment_name: str, + node_type: str, + ) -> tuple[VirtualEnvironmentNodeRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, node_type, node_name, version_hash " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " + f"WHERE virtual_environment_name = {p} AND node_type = {p} ORDER BY node_name", + params=[virtual_environment_name, node_type], + ) + return tuple( + VirtualEnvironmentNodeRefRecord( + virtual_environment_name=row[0], + node_type=row[1], + node_name=row[2], + version_hash=row[3], + ) + for row in rows + ) + + def get_virtual_environment_model_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentModelRefRecord, ...]: + refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="model", + ) + return tuple( + VirtualEnvironmentModelRefRecord( + virtual_environment_name=ref.virtual_environment_name, + model_name=ref.node_name, + version_hash=ref.version_hash, + ) + for ref in refs + ) + + def get_virtual_environment_function_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentFunctionRefRecord, ...]: + refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + *self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="udf", + ), + *self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="table_fn", + ), + ) + return tuple( + VirtualEnvironmentFunctionRefRecord( + virtual_environment_name=ref.virtual_environment_name, + node_type=ref.node_type, + function_name=ref.node_name, + version_hash=ref.version_hash, + ) + for ref in sorted(refs, key=lambda item: (item.node_type, item.node_name)) + ) + + def get_virtual_environment_seed_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentSeedRefRecord, ...]: + refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name=virtual_environment_name, + node_type="seed", + ) + return tuple( + VirtualEnvironmentSeedRefRecord( + virtual_environment_name=ref.virtual_environment_name, + seed_name=ref.node_name, + version_hash=ref.version_hash, + ) + for ref in refs + ) + + def get_virtual_environment_python_node_refs( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentPythonNodeRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, node_type, node_name, version_hash " + f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " + f"WHERE virtual_environment_name = {p} " + "AND node_type IN ('task', 'loader', 'asset', 'check', 'hook') " + "ORDER BY node_type, node_name", + params=[virtual_environment_name], + ) + return tuple( + VirtualEnvironmentPythonNodeRefRecord( + virtual_environment_name=row[0], + node_type=row[1], + node_name=row[2], + version_hash=row[3], + ) + for row in rows + ) + + def get_virtual_environment_source_freshness( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[SourceFreshnessRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT virtual_environment_name, source_name, strategy, value_kind, " + "data_version, data_version_hash, observed_at " + f"FROM {self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " + f"WHERE virtual_environment_name = {p} ORDER BY source_name", + params=[virtual_environment_name], + ) + return tuple( + SourceFreshnessRecord( + virtual_environment_name=row[0], + source_name=row[1], + strategy=row[2], + value_kind=row[3], + data_version=row[4], + data_version_hash=row[5], + observed_at=from_naive_utc_wall_clock(row[6]), + ) + for row in rows + ) + + def list_virtual_environment_checkpoints( + self, *, connection: Any, schema: str, virtual_environment_name: str + ) -> tuple[VirtualEnvironmentCheckpointRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT checkpoint_id, virtual_environment_name, created_at " + "FROM " + f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE)} " + f"WHERE virtual_environment_name = {p} ORDER BY created_at DESC, checkpoint_id DESC", + params=[virtual_environment_name], + ) + return tuple( + VirtualEnvironmentCheckpointRecord( + checkpoint_id=row[0], + virtual_environment_name=row[1], + created_at=row[2], + ) + for row in rows + ) + + def get_virtual_environment_checkpoint_model_refs( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> tuple[VirtualEnvironmentCheckpointModelRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT checkpoint_id, model_name, version_hash FROM " + + self._qualified_name( + schema=schema, + table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, + ) + + " " + f"WHERE checkpoint_id = {p} ORDER BY model_name", + params=[checkpoint_id], + ) + return tuple( + VirtualEnvironmentCheckpointModelRefRecord( + checkpoint_id=row[0], + model_name=row[1], + version_hash=row[2], + ) + for row in rows + ) + + def get_virtual_environment_checkpoint_function_refs( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...]: + p: str = self._placeholder + checkpoint_function_ref_table: str = self._qualified_name( + schema=schema, + table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, + ) + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql=f"SELECT checkpoint_id, function_name, version_hash " + f"FROM {checkpoint_function_ref_table} " + f"WHERE checkpoint_id = {p} ORDER BY function_name", + params=[checkpoint_id], + ) + return tuple( + VirtualEnvironmentCheckpointFunctionRefRecord( + checkpoint_id=row[0], + function_name=row[1], + version_hash=row[2], + ) + for row in rows + ) + + def get_virtual_environment_checkpoint_seed_refs( + self, *, connection: Any, schema: str, checkpoint_id: str + ) -> tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...]: + p: str = self._placeholder + rows: list[tuple[Any, ...]] = self._fetch_all( + connection=connection, + sql="SELECT checkpoint_id, seed_name, version_hash FROM " + + self._qualified_name( + schema=schema, + table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, + ) + + " " + f"WHERE checkpoint_id = {p} ORDER BY seed_name", + params=[checkpoint_id], + ) + return tuple( + VirtualEnvironmentCheckpointSeedRefRecord( + checkpoint_id=row[0], + seed_name=row[1], + version_hash=row[2], + ) + for row in rows + ) + + def get_state_operation( + self, *, connection: Any, schema: str, operation_id: str + ) -> StateOperationRecord | None: + p: str = self._placeholder + row: tuple[Any, ...] | None = self._fetch_one( + connection=connection, + sql="SELECT operation_id, operation_type, status, virtual_environment_name " + f"FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " + f"WHERE operation_id = {p}", + params=[operation_id], + ) + if row is None: + return None + return StateOperationRecord( + operation_id=row[0], + operation_type=StateOperationType(row[1]), + status=StateOperationStatus(row[2]), + virtual_environment_name=row[3], + ) diff --git a/src/sqlbuild/virtual/state/classes/duckdb.py b/src/sqlbuild/virtual/state/classes/duckdb.py index 9c1405d714..d4edb7945f 100644 --- a/src/sqlbuild/virtual/state/classes/duckdb.py +++ b/src/sqlbuild/virtual/state/classes/duckdb.py @@ -2,24 +2,16 @@ from __future__ import annotations +from collections.abc import Iterator, Sequence +from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import datetime -from typing import Any +from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection -from sqlbuild.compiler.compile.types import CompiledResourceType -from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json -from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json -from sqlbuild.executor.node_results.models import ( - NodeResultEnvelope, - NodeResultQuery, - NodeResultRecord, -) from sqlbuild.microbatches.models import MicrobatchEvent, MicrobatchScope, MicrobatchWriteResult from sqlbuild.virtual.state._helpers.state_storage.datetime import ( - from_naive_utc_wall_clock, to_naive_utc_wall_clock, ) -from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id from sqlbuild.virtual.state._helpers.state_storage.microbatch_events import ( append_duckdb_microbatch_event, append_duckdb_microbatch_events, @@ -28,88 +20,38 @@ read_duckdb_microbatch_scope_history, ) from sqlbuild.virtual.state._helpers.state_storage.validation import build_validation_result -from sqlbuild.virtual.state.classes._duckdb_conditional_publish import ( - DuckDbConditionalPublishMixin, -) -from sqlbuild.virtual.state.classes.state_backend import StateBackend +from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( - CURRENT_STATE_SCHEMA_VERSION, DUCKDB_DATETIME_TYPE_TOKEN, DUCKDB_INTEGER_TYPE_TOKEN, DUCKDB_TIMESTAMP_TYPE_TOKEN, - FUNCTION_VERSION_TABLE, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - MODEL_VERSION_TABLE, - NODE_RESULTS_TABLE, - NON_UNIQUE_STATE_INDEXES, - PHYSICAL_RELATION_ANCESTRY_TABLE, - PHYSICAL_RELATION_TABLE, PYTHON_NODE_VERSION_TABLE, - RECONCILE_EVENT_TABLE, - SEED_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, - STATE_BOOLEAN_TRUE, - STATE_MIGRATION_EVENTS_TABLE, - STATE_OPERATION_EVENT_TABLE, - STATE_OPERATION_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, - STATE_TABLES, - STATE_VERSION_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, - VIRTUAL_ENVIRONMENT_TABLE, ) from sqlbuild.virtual.state.exceptions import ( StateBackendConfigError, - StateBackupNotFoundError, - StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( - FunctionVersionRecord, - ModelVersionRecord, - PhysicalRelationAncestryRecord, - PhysicalRelationRecord, - PythonNodeVersionRecord, - ReconcileEventRecord, - SeedVersionRecord, SourceFreshnessRecord, - StateBackupRecord, - StateLockRecord, - StateOperationEventRecord, - StateOperationRecord, + StateLockLease, StateSchemaValidationResult, - VirtualEnvironmentCheckpointFunctionRefRecord, - VirtualEnvironmentCheckpointModelRefRecord, - VirtualEnvironmentCheckpointRecord, - VirtualEnvironmentCheckpointSeedRefRecord, - VirtualEnvironmentFunctionRefRecord, - VirtualEnvironmentModelRefRecord, VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentPythonNodeRefRecord, - VirtualEnvironmentRecord, - VirtualEnvironmentRetentionRecord, - VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( - ModelVersionStatus, - PhysicalArtifactType, StateColumnType, - StateMigrationAction, - StateMigrationStatus, - StateOperationStatus, - StateOperationType, - VirtualEnvironmentStatus, ) -class DuckDbStateBackend(DuckDbConditionalPublishMixin, StateBackend): +class DuckDbStateBackend(SqlStateBackend): """DuckDB implementation for virtual state.""" + _placeholder: ClassVar[str] = "?" + def connect(self, config: dict[str, object]) -> Any: import duckdb @@ -118,54 +60,56 @@ def connect(self, config: dict[str, object]) -> Any: raise StateBackendConfigError("DuckDB state backend requires state.connection.database") return ObservedConnection(raw_connection=duckdb.connect(database), adapter="duckdb") - def close(self, connection: Any) -> None: - connection.close() - - def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: + def _fetch_one( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> tuple[Any, ...] | None: + if params is None: + return connection.execute(sql).fetchone() + return connection.execute(sql, params).fetchone() + + def _fetch_all( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> list[tuple[Any, ...]]: + if params is None: + return connection.execute(sql).fetchall() + return connection.execute(sql, params).fetchall() + + @contextmanager + def _write_transaction(self, *, connection: Any) -> Iterator[Any]: connection.execute("BEGIN") try: - connection.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - connection.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} (" - "schema_version INTEGER NOT NULL, " - "sqlbuild_version TEXT NOT NULL, " - "updated_at TIMESTAMP NOT NULL" - ")" - ) - connection.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} (" - "event_id TEXT NOT NULL, " - "action TEXT NOT NULL, " - "backup_id TEXT, " - "status TEXT NOT NULL, " - "message TEXT, " - "created_at TIMESTAMP NOT NULL" - ")" - ) - self._create_additional_state_tables(connection=connection, schema=schema) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)}" - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} " - "(schema_version, sqlbuild_version, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)", - [CURRENT_STATE_SCHEMA_VERSION, sqlbuild_version], - ) - self._record_event( - connection=connection, - schema=schema, - action=StateMigrationAction.INIT, - backup_id_value=None, - status=StateMigrationStatus.SUCCESS, - message=None, - ) + yield connection connection.execute("COMMIT") except BaseException: connection.execute("ROLLBACK") raise + def _statement_executor(self, *, connection: Any) -> AbstractContextManager[Any]: + return nullcontext(connection) + + def _execute_in( + self, *, executor: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + if params is None: + executor.execute(sql) + return + executor.execute(sql, params) + + def _lease_is_owned(self, *, executor: Any, schema: str, lease: StateLockLease) -> bool: + owned: tuple[Any, ...] | None = executor.execute( + f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " + "SET updated_at = CURRENT_TIMESTAMP " + "WHERE lock_key = ? AND owner_id = ? AND expires_at > CURRENT_TIMESTAMP " + "RETURNING lock_key", + [lease.lock_key, lease.owner_id], + ).fetchone() + return owned is not None + + def _fetch_one_in( + self, *, executor: Any, sql: str, params: Sequence[object] + ) -> tuple[Any, ...] | None: + return executor.execute(sql, params).fetchone() + def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidationResult: tables: set[str] = { row[0] @@ -193,1017 +137,69 @@ def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidati expected_columns=STATE_TABLE_COLUMNS, type_matches=self._state_type_matches, expected_indexes=STATE_TABLE_INDEXES, - existing_indexes_by_table=indexes_by_table, - ) - - def create_backup(self, *, connection: Any, schema: str) -> str: - validation: StateSchemaValidationResult = self.inspect_schema( - connection=connection, schema=schema - ) - if not validation.valid: - raise StateSchemaInvalidError("Cannot backup invalid state schema") - backup_id_value: str = backup_id() - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - connection.execute("BEGIN") - try: - connection.execute(f"CREATE SCHEMA {self._quote_identifier(backup_schema)}") - table_name: str - for table_name in STATE_TABLES: - connection.execute( - f"CREATE TABLE {self._qualified_name(schema=backup_schema, table=table_name)} " - "AS " - f"SELECT * FROM {self._qualified_name(schema=schema, table=table_name)}" - ) - self._record_event( - connection=connection, - schema=schema, - action=StateMigrationAction.BACKUP, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - return backup_id_value - - def rollback(self, *, connection: Any, schema: str, backup_id: str | None = None) -> str: - backup_id_value: str = backup_id or self._latest_backup_id( - connection=connection, schema=schema - ) - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - if not self._schema_exists(connection=connection, schema=backup_schema): - raise StateBackupNotFoundError(f"State backup schema '{backup_schema}' does not exist") - connection.execute("BEGIN") - try: - table_name: str - for table_name in STATE_TABLES: - connection.execute( - f"DROP TABLE IF EXISTS {self._qualified_name(schema=schema, table=table_name)}" - ) - connection.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - for table_name in STATE_TABLES: - connection.execute( - f"CREATE TABLE {self._qualified_name(schema=schema, table=table_name)} AS " - f"SELECT * FROM {self._qualified_name(schema=backup_schema, table=table_name)}" - ) - self._create_state_indexes(connection=connection, schema=schema) - self._record_event( - connection=connection, - schema=schema, - action=StateMigrationAction.ROLLBACK, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - return backup_id_value - - def reset(self, *, connection: Any, schema: str) -> None: - connection.execute("BEGIN") - try: - for table_name in STATE_TABLES: - connection.execute( - f"DROP TABLE IF EXISTS {self._qualified_name(schema=schema, table=table_name)}" - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_model_version( - self, *, connection: Any, schema: str, record: ModelVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=MODEL_VERSION_TABLE, - where_sql="model_name = ? AND version_hash = ?", - params=[record.model_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [record.model_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "(model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status, " - "created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, " - "COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.definition_identity_hash, - record.identity_metadata_hash, - record.definition_text_b64, - record.identity_metadata_json_b64, - record.compiled_sql_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_model_version( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> ModelVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status " - f"FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [model_name, version_hash], - ).fetchone() - if row is None: - return None - return ModelVersionRecord( - model_name=row[0], - version_hash=row[1], - definition_identity_hash=row[2], - identity_metadata_hash=row[3], - definition_text_b64=row[4], - identity_metadata_json_b64=row[5], - compiled_sql_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - - def upsert_function_version( - self, *, connection: Any, schema: str, record: FunctionVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=FUNCTION_VERSION_TABLE, - where_sql="function_name = ? AND version_hash = ?", - params=[record.function_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = ? AND version_hash = ?", - [record.function_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "(function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, " - "COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.function_name, - record.version_hash, - record.language, - record.returns, - record.arguments_json_b64, - record.return_columns_json_b64, - record.packages_json_b64, - record.runtime_version, - record.entry_point, - record.body_sql_b64, - record.definition_text_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_function_version( - self, *, connection: Any, schema: str, function_name: str, version_hash: str - ) -> FunctionVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status " - f"FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = ? AND version_hash = ?", - [function_name, version_hash], - ).fetchone() - if row is None: - return None - return FunctionVersionRecord( - function_name=row[0], - version_hash=row[1], - language=row[2], - returns=row[3], - arguments_json_b64=row[4], - return_columns_json_b64=row[5], - packages_json_b64=row[6], - runtime_version=row[7], - entry_point=row[8], - body_sql_b64=row[9], - definition_text_b64=row[10], - status=ModelVersionStatus(row[11]), - ) - - def upsert_seed_version( - self, *, connection: Any, schema: str, record: SeedVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=SEED_VERSION_TABLE, - where_sql="seed_name = ? AND version_hash = ?", - params=[record.seed_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = ? AND version_hash = ?", - [record.seed_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "(seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.seed_name, - record.version_hash, - record.identity_metadata_hash, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_seed_version( - self, *, connection: Any, schema: str, seed_name: str, version_hash: str - ) -> SeedVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = ? AND version_hash = ?", - [seed_name, version_hash], - ).fetchone() - if row is None: - return None - return SeedVersionRecord( - seed_name=row[0], - version_hash=row[1], - identity_metadata_hash=row[2], - identity_metadata_json_b64=row[3], - status=ModelVersionStatus(row[4]), - ) - - def upsert_python_node_version( - self, *, connection: Any, schema: str, record: PythonNodeVersionRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=PYTHON_NODE_VERSION_TABLE, - where_sql="node_type = ? AND node_name = ? AND version_hash = ?", - params=[record.node_type, record.node_name, record.version_hash], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = ? AND node_name = ? AND version_hash = ?", - [record.node_type, record.node_name, record.version_hash], - ) - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "(node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, " - "status, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), " - "CURRENT_TIMESTAMP)", - [ - record.node_type, - record.node_name, - record.version_hash, - record.definition_hash, - record.identity_metadata_hash, - record.definition_json_b64, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_python_node_version( - self, - *, - connection: Any, - schema: str, - node_type: str, - node_name: str, - version_hash: str, - ) -> PythonNodeVersionRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = ? AND node_name = ? AND version_hash = ?", - [node_type, node_name, version_hash], - ).fetchone() - if row is None: - return None - return PythonNodeVersionRecord( - node_type=row[0], - node_name=row[1], - version_hash=row[2], - definition_hash=row[3], - identity_metadata_hash=row[4], - definition_json_b64=row[5], - identity_metadata_json_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - - def insert_node_result( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - record: NodeResultRecord, - ) -> None: - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - "(virtual_environment_name, node_type, node_name, target_database, target_schema, " - "target_name, run_id, status, payload_json_b64, metadata_json_b64, error_message, " - "materialized, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - [ - virtual_environment_name, - record.node_type, - record.node_name, - record.target_database, - record.target_schema, - record.target_name, - record.run_id, - record.status, - encode_node_result_json( - value=record.payload, label="payload", node_name=record.node_name - ), - encode_node_result_json( - value=record.metadata, label="metadata", node_name=record.node_name - ), - record.error_message, - self._materialized_storage(record.materialized), - record.ts, - ], - ) - - def read_node_results( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - query: NodeResultQuery, - ) -> tuple[NodeResultEnvelope, ...]: - if query.limit < 1: - return () - predicates: list[str] = [ - "virtual_environment_name = ?", - "node_type = ?", - "node_name = ?", - self._optional_equality_sql( - column="target_database", value=query.target_database, placeholder="?" - ), - self._optional_equality_sql( - column="target_schema", value=query.target_schema, placeholder="?" - ), - self._optional_equality_sql( - column="target_name", value=query.target_name, placeholder="?" - ), - ] - params: list[object] = [virtual_environment_name, query.node_type, query.node_name] - for value in (query.target_database, query.target_schema, query.target_name): - if value is not None: - params.append(value) - if query.statuses is not None: - placeholders: str = ", ".join("?" for _ in query.statuses) - predicates.append(f"status IN ({placeholders})") - params.extend(query.statuses) - if query.run_id is not None: - predicates.append("run_id = ?") - params.append(query.run_id) - params.append(query.limit) - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT node_type, node_name, run_id, status, payload_json_b64, metadata_json_b64, " - "error_message, materialized, created_at " - f"FROM {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - f"WHERE {' AND '.join(predicates)} " - "ORDER BY created_at DESC, run_id DESC LIMIT ?", - params, - ).fetchall() - return tuple(self._node_result_row_to_envelope(row) for row in rows) - - def append_microbatch_event( - self, *, connection: Any, schema: str, event: MicrobatchEvent - ) -> None: - append_duckdb_microbatch_event( - connection=connection, - qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), - event=event, - ) - - def append_microbatch_events( - self, *, connection: Any, schema: str, events: tuple[MicrobatchEvent, ...] - ) -> MicrobatchWriteResult: - return append_duckdb_microbatch_events( - connection=connection, - qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), - events=events, - ) - - def read_microbatch_scope_history( - self, *, connection: Any, schema: str, scope: MicrobatchScope - ) -> tuple[MicrobatchEvent, ...]: - return read_duckdb_microbatch_scope_history( - connection=connection, - qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), - scope=scope, - ) - - def read_microbatch_retention_history( - self, *, connection: Any, schema: str - ) -> tuple[MicrobatchEvent, ...]: - return read_duckdb_microbatch_retention_history( - connection=connection, - qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), - ) - - def read_microbatch_model_history( - self, *, connection: Any, schema: str, scope: MicrobatchScope - ) -> tuple[MicrobatchEvent, ...]: - return read_duckdb_microbatch_model_history( - connection=connection, - qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), - scope=scope, - ) - - def upsert_physical_relation( - self, *, connection: Any, schema: str, record: PhysicalRelationRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=PHYSICAL_RELATION_TABLE, - where_sql="artifact_type = ? AND artifact_name = ? AND version_hash = ?", - params=[record.artifact_type.value, record.artifact_name, record.version_hash], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = ? AND artifact_name = ? AND version_hash = ?", - [record.artifact_type.value, record.artifact_name, record.version_hash], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "(artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.artifact_type.value, - record.artifact_name, - record.version_hash, - record.database_name, - record.schema_name, - record.relation_name, - record.relation_type, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_physical_relation_for_artifact( - self, - *, - connection: Any, - schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - version_hash: str, - ) -> PhysicalRelationRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = ? AND artifact_name = ? AND version_hash = ?", - [artifact_type.value, artifact_name, version_hash], - ).fetchone() - if row is None: - return None - return PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - - def list_physical_relations_for_artifact( - self, - *, - connection: Any, - schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - ) -> tuple[PhysicalRelationRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = ? AND artifact_name = ? " - "ORDER BY updated_at DESC, version_hash DESC", - [artifact_type.value, artifact_name], - ).fetchall() - return tuple( - PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - for row in rows - ) - - def upsert_physical_relation_ancestry( - self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=PHYSICAL_RELATION_ANCESTRY_TABLE, - where_sql="model_name = ? AND version_hash = ?", - params=[record.model_name, record.version_hash], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [record.model_name, record.version_hash], - ) - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "(model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy, created_at, updated_at) " - "VALUES (?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.parent_model_name, - record.parent_version_hash, - record.seed_strategy, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_physical_relation_ancestry( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> PhysicalRelationAncestryRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "WHERE model_name = ? AND version_hash = ?", - [model_name, version_hash], - ).fetchone() - if row is None: - return None - return PhysicalRelationAncestryRecord( - model_name=row[0], - version_hash=row[1], - parent_model_name=row[2], - parent_version_hash=row[3], - seed_strategy=row[4], - ) - - def upsert_virtual_environment( - self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - connection.execute("BEGIN") - try: - self._upsert_virtual_environment_record( - connection=connection, schema=schema, record=record - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> VirtualEnvironmentRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT virtual_environment_name, status, baseline_virtual_environment_name, " - "finalized_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ).fetchone() - if row is None: - return None - return VirtualEnvironmentRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - baseline_virtual_environment_name=row[2], - finalized_at=row[3], - ) - - def list_virtual_environments( - self, *, connection: Any, schema: str - ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, status, updated_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "ORDER BY updated_at DESC, virtual_environment_name DESC" - ).fetchall() - return tuple( - VirtualEnvironmentRetentionRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - updated_at=row[2], - ) - for row in rows - ) - - def delete_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> None: - connection.execute("BEGIN") - try: - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = ?", - [virtual_environment_name], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def replace_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type={node_type: refs}, - ) - - def replace_virtual_environment_node_ref_groups( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - connection.execute("BEGIN") - try: - self._replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_virtual_environment_and_replace_node_ref_groups( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - connection.execute("BEGIN") - try: - self._upsert_virtual_environment_record( - connection=connection, schema=schema, record=record - ) - self._replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - ) -> tuple[VirtualEnvironmentNodeRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = ? AND node_type = ? ORDER BY node_name", - [virtual_environment_name, node_type], - ).fetchall() - return tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - - def upsert_virtual_environment_node_ref( - self, - *, - connection: Any, - schema: str, - ref: VirtualEnvironmentNodeRefRecord, - ) -> None: - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "(virtual_environment_name, node_type, node_name, version_hash, updated_at) " - "VALUES (?, ?, ?, ?, now()) " - "ON CONFLICT (virtual_environment_name, node_type, node_name) " - "DO UPDATE SET version_hash = excluded.version_hash, updated_at = now()", - [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], + existing_indexes_by_table=indexes_by_table, ) - def replace_virtual_environment_model_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentModelRefRecord, ...], + def append_microbatch_event( + self, *, connection: Any, schema: str, event: MicrobatchEvent ) -> None: - self.replace_virtual_environment_node_refs( + append_duckdb_microbatch_event( connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="model", - node_name=ref.model_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), + qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), + event=event, ) - def get_virtual_environment_model_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentModelRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( + def append_microbatch_events( + self, *, connection: Any, schema: str, events: tuple[MicrobatchEvent, ...] + ) -> MicrobatchWriteResult: + return append_duckdb_microbatch_events( connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - ) - return tuple( - VirtualEnvironmentModelRefRecord( - virtual_environment_name=ref.virtual_environment_name, - model_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs + qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), + events=events, ) - def replace_virtual_environment_function_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentFunctionRefRecord, ...], - ) -> None: - ref: VirtualEnvironmentFunctionRefRecord - for ref in refs: - if ref.node_type not in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - raise StateBackendConfigError("Function ref node_type must be 'udf' or 'table_fn'") - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]] = {} - for node_type in ("udf", "table_fn"): - node_refs: list[VirtualEnvironmentNodeRefRecord] = [] - for ref in refs: - if ref.node_type == node_type: - node_refs.append( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.function_name, - version_hash=ref.version_hash, - ) - ) - refs_by_node_type[node_type] = tuple(node_refs) - self.replace_virtual_environment_node_ref_groups( + def read_microbatch_scope_history( + self, *, connection: Any, schema: str, scope: MicrobatchScope + ) -> tuple[MicrobatchEvent, ...]: + return read_duckdb_microbatch_scope_history( connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - - def get_virtual_environment_function_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentFunctionRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="udf", - ), - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="table_fn", - ), - ) - return tuple( - VirtualEnvironmentFunctionRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - function_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in sorted(refs, key=lambda item: (item.node_type, item.node_name)) + qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), + scope=scope, ) - def replace_virtual_environment_seed_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentSeedRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( + def read_microbatch_retention_history( + self, *, connection: Any, schema: str + ) -> tuple[MicrobatchEvent, ...]: + return read_duckdb_microbatch_retention_history( connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="seed", - node_name=ref.seed_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), + qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), ) - def get_virtual_environment_seed_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentSeedRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( + def read_microbatch_model_history( + self, *, connection: Any, schema: str, scope: MicrobatchScope + ) -> tuple[MicrobatchEvent, ...]: + return read_duckdb_microbatch_model_history( connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - ) - return tuple( - VirtualEnvironmentSeedRefRecord( - virtual_environment_name=ref.virtual_environment_name, - seed_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs + qualified_table=self._qualified_name(schema=schema, table=MICROBATCH_EVENT_TABLE), + scope=scope, ) - def upsert_virtual_environment_python_node_ref( + def upsert_virtual_environment_node_ref( self, *, connection: Any, schema: str, - ref: VirtualEnvironmentPythonNodeRefRecord, + ref: VirtualEnvironmentNodeRefRecord, ) -> None: - self.upsert_virtual_environment_node_ref( - connection=connection, - schema=schema, - ref=VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.node_name, - version_hash=ref.version_hash, - ), - ) - - def get_virtual_environment_python_node_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentPythonNodeRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = ? " - "AND node_type IN ('task', 'loader', 'asset', 'check', 'hook') " - "ORDER BY node_type, node_name", - [virtual_environment_name], - ).fetchall() - return tuple( - VirtualEnvironmentPythonNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: - row: tuple[Any, ...] = connection.execute( - "SELECT COUNT(*) " - f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} versions " - "WHERE NOT EXISTS (" - "SELECT 1 " - "FROM " + connection.execute( + "INSERT INTO " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "refs " - "WHERE refs.node_type = versions.node_type " - "AND refs.node_name = versions.node_name " - "AND refs.version_hash = versions.version_hash)" - ).fetchone() - return int(row[0]) + "(virtual_environment_name, node_type, node_name, version_hash, updated_at) " + "VALUES (?, ?, ?, ?, now()) " + "ON CONFLICT (virtual_environment_name, node_type, node_name) " + "DO UPDATE SET version_hash = excluded.version_hash, updated_at = now()", + [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], + ) def prune_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: before_count: int = self.count_unreferenced_python_node_versions( @@ -1296,259 +292,6 @@ def replace_virtual_environment_source_freshness( connection.execute("ROLLBACK") raise - def get_virtual_environment_source_freshness( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[SourceFreshnessRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT virtual_environment_name, source_name, strategy, value_kind, " - "data_version, data_version_hash, observed_at " - f"FROM {self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " - "WHERE virtual_environment_name = ? ORDER BY source_name", - [virtual_environment_name], - ).fetchall() - return tuple( - SourceFreshnessRecord( - virtual_environment_name=row[0], - source_name=row[1], - strategy=row[2], - value_kind=row[3], - data_version=row[4], - data_version_hash=row[5], - observed_at=from_naive_utc_wall_clock(row[6]), - ) - for row in rows - ) - - def create_virtual_environment_checkpoint( - self, - *, - connection: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> None: - connection.execute("BEGIN") - try: - self._insert_virtual_environment_checkpoint_rows( - connection=connection, - schema=schema, - checkpoint=checkpoint, - refs=refs, - function_refs=function_refs, - seed_refs=seed_refs, - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def list_virtual_environment_checkpoints( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentCheckpointRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT checkpoint_id, virtual_environment_name, created_at " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE)} " - "WHERE virtual_environment_name = ? ORDER BY created_at DESC, checkpoint_id DESC", - [virtual_environment_name], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointRecord( - checkpoint_id=row[0], - virtual_environment_name=row[1], - created_at=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_model_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointModelRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT checkpoint_id, model_name, version_hash FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - + " " - "WHERE checkpoint_id = ? ORDER BY model_name", - [checkpoint_id], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointModelRefRecord( - checkpoint_id=row[0], - model_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_function_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...]: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - rows: list[tuple[Any, ...]] = connection.execute( - f"SELECT checkpoint_id, function_name, version_hash " - f"FROM {checkpoint_function_ref_table} " - "WHERE checkpoint_id = ? ORDER BY function_name", - [checkpoint_id], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointFunctionRefRecord( - checkpoint_id=row[0], - function_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_seed_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT checkpoint_id, seed_name, version_hash FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - + " " - "WHERE checkpoint_id = ? ORDER BY seed_name", - [checkpoint_id], - ).fetchall() - return tuple( - VirtualEnvironmentCheckpointSeedRefRecord( - checkpoint_id=row[0], - seed_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def delete_virtual_environment_checkpoint( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> None: - connection.execute("BEGIN") - try: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - checkpoint_seed_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - checkpoint_model_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - connection.execute( - f"DELETE FROM {checkpoint_seed_ref_table} WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute( - f"DELETE FROM {checkpoint_function_ref_table} WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute( - f"DELETE FROM {checkpoint_model_ref_table} WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE)}" - " " - "WHERE checkpoint_id = ?", - [checkpoint_id], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def upsert_state_operation( - self, *, connection: Any, schema: str, record: StateOperationRecord - ) -> None: - connection.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=STATE_OPERATION_TABLE, - where_sql="operation_id = ?", - params=[record.operation_id], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = ?", - [record.operation_id], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "(operation_id, operation_type, status, virtual_environment_name, " - "created_at, updated_at) " - "VALUES (?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.operation_id, - record.operation_type.value, - record.status.value, - record.virtual_environment_name, - existing_created_at, - ], - ) - connection.execute("COMMIT") - except BaseException: - connection.execute("ROLLBACK") - raise - - def get_state_operation( - self, *, connection: Any, schema: str, operation_id: str - ) -> StateOperationRecord | None: - row: tuple[Any, ...] | None = connection.execute( - "SELECT operation_id, operation_type, status, virtual_environment_name " - f"FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = ?", - [operation_id], - ).fetchone() - if row is None: - return None - return StateOperationRecord( - operation_id=row[0], - operation_type=StateOperationType(row[1]), - status=StateOperationStatus(row[2]), - virtual_environment_name=row[3], - ) - - def create_state_operation_event( - self, *, connection: Any, schema: str, record: StateOperationEventRecord - ) -> None: - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_OPERATION_EVENT_TABLE)} " - "(event_id, operation_id, action, status, message, created_at) " - "VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)", - [ - record.event_id, - record.operation_id, - record.action, - record.status.value, - record.message, - ], - ) - - def create_reconcile_event( - self, *, connection: Any, schema: str, record: ReconcileEventRecord - ) -> None: - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=RECONCILE_EVENT_TABLE)} " - "(event_id, action, status, message, created_at) " - "VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)", - [record.event_id, record.action.value, record.status.value, record.message], - ) - def acquire_lock( self, *, @@ -1609,44 +352,6 @@ def release_lock(self, *, connection: Any, schema: str, lock_key: str, owner_id: connection.execute("ROLLBACK") raise - def renew_lock( - self, - *, - connection: Any, - schema: str, - lock_key: str, - owner_id: str, - expires_at: datetime, - ) -> bool: - row: tuple[Any, ...] | None = connection.execute( - f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "SET expires_at = ?, updated_at = CURRENT_TIMESTAMP " - "WHERE lock_key = ? AND owner_id = ? AND expires_at > CURRENT_TIMESTAMP " - "RETURNING lock_key", - [expires_at, lock_key, owner_id], - ).fetchone() - return row is not None - - def list_active_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at > CURRENT_TIMESTAMP ORDER BY lock_key" - ).fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - - def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY lock_key" - ).fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: connection.execute( "DELETE FROM " @@ -1654,255 +359,14 @@ def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: [lock_key], ) - def list_state_backups(self, *, connection: Any, schema: str) -> tuple[StateBackupRecord, ...]: - prefix: str = f"{schema}__backup_%" - rows: list[tuple[Any, ...]] = connection.execute( - "SELECT s.schema_name, e.backup_id, MAX(e.created_at) " - "FROM information_schema.schemata s " - "LEFT JOIN " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} e " - "ON s.schema_name = ? || e.backup_id " - "WHERE s.schema_name LIKE ? " - "GROUP BY s.schema_name, e.backup_id ORDER BY s.schema_name DESC", - [f"{schema}__backup_", prefix], - ).fetchall() - return tuple( - StateBackupRecord( - backup_id=row[1] or str(row[0]).removeprefix(f"{schema}__backup_"), - schema_name=row[0], - created_at=row[2], - ) - for row in rows - ) - def delete_state_backup(self, *, connection: Any, schema: str, backup_id: str) -> None: backup_schema: str = self._backup_schema_name(schema=schema, backup_id_value=backup_id) connection.execute(f"DROP SCHEMA IF EXISTS {self._quote_identifier(backup_schema)} CASCADE") - def _latest_backup_id(self, *, connection: Any, schema: str) -> str: - prefix: str = f"{schema}__backup_%" - rows: list[tuple[str]] = connection.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE ? " - "ORDER BY schema_name DESC LIMIT 1", - [prefix], - ).fetchall() - if not rows: - raise StateBackupNotFoundError("No state backup is available for rollback") - return rows[0][0].removeprefix(f"{schema}__backup_") - - def _schema_exists(self, *, connection: Any, schema: str) -> bool: - rows: list[tuple[str]] = connection.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name = ?", - [schema], - ).fetchall() - return bool(rows) - - def _record_event( - self, - *, - connection: Any, - schema: str, - action: StateMigrationAction, - backup_id_value: str | None, - status: StateMigrationStatus, - message: str | None, - ) -> None: - connection.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} " - "(event_id, action, backup_id, status, message, created_at) " - "VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)", - [event_id(), action.value, backup_id_value, status.value, message], - ) - - def _created_at_for_key( - self, - *, - connection: Any, - schema: str, - table_name: str, - where_sql: str, - params: list[object], - ) -> datetime | None: - row: tuple[Any, ...] | None = connection.execute( - f"SELECT created_at FROM {self._qualified_name(schema=schema, table=table_name)} " - f"WHERE {where_sql}", - params, - ).fetchone() - if row is None: - return None - return row[0] - - def _create_additional_state_tables(self, *, connection: Any, schema: str) -> None: - table_name: str - columns: dict[str, StateColumnType] - for table_name, columns in STATE_TABLE_COLUMNS.items(): - if table_name in {STATE_VERSION_TABLE, STATE_MIGRATION_EVENTS_TABLE}: - continue - column_sql: str = ", ".join( - f"{self._quote_identifier(column_name)} {self._state_column_sql_type(column_type)}" - for column_name, column_type in columns.items() - ) - connection.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)} " - f"({column_sql})" - ) - column_name: str - column_type: StateColumnType - for column_name, column_type in columns.items(): - connection.execute( - f"ALTER TABLE {self._qualified_name(schema=schema, table=table_name)} " - f"ADD COLUMN IF NOT EXISTS {self._quote_identifier(column_name)} " - f"{self._state_column_sql_type(column_type)}" - ) - self._create_state_indexes(connection=connection, schema=schema) - - def _create_state_indexes(self, *, connection: Any, schema: str) -> None: - table_name: str - indexes: dict[str, tuple[str, ...]] - for table_name, indexes in STATE_TABLE_INDEXES.items(): - index_name: str - columns: tuple[str, ...] - for index_name, columns in indexes.items(): - column_sql: str = ", ".join(self._quote_identifier(column) for column in columns) - unique_sql: str = "" if index_name in NON_UNIQUE_STATE_INDEXES else "UNIQUE " - connection.execute( - f"CREATE {unique_sql}INDEX IF NOT EXISTS {self._quote_identifier(index_name)} " - f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})" - ) - - def _state_column_sql_type(self, column_type: StateColumnType) -> str: - match column_type: - case StateColumnType.INTEGER: - return "INTEGER" - case StateColumnType.TEXT: - return "TEXT" - case StateColumnType.TIMESTAMP: - return "TIMESTAMP" - raise StateBackendConfigError(f"Unsupported state column type: {column_type}") - - def _node_result_row_to_envelope(self, row: tuple[Any, ...]) -> NodeResultEnvelope: - node_name: str = str(row[1]) - metadata: object = decode_node_result_json( - value=str(row[5]), label="metadata", node_name=node_name - ) - normalized_metadata: dict[str, object] = ( - {str(key): value for key, value in metadata.items()} - if isinstance(metadata, dict) - else {} - ) - return NodeResultEnvelope( - node_type=str(row[0]), - node_name=node_name, - run_id=str(row[2]), - status=str(row[3]), - payload=decode_node_result_json( - value=str(row[4]), label="payload", node_name=node_name - ), - metadata=normalized_metadata, - error_message=str(row[6]) if row[6] is not None else None, - materialized=self._parse_materialized(row[7]), - ts=row[8], - ) - - def _optional_equality_sql(self, *, column: str, value: object | None, placeholder: str) -> str: - if value is None: - return f"{column} IS NULL" - return f"{column} = {placeholder}" - - def _materialized_storage(self, value: bool | None) -> str | None: - if value is None: - return None - return "true" if value else "false" - - def _parse_materialized(self, value: object) -> bool | None: - if value is None: - return None - return str(value).lower() == STATE_BOOLEAN_TRUE - - def _quote_identifier(self, identifier: str) -> str: - return '"' + identifier.replace('"', '""') + '"' - - def _qualified_name(self, *, schema: str, table: str) -> str: - return f"{self._quote_identifier(schema)}.{self._quote_identifier(table)}" - - def _validate_source_freshness_records( - self, - *, - virtual_environment_name: str, - records: tuple[SourceFreshnessRecord, ...], - ) -> None: - seen_source_names: set[str] = set() - record: SourceFreshnessRecord - for record in records: - if record.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Source freshness record virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if record.source_name in seen_source_names: - raise StateBackendConfigError( - f"Duplicate source freshness record for source '{record.source_name}'" - ) - seen_source_names.add(record.source_name) - - def _validate_node_ref_replacement( - self, - *, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - seen_node_names: set[str] = set() - ref: VirtualEnvironmentNodeRefRecord - for ref in refs: - if ref.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Node ref virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if ref.node_type != node_type: - raise StateBackendConfigError("Node ref node_type must match replacement node_type") - if ref.node_name in seen_node_names: - raise StateBackendConfigError( - f"Duplicate node ref for node type '{node_type}' and name '{ref.node_name}'" - ) - seen_node_names.add(ref.node_name) - - def _upsert_virtual_environment_record( - self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - existing_created_at: datetime | None = self._created_at_for_key( - connection=connection, - schema=schema, - table_name=VIRTUAL_ENVIRONMENT_TABLE, - where_sql="virtual_environment_name = ?", - params=[record.virtual_environment_name], - ) - connection.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = ?", - [record.virtual_environment_name], - ) - connection.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "(virtual_environment_name, status, baseline_virtual_environment_name, " - "created_at, updated_at, finalized_at) " - "VALUES (?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP, ?)", - [ - record.virtual_environment_name, - record.status.value, - record.baseline_virtual_environment_name, - existing_created_at, - record.finalized_at, - ], - ) - def _replace_virtual_environment_node_ref_groups( self, *, - connection: Any, + executor: Any, schema: str, virtual_environment_name: str, refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], @@ -1916,8 +380,8 @@ def _replace_virtual_environment_node_ref_groups( node_type=node_type, refs=refs, ) - connection.execute(f"DROP TABLE IF EXISTS {temp_table_name}") - connection.execute( + executor.execute(f"DROP TABLE IF EXISTS {temp_table_name}") + executor.execute( f"CREATE TEMP TABLE {temp_table_name} (" "virtual_environment_name TEXT NOT NULL, " "node_type TEXT NOT NULL, " @@ -1927,7 +391,7 @@ def _replace_virtual_environment_node_ref_groups( ) ref: VirtualEnvironmentNodeRefRecord for ref in refs: - connection.execute( + executor.execute( f"INSERT INTO {temp_table_name} " "(virtual_environment_name, node_type, node_name, version_hash) " "VALUES (?, ?, ?, ?)", @@ -1938,7 +402,7 @@ def _replace_virtual_environment_node_ref_groups( ref.version_hash, ], ) - connection.execute( + executor.execute( "DELETE FROM " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " "WHERE virtual_environment_name = ? AND node_type = ? " @@ -1949,7 +413,7 @@ def _replace_virtual_environment_node_ref_groups( ".node_name)", [virtual_environment_name, node_type], ) - connection.execute( + executor.execute( "INSERT INTO " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " "(virtual_environment_name, node_type, node_name, version_hash, updated_at) " @@ -1958,10 +422,7 @@ def _replace_virtual_environment_node_ref_groups( "ON CONFLICT (virtual_environment_name, node_type, node_name) " "DO UPDATE SET version_hash = excluded.version_hash, updated_at = now()" ) - connection.execute(f"DROP TABLE IF EXISTS {temp_table_name}") - - def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: - return f"{schema}__backup_{backup_id_value}" + executor.execute(f"DROP TABLE IF EXISTS {temp_table_name}") def _state_type_matches(self, *, actual_type: str, expected_type: StateColumnType) -> bool: actual: str = actual_type.lower() diff --git a/src/sqlbuild/virtual/state/classes/postgres.py b/src/sqlbuild/virtual/state/classes/postgres.py index 48e910d476..867e80760e 100644 --- a/src/sqlbuild/virtual/state/classes/postgres.py +++ b/src/sqlbuild/virtual/state/classes/postgres.py @@ -2,18 +2,12 @@ from __future__ import annotations +from collections.abc import Iterator, Sequence +from contextlib import AbstractContextManager, contextmanager from datetime import datetime -from typing import Any +from typing import Any, ClassVar from sqlbuild.adapter.contract.classes.observed_connection import ObservedConnection -from sqlbuild.compiler.compile.types import CompiledResourceType -from sqlbuild.executor.node_results.main.decode_json import decode_node_result_json -from sqlbuild.executor.node_results.main.encode_json import encode_node_result_json -from sqlbuild.executor.node_results.models import ( - NodeResultEnvelope, - NodeResultQuery, - NodeResultRecord, -) from sqlbuild.microbatches.classes.event_codec import MicrobatchEventCodec from sqlbuild.microbatches.constants import ( MICROBATCH_COLUMNS, @@ -22,93 +16,42 @@ ) from sqlbuild.microbatches.models import MicrobatchEvent, MicrobatchScope, MicrobatchWriteResult from sqlbuild.virtual.state._helpers.state_storage.datetime import ( - from_naive_utc_wall_clock, to_naive_utc_wall_clock, ) -from sqlbuild.virtual.state._helpers.state_storage.events import backup_id, event_id from sqlbuild.virtual.state._helpers.state_storage.validation import ( build_validation_result, - validate_conditional_virtual_environment_publication, ) -from sqlbuild.virtual.state.classes.state_backend import StateBackend +from sqlbuild.virtual.state.classes._sql_state_backend import SqlStateBackend from sqlbuild.virtual.state.constants import ( - CURRENT_STATE_SCHEMA_VERSION, - FUNCTION_VERSION_TABLE, LOCK_TABLE, MICROBATCH_EVENT_TABLE, - MODEL_VERSION_TABLE, - NODE_RESULTS_TABLE, - NON_UNIQUE_STATE_INDEXES, - PHYSICAL_RELATION_ANCESTRY_TABLE, - PHYSICAL_RELATION_TABLE, POSTGRES_INTEGER_TYPES, POSTGRES_TEXT_TYPES, PYTHON_NODE_VERSION_TABLE, - RECONCILE_EVENT_TABLE, - SEED_VERSION_TABLE, SOURCE_FRESHNESS_OBSERVATION_TABLE, - STATE_BOOLEAN_TRUE, - STATE_MIGRATION_EVENTS_TABLE, - STATE_OPERATION_EVENT_TABLE, - STATE_OPERATION_TABLE, STATE_TABLE_COLUMNS, STATE_TABLE_INDEXES, - STATE_TABLES, - STATE_VERSION_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, - VIRTUAL_ENVIRONMENT_TABLE, ) from sqlbuild.virtual.state.exceptions import ( StateBackendConfigError, - StateBackupNotFoundError, - StateSchemaInvalidError, ) from sqlbuild.virtual.state.models import ( - FunctionVersionRecord, - ModelVersionRecord, - PhysicalRelationAncestryRecord, - PhysicalRelationRecord, - PythonNodeVersionRecord, - ReconcileEventRecord, - SeedVersionRecord, SourceFreshnessRecord, - StateBackupRecord, StateLockLease, - StateLockRecord, - StateOperationEventRecord, - StateOperationRecord, StateSchemaValidationResult, - VirtualEnvironmentCheckpointFunctionRefRecord, - VirtualEnvironmentCheckpointModelRefRecord, - VirtualEnvironmentCheckpointRecord, - VirtualEnvironmentCheckpointSeedRefRecord, - VirtualEnvironmentFunctionRefRecord, - VirtualEnvironmentModelRefRecord, VirtualEnvironmentNodeRefRecord, - VirtualEnvironmentPythonNodeRefRecord, - VirtualEnvironmentRecord, - VirtualEnvironmentRetentionRecord, - VirtualEnvironmentSeedRefRecord, ) from sqlbuild.virtual.state.types import ( - ModelVersionStatus, - PhysicalArtifactType, StateColumnType, - StateMigrationAction, - StateMigrationStatus, - StateOperationStatus, - StateOperationType, - VirtualEnvironmentStatus, ) -class PostgresStateBackend(StateBackend): +class PostgresStateBackend(SqlStateBackend): """Postgres implementation for virtual state.""" + _placeholder: ClassVar[str] = "%s" + def connect(self, config: dict[str, object]) -> Any: try: import psycopg @@ -133,56 +76,64 @@ def connect(self, config: dict[str, object]) -> Any: except Exception as error: raise StateBackendConfigError("Could not connect to Postgres state backend") from error - def close(self, connection: Any) -> None: - connection.close() + def _fetch_one( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> tuple[Any, ...] | None: + with connection.cursor() as cursor: + if params is None: + cursor.execute(sql) + else: + cursor.execute(sql, params) + return cursor.fetchone() + + def _fetch_all( + self, *, connection: Any, sql: str, params: Sequence[object] | None = None + ) -> list[tuple[Any, ...]]: + with connection.cursor() as cursor: + if params is None: + cursor.execute(sql) + else: + cursor.execute(sql, params) + return cursor.fetchall() - def initialize(self, *, connection: Any, schema: str, sqlbuild_version: str) -> None: + @contextmanager + def _write_transaction(self, *, connection: Any) -> Iterator[Any]: with connection.cursor() as cursor: cursor.execute("BEGIN") try: - cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - cursor.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} (" - "schema_version INTEGER NOT NULL, " - "sqlbuild_version TEXT NOT NULL, " - "updated_at TIMESTAMP NOT NULL" - ")" - ) - cursor.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} (" - "event_id TEXT NOT NULL, " - "action TEXT NOT NULL, " - "backup_id TEXT, " - "status TEXT NOT NULL, " - "message TEXT, " - "created_at TIMESTAMP NOT NULL" - ")" - ) - self._create_additional_state_tables(cursor=cursor, schema=schema) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)}" - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=STATE_VERSION_TABLE)} " - "(schema_version, sqlbuild_version, updated_at) " - "VALUES (%s, %s, CURRENT_TIMESTAMP)", - [CURRENT_STATE_SCHEMA_VERSION, sqlbuild_version], - ) - self._record_event( - cursor=cursor, - schema=schema, - action=StateMigrationAction.INIT, - backup_id_value=None, - status=StateMigrationStatus.SUCCESS, - message=None, - ) + yield cursor cursor.execute("COMMIT") except BaseException: cursor.execute("ROLLBACK") raise + def _statement_executor(self, *, connection: Any) -> AbstractContextManager[Any]: + return connection.cursor() + + def _execute_in( + self, *, executor: Any, sql: str, params: Sequence[object] | None = None + ) -> None: + if params is None: + executor.execute(sql) + return + executor.execute(sql, params) + + def _lease_is_owned(self, *, executor: Any, schema: str, lease: StateLockLease) -> bool: + lock_table: str = self._qualified_name(schema=schema, table=LOCK_TABLE) + executor.execute( + f"SELECT lock_key FROM {lock_table} " + "WHERE lock_key = %s AND owner_id = %s " + "AND expires_at > CURRENT_TIMESTAMP FOR UPDATE", + [lease.lock_key, lease.owner_id], + ) + return executor.fetchone() is not None + + def _fetch_one_in( + self, *, executor: Any, sql: str, params: Sequence[object] + ) -> tuple[Any, ...] | None: + executor.execute(sql, params) + return executor.fetchone() + def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidationResult: with connection.cursor() as cursor: cursor.execute( @@ -214,465 +165,6 @@ def inspect_schema(self, *, connection: Any, schema: str) -> StateSchemaValidati existing_indexes_by_table=indexes_by_table, ) - def create_backup(self, *, connection: Any, schema: str) -> str: - validation: StateSchemaValidationResult = self.inspect_schema( - connection=connection, schema=schema - ) - if not validation.valid: - raise StateSchemaInvalidError("Cannot backup invalid state schema") - backup_id_value: str = backup_id() - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - cursor.execute(f"CREATE SCHEMA {self._quote_identifier(backup_schema)}") - table_name: str - for table_name in STATE_TABLES: - cursor.execute( - "CREATE TABLE " - + self._qualified_name(schema=backup_schema, table=table_name) - + " AS " - f"SELECT * FROM {self._qualified_name(schema=schema, table=table_name)}" - ) - self._record_event( - cursor=cursor, - schema=schema, - action=StateMigrationAction.BACKUP, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - return backup_id_value - - def rollback(self, *, connection: Any, schema: str, backup_id: str | None = None) -> str: - backup_id_value: str = backup_id or self._latest_backup_id( - connection=connection, schema=schema - ) - backup_schema: str = self._backup_schema_name( - schema=schema, - backup_id_value=backup_id_value, - ) - if not self._schema_exists(connection=connection, schema=backup_schema): - raise StateBackupNotFoundError(f"State backup schema '{backup_schema}' does not exist") - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - table_name: str - for table_name in STATE_TABLES: - cursor.execute( - "DROP TABLE IF EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)}" - ) - cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {self._quote_identifier(schema)}") - for table_name in STATE_TABLES: - cursor.execute( - f"CREATE TABLE {self._qualified_name(schema=schema, table=table_name)} AS " - "SELECT * FROM " - f"{self._qualified_name(schema=backup_schema, table=table_name)}" - ) - self._create_state_indexes(cursor=cursor, schema=schema) - self._record_event( - cursor=cursor, - schema=schema, - action=StateMigrationAction.ROLLBACK, - backup_id_value=backup_id_value, - status=StateMigrationStatus.SUCCESS, - message=None, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - return backup_id_value - - def reset(self, *, connection: Any, schema: str) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - table_name: str - for table_name in STATE_TABLES: - cursor.execute( - "DROP TABLE IF EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)}" - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_model_version( - self, *, connection: Any, schema: str, record: ModelVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=MODEL_VERSION_TABLE, - where_sql="model_name = %s AND version_hash = %s", - params=[record.model_name, record.version_hash], - ) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = %s AND version_hash = %s", - [record.model_name, record.version_hash], - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "(model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status, " - "created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.definition_identity_hash, - record.identity_metadata_hash, - record.definition_text_b64, - record.identity_metadata_json_b64, - record.compiled_sql_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_model_version( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> ModelVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT model_name, version_hash, definition_identity_hash, " - "identity_metadata_hash, definition_text_b64, identity_metadata_json_b64, " - "compiled_sql_b64, status " - f"FROM {self._qualified_name(schema=schema, table=MODEL_VERSION_TABLE)} " - "WHERE model_name = %s AND version_hash = %s", - [model_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return ModelVersionRecord( - model_name=row[0], - version_hash=row[1], - definition_identity_hash=row[2], - identity_metadata_hash=row[3], - definition_text_b64=row[4], - identity_metadata_json_b64=row[5], - compiled_sql_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - - def upsert_function_version( - self, *, connection: Any, schema: str, record: FunctionVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=FUNCTION_VERSION_TABLE, - where_sql="function_name = %s AND version_hash = %s", - params=[record.function_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = %s AND version_hash = %s", - [record.function_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "(function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.function_name, - record.version_hash, - record.language, - record.returns, - record.arguments_json_b64, - record.return_columns_json_b64, - record.packages_json_b64, - record.runtime_version, - record.entry_point, - record.body_sql_b64, - record.definition_text_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_function_version( - self, *, connection: Any, schema: str, function_name: str, version_hash: str - ) -> FunctionVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT function_name, version_hash, language, returns, arguments_json_b64, " - "return_columns_json_b64, packages_json_b64, runtime_version, entry_point, " - "body_sql_b64, definition_text_b64, status " - f"FROM {self._qualified_name(schema=schema, table=FUNCTION_VERSION_TABLE)} " - "WHERE function_name = %s AND version_hash = %s", - [function_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return FunctionVersionRecord( - function_name=row[0], - version_hash=row[1], - language=row[2], - returns=row[3], - arguments_json_b64=row[4], - return_columns_json_b64=row[5], - packages_json_b64=row[6], - runtime_version=row[7], - entry_point=row[8], - body_sql_b64=row[9], - definition_text_b64=row[10], - status=ModelVersionStatus(row[11]), - ) - - def upsert_seed_version( - self, *, connection: Any, schema: str, record: SeedVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=SEED_VERSION_TABLE, - where_sql="seed_name = %s AND version_hash = %s", - params=[record.seed_name, record.version_hash], - ) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = %s AND version_hash = %s", - [record.seed_name, record.version_hash], - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "(seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.seed_name, - record.version_hash, - record.identity_metadata_hash, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_seed_version( - self, *, connection: Any, schema: str, seed_name: str, version_hash: str - ) -> SeedVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT seed_name, version_hash, identity_metadata_hash, " - "identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=SEED_VERSION_TABLE)} " - "WHERE seed_name = %s AND version_hash = %s", - [seed_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return SeedVersionRecord( - seed_name=row[0], - version_hash=row[1], - identity_metadata_hash=row[2], - identity_metadata_json_b64=row[3], - status=ModelVersionStatus(row[4]), - ) - - def upsert_python_node_version( - self, *, connection: Any, schema: str, record: PythonNodeVersionRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=PYTHON_NODE_VERSION_TABLE, - where_sql="node_type = %s AND node_name = %s AND version_hash = %s", - params=[record.node_type, record.node_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = %s AND node_name = %s AND version_hash = %s", - [record.node_type, record.node_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "(node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, " - "status, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.node_type, - record.node_name, - record.version_hash, - record.definition_hash, - record.identity_metadata_hash, - record.definition_json_b64, - record.identity_metadata_json_b64, - record.status.value, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_python_node_version( - self, - *, - connection: Any, - schema: str, - node_type: str, - node_name: str, - version_hash: str, - ) -> PythonNodeVersionRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT node_type, node_name, version_hash, definition_hash, " - "identity_metadata_hash, definition_json_b64, identity_metadata_json_b64, status " - f"FROM {self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "WHERE node_type = %s AND node_name = %s AND version_hash = %s", - [node_type, node_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return PythonNodeVersionRecord( - node_type=row[0], - node_name=row[1], - version_hash=row[2], - definition_hash=row[3], - identity_metadata_hash=row[4], - definition_json_b64=row[5], - identity_metadata_json_b64=row[6], - status=ModelVersionStatus(row[7]), - ) - - def insert_node_result( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - record: NodeResultRecord, - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - "(virtual_environment_name, node_type, node_name, target_database, target_schema, " - "target_name, run_id, status, payload_json_b64, metadata_json_b64, error_message, " - "materialized, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, " - "%s, %s)", - [ - virtual_environment_name, - record.node_type, - record.node_name, - record.target_database, - record.target_schema, - record.target_name, - record.run_id, - record.status, - encode_node_result_json( - value=record.payload, label="payload", node_name=record.node_name - ), - encode_node_result_json( - value=record.metadata, label="metadata", node_name=record.node_name - ), - record.error_message, - self._materialized_storage(record.materialized), - record.ts, - ], - ) - - def read_node_results( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - query: NodeResultQuery, - ) -> tuple[NodeResultEnvelope, ...]: - if query.limit < 1: - return () - predicates: list[str] = [ - "virtual_environment_name = %s", - "node_type = %s", - "node_name = %s", - self._optional_equality_sql( - column="target_database", value=query.target_database, placeholder="%s" - ), - self._optional_equality_sql( - column="target_schema", value=query.target_schema, placeholder="%s" - ), - self._optional_equality_sql( - column="target_name", value=query.target_name, placeholder="%s" - ), - ] - params: list[object] = [virtual_environment_name, query.node_type, query.node_name] - for value in (query.target_database, query.target_schema, query.target_name): - if value is not None: - params.append(value) - if query.statuses is not None: - placeholders: str = ", ".join("%s" for _ in query.statuses) - predicates.append(f"status IN ({placeholders})") - params.extend(query.statuses) - if query.run_id is not None: - predicates.append("run_id = %s") - params.append(query.run_id) - params.append(query.limit) - with connection.cursor() as cursor: - cursor.execute( - "SELECT node_type, node_name, run_id, status, payload_json_b64, " - "metadata_json_b64, error_message, materialized, created_at " - f"FROM {self._qualified_name(schema=schema, table=NODE_RESULTS_TABLE)} " - f"WHERE {' AND '.join(predicates)} " - "ORDER BY created_at DESC, run_id DESC LIMIT %s", - params, - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple(self._node_result_row_to_envelope(row) for row in rows) - def append_microbatch_event( self, *, connection: Any, schema: str, event: MicrobatchEvent ) -> None: @@ -771,679 +263,69 @@ def read_microbatch_model_history( rows: list[tuple[Any, ...]] = cursor.fetchall() return MicrobatchEventCodec.from_rows(tuple(row) for row in rows) - def upsert_physical_relation( - self, *, connection: Any, schema: str, record: PhysicalRelationRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=PHYSICAL_RELATION_TABLE, - where_sql="artifact_type = %s AND artifact_name = %s AND version_hash = %s", - params=[record.artifact_type.value, record.artifact_name, record.version_hash], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = %s AND artifact_name = %s AND version_hash = %s", - [record.artifact_type.value, record.artifact_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "(artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.artifact_type.value, - record.artifact_name, - record.version_hash, - record.database_name, - record.schema_name, - record.relation_name, - record.relation_type, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_physical_relation_for_artifact( + def upsert_virtual_environment_node_ref( self, *, connection: Any, schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - version_hash: str, - ) -> PhysicalRelationRecord | None: + ref: VirtualEnvironmentNodeRefRecord, + ) -> None: with connection.cursor() as cursor: cursor.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = %s AND artifact_name = %s AND version_hash = %s", - [artifact_type.value, artifact_name, version_hash], + "INSERT INTO " + f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " + "(virtual_environment_name, node_type, node_name, version_hash, updated_at) " + "VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP) " + "ON CONFLICT (virtual_environment_name, node_type, node_name) " + "DO UPDATE SET version_hash = excluded.version_hash, " + "updated_at = CURRENT_TIMESTAMP", + [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], + + def prune_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: + before_count: int = self.count_unreferenced_python_node_versions( + connection=connection, schema=schema ) + with connection.cursor() as cursor: + cursor.execute( + "DELETE FROM " + f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " + "AS versions " + "WHERE NOT EXISTS (" + "SELECT 1 " + "FROM " + f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " + "refs " + "WHERE refs.node_type = versions.node_type " + "AND refs.node_name = versions.node_name " + "AND refs.version_hash = versions.version_hash)" + ) + return before_count - def list_physical_relations_for_artifact( + def replace_virtual_environment_source_freshness( self, *, connection: Any, schema: str, - artifact_type: PhysicalArtifactType, - artifact_name: str, - ) -> tuple[PhysicalRelationRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT artifact_type, artifact_name, version_hash, database_name, schema_name, " - "relation_name, relation_type " - f"FROM {self._qualified_name(schema=schema, table=PHYSICAL_RELATION_TABLE)} " - "WHERE artifact_type = %s AND artifact_name = %s " - "ORDER BY updated_at DESC, version_hash DESC", - [artifact_type.value, artifact_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - PhysicalRelationRecord( - artifact_type=PhysicalArtifactType(row[0]), - artifact_name=row[1], - version_hash=row[2], - database_name=row[3], - schema_name=row[4], - relation_name=row[5], - relation_type=row[6], - ) - for row in rows - ) - - def upsert_physical_relation_ancestry( - self, *, connection: Any, schema: str, record: PhysicalRelationAncestryRecord + virtual_environment_name: str, + records: tuple[SourceFreshnessRecord, ...], ) -> None: + self._validate_source_freshness_records( + virtual_environment_name=virtual_environment_name, + records=records, + ) with connection.cursor() as cursor: cursor.execute("BEGIN") try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=PHYSICAL_RELATION_ANCESTRY_TABLE, - where_sql="model_name = %s AND version_hash = %s", - params=[record.model_name, record.version_hash], - ) cursor.execute( "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)}" - " " - "WHERE model_name = %s AND version_hash = %s", - [record.model_name, record.version_hash], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)}" - " " - "(model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy, created_at, updated_at) " - "VALUES (%s, %s, %s, %s, %s, " - "COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.model_name, - record.version_hash, - record.parent_model_name, - record.parent_version_hash, - record.seed_strategy, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_physical_relation_ancestry( - self, *, connection: Any, schema: str, model_name: str, version_hash: str - ) -> PhysicalRelationAncestryRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT model_name, version_hash, parent_model_name, parent_version_hash, " - "seed_strategy " - "FROM " - f"{self._qualified_name(schema=schema, table=PHYSICAL_RELATION_ANCESTRY_TABLE)} " - "WHERE model_name = %s AND version_hash = %s", - [model_name, version_hash], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return PhysicalRelationAncestryRecord( - model_name=row[0], - version_hash=row[1], - parent_model_name=row[2], - parent_version_hash=row[3], - seed_strategy=row[4], - ) - - def upsert_virtual_environment( - self, *, connection: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._upsert_virtual_environment_record(cursor=cursor, schema=schema, record=record) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> VirtualEnvironmentRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, status, baseline_virtual_environment_name, " - "finalized_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return VirtualEnvironmentRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - baseline_virtual_environment_name=row[2], - finalized_at=row[3], - ) - - def list_virtual_environments( - self, *, connection: Any, schema: str - ) -> tuple[VirtualEnvironmentRetentionRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, status, updated_at " - f"FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "ORDER BY updated_at DESC, virtual_environment_name DESC" - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentRetentionRecord( - virtual_environment_name=row[0], - status=VirtualEnvironmentStatus(row[1]), - updated_at=row[2], - ) - for row in rows - ) - - def delete_virtual_environment( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=SOURCE_FRESHNESS_OBSERVATION_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def replace_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type={node_type: refs}, - ) - - def replace_virtual_environment_node_ref_groups( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._replace_virtual_environment_node_ref_groups( - cursor=cursor, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_virtual_environment_and_replace_node_ref_groups( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._upsert_virtual_environment_record(cursor=cursor, schema=schema, record=record) - self._replace_virtual_environment_node_ref_groups( - cursor=cursor, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_virtual_environment_and_replace_node_ref_groups_if_locks_owned( - self, - *, - connection: Any, - schema: str, - record: VirtualEnvironmentRecord, - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], - leases: tuple[StateLockLease, ...], - checkpoint: VirtualEnvironmentCheckpointRecord | None = None, - checkpoint_refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...] = (), - checkpoint_function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - checkpoint_seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> bool: - validate_conditional_virtual_environment_publication( - record=record, - refs_by_node_type=refs_by_node_type, - checkpoint=checkpoint, - checkpoint_refs=checkpoint_refs, - checkpoint_function_refs=checkpoint_function_refs, - checkpoint_seed_refs=checkpoint_seed_refs, - ) - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - lease: StateLockLease - for lease in leases: - lock_table: str = self._qualified_name(schema=schema, table=LOCK_TABLE) - cursor.execute( - f"SELECT lock_key FROM {lock_table} " - "WHERE lock_key = %s AND owner_id = %s " - "AND expires_at > CURRENT_TIMESTAMP FOR UPDATE", - [lease.lock_key, lease.owner_id], - ) - if cursor.fetchone() is None: - cursor.execute("ROLLBACK") - return False - self._upsert_virtual_environment_record(cursor=cursor, schema=schema, record=record) - self._replace_virtual_environment_node_ref_groups( - cursor=cursor, - schema=schema, - virtual_environment_name=record.virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - if checkpoint is not None: - self._insert_virtual_environment_checkpoint_rows( - cursor=cursor, - schema=schema, - checkpoint=checkpoint, - refs=checkpoint_refs, - function_refs=checkpoint_function_refs, - seed_refs=checkpoint_seed_refs, - ) - cursor.execute("COMMIT") - return True - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_virtual_environment_node_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - node_type: str, - ) -> tuple[VirtualEnvironmentNodeRefRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = %s AND node_type = %s ORDER BY node_name", - [virtual_environment_name, node_type], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - - def upsert_virtual_environment_node_ref( - self, - *, - connection: Any, - schema: str, - ref: VirtualEnvironmentNodeRefRecord, - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "(virtual_environment_name, node_type, node_name, version_hash, updated_at) " - "VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP) " - "ON CONFLICT (virtual_environment_name, node_type, node_name) " - "DO UPDATE SET version_hash = excluded.version_hash, " - "updated_at = CURRENT_TIMESTAMP", - [ref.virtual_environment_name, ref.node_type, ref.node_name, ref.version_hash], - ) - - def replace_virtual_environment_model_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentModelRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="model", - node_name=ref.model_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), - ) - - def get_virtual_environment_model_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentModelRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="model", - ) - return tuple( - VirtualEnvironmentModelRefRecord( - virtual_environment_name=ref.virtual_environment_name, - model_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs - ) - - def replace_virtual_environment_function_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentFunctionRefRecord, ...], - ) -> None: - ref: VirtualEnvironmentFunctionRefRecord - for ref in refs: - if ref.node_type not in { - CompiledResourceType.UDF, - CompiledResourceType.TABLE_FN, - }: - raise StateBackendConfigError("Function ref node_type must be 'udf' or 'table_fn'") - refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]] = {} - for node_type in ("udf", "table_fn"): - node_refs: list[VirtualEnvironmentNodeRefRecord] = [] - for ref in refs: - if ref.node_type == node_type: - node_refs.append( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.function_name, - version_hash=ref.version_hash, - ) - ) - refs_by_node_type[node_type] = tuple(node_refs) - self.replace_virtual_environment_node_ref_groups( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - refs_by_node_type=refs_by_node_type, - ) - - def get_virtual_environment_function_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentFunctionRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="udf", - ), - *self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="table_fn", - ), - ) - return tuple( - VirtualEnvironmentFunctionRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - function_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in sorted(refs, key=lambda item: (item.node_type, item.node_name)) - ) - - def replace_virtual_environment_seed_refs( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - refs: tuple[VirtualEnvironmentSeedRefRecord, ...], - ) -> None: - self.replace_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - refs=tuple( - VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type="seed", - node_name=ref.seed_name, - version_hash=ref.version_hash, - ) - for ref in refs - ), - ) - - def get_virtual_environment_seed_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentSeedRefRecord, ...]: - refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = self.get_virtual_environment_node_refs( - connection=connection, - schema=schema, - virtual_environment_name=virtual_environment_name, - node_type="seed", - ) - return tuple( - VirtualEnvironmentSeedRefRecord( - virtual_environment_name=ref.virtual_environment_name, - seed_name=ref.node_name, - version_hash=ref.version_hash, - ) - for ref in refs - ) - - def upsert_virtual_environment_python_node_ref( - self, - *, - connection: Any, - schema: str, - ref: VirtualEnvironmentPythonNodeRefRecord, - ) -> None: - self.upsert_virtual_environment_node_ref( - connection=connection, - schema=schema, - ref=VirtualEnvironmentNodeRefRecord( - virtual_environment_name=ref.virtual_environment_name, - node_type=ref.node_type, - node_name=ref.node_name, - version_hash=ref.version_hash, - ), - ) - - def get_virtual_environment_python_node_refs( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentPythonNodeRefRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, node_type, node_name, version_hash " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "WHERE virtual_environment_name = %s " - "AND node_type IN ('task', 'loader', 'asset', 'check', 'hook') " - "ORDER BY node_type, node_name", - [virtual_environment_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentPythonNodeRefRecord( - virtual_environment_name=row[0], - node_type=row[1], - node_name=row[2], - version_hash=row[3], - ) - for row in rows - ) - - def count_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: - with connection.cursor() as cursor: - cursor.execute( - "SELECT COUNT(*) " - "FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} versions " - "WHERE NOT EXISTS (" - "SELECT 1 " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "refs " - "WHERE refs.node_type = versions.node_type " - "AND refs.node_name = versions.node_name " - "AND refs.version_hash = versions.version_hash)" - ) - row: tuple[Any, ...] = cursor.fetchone() - return int(row[0]) - - def prune_unreferenced_python_node_versions(self, *, connection: Any, schema: str) -> int: - before_count: int = self.count_unreferenced_python_node_versions( - connection=connection, schema=schema - ) - with connection.cursor() as cursor: - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=PYTHON_NODE_VERSION_TABLE)} " - "AS versions " - "WHERE NOT EXISTS (" - "SELECT 1 " - "FROM " - f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " - "refs " - "WHERE refs.node_type = versions.node_type " - "AND refs.node_name = versions.node_name " - "AND refs.version_hash = versions.version_hash)" - ) - return before_count - - def replace_virtual_environment_source_freshness( - self, - *, - connection: Any, - schema: str, - virtual_environment_name: str, - records: tuple[SourceFreshnessRecord, ...], - ) -> None: - self._validate_source_freshness_records( - virtual_environment_name=virtual_environment_name, - records=records, - ) - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=SOURCE_FRESHNESS_OBSERVATION_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s", - [virtual_environment_name], + + self._qualified_name( + schema=schema, + table=SOURCE_FRESHNESS_OBSERVATION_TABLE, + ) + + " " + "WHERE virtual_environment_name = %s", + [virtual_environment_name], ) record: SourceFreshnessRecord for record in records: @@ -1472,337 +354,6 @@ def replace_virtual_environment_source_freshness( cursor.execute("ROLLBACK") raise - def get_virtual_environment_source_freshness( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[SourceFreshnessRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT virtual_environment_name, source_name, strategy, value_kind, " - "data_version, data_version_hash, observed_at " - "FROM " - f"{self._qualified_name(schema=schema, table=SOURCE_FRESHNESS_OBSERVATION_TABLE)} " - "WHERE virtual_environment_name = %s ORDER BY source_name", - [virtual_environment_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - SourceFreshnessRecord( - virtual_environment_name=row[0], - source_name=row[1], - strategy=row[2], - value_kind=row[3], - data_version=row[4], - data_version_hash=row[5], - observed_at=from_naive_utc_wall_clock(row[6]), - ) - for row in rows - ) - - def create_virtual_environment_checkpoint( - self, - *, - connection: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...] = (), - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...] = (), - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - self._insert_virtual_environment_checkpoint_rows( - cursor=cursor, - schema=schema, - checkpoint=checkpoint, - refs=refs, - function_refs=function_refs, - seed_refs=seed_refs, - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def _insert_virtual_environment_checkpoint_rows( - self, - *, - cursor: Any, - schema: str, - checkpoint: VirtualEnvironmentCheckpointRecord, - refs: tuple[VirtualEnvironmentCheckpointModelRefRecord, ...], - function_refs: tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...], - seed_refs: tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...], - ) -> None: - checkpoint_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE - ) - model_ref_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE - ) - function_ref_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE - ) - seed_ref_table: str = self._qualified_name( - schema=schema, table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE - ) - cursor.execute( - f"INSERT INTO {checkpoint_table} " - "(checkpoint_id, virtual_environment_name, created_at) " - "VALUES (%s, %s, CURRENT_TIMESTAMP)", - [checkpoint.checkpoint_id, checkpoint.virtual_environment_name], - ) - ref: VirtualEnvironmentCheckpointModelRefRecord - for ref in refs: - cursor.execute( - f"INSERT INTO {model_ref_table} " - "(checkpoint_id, model_name, version_hash) VALUES (%s, %s, %s)", - [ref.checkpoint_id, ref.model_name, ref.version_hash], - ) - function_ref: VirtualEnvironmentCheckpointFunctionRefRecord - for function_ref in function_refs: - cursor.execute( - f"INSERT INTO {function_ref_table} " - "(checkpoint_id, function_name, version_hash) VALUES (%s, %s, %s)", - [function_ref.checkpoint_id, function_ref.function_name, function_ref.version_hash], - ) - seed_ref: VirtualEnvironmentCheckpointSeedRefRecord - for seed_ref in seed_refs: - cursor.execute( - f"INSERT INTO {seed_ref_table} " - "(checkpoint_id, seed_name, version_hash) VALUES (%s, %s, %s)", - [seed_ref.checkpoint_id, seed_ref.seed_name, seed_ref.version_hash], - ) - - def list_virtual_environment_checkpoints( - self, *, connection: Any, schema: str, virtual_environment_name: str - ) -> tuple[VirtualEnvironmentCheckpointRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT checkpoint_id, virtual_environment_name, created_at " - "FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, - ) - + " " - "WHERE virtual_environment_name = %s ORDER BY created_at DESC, checkpoint_id DESC", - [virtual_environment_name], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointRecord( - checkpoint_id=row[0], - virtual_environment_name=row[1], - created_at=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_model_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointModelRefRecord, ...]: - with connection.cursor() as cursor: - checkpoint_model_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - cursor.execute( - f"SELECT checkpoint_id, model_name, version_hash " - f"FROM {checkpoint_model_ref_table} " - "WHERE checkpoint_id = %s ORDER BY model_name", - [checkpoint_id], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointModelRefRecord( - checkpoint_id=row[0], - model_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_function_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointFunctionRefRecord, ...]: - with connection.cursor() as cursor: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - cursor.execute( - f"SELECT checkpoint_id, function_name, version_hash " - f"FROM {checkpoint_function_ref_table} " - "WHERE checkpoint_id = %s ORDER BY function_name", - [checkpoint_id], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointFunctionRefRecord( - checkpoint_id=row[0], - function_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def get_virtual_environment_checkpoint_seed_refs( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> tuple[VirtualEnvironmentCheckpointSeedRefRecord, ...]: - with connection.cursor() as cursor: - checkpoint_seed_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - cursor.execute( - f"SELECT checkpoint_id, seed_name, version_hash " - f"FROM {checkpoint_seed_ref_table} " - "WHERE checkpoint_id = %s ORDER BY seed_name", - [checkpoint_id], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - VirtualEnvironmentCheckpointSeedRefRecord( - checkpoint_id=row[0], - seed_name=row[1], - version_hash=row[2], - ) - for row in rows - ) - - def delete_virtual_environment_checkpoint( - self, *, connection: Any, schema: str, checkpoint_id: str - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - checkpoint_function_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_FUNCTION_REF_TABLE, - ) - checkpoint_model_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_MODEL_REF_TABLE, - ) - checkpoint_seed_ref_table: str = self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_SEED_REF_TABLE, - ) - cursor.execute( - f"DELETE FROM {checkpoint_seed_ref_table} WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute( - f"DELETE FROM {checkpoint_function_ref_table} WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute( - f"DELETE FROM {checkpoint_model_ref_table} WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute( - "DELETE FROM " - + self._qualified_name( - schema=schema, - table=VIRTUAL_ENVIRONMENT_CHECKPOINT_TABLE, - ) - + " " - "WHERE checkpoint_id = %s", - [checkpoint_id], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def upsert_state_operation( - self, *, connection: Any, schema: str, record: StateOperationRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute("BEGIN") - try: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=STATE_OPERATION_TABLE, - where_sql="operation_id = %s", - params=[record.operation_id], - ) - cursor.execute( - "DELETE FROM " - f"{self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = %s", - [record.operation_id], - ) - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "(operation_id, operation_type, status, virtual_environment_name, " - "created_at, updated_at) " - "VALUES (%s, %s, %s, %s, COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP)", - [ - record.operation_id, - record.operation_type.value, - record.status.value, - record.virtual_environment_name, - existing_created_at, - ], - ) - cursor.execute("COMMIT") - except BaseException: - cursor.execute("ROLLBACK") - raise - - def get_state_operation( - self, *, connection: Any, schema: str, operation_id: str - ) -> StateOperationRecord | None: - with connection.cursor() as cursor: - cursor.execute( - "SELECT operation_id, operation_type, status, virtual_environment_name " - f"FROM {self._qualified_name(schema=schema, table=STATE_OPERATION_TABLE)} " - "WHERE operation_id = %s", - [operation_id], - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return StateOperationRecord( - operation_id=row[0], - operation_type=StateOperationType(row[1]), - status=StateOperationStatus(row[2]), - virtual_environment_name=row[3], - ) - - def create_state_operation_event( - self, *, connection: Any, schema: str, record: StateOperationEventRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_OPERATION_EVENT_TABLE)} " - "(event_id, operation_id, action, status, message, created_at) " - "VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP)", - [ - record.event_id, - record.operation_id, - record.action, - record.status.value, - record.message, - ], - ) - - def create_reconcile_event( - self, *, connection: Any, schema: str, record: ReconcileEventRecord - ) -> None: - with connection.cursor() as cursor: - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=RECONCILE_EVENT_TABLE)} " - "(event_id, action, status, message, created_at) " - "VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP)", - [record.event_id, record.action.value, record.status.value, record.message], - ) - def acquire_lock( self, *, @@ -1862,49 +413,6 @@ def release_lock(self, *, connection: Any, schema: str, lock_key: str, owner_id: cursor.execute("ROLLBACK") raise - def renew_lock( - self, - *, - connection: Any, - schema: str, - lock_key: str, - owner_id: str, - expires_at: datetime, - ) -> bool: - with connection.cursor() as cursor: - cursor.execute( - f"UPDATE {self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "SET expires_at = %s, updated_at = CURRENT_TIMESTAMP " - "WHERE lock_key = %s AND owner_id = %s AND expires_at > CURRENT_TIMESTAMP " - "RETURNING lock_key", - [expires_at, lock_key, owner_id], - ) - return cursor.fetchone() is not None - - def list_active_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at > CURRENT_TIMESTAMP ORDER BY lock_key" - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - - def list_expired_locks(self, *, connection: Any, schema: str) -> tuple[StateLockRecord, ...]: - with connection.cursor() as cursor: - cursor.execute( - "SELECT lock_key, owner_id, expires_at FROM " - f"{self._qualified_name(schema=schema, table=LOCK_TABLE)} " - "WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY lock_key" - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - StateLockRecord(lock_key=row[0], owner_id=row[1], expires_at=row[2]) for row in rows - ) - def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: with connection.cursor() as cursor: cursor.execute( @@ -1914,264 +422,16 @@ def delete_lock(self, *, connection: Any, schema: str, lock_key: str) -> None: ) connection.commit() - def list_state_backups(self, *, connection: Any, schema: str) -> tuple[StateBackupRecord, ...]: - prefix: str = f"{schema}__backup_%" - with connection.cursor() as cursor: - cursor.execute( - "SELECT s.schema_name, e.backup_id, MAX(e.created_at) " - "FROM information_schema.schemata s " - "LEFT JOIN " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} e " - "ON s.schema_name = %s || e.backup_id " - "WHERE s.schema_name LIKE %s " - "GROUP BY s.schema_name, e.backup_id ORDER BY s.schema_name DESC", - [f"{schema}__backup_", prefix], - ) - rows: list[tuple[Any, ...]] = cursor.fetchall() - return tuple( - StateBackupRecord( - backup_id=row[1] or str(row[0]).removeprefix(f"{schema}__backup_"), - schema_name=row[0], - created_at=row[2], - ) - for row in rows - ) - def delete_state_backup(self, *, connection: Any, schema: str, backup_id: str) -> None: backup_schema: str = self._backup_schema_name(schema=schema, backup_id_value=backup_id) with connection.cursor() as cursor: cursor.execute(f"DROP SCHEMA IF EXISTS {self._quote_identifier(backup_schema)} CASCADE") connection.commit() - def _latest_backup_id(self, *, connection: Any, schema: str) -> str: - prefix: str = f"{schema}__backup_%" - with connection.cursor() as cursor: - cursor.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE %s " - "ORDER BY schema_name DESC LIMIT 1", - [prefix], - ) - row: tuple[str] | None = cursor.fetchone() - if row is None: - raise StateBackupNotFoundError("No state backup is available for rollback") - return row[0].removeprefix(f"{schema}__backup_") - - def _schema_exists(self, *, connection: Any, schema: str) -> bool: - with connection.cursor() as cursor: - cursor.execute( - "SELECT schema_name FROM information_schema.schemata WHERE schema_name = %s", - [schema], - ) - row: tuple[str] | None = cursor.fetchone() - return row is not None - - def _record_event( - self, - *, - cursor: Any, - schema: str, - action: StateMigrationAction, - backup_id_value: str | None, - status: StateMigrationStatus, - message: str | None, - ) -> None: - cursor.execute( - "INSERT INTO " - f"{self._qualified_name(schema=schema, table=STATE_MIGRATION_EVENTS_TABLE)} " - "(event_id, action, backup_id, status, message, created_at) " - "VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP)", - [event_id(), action.value, backup_id_value, status.value, message], - ) - - def _create_additional_state_tables(self, *, cursor: Any, schema: str) -> None: - table_name: str - columns: dict[str, StateColumnType] - for table_name, columns in STATE_TABLE_COLUMNS.items(): - if table_name in {STATE_VERSION_TABLE, STATE_MIGRATION_EVENTS_TABLE}: - continue - column_sql: str = ", ".join( - f"{self._quote_identifier(column_name)} {self._state_column_sql_type(column_type)}" - for column_name, column_type in columns.items() - ) - cursor.execute( - "CREATE TABLE IF NOT EXISTS " - f"{self._qualified_name(schema=schema, table=table_name)} " - f"({column_sql})" - ) - column_name: str - column_type: StateColumnType - for column_name, column_type in columns.items(): - cursor.execute( - f"ALTER TABLE {self._qualified_name(schema=schema, table=table_name)} " - f"ADD COLUMN IF NOT EXISTS {self._quote_identifier(column_name)} " - f"{self._state_column_sql_type(column_type)}" - ) - self._create_state_indexes(cursor=cursor, schema=schema) - - def _create_state_indexes(self, *, cursor: Any, schema: str) -> None: - table_name: str - indexes: dict[str, tuple[str, ...]] - for table_name, indexes in STATE_TABLE_INDEXES.items(): - index_name: str - columns: tuple[str, ...] - for index_name, columns in indexes.items(): - column_sql: str = ", ".join(self._quote_identifier(column) for column in columns) - unique_sql: str = "" if index_name in NON_UNIQUE_STATE_INDEXES else "UNIQUE " - cursor.execute( - f"CREATE {unique_sql}INDEX IF NOT EXISTS {self._quote_identifier(index_name)} " - f"ON {self._qualified_name(schema=schema, table=table_name)} ({column_sql})" - ) - - def _created_at_for_key( - self, - *, - cursor: Any, - schema: str, - table_name: str, - where_sql: str, - params: list[object], - ) -> datetime | None: - cursor.execute( - f"SELECT created_at FROM {self._qualified_name(schema=schema, table=table_name)} " - f"WHERE {where_sql}", - params, - ) - row: tuple[Any, ...] | None = cursor.fetchone() - if row is None: - return None - return row[0] - - def _state_column_sql_type(self, column_type: StateColumnType) -> str: - match column_type: - case StateColumnType.INTEGER: - return "INTEGER" - case StateColumnType.TEXT: - return "TEXT" - case StateColumnType.TIMESTAMP: - return "TIMESTAMP" - raise StateBackendConfigError(f"Unsupported state column type: {column_type}") - - def _node_result_row_to_envelope(self, row: tuple[Any, ...]) -> NodeResultEnvelope: - node_name: str = str(row[1]) - metadata: object = decode_node_result_json( - value=str(row[5]), label="metadata", node_name=node_name - ) - normalized_metadata: dict[str, object] = ( - {str(key): value for key, value in metadata.items()} - if isinstance(metadata, dict) - else {} - ) - return NodeResultEnvelope( - node_type=str(row[0]), - node_name=node_name, - run_id=str(row[2]), - status=str(row[3]), - payload=decode_node_result_json( - value=str(row[4]), label="payload", node_name=node_name - ), - metadata=normalized_metadata, - error_message=str(row[6]) if row[6] is not None else None, - materialized=self._parse_materialized(row[7]), - ts=row[8], - ) - - def _optional_equality_sql(self, *, column: str, value: object | None, placeholder: str) -> str: - if value is None: - return f"{column} IS NULL" - return f"{column} = {placeholder}" - - def _materialized_storage(self, value: bool | None) -> str | None: - if value is None: - return None - return "true" if value else "false" - - def _parse_materialized(self, value: object) -> bool | None: - if value is None: - return None - return str(value).lower() == STATE_BOOLEAN_TRUE - - def _quote_identifier(self, identifier: str) -> str: - return '"' + identifier.replace('"', '""') + '"' - - def _qualified_name(self, *, schema: str, table: str) -> str: - return f"{self._quote_identifier(schema)}.{self._quote_identifier(table)}" - - def _validate_source_freshness_records( - self, - *, - virtual_environment_name: str, - records: tuple[SourceFreshnessRecord, ...], - ) -> None: - seen_source_names: set[str] = set() - record: SourceFreshnessRecord - for record in records: - if record.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Source freshness record virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if record.source_name in seen_source_names: - raise StateBackendConfigError( - f"Duplicate source freshness record for source '{record.source_name}'" - ) - seen_source_names.add(record.source_name) - - def _validate_node_ref_replacement( - self, - *, - virtual_environment_name: str, - node_type: str, - refs: tuple[VirtualEnvironmentNodeRefRecord, ...], - ) -> None: - seen_node_names: set[str] = set() - ref: VirtualEnvironmentNodeRefRecord - for ref in refs: - if ref.virtual_environment_name != virtual_environment_name: - raise StateBackendConfigError( - "Node ref virtual_environment_name must match replacement " - "virtual_environment_name" - ) - if ref.node_type != node_type: - raise StateBackendConfigError("Node ref node_type must match replacement node_type") - if ref.node_name in seen_node_names: - raise StateBackendConfigError( - f"Duplicate node ref for node type '{node_type}' and name '{ref.node_name}'" - ) - seen_node_names.add(ref.node_name) - - def _upsert_virtual_environment_record( - self, *, cursor: Any, schema: str, record: VirtualEnvironmentRecord - ) -> None: - existing_created_at: datetime | None = self._created_at_for_key( - cursor=cursor, - schema=schema, - table_name=VIRTUAL_ENVIRONMENT_TABLE, - where_sql="virtual_environment_name = %s", - params=[record.virtual_environment_name], - ) - cursor.execute( - f"DELETE FROM {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "WHERE virtual_environment_name = %s", - [record.virtual_environment_name], - ) - cursor.execute( - f"INSERT INTO {self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_TABLE)} " - "(virtual_environment_name, status, baseline_virtual_environment_name, " - "created_at, updated_at, finalized_at) " - "VALUES (%s, %s, %s, COALESCE(%s, CURRENT_TIMESTAMP), CURRENT_TIMESTAMP, %s)", - [ - record.virtual_environment_name, - record.status.value, - record.baseline_virtual_environment_name, - existing_created_at, - record.finalized_at, - ], - ) - def _replace_virtual_environment_node_ref_groups( self, *, - cursor: Any, + executor: Any, schema: str, virtual_environment_name: str, refs_by_node_type: dict[str, tuple[VirtualEnvironmentNodeRefRecord, ...]], @@ -2184,7 +444,7 @@ def _replace_virtual_environment_node_ref_groups( node_type=node_type, refs=refs, ) - cursor.execute( + executor.execute( "DELETE FROM " f"{self._qualified_name(schema=schema, table=VIRTUAL_ENVIRONMENT_NODE_REF_TABLE)} " "WHERE virtual_environment_name = %s AND node_type = %s", @@ -2192,7 +452,7 @@ def _replace_virtual_environment_node_ref_groups( ) ref: VirtualEnvironmentNodeRefRecord for ref in refs: - cursor.execute( + executor.execute( "INSERT INTO " + self._qualified_name( schema=schema, @@ -2212,9 +472,6 @@ def _replace_virtual_environment_node_ref_groups( ], ) - def _backup_schema_name(self, *, schema: str, backup_id_value: str) -> str: - return f"{schema}__backup_{backup_id_value}" - def _state_type_matches(self, *, actual_type: str, expected_type: StateColumnType) -> bool: actual: str = actual_type.lower() match expected_type: diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py index 7687fc89aa..5ae694aeed 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/_test_types.py @@ -1055,6 +1055,19 @@ class RemoveColumnSemanticsBuildE2ETestCase: expected_warning_count: int +@dataclass(frozen=True) +class SnapshotReappearingKeyBuildE2ETestCase: + """Test case for hard-deleted keys that reappear in historical check snapshots.""" + + description: str + repo_files: dict[str, str] + initial_seed_sql: str + observation_sql_by_round: tuple[tuple[str, ...], ...] + history_query_template: str + expected_incremental_rows_by_round: tuple[tuple[tuple[object, ...], ...], ...] + expected_full_history_rows: tuple[tuple[object, ...], ...] + + @dataclass(frozen=True) class DynamicColumnContractE2ETestCase: description: str @@ -1085,3 +1098,14 @@ class WindowLimitedReplayE2ETestCase: changed_rows_sql: str expected_changed_rows: int expected_total_rows: int + + +@dataclass(frozen=True) +class SqlHookIdentityBuildE2ETestCase: + """Test case for model version identity after a SQL hook resource edit.""" + + description: str + removed_hook_path: str + added_hook_path: str + added_hook_contents: str + expected_reason: str diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py index de079fee7d..3268604f80 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_snapshot_build.py @@ -18,6 +18,7 @@ SnapshotHistoricalCheckBuildE2ETestCase, SnapshotHistoricalTimestampBuildE2ETestCase, SnapshotHookBuildE2ETestCase, + SnapshotReappearingKeyBuildE2ETestCase, SnapshotSelectorBuildE2ETestCase, SnapshotTimestampBuildE2ETestCase, SnapshotTimestampFailureBuildE2ETestCase, @@ -31,6 +32,56 @@ table_exists, ) +_TIMESTAMP_REAPPEARING_KEY_FINAL_HISTORY: tuple[tuple[object, ...], ...] = ( + (1, "basic", "2024-01-15", "2024-02-03"), + (1, "pro", "2024-02-03", None), + (2, "basic", "2024-01-15", "2024-02-02"), + (2, "basic", "2024-02-03", None), + (3, "basic", "2024-01-15", "2024-02-04"), + (3, "basic", "2024-02-05", None), +) +_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS: tuple[str, ...] = ( + "(1, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-01'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-01'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-01')", + "(1, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-02'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-02')", + "(1, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-03'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-03'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-03')", + "(1, 'pro', TIMESTAMP '2024-02-03', TIMESTAMP '2024-02-04'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-04'), " + "(1, 'pro', TIMESTAMP '2024-02-03', TIMESTAMP '2024-02-05'), " + "(2, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-05'), " + "(3, 'basic', TIMESTAMP '2024-01-15', TIMESTAMP '2024-02-05')", +) +_TIMESTAMP_REAPPEARING_KEY_TABLE_SQL: str = ( + "(customer_id INTEGER, status VARCHAR, updated_at TIMESTAMP, observed_at TIMESTAMP)" +) +_REAPPEARING_KEY_FINAL_HISTORY: tuple[tuple[object, ...], ...] = ( + (1, "active", "2024-02-01", "2024-02-04"), + (1, "paused", "2024-02-04", None), + (2, "active", "2024-02-01", "2024-02-02"), + (2, "active", "2024-02-03", None), + (3, "active", "2024-02-01", "2024-02-04"), + (3, "active", "2024-02-05", None), +) +_REAPPEARING_KEY_OBSERVATIONS_SQL: str = ( + "SELECT 1 AS customer_id, 'active' AS status, TIMESTAMP '2024-02-01' AS observed_at " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-01' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-01' " + "UNION ALL SELECT 1, 'active', TIMESTAMP '2024-02-02' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-02' " + "UNION ALL SELECT 1, 'active', TIMESTAMP '2024-02-03' " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-03' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-03' " + "UNION ALL SELECT 1, 'paused', TIMESTAMP '2024-02-04' " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-04' " + "UNION ALL SELECT 1, 'paused', TIMESTAMP '2024-02-05' " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-05' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-05'" +) + @pytest.mark.parametrize( "test_case", @@ -1927,6 +1978,294 @@ def test_given_shallow_waffle_shop_snapshot_edges_when_sources_change_then_cli_t ) +@pytest.mark.parametrize( + "test_case", + [ + SnapshotReappearingKeyBuildE2ETestCase( + description=( + "historical check snapshot reopens hard-deleted keys that reappear unchanged" + ), + repo_files={ + "sqlbuild_project.toml": dedent( + """ + name = "snapshot_reappearing_keys" + adapter = "duckdb" + + [connection] + database = "snapshot_reappearing_keys.duckdb" + """ + ).strip() + + "\n", + "sources/raw.yml": dedent( + """ + sources: + - name: raw_membership_observations + schema: main + table: raw_membership_observations + - name: raw_membership_full_history + schema: main + table: raw_membership_full_history + """ + ).strip() + + "\n", + "models/membership_incremental_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy check, + check_columns [status], + observed_at observed_at, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, observed_at + FROM __source("raw_membership_observations") + """ + ).strip() + + "\n", + "models/membership_full_history_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy check, + check_columns [status], + observed_at observed_at, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, observed_at + FROM __source("raw_membership_full_history") + """ + ).strip() + + "\n", + }, + initial_seed_sql=( + "CREATE TABLE main.raw_membership_observations AS " + "SELECT 1 AS customer_id, 'active' AS status, TIMESTAMP '2024-02-01' AS observed_at " + "UNION ALL SELECT 2, 'active', TIMESTAMP '2024-02-01' " + "UNION ALL SELECT 3, 'active', TIMESTAMP '2024-02-01'; " + f"CREATE TABLE main.raw_membership_full_history AS {_REAPPEARING_KEY_OBSERVATIONS_SQL}" + ), + observation_sql_by_round=( + (), + ( + "INSERT INTO main.raw_membership_observations VALUES " + "(1, 'active', TIMESTAMP '2024-02-02'), (3, 'active', TIMESTAMP '2024-02-02')", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + "(1, 'active', TIMESTAMP '2024-02-03'), (2, 'active', TIMESTAMP '2024-02-03'), " + "(3, 'active', TIMESTAMP '2024-02-03')", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + "(1, 'paused', TIMESTAMP '2024-02-04'), (2, 'active', TIMESTAMP '2024-02-04'), " + "(1, 'paused', TIMESTAMP '2024-02-05'), (2, 'active', TIMESTAMP '2024-02-05'), " + "(3, 'active', TIMESTAMP '2024-02-05')", + ), + ), + history_query_template=( + "SELECT customer_id, status, strftime(valid_from, '%Y-%m-%d'), " + "strftime(valid_to, '%Y-%m-%d') FROM main.{table_name} " + "ORDER BY customer_id, valid_from" + ), + expected_incremental_rows_by_round=( + ( + (1, "active", "2024-02-01", None), + (2, "active", "2024-02-01", None), + (3, "active", "2024-02-01", None), + ), + ( + (1, "active", "2024-02-01", None), + (2, "active", "2024-02-01", "2024-02-02"), + (3, "active", "2024-02-01", None), + ), + ( + (1, "active", "2024-02-01", None), + (2, "active", "2024-02-01", "2024-02-02"), + (2, "active", "2024-02-03", None), + (3, "active", "2024-02-01", None), + ), + _REAPPEARING_KEY_FINAL_HISTORY, + ), + expected_full_history_rows=_REAPPEARING_KEY_FINAL_HISTORY, + ), + SnapshotReappearingKeyBuildE2ETestCase( + description=( + "historical timestamp snapshot reopens hard-deleted keys at the observation time" + ), + repo_files={ + "sqlbuild_project.toml": dedent( + """ + name = "snapshot_reappearing_keys" + adapter = "duckdb" + + [connection] + database = "snapshot_reappearing_keys.duckdb" + """ + ).strip() + + "\n", + "sources/raw.yml": dedent( + """ + sources: + - name: raw_membership_observations + schema: main + table: raw_membership_observations + - name: raw_membership_full_history + schema: main + table: raw_membership_full_history + """ + ).strip() + + "\n", + "models/membership_incremental_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy timestamp, + updated_at updated_at, + observed_at observed_at, + historical_input snapshot, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, updated_at, observed_at + FROM __source("raw_membership_observations") + """ + ).strip() + + "\n", + "models/membership_full_history_snapshot.sql": dedent( + """ + MODEL ( + materialized snapshot, + unique_key [customer_id], + snapshot_strategy timestamp, + updated_at updated_at, + observed_at observed_at, + historical_input snapshot, + invalidate_hard_deletes true + ); + + SELECT customer_id, status, updated_at, observed_at + FROM __source("raw_membership_full_history") + """ + ).strip() + + "\n", + }, + initial_seed_sql=( + f"CREATE TABLE main.raw_membership_observations {_TIMESTAMP_REAPPEARING_KEY_TABLE_SQL}; " + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[0]}; " + f"CREATE TABLE main.raw_membership_full_history {_TIMESTAMP_REAPPEARING_KEY_TABLE_SQL}; " + "INSERT INTO main.raw_membership_full_history VALUES " + f"{', '.join(_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS)}" + ), + observation_sql_by_round=( + (), + ( + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[1]}", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[2]}", + ), + ( + "INSERT INTO main.raw_membership_observations VALUES " + f"{_TIMESTAMP_REAPPEARING_KEY_ROUND_ROWS[3]}", + ), + ), + history_query_template=( + "SELECT customer_id, status, strftime(valid_from, '%Y-%m-%d'), " + "strftime(valid_to, '%Y-%m-%d') FROM main.{table_name} " + "ORDER BY customer_id, valid_from" + ), + expected_incremental_rows_by_round=( + ( + (1, "basic", "2024-01-15", None), + (2, "basic", "2024-01-15", None), + (3, "basic", "2024-01-15", None), + ), + ( + (1, "basic", "2024-01-15", None), + (2, "basic", "2024-01-15", "2024-02-02"), + (3, "basic", "2024-01-15", None), + ), + ( + (1, "basic", "2024-01-15", None), + (2, "basic", "2024-01-15", "2024-02-02"), + (2, "basic", "2024-02-03", None), + (3, "basic", "2024-01-15", None), + ), + _TIMESTAMP_REAPPEARING_KEY_FINAL_HISTORY, + ), + expected_full_history_rows=_TIMESTAMP_REAPPEARING_KEY_FINAL_HISTORY, + ), + ], + ids=lambda case: case.description, +) +def test_given_hard_deleted_key_when_it_reappears_unchanged_then_snapshot_opens_new_version( + test_case: SnapshotReappearingKeyBuildE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="snapshot_reappearing_keys", + repo_files=test_case.repo_files, + ) + db_path: Path = project_dir / "snapshot_reappearing_keys.duckdb" + + import duckdb + + connection: duckdb.DuckDBPyConnection = duckdb.connect(str(db_path)) + connection.execute(test_case.initial_seed_sql) + connection.close() + + observation_sql_round: tuple[str, ...] + expected_incremental_rows: tuple[tuple[object, ...], ...] + for observation_sql_round, expected_incremental_rows in zip( + test_case.observation_sql_by_round, + test_case.expected_incremental_rows_by_round, + strict=True, + ): + connection = duckdb.connect(str(db_path)) + observation_sql: str + for observation_sql in observation_sql_round: + connection.execute(observation_sql) + connection.close() + + result: object = run_sqb(command=("--no-color", "build"), project_dir=project_dir) + assert result.returncode == 0, result.stdout + result.stderr + + incremental_rows: tuple[tuple[object, ...], ...] = tuple( + tuple(row) + for row in query_duckdb( + db_path=db_path, + sql=test_case.history_query_template.format( + table_name="membership_incremental_snapshot" + ), + ) + ) + full_history_rows: tuple[tuple[object, ...], ...] = tuple( + tuple(row) + for row in query_duckdb( + db_path=db_path, + sql=test_case.history_query_template.format( + table_name="membership_full_history_snapshot" + ), + ) + ) + assert incremental_rows == expected_incremental_rows + assert full_history_rows == test_case.expected_full_history_rows + table_name: str + for table_name in ("membership_incremental_snapshot", "membership_full_history_snapshot"): + assert_snapshot_scd2_invariants( + db_path=db_path, table_name=table_name, key_columns=("customer_id",) + ) + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py new file mode 100644 index 0000000000..86f73a6061 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/build/test_sql_hook_identity.py @@ -0,0 +1,111 @@ +"""E2E tests for model version identity across SQL hook resource edits.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from textwrap import dedent + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.build._test_types import ( + SqlHookIdentityBuildE2ETestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project, run_sqb + +_HOOK_PATH: str = "hooks/sql/record_order_access.sql" +_HOOK_CONTENTS: str = ( + dedent( + """ + HOOK (description "Record order table access"); + + CREATE OR REPLACE TABLE main.order_access AS SELECT @label AS label + """ + ).strip() + + "\n" +) + + +@pytest.mark.parametrize( + "test_case", + [ + SqlHookIdentityBuildE2ETestCase( + description="hook description edit keeps the model unchanged", + removed_hook_path=_HOOK_PATH, + added_hook_path=_HOOK_PATH, + added_hook_contents=_HOOK_CONTENTS.replace( + "Record order table access", "Audit order reads" + ), + expected_reason="no_change", + ), + SqlHookIdentityBuildE2ETestCase( + description="hook file move keeps the model unchanged", + removed_hook_path=_HOOK_PATH, + added_hook_path="hooks/sql/audit/record_order_access.sql", + added_hook_contents=_HOOK_CONTENTS, + expected_reason="no_change", + ), + SqlHookIdentityBuildE2ETestCase( + description="hook SQL edit changes the model", + removed_hook_path=_HOOK_PATH, + added_hook_path=_HOOK_PATH, + added_hook_contents=_HOOK_CONTENTS.replace("@label AS label", "@label AS access"), + expected_reason="config_changed", + ), + ], + ids=lambda case: case.description, +) +def test_given_built_model_with_sql_hook_when_editing_hook_then_plans_expected_reason( + test_case: SqlHookIdentityBuildE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, + project_name="sql_hook_identity_project", + repo_files={ + "sqlbuild_project.toml": dedent( + """ + name = "sql_hook_identity_project" + adapter = "duckdb" + + [connection] + database = "sql_hook_identity_project.duckdb" + """ + ).strip() + + "\n", + _HOOK_PATH: _HOOK_CONTENTS, + "models/orders.sql": dedent( + """ + MODEL ( + materialized table, + post_hooks [sql("record_order_access", label: 1)] + ); + + SELECT 1 AS order_id + """ + ).strip() + + "\n", + }, + ) + first_build: subprocess.CompletedProcess[str] = run_sqb( + command=("--no-color", "build"), + project_dir=project_dir, + ) + assert first_build.returncode == 0, first_build.stdout + first_build.stderr + + (project_dir / test_case.removed_hook_path).unlink() + added_hook: Path = project_dir / test_case.added_hook_path + added_hook.parent.mkdir(parents=True, exist_ok=True) + added_hook.write_text(test_case.added_hook_contents) + plan: subprocess.CompletedProcess[str] = run_sqb( + command=("plan", "--json"), + project_dir=project_dir, + ) + + assert plan.returncode == 0, plan.stdout + plan.stderr + payload: dict[str, object] = json.loads(plan.stdout) + reasons_by_name: dict[str, str] = { + str(entry["name"]): str(entry["reason"]) for entry in payload["models"] + } + assert reasons_by_name["orders"] == test_case.expected_reason diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py b/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py index fde9573d92..0fdf5f26e5 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/postgres/test_postgres.py @@ -1116,6 +1116,7 @@ def test_given_intermediate_strategy_project_when_loading_twice_on_postgres_then ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), expected_failure_fragments=( "current_customer_snapshot", @@ -1231,6 +1232,7 @@ def test_given_snapshot_project_when_building_on_postgres_then_scd2_history_is_v ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), ) ], diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py b/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py index 07793fe235..2907149c44 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/scenario/test_scenario.py @@ -42,6 +42,36 @@ ) +@pytest.mark.parametrize( + "test_case", + ( + ScenarioCliE2ETestCase( + description="unclosed scenario comment names the scenario", + command=("--no-color", "scenario", "test"), + expected_exit_code=1, + expected_stderr_fragments=("SQL scenario contains an unclosed block comment",), + ), + ), + ids=lambda case: case.description, +) +def test_given_unclosed_scenario_comment_when_running_cli_then_error_names_scenario( + test_case: ScenarioCliE2ETestCase, + tmp_path: Path, +) -> None: + repo_files: dict[str, str] = build_scenario_project_files() + repo_files["tests/scenarios/order_totals_pass.sql"] = "SCENARIO ();\nWITH /* unfinished" + project_dir: Path = prepare_inline_project( + tmp_path=tmp_path, project_name="scenario_comment", repo_files=repo_files + ) + result: subprocess.CompletedProcess[str] = run_sqb( + command=test_case.command, project_dir=project_dir + ) + + assert result.returncode == test_case.expected_exit_code + for fragment in test_case.expected_stderr_fragments: + assert fragment in result.stderr + + @pytest.mark.parametrize( "test_case", ( diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py b/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py index 346f8e23ff..7c0290f4ec 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/sqlserver/test_sqlserver.py @@ -1764,6 +1764,7 @@ def test_given_intermediate_strategy_project_when_loading_twice_on_sqlserver_the ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), expected_failure_fragments=( "current_customer_snapshot", @@ -1885,6 +1886,7 @@ def test_given_snapshot_project_when_building_on_sqlserver_then_scd2_history_is_ ("1", "active", "2026-01-01", "2026-01-03"), ("1", "paused", "2026-01-03", None), ("2", "active", "2026-01-01", "2026-01-02"), + ("2", "active", "2026-01-03", None), ), ) ], 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 e196fd3377..45639506f2 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/helpers.py @@ -870,6 +870,43 @@ def build_star_mock_fixture_project_files() -> dict[str, str]: return files +def build_union_distinct_expected_project_files() -> dict[str, str]: + """Build a SQL test whose expected rows combine branches with UNION DISTINCT.""" + + return { + "sqlbuild_project.toml": ( + 'name = "union_distinct_expected"\n' + 'adapter = "duckdb"\n\n' + "[connection]\n" + 'database = "union_distinct_expected.duckdb"\n' + ), + "models/orders.sql": ( + 'MODEL ();\n\nSELECT DISTINCT order_id FROM __source("raw_orders")\n' + ), + "sources/raw_orders.yml": ( + "sources:\n" + " - name: raw_orders\n" + " schema: main\n" + " table: raw_orders\n" + " columns:\n" + " - name: order_id\n type: INTEGER\n" + ), + "tests/unit/test_orders.sql": ( + "TEST();\n\n" + "WITH\n" + "__source__raw_orders AS (\n" + " SELECT 1 AS order_id UNION ALL SELECT 1 AS order_id UNION ALL SELECT 2 AS order_id\n" + "),\n" + "__expected__orders AS (\n" + " SELECT 1 AS order_id\n" + " UNION DISTINCT\n" + " SELECT 2 AS order_id\n" + ")\n" + "SELECT 1\n" + ), + } + + def build_transformed_collection_project_files() -> dict[str, str]: """Build a valid aggregation that changes a scalar input into a collection output.""" diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py index 2ca60ccc35..3c4320c7fb 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/test/test_test.py @@ -62,6 +62,7 @@ build_star_partial_fixture_project_files, build_table_function_fixture_project_files, build_transformed_collection_project_files, + build_union_distinct_expected_project_files, build_unsatisfied_leaf_test_project_files, build_unspecified_nullability_fixture_project_files, build_untyped_null_fixture_project_files, @@ -316,6 +317,11 @@ def test_given_inspection_when_requesting_json_file_then_parser_rejects_conflict repo_files=build_star_mock_fixture_project_files(), expected_stdout_fragment="PASS=1", ), + FixtureCompatibilityE2ETestCase( + description="union distinct expected branches are validated per branch", + repo_files=build_union_distinct_expected_project_files(), + expected_stdout_fragment="PASS=1", + ), FixtureCompatibilityE2ETestCase( description="aggregation output does not inherit scalar input type", repo_files=build_transformed_collection_project_files(), diff --git a/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py b/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py index f5edc37879..713f303186 100644 --- a/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py +++ b/tests/e2e/src/sqlbuild/integrations/dagster/_test_types.py @@ -126,3 +126,12 @@ class DagsterComputeLogE2ETestCase: expected_stdout_fragments: tuple[str, ...] unexpected_structured_log_fragment: str expected_materialization_fragment: str + + +@dataclass(frozen=True) +class DagsterSelectorParityE2ETestCase: + """Test case comparing Dagster asset selection with the SQLBuild planner.""" + + description: str + select: str + expected_model_names: tuple[str, ...] diff --git a/tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py b/tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py new file mode 100644 index 0000000000..2ce58984c4 --- /dev/null +++ b/tests/e2e/src/sqlbuild/integrations/dagster/test_selection_parity.py @@ -0,0 +1,84 @@ +"""E2E parity between Dagster asset selection and SQLBuild planner selection.""" + +import json +import subprocess +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import dagster as dg +import pytest + +from sqlbuild.integrations.dagster import build_sqlbuild_asset_selection, sqlbuild_assets +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_waffle_shop, run_sqb +from tests.e2e.src.sqlbuild.integrations.dagster._test_types import ( + DagsterSelectorParityE2ETestCase, +) + + +@pytest.mark.parametrize( + "test_case", + ( + DagsterSelectorParityE2ETestCase( + description="both-direction expansion matches the planner", + select="+daily_revenue+", + expected_model_names=("daily_revenue", "stg_orders", "stg_payments"), + ), + DagsterSelectorParityE2ETestCase( + description="path selector keeps every path like the planner", + select="+fact_orders~hourly_activity_with_daily_context", + expected_model_names=( + "daily_activity_rollup", + "fact_orders", + "hourly_activity_with_daily_context", + "hourly_order_activity", + "stg_orders", + "stg_payments", + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_waffle_shop_selector_when_resolving_in_dagster_then_models_match_planner( + test_case: DagsterSelectorParityE2ETestCase, + tmp_path: Path, +) -> None: + project_dir: Path = prepare_waffle_shop(tmp_path) + dag_result: subprocess.CompletedProcess[str] = run_sqb( + command=("dag", "--json"), project_dir=project_dir + ) + assert dag_result.returncode == 0, dag_result.stdout + dag_result.stderr + dag: Mapping[str, Any] = json.loads(dag_result.stdout) + plan_result: subprocess.CompletedProcess[str] = run_sqb( + command=("plan", "--select", test_case.select, "--json"), project_dir=project_dir + ) + assert plan_result.returncode == 0, plan_result.stdout + plan_result.stderr + planned_model_names: tuple[str, ...] = tuple( + sorted(str(entry["name"]) for entry in json.loads(plan_result.stdout)["models"]) + ) + + @sqlbuild_assets(dag=dag) + def waffle_shop_assets() -> dg.MaterializeResult: + return dg.MaterializeResult() + + selection: Any = build_sqlbuild_asset_selection( + sqlbuild_assets=[waffle_shop_assets], + dag=dag, + sqlbuild_select=test_case.select, + ) + kind_and_name_by_asset_key: dict[tuple[str, ...], tuple[str, str]] = { + tuple(node["asset_key"]): (str(node["kind"]), str(node["name"])) for node in dag["nodes"] + } + selected_kind_and_names: set[tuple[str, str]] = { + kind_and_name_by_asset_key[tuple(key.path)] + for key in selection.resolve([waffle_shop_assets]) + } + model_kind_and_names: set[tuple[str, str]] = { + ("model", name) for _, name in kind_and_name_by_asset_key.values() + } + dagster_model_names: tuple[str, ...] = tuple( + sorted(name for _, name in selected_kind_and_names & model_kind_and_names) + ) + + assert planned_model_names == test_case.expected_model_names + assert dagster_model_names == test_case.expected_model_names diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py index cec73549b2..4e5bfe9e98 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/_test_types.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from dataclasses import dataclass from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -13,3 +14,38 @@ class SnapshotSqlRenderingAdapterTestCase: expected_historical_timestamp_initial_hard_delete_fragments: tuple[str, ...] expected_historical_timestamp_apply_hard_delete_fragments: tuple[str, ...] expected_historical_check_apply_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class SnapshotReappearanceRenderingTestCase: + description: str + adapter: BaseAdapter + expected_initial_fragments: tuple[str, ...] + expected_apply_fragments: tuple[str, ...] + unexpected_apply_fragments: tuple[str, ...] + + +@dataclass(frozen=True) +class SnapshotExecutionTestCase: + """One adapter executing one snapshot scenario through one build path on DuckDB.""" + + description: str + adapter_type: type[BaseAdapter] + normalize_sql: Callable[[str], str] + source_select_sql: Callable[[tuple[object, ...]], str] + render_initial: Callable[[BaseAdapter], tuple[str, ...]] + render_apply: Callable[[BaseAdapter], tuple[str, ...]] + history_sql: str + builds: tuple[tuple[tuple[object, ...], ...], ...] + expected_history: tuple[tuple[object, ...], ...] + retry_after_insert: bool = False + + +@dataclass(frozen=True) +class SqlServerSnapshotRenderingTestCase: + """SQL Server snapshot renders that must stay valid T-SQL.""" + + description: str + unexpected_fragments: tuple[str, ...] + expected_changes_apply_prefix: str + expected_changes_apply_fragments: tuple[str, ...] diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py new file mode 100644 index 0000000000..a138efe157 --- /dev/null +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/helpers.py @@ -0,0 +1,576 @@ +"""Helpers that execute adapter-rendered snapshot SQL on DuckDB over shared scenarios. + +DuckDB runs the SQL rendered by the DuckDB, MotherDuck, Postgres, BigQuery, Snowflake and +Databricks adapters. The only shims are identifier quoting for BigQuery and a pinned execution +clock so current-state histories are deterministic; neither changes statement semantics. SQL +Server renders T-SQL that DuckDB cannot parse, so it is covered by rendering assertions only. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import duckdb + +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.contract.models import SnapshotChangeTarget +from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter +from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.adapters.motherduck.classes.motherduck_adapter import MotherDuckAdapter +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter +from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.spec.contracts.types import TableType +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( + SnapshotExecutionTestCase, +) + + +@dataclass(frozen=True) +class SnapshotExecutionAdapter: + """Adapter whose rendered snapshot SQL is executed on DuckDB.""" + + name: str + adapter_type: type[BaseAdapter] + normalize_sql: Callable[[str], str] + + +@dataclass(frozen=True) +class SnapshotExecutionKind: + """Snapshot configuration family shared by execution scenarios.""" + + source_select_sql: Callable[[tuple[object, ...]], str] + render_initial: Callable[[BaseAdapter], tuple[str, ...]] + render_apply: Callable[[BaseAdapter], tuple[str, ...]] + history_sql: str + + +@dataclass(frozen=True) +class SnapshotExecutionScenario: + """Source rows per build and the history every adapter must produce.""" + + description: str + kind: SnapshotExecutionKind + builds: tuple[tuple[tuple[object, ...], ...], ...] + expected_history: tuple[tuple[object, ...], ...] + + +@dataclass(frozen=True) +class SnapshotExecutionRun: + """One adapter executing one scenario through incremental builds or one initial build.""" + + adapter: SnapshotExecutionAdapter + scenario: SnapshotExecutionScenario + path: str + builds: tuple[tuple[tuple[object, ...], ...], ...] + + +_TARGET: str = "snapshot_target" +_SOURCE: str = "snapshot_source" +FIRST_CLOCK_DAY: int = 20 +_KEY: tuple[str, ...] = ("customer_id",) + + +def _timestamp(day: object) -> str: + return f"TIMESTAMP '2026-01-{int(str(day)):02d}'" + + +def _check_history_row_sql(row: tuple[object, ...]) -> str: + return ( + f"SELECT {row[0]} AS customer_id, '{row[1]}' AS status, {_timestamp(row[2])} AS observed_at" + ) + + +def _timestamp_history_row_sql(row: tuple[object, ...]) -> str: + return ( + f"SELECT {row[0]} AS customer_id, '{row[1]}' AS plan, " + f"{_timestamp(row[2])} AS updated_at, {_timestamp(row[3])} AS observed_at" + ) + + +def _current_check_row_sql(row: tuple[object, ...]) -> str: + return f"SELECT {row[0]} AS customer_id, '{row[1]}' AS status" + + +def _current_timestamp_row_sql(row: tuple[object, ...]) -> str: + return f"SELECT {row[0]} AS customer_id, '{row[1]}' AS plan, {_timestamp(row[2])} AS updated_at" + + +def _historical_check_initial(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_create_initial_historical_check_snapshot_destination( + table_type=TableType.PERMANENT, + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + check_columns=("status",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "status", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _historical_check_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_historical_check_snapshot_changes( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + check_columns=("status",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "status", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _historical_timestamp_initial(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_create_initial_historical_timestamp_snapshot_destination( + table_type=TableType.PERMANENT, + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "updated_at", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _historical_timestamp_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_historical_timestamp_snapshot_changes( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "updated_at", "observed_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _current_initial( + adapter: BaseAdapter, strategy: str, updated_at_column: str | None +) -> tuple[str, ...]: + return adapter.render_create_initial_snapshot_destination( + table_type=TableType.PERMANENT, + destination=_TARGET, + origin=_SOURCE, + snapshot_strategy=strategy, + updated_at_column=updated_at_column, + observed_at_column=None, + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + ) + + +def _current_timestamp_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_timestamp_snapshot_changes( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + updated_at_column="updated_at", + observed_at_column=None, + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + output_columns=("customer_id", "plan", "updated_at"), + invalidate_hard_deletes=hard_deletes, + ) + + +def _current_check_apply(adapter: BaseAdapter, hard_deletes: bool) -> tuple[str, ...]: + return adapter.render_apply_check_snapshot_changes( + target=SnapshotChangeTarget( + destination=_TARGET, + origin=_SOURCE, + unique_key=_KEY, + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "status"), + ), + check_columns=("status",), + updated_at_column=None, + observed_at_column=None, + initial_valid_from=None, + invalidate_hard_deletes=hard_deletes, + ) + + +def _history_sql(value_column: str) -> str: + return ( + f"SELECT customer_id, {value_column}, day(valid_from), day(valid_to) " + f"FROM {_TARGET} ORDER BY customer_id, valid_from" + ) + + +def _unquote_backticks(sql: str) -> str: + return sql.replace("`", '"') + + +def _unchanged(sql: str) -> str: + return sql + + +HISTORICAL_CHECK: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_check_history_row_sql, + render_initial=lambda adapter: _historical_check_initial(adapter, False), + render_apply=lambda adapter: _historical_check_apply(adapter, False), + history_sql=_history_sql("status"), +) +HISTORICAL_CHECK_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_check_history_row_sql, + render_initial=lambda adapter: _historical_check_initial(adapter, True), + render_apply=lambda adapter: _historical_check_apply(adapter, True), + history_sql=_history_sql("status"), +) +HISTORICAL_TIMESTAMP: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_timestamp_history_row_sql, + render_initial=lambda adapter: _historical_timestamp_initial(adapter, False), + render_apply=lambda adapter: _historical_timestamp_apply(adapter, False), + history_sql=_history_sql("plan"), +) +HISTORICAL_TIMESTAMP_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_timestamp_history_row_sql, + render_initial=lambda adapter: _historical_timestamp_initial(adapter, True), + render_apply=lambda adapter: _historical_timestamp_apply(adapter, True), + history_sql=_history_sql("plan"), +) +CURRENT_TIMESTAMP: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_current_timestamp_row_sql, + render_initial=lambda adapter: _current_initial(adapter, "timestamp", "updated_at"), + render_apply=lambda adapter: _current_timestamp_apply(adapter, False), + history_sql=_history_sql("plan"), +) +CURRENT_TIMESTAMP_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_current_timestamp_row_sql, + render_initial=lambda adapter: _current_initial(adapter, "timestamp", "updated_at"), + render_apply=lambda adapter: _current_timestamp_apply(adapter, True), + history_sql=_history_sql("plan"), +) +CURRENT_CHECK_HARD_DELETES: SnapshotExecutionKind = SnapshotExecutionKind( + source_select_sql=_current_check_row_sql, + render_initial=lambda adapter: _current_initial(adapter, "check", None), + render_apply=lambda adapter: _current_check_apply(adapter, True), + history_sql=_history_sql("status"), +) + +ADAPTERS: tuple[SnapshotExecutionAdapter, ...] = ( + SnapshotExecutionAdapter(name="duckdb", adapter_type=DuckDbAdapter, normalize_sql=_unchanged), + SnapshotExecutionAdapter( + name="motherduck", adapter_type=MotherDuckAdapter, normalize_sql=_unchanged + ), + SnapshotExecutionAdapter( + name="postgres", adapter_type=PostgresAdapter, normalize_sql=_unchanged + ), + SnapshotExecutionAdapter( + name="bigquery", adapter_type=BigQueryAdapter, normalize_sql=_unquote_backticks + ), + SnapshotExecutionAdapter( + name="snowflake", adapter_type=SnowflakeAdapter, normalize_sql=_unchanged + ), + SnapshotExecutionAdapter( + name="databricks", adapter_type=DatabricksAdapter, normalize_sql=_unchanged + ), +) + +INVARIANT_VIOLATIONS_SQL: str = ( + "SELECT " + f"(SELECT COUNT(*) FROM (SELECT customer_id FROM {_TARGET} WHERE valid_to IS NULL " + "GROUP BY customer_id HAVING COUNT(*) > 1)), " + f"(SELECT COUNT(*) FROM {_TARGET} WHERE valid_to IS NOT NULL AND valid_from >= valid_to), " + f"(SELECT COUNT(*) FROM {_TARGET} AS left_row JOIN {_TARGET} AS right_row " + "ON left_row.customer_id = right_row.customer_id " + "AND left_row.valid_from < right_row.valid_from " + "AND right_row.valid_from < COALESCE(left_row.valid_to, TIMESTAMP '9999-12-31'))" +) + + +def build_execution_runs( + *, + incremental_scenarios: tuple[SnapshotExecutionScenario, ...], + full_history_scenarios: tuple[SnapshotExecutionScenario, ...], +) -> tuple[SnapshotExecutionRun, ...]: + """Run every scenario on every adapter, and historical scenarios as one initial build.""" + + runs: list[SnapshotExecutionRun] = [] + adapter: SnapshotExecutionAdapter + scenario: SnapshotExecutionScenario + for adapter in ADAPTERS: + for scenario in incremental_scenarios: + runs.append( + SnapshotExecutionRun( + adapter=adapter, + scenario=scenario, + path="incremental builds", + builds=scenario.builds, + ) + ) + for scenario in full_history_scenarios: + runs.append( + SnapshotExecutionRun( + adapter=adapter, + scenario=scenario, + path="initial full history", + builds=(scenario.builds[-1],), + ) + ) + return tuple(runs) + + +def run_snapshot_builds( + test_case: SnapshotExecutionTestCase, +) -> tuple[tuple[tuple[object, ...], ...], tuple[object, ...] | None]: + """Execute each build's rendered statements and return the history and invariant counts.""" + + connection: duckdb.DuckDBPyConnection = duckdb.connect() + _run_build( + connection=connection, + test_case=test_case, + rows=test_case.builds[0], + render=test_case.render_initial, + day=FIRST_CLOCK_DAY, + ) + build_index: int + rows: tuple[tuple[object, ...], ...] + for build_index, rows in enumerate(test_case.builds[1:], start=1): + _run_build( + connection=connection, + test_case=test_case, + rows=rows, + render=test_case.render_apply, + day=FIRST_CLOCK_DAY + build_index, + retry_after_insert=test_case.retry_after_insert, + ) + history: tuple[tuple[object, ...], ...] = tuple( + tuple(row) for row in connection.execute(test_case.history_sql).fetchall() + ) + violations: tuple[object, ...] | None = connection.execute(INVARIANT_VIOLATIONS_SQL).fetchone() + connection.close() + return history, violations + + +def _pinned_clock_adapter(adapter_type: type[BaseAdapter], day: int) -> BaseAdapter: + clock_sql: str = _timestamp(day) + pinned_type: type[BaseAdapter] = type( + f"PinnedClock{adapter_type.__name__}", + (adapter_type,), + {"render_current_timestamp": lambda _self: clock_sql}, + ) + return pinned_type() + + +def _run_build( + *, + connection: duckdb.DuckDBPyConnection, + test_case: SnapshotExecutionTestCase, + rows: tuple[tuple[object, ...], ...], + render: Callable[[BaseAdapter], tuple[str, ...]], + day: int, + retry_after_insert: bool = False, +) -> None: + source_sql: str = " UNION ALL ".join(test_case.source_select_sql(row) for row in rows) + connection.execute(f"CREATE OR REPLACE TABLE {_SOURCE} AS {source_sql}") + adapter: BaseAdapter = _pinned_clock_adapter(test_case.adapter_type, day) + statements: tuple[str, ...] = render(adapter) + statement: str + for statement in (*statements[: int(retry_after_insert)], *statements): + connection.execute(test_case.normalize_sql(statement)) + + +def render_snapshot_sql_matrix(adapter: BaseAdapter) -> dict[str, str]: + """Render every snapshot contract method over keys, strategies, hard deletes and paths.""" + + statements_by_label: dict[str, str] = {} + unique_key: tuple[str, ...] + for unique_key in (("customer_id",), ("customer_id", "region")): + description: str + statements: tuple[str, ...] + for description, statements in ( + *_current_state_renders(adapter=adapter, unique_key=unique_key), + *_historical_renders(adapter=adapter, unique_key=unique_key), + ): + index: int + statement: str + for index, statement in enumerate(statements): + statements_by_label[f"{description} key {unique_key} #{index}"] = statement + return statements_by_label + + +def _matrix_output_columns(unique_key: tuple[str, ...]) -> tuple[str, ...]: + return (*unique_key, "plan", "status", "updated_at", "observed_at") + + +def _current_state_renders( + *, adapter: BaseAdapter, unique_key: tuple[str, ...] +) -> list[tuple[str, tuple[str, ...]]]: + renders: list[tuple[str, tuple[str, ...]]] = [] + target: SnapshotChangeTarget = SnapshotChangeTarget( + destination="target_table", + origin="source_table", + unique_key=unique_key, + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=_matrix_output_columns(unique_key), + ) + strategy: str + for strategy in ("timestamp", "check"): + renders.append( + ( + f"current-state {strategy} initial", + adapter.render_create_initial_snapshot_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + snapshot_strategy=strategy, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + ), + ) + ) + hard_deletes: bool + for hard_deletes in (False, True): + renders.append( + ( + f"current-state timestamp apply (hard deletes {hard_deletes})", + adapter.render_apply_timestamp_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + initial_valid_from=None, + output_columns=target.output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"current-state check apply (hard deletes {hard_deletes})", + adapter.render_apply_check_snapshot_changes( + target=target, + check_columns=("plan", "status"), + updated_at_column="updated_at", + observed_at_column="observed_at", + initial_valid_from=None, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + return renders + + +def _historical_renders( + *, adapter: BaseAdapter, unique_key: tuple[str, ...] +) -> list[tuple[str, tuple[str, ...]]]: + output_columns: tuple[str, ...] = _matrix_output_columns(unique_key) + renders: list[tuple[str, tuple[str, ...]]] = [ + ( + "historical changes initial", + adapter.render_create_initial_historical_timestamp_changes_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + ), + ), + ( + "historical changes apply", + adapter.render_apply_historical_timestamp_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + ), + ), + ] + hard_deletes: bool + for hard_deletes in (False, True): + suffix: str = f"(hard deletes {hard_deletes})" + renders.append( + ( + f"historical timestamp initial {suffix}", + adapter.render_create_initial_historical_timestamp_snapshot_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"historical timestamp apply {suffix}", + adapter.render_apply_historical_timestamp_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + updated_at_column="updated_at", + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"historical check initial {suffix}", + adapter.render_create_initial_historical_check_snapshot_destination( + table_type=TableType.PERMANENT, + destination="target_table", + origin="source_table", + unique_key=unique_key, + check_columns=("plan", "status"), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + renders.append( + ( + f"historical check apply {suffix}", + adapter.render_apply_historical_check_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=unique_key, + check_columns=("plan", "status"), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=output_columns, + invalidate_hard_deletes=hard_deletes, + ), + ) + ) + return renders diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py new file mode 100644 index 0000000000..a1bd95630e --- /dev/null +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_duckdb_execution.py @@ -0,0 +1,376 @@ +"""Cross-adapter snapshot SQL executed on DuckDB; see helpers for the adapter shims.""" + +from __future__ import annotations + +import pytest + +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( + SnapshotExecutionTestCase, +) +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering.helpers import ( + CURRENT_CHECK_HARD_DELETES, + CURRENT_TIMESTAMP, + CURRENT_TIMESTAMP_HARD_DELETES, + FIRST_CLOCK_DAY, + HISTORICAL_CHECK, + HISTORICAL_CHECK_HARD_DELETES, + HISTORICAL_TIMESTAMP, + HISTORICAL_TIMESTAMP_HARD_DELETES, + SnapshotExecutionRun, + SnapshotExecutionScenario, + build_execution_runs, + run_snapshot_builds, +) + +_CHECK_DELETE_AND_REAPPEAR_ROWS: tuple[tuple[object, ...], ...] = ( + (1, "active", 1), + (2, "active", 1), + (1, "active", 2), + (1, "active", 3), + (2, "active", 3), +) +_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS: tuple[tuple[object, ...], ...] = ( + (1, "basic", 1, 1), + (2, "pro", 1, 1), + (1, "basic", 1, 2), + (1, "basic", 1, 3), + (2, "pro", 1, 3), +) + +_HISTORICAL_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="check without hard deletes records changes and ignores missing keys", + kind=HISTORICAL_CHECK, + builds=( + ((1, "active", 1), (2, "active", 1)), + ((1, "active", 1), (2, "active", 1), (1, "paused", 2), (1, "paused", 3)), + ), + expected_history=((1, "active", 1, 2), (1, "paused", 2, None), (2, "active", 1, None)), + ), + SnapshotExecutionScenario( + description="check initial observation", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=(((1, "active", 1), (2, "active", 1)),), + expected_history=((1, "active", 1, None), (2, "active", 1, None)), + ), + SnapshotExecutionScenario( + description="check hard delete", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=(_CHECK_DELETE_AND_REAPPEAR_ROWS[:2], _CHECK_DELETE_AND_REAPPEAR_ROWS[:3]), + expected_history=((1, "active", 1, None), (2, "active", 1, 2)), + ), + SnapshotExecutionScenario( + description="check unchanged reappearance across builds", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + _CHECK_DELETE_AND_REAPPEAR_ROWS[:3], + _CHECK_DELETE_AND_REAPPEAR_ROWS, + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, None), + ), + ), + SnapshotExecutionScenario( + description="check unchanged reappearance in the same build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=(_CHECK_DELETE_AND_REAPPEAR_ROWS[:2], _CHECK_DELETE_AND_REAPPEAR_ROWS), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp without hard deletes records updated_at changes", + kind=HISTORICAL_TIMESTAMP, + builds=( + ((1, "basic", 1, 1),), + ((1, "basic", 1, 1), (1, "pro", 2, 3), (1, "pro", 2, 4)), + ), + expected_history=((1, "basic", 1, 2), (1, "pro", 2, None)), + ), + SnapshotExecutionScenario( + description="timestamp initial observation", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=(_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2],), + expected_history=((1, "basic", 1, None), (2, "pro", 1, None)), + ), + SnapshotExecutionScenario( + description="timestamp hard delete", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2)), + ), + SnapshotExecutionScenario( + description="timestamp change followed by an unchanged repeat in one build", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1, 1),), + ((1, "basic", 1, 1), (1, "pro", 2, 3), (1, "pro", 2, 4)), + ), + expected_history=((1, "basic", 1, 2), (1, "pro", 2, None)), + ), + SnapshotExecutionScenario( + description="check change followed by an unchanged repeat in one build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + ((1, "active", 1),), + ((1, "active", 1), (1, "paused", 2), (1, "paused", 3)), + ), + expected_history=((1, "active", 1, 2), (1, "paused", 2, None)), + ), + SnapshotExecutionScenario( + description="check multiple delete and reappear cycles in one build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + ( + *_CHECK_DELETE_AND_REAPPEAR_ROWS, + (1, "active", 4), + (1, "active", 5), + (2, "active", 5), + ), + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, 4), + (2, "active", 5, None), + ), + ), + SnapshotExecutionScenario( + description="check reappearance, change and change back in one build", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + ( + *_CHECK_DELETE_AND_REAPPEAR_ROWS, + (1, "active", 4), + (2, "paused", 4), + (1, "active", 5), + (2, "active", 5), + ), + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, 4), + (2, "paused", 4, 5), + (2, "active", 5, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp unchanged updated_at reappearance across builds", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS, + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "pro", 3, None)), + ), + SnapshotExecutionScenario( + description="timestamp unchanged updated_at reappearance in the same build", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=(_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "pro", 3, None)), + ), + SnapshotExecutionScenario( + description="timestamp multiple delete and reappear cycles in one build", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + ( + *_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS, + (1, "basic", 1, 4), + (1, "basic", 1, 5), + (2, "pro", 1, 5), + ), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, 2), + (2, "pro", 3, 4), + (2, "pro", 5, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp changed reappearance starts at the new updated_at", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + (*_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], (1, "basic", 1, 4), (2, "team", 3, 4)), + (*_TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], (1, "basic", 1, 4), (2, "team", 3, 4)), + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "team", 3, None)), + ), +) + +_WINDOWED_SOURCE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="check reappearance after the deletion observation left the source window", + kind=HISTORICAL_CHECK_HARD_DELETES, + builds=( + _CHECK_DELETE_AND_REAPPEAR_ROWS[:2], + _CHECK_DELETE_AND_REAPPEAR_ROWS[:3], + _CHECK_DELETE_AND_REAPPEAR_ROWS[3:], + ), + expected_history=( + (1, "active", 1, None), + (2, "active", 1, 2), + (2, "active", 3, None), + ), + ), + SnapshotExecutionScenario( + description="timestamp reappearance after the deletion observation left the source window", + kind=HISTORICAL_TIMESTAMP_HARD_DELETES, + builds=( + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:2], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[:3], + _TIMESTAMP_DELETE_AND_REAPPEAR_ROWS[3:], + ), + expected_history=((1, "basic", 1, None), (2, "pro", 1, 2), (2, "pro", 3, None)), + ), +) + +_CURRENT_STATE_SCENARIOS: tuple[SnapshotExecutionScenario, ...] = ( + SnapshotExecutionScenario( + description="current timestamp reappearance updated_at equals deletion time", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1), (2, "pro", 1)), + ((1, "basic", 1),), + ((1, "basic", 1), (2, "team", FIRST_CLOCK_DAY + 1)), + ((1, "basic", 1), (2, "team", FIRST_CLOCK_DAY + 1)), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, FIRST_CLOCK_DAY + 1), + (2, "team", FIRST_CLOCK_DAY + 2, None), + ), + ), + SnapshotExecutionScenario( + description="current timestamp without hard deletes records updated_at changes", + kind=CURRENT_TIMESTAMP, + builds=(((1, "basic", 1),), ((1, "pro", 5),)), + expected_history=((1, "basic", 1, 5), (1, "pro", 5, None)), + ), + SnapshotExecutionScenario( + description="current timestamp hard deletes keep updated_at for changes", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=(((1, "basic", 1),), ((1, "pro", 5),)), + expected_history=((1, "basic", 1, 5), (1, "pro", 5, None)), + ), + SnapshotExecutionScenario( + description="current check hard delete then unchanged reappearance", + kind=CURRENT_CHECK_HARD_DELETES, + builds=(((1, "active"), (2, "active")), ((1, "active"),), ((1, "active"), (2, "active"))), + expected_history=( + (1, "active", FIRST_CLOCK_DAY, None), + (2, "active", FIRST_CLOCK_DAY, FIRST_CLOCK_DAY + 1), + (2, "active", FIRST_CLOCK_DAY + 2, None), + ), + ), + SnapshotExecutionScenario( + description="current timestamp hard delete then unchanged reappearance", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1), (2, "pro", 1)), + ((1, "basic", 1),), + ((1, "basic", 1), (2, "pro", 1)), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, FIRST_CLOCK_DAY + 1), + (2, "pro", FIRST_CLOCK_DAY + 2, None), + ), + ), + SnapshotExecutionScenario( + description="current timestamp hard delete then reappearance with a newer updated_at", + kind=CURRENT_TIMESTAMP_HARD_DELETES, + builds=( + ((1, "basic", 1), (2, "pro", 1)), + ((1, "basic", 1),), + ((1, "basic", 1), (2, "team", 5)), + ), + expected_history=( + (1, "basic", 1, None), + (2, "pro", 1, FIRST_CLOCK_DAY + 1), + (2, "team", FIRST_CLOCK_DAY + 2, None), + ), + ), +) + +_RUNS: tuple[SnapshotExecutionRun, ...] = build_execution_runs( + incremental_scenarios=( + *_HISTORICAL_SCENARIOS, + *_WINDOWED_SOURCE_SCENARIOS, + *_CURRENT_STATE_SCENARIOS, + ), + full_history_scenarios=_HISTORICAL_SCENARIOS, +) + + +@pytest.mark.parametrize( + "test_case", + [ + SnapshotExecutionTestCase( + description=f"{run.adapter.name}: {run.scenario.description} ({run.path})", + adapter_type=run.adapter.adapter_type, + normalize_sql=run.adapter.normalize_sql, + source_select_sql=run.scenario.kind.source_select_sql, + render_initial=run.scenario.kind.render_initial, + render_apply=run.scenario.kind.render_apply, + history_sql=run.scenario.kind.history_sql, + builds=run.builds, + expected_history=run.scenario.expected_history, + ) + for run in _RUNS + ], + ids=lambda case: case.description, +) +def test_given_snapshot_scenario_when_executing_adapter_sql_on_duckdb_then_history_matches( + test_case: SnapshotExecutionTestCase, +) -> None: + history: tuple[tuple[object, ...], ...] + violations: tuple[object, ...] | None + history, violations = run_snapshot_builds(test_case) + + assert history == test_case.expected_history + assert violations == (0, 0, 0) + + +@pytest.mark.parametrize( + "test_case", + [ + SnapshotExecutionTestCase( + description=f"{run.adapter.name}: {run.scenario.description} interrupted", + adapter_type=run.adapter.adapter_type, + normalize_sql=run.adapter.normalize_sql, + source_select_sql=run.scenario.kind.source_select_sql, + render_initial=run.scenario.kind.render_initial, + render_apply=run.scenario.kind.render_apply, + history_sql=run.scenario.kind.history_sql, + builds=run.builds, + expected_history=run.scenario.expected_history, + retry_after_insert=True, + ) + for run in _RUNS + if run.scenario.kind is CURRENT_TIMESTAMP_HARD_DELETES + ], + ids=lambda case: case.description, +) +def test_given_interrupted_timestamp_insert_when_retrying_then_history_is_idempotent( + test_case: SnapshotExecutionTestCase, +) -> None: + history, violations = run_snapshot_builds(test_case) + + assert history == test_case.expected_history + assert violations == (0, 0, 0) diff --git a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py index 10d33e17db..f297d819fb 100644 --- a/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py +++ b/tests/integration/src/sqlbuild/adapters/snapshot_sql_rendering/test_rendering.py @@ -7,10 +7,40 @@ from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter from sqlbuild.adapters.databricks.classes.databricks_adapter import DatabricksAdapter from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.adapters.motherduck.classes.motherduck_adapter import MotherDuckAdapter +from sqlbuild.adapters.postgres.classes.postgres_adapter import PostgresAdapter from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.adapters.sqlserver.classes.sqlserver_adapter import SqlServerAdapter from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering._test_types import ( + SnapshotReappearanceRenderingTestCase, SnapshotSqlRenderingAdapterTestCase, + SqlServerSnapshotRenderingTestCase, ) +from tests.integration.src.sqlbuild.adapters.snapshot_sql_rendering.helpers import ( + render_snapshot_sql_matrix, +) + +_REAPPEARANCE_INITIAL_FRAGMENTS: tuple[str, ...] = ( + "__observed_group_sequence AS (SELECT __observed_at, " + "LAG(__observed_at) OVER (ORDER BY __observed_at) AS __prev_group_observed_at, " + "LEAD(__observed_at) OVER (ORDER BY __observed_at) AS __next_group_observed_at " + "FROM (SELECT DISTINCT observed_at AS __observed_at FROM source_table)", + "(__prev_observed_at IS NOT NULL AND __prev_observed_at <> __prev_group_observed_at)", + "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING) AS __next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", +) +_REAPPEARANCE_APPLY_FRAGMENTS: tuple[str, ...] = ( + "AS __latest_rn FROM target_table) AS __latest_ranked WHERE __latest_rn = 1", + "__reappearances AS (SELECT __classified.customer_id, " + "MIN(__classified.observed_at) AS __reappeared_at", + "WHERE __latest.valid_to IS NOT NULL AND __closing_group.__observed_at IS NULL", + "OR __classified.observed_at < __first_new_starts.__first_version_start", + "__new_changes AS (SELECT *, LEAD(__version_start) OVER", + "customer_id, __version_start AS __close_at FROM __new_changes", + "__new_changes.__version_start, CASE WHEN __new_changes.__next_version_start IS NULL " + "THEN __new_changes.__next_absence_at", +) +_PORTABLE_APPLY_EXCLUSIONS: tuple[str, ...] = ("QUALIFY", "UNION DISTINCT", "NOT EXISTS") @pytest.mark.parametrize( @@ -29,16 +59,17 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "NOT EXISTS", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "NOT EXISTS", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -66,16 +97,17 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -103,16 +135,17 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -140,16 +173,17 @@ "__target.customer_id = __source.customer_id", "__target.region = __source.region", "WHEN MATCHED THEN UPDATE SET effective_to = __source.updated_at", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at IS NOT NULL THEN", ), expected_historical_check_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_initial_hard_delete_fragments=( - "__hard_deleted_at", - "WHEN __hard_deleted_at < __next_change_at THEN __hard_deleted_at", - "__hard_delete_candidates AS (", + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", ), expected_historical_timestamp_apply_hard_delete_fragments=( "__hard_deletes AS (", @@ -164,6 +198,46 @@ "UNION ALL", ), ), + SnapshotSqlRenderingAdapterTestCase( + description="sqlserver renders complete snapshot SQL matrix", + adapter=SqlServerAdapter(), + expected_create_initial_fragments=( + "SELECT * INTO target_table FROM (SELECT *", + "updated_at AS valid_from", + "CAST(NULL AS DATETIME2) AS valid_to", + ), + expected_timestamp_hard_delete_fragments=( + "UPDATE __target SET effective_to = __source.updated_at", + "__target.customer_id = __source.customer_id", + "__target.region = __source.region", + "WHERE __target.effective_to IS NULL AND NOT EXISTS", + "MAX(effective_to) AS __closed_at FROM target_table " + "GROUP BY customer_id, region) AS __history", + "AND __history.__closed_at IS NOT NULL THEN", + ), + expected_historical_check_initial_hard_delete_fragments=( + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", + "AS valid_to INTO target_table FROM __versions", + ), + expected_historical_timestamp_initial_hard_delete_fragments=( + "__next_absence_at", + "WHEN __next_absence_at < __next_version_start THEN __next_absence_at", + "AS valid_to INTO target_table FROM __versions", + ), + expected_historical_timestamp_apply_hard_delete_fragments=( + "__hard_deletes AS (", + "AS __close_at FROM target_table AS __target", + "UNION ALL", + ), + expected_historical_check_apply_fragments=( + "LAG(plan) OVER (PARTITION BY customer_id ORDER BY observed_at)", + "__hard_deletes AS (", + "INSERT INTO target_table (customer_id, plan, observed_at, valid_from, valid_to)", + "UPDATE __target SET valid_to = (SELECT MIN(__close_candidates.__close_at)", + "UNION ALL", + ), + ), ], ids=lambda case: case.description, ) @@ -265,3 +339,181 @@ def test_given_builtin_adapter_when_rendering_snapshot_sql_then_covers_snapshot_ assert expected_fragment in historical_timestamp_apply_sql for expected_fragment in test_case.expected_historical_check_apply_fragments: assert expected_fragment in historical_check_apply_sql + + +@pytest.mark.parametrize( + "test_case", + [ + SnapshotReappearanceRenderingTestCase( + description="duckdb renders portable hard-delete history SQL", + adapter=DuckDbAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="motherduck renders portable hard-delete history SQL", + adapter=MotherDuckAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="postgres renders portable hard-delete history SQL", + adapter=PostgresAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="bigquery renders portable hard-delete history SQL", + adapter=BigQueryAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="snowflake renders portable hard-delete history SQL", + adapter=SnowflakeAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="databricks renders portable hard-delete history SQL", + adapter=DatabricksAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=_PORTABLE_APPLY_EXCLUSIONS, + ), + SnapshotReappearanceRenderingTestCase( + description="sqlserver renders portable hard-delete history SQL", + adapter=SqlServerAdapter(), + expected_initial_fragments=_REAPPEARANCE_INITIAL_FRAGMENTS, + expected_apply_fragments=_REAPPEARANCE_APPLY_FRAGMENTS, + unexpected_apply_fragments=(*_PORTABLE_APPLY_EXCLUSIONS, "IS DISTINCT FROM"), + ), + ], + ids=lambda case: case.description, +) +def test_given_hard_deletes_when_rendering_historical_check_snapshot_then_reopens_reappearing_keys( + test_case: SnapshotReappearanceRenderingTestCase, +) -> None: + initial_sql: str = "\n".join( + test_case.adapter.render_create_initial_historical_check_snapshot_destination( + table_type="table", + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=True, + ) + ) + apply_statements: tuple[str, ...] = ( + test_case.adapter.render_apply_historical_check_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=True, + ) + ) + without_hard_deletes_sql: str = "\n".join( + ( + *test_case.adapter.render_create_initial_historical_check_snapshot_destination( + table_type="table", + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=False, + ), + *test_case.adapter.render_apply_historical_check_snapshot_changes( + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + check_columns=("plan",), + observed_at_column="observed_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "observed_at"), + invalidate_hard_deletes=False, + ), + ) + ) + + expected_fragment: str + for expected_fragment in test_case.expected_initial_fragments: + assert expected_fragment in initial_sql + apply_sql: str = "\n".join(apply_statements) + for expected_fragment in test_case.expected_apply_fragments: + assert expected_fragment in apply_sql + unexpected_fragment: str + for unexpected_fragment in test_case.unexpected_apply_fragments: + assert unexpected_fragment not in apply_sql + assert "__prev_group_observed_at" not in without_hard_deletes_sql + assert "__reappearances" not in without_hard_deletes_sql + assert "__next_absence_at" not in without_hard_deletes_sql + + +@pytest.mark.parametrize( + "test_case", + [ + SqlServerSnapshotRenderingTestCase( + description="sqlserver renders every snapshot statement without QUALIFY", + unexpected_fragments=("QUALIFY",), + expected_changes_apply_prefix=( + ";WITH __latest_ordered AS (SELECT *, ROW_NUMBER() OVER (" + "PARTITION BY customer_id ORDER BY updated_at DESC) AS __rn FROM target_table), " + "__latest AS (SELECT * FROM __latest_ordered WHERE __rn = 1), __new_changes AS (" + ), + expected_changes_apply_fragments=( + "UPDATE __target SET valid_to = (", + "FROM target_table AS __target WHERE __target.valid_to IS NULL", + "INSERT INTO target_table (", + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_sqlserver_when_rendering_snapshot_matrix_then_sql_is_tsql_compatible( + test_case: SqlServerSnapshotRenderingTestCase, +) -> None: + statements_by_label: dict[str, str] = render_snapshot_sql_matrix(SqlServerAdapter()) + changes_apply_statements: tuple[str, ...] = ( + SqlServerAdapter().render_apply_historical_timestamp_changes( + destination="target_table", + origin="source_table", + unique_key=("customer_id",), + updated_at_column="updated_at", + valid_from_column="valid_from", + valid_to_column="valid_to", + output_columns=("customer_id", "plan", "updated_at"), + ) + ) + + fragment: str + for fragment in test_case.unexpected_fragments: + assert { + label: fragment in statement for label, statement in statements_by_label.items() + } == dict.fromkeys(statements_by_label, False) + assert len(changes_apply_statements) == 2 + statement: str + for statement in changes_apply_statements: + assert statement.startswith(test_case.expected_changes_apply_prefix) + changes_apply_sql: str = "\n".join(changes_apply_statements) + expected_fragment: str + for expected_fragment in test_case.expected_changes_apply_fragments: + assert expected_fragment in changes_apply_sql diff --git a/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py new file mode 100644 index 0000000000..e0de81c887 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/__init__.py @@ -0,0 +1 @@ +"""Integration tests for shared CLI command runtime helpers.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py new file mode 100644 index 0000000000..8d5e37e5b1 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/_test_types.py @@ -0,0 +1,20 @@ +from collections.abc import Callable +from dataclasses import dataclass + +from sqlbuild.cli.commands.models import ( + AuditInvocation, + CheckInvocation, + SeedInvocation, + TestInvocation, +) + + +@dataclass(frozen=True) +class CommandInvocationTestCase: + description: str + resolve: Callable[..., AuditInvocation | CheckInvocation | SeedInvocation | TestInvocation] + request_type: Callable[..., object] + json_output: bool + expected_invocation_type_name: str + expected_stream_name: str + expected_has_progress_reporters: bool diff --git a/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py new file mode 100644 index 0000000000..1f6b7b4de5 --- /dev/null +++ b/tests/integration/src/sqlbuild/cli/commands/_helpers/runtime/test_command_invocation.py @@ -0,0 +1,89 @@ +"""Integration tests for resolving CLI command invocations against a real DuckDB project.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from sqlbuild.cli.commands import models as command_models +from sqlbuild.cli.commands._helpers.audit.invocation import resolve_audit_invocation +from sqlbuild.cli.commands._helpers.check.invocation import resolve_check_invocation +from sqlbuild.cli.commands._helpers.seed.invocation import resolve_seed_invocation +from sqlbuild.cli.commands._helpers.test.invocation import resolve_test_invocation +from tests.integration.src.sqlbuild.cli.commands._helpers.runtime._test_types import ( + CommandInvocationTestCase, +) + +_PROJECT_TOML: str = ( + 'name = "orders_shop"\nadapter = "duckdb"\n\n[connection]\ndatabase = "orders.duckdb"\n' +) + + +@pytest.mark.parametrize( + "test_case", + ( + CommandInvocationTestCase( + description="audit json output reports progress on stderr", + resolve=resolve_audit_invocation, + request_type=command_models.AuditCommandRequest, + json_output=True, + expected_invocation_type_name="AuditInvocation", + expected_stream_name="stderr", + expected_has_progress_reporters=True, + ), + CommandInvocationTestCase( + description="check text output reports progress on stdout", + resolve=resolve_check_invocation, + request_type=command_models.CheckCommandRequest, + json_output=False, + expected_invocation_type_name="CheckInvocation", + expected_stream_name="stdout", + expected_has_progress_reporters=True, + ), + CommandInvocationTestCase( + description="test json output reports progress on stderr", + resolve=resolve_test_invocation, + request_type=command_models.TestCommandRequest, + json_output=True, + expected_invocation_type_name="TestInvocation", + expected_stream_name="stderr", + expected_has_progress_reporters=True, + ), + CommandInvocationTestCase( + description="seed text output omits shared progress reporters", + resolve=resolve_seed_invocation, + request_type=command_models.SeedCommandRequest, + json_output=False, + expected_invocation_type_name="SeedInvocation", + expected_stream_name="stdout", + expected_has_progress_reporters=False, + ), + ), + ids=lambda case: case.description, +) +def test_given_duckdb_project_when_resolving_command_invocation_then_returns_command_context( + tmp_path: Path, + test_case: CommandInvocationTestCase, +) -> None: + _ = (tmp_path / "sqlbuild_project.toml").write_text(_PROJECT_TOML, encoding="utf-8") + request: object = test_case.request_type( + project_dir=tmp_path, no_color=True, json_output=test_case.json_output + ) + + invocation: ( + command_models.AuditInvocation + | command_models.CheckInvocation + | command_models.SeedInvocation + | command_models.TestInvocation + ) = test_case.resolve(request=request) + + assert type(invocation).__name__ == test_case.expected_invocation_type_name + assert invocation.effective_project_dir == tmp_path + assert invocation.adapter_name == "duckdb" + assert invocation.discovered_inputs.project_config.name == "orders_shop" + assert invocation.use_color is False + assert invocation.progress_stream is getattr(sys, test_case.expected_stream_name) + assert hasattr(invocation, "connection_progress") is test_case.expected_has_progress_reporters + assert hasattr(invocation, "planning_progress") is test_case.expected_has_progress_reporters diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py b/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py index d52b2139a0..f1ae320a32 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/_test_types.py @@ -2,8 +2,10 @@ from dataclasses import dataclass +from sqlbuild.virtual.state.types import VirtualEnvironmentStatus from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ConditionalPublicationPayload, + StateRefContractObservation, ) @@ -252,3 +254,21 @@ class MicrobatchStateRoundTripTestCase: description: str expected_event_count: int + + +@dataclass(frozen=True) +class StateReadContractTestCase: + description: str + schema: str + sqlbuild_version: str + expected_environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] + expected_active_lock_keys: tuple[str, ...] + expected_expired_lock_keys: tuple[str, ...] + + +@dataclass(frozen=True) +class DuckDbStateBackendRefContractTestCase: + description: str + expected_observation: StateRefContractObservation + schema: str + sqlbuild_version: str diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py index 5cc4b6c320..34300313bd 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/helpers.py @@ -1,19 +1,28 @@ from __future__ import annotations from dataclasses import replace +from datetime import datetime, timedelta from pathlib import Path from typing import Any, NamedTuple from sqlbuild.virtual.state.classes.duckdb import DuckDbStateBackend +from sqlbuild.virtual.state.classes.state_backend import StateBackend from sqlbuild.virtual.state.models import ( + FunctionVersionRecord, + PhysicalRelationRecord, VirtualEnvironmentCheckpointFunctionRefRecord, VirtualEnvironmentCheckpointModelRefRecord, VirtualEnvironmentCheckpointRecord, VirtualEnvironmentCheckpointSeedRefRecord, + VirtualEnvironmentFunctionRefRecord, VirtualEnvironmentNodeRefRecord, VirtualEnvironmentRecord, ) -from sqlbuild.virtual.state.types import VirtualEnvironmentStatus +from sqlbuild.virtual.state.types import ( + ModelVersionStatus, + PhysicalArtifactType, + VirtualEnvironmentStatus, +) _ENVIRONMENT: str = "dev" _CHECKPOINT_ID: str = "checkpoint-1" @@ -135,3 +144,308 @@ def open_duckdb_state_backend(*, db_path: Path) -> tuple[DuckDbStateBackend, Any def fetch_all(connection: Any, sql: str) -> list[tuple[Any, ...]]: return connection.execute(sql).fetchall() + + +STATE_READ_CONTRACT_FUNCTION: FunctionVersionRecord = FunctionVersionRecord( + function_name="is_large_order", + version_hash="function-v1", + language="sql", + returns="BOOLEAN", + arguments_json_b64="W3sibmFtZSI6ImFtb3VudCIsInR5cGUiOiJJTlRFR0VSIn1d", + return_columns_json_b64="W10=", + packages_json_b64="W10=", + runtime_version=None, + entry_point=None, + body_sql_b64="YW1vdW50ID4gOQ==", + definition_text_b64="YW1vdW50ID4gOQ==", + status=ModelVersionStatus.READY, +) +STATE_READ_CONTRACT_RELATION_V1: PhysicalRelationRecord = PhysicalRelationRecord( + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + version_hash="orders-v1", + database_name=None, + schema_name="sqb_orders", + relation_name="orders__v1", + relation_type="table", +) +STATE_READ_CONTRACT_RELATION_V2: PhysicalRelationRecord = replace( + STATE_READ_CONTRACT_RELATION_V1, version_hash="orders-v2", relation_name="orders__v2" +) + + +class StateReadContractObservation(NamedTuple): + """Values read back after writing the shared state read contract fixture.""" + + function_before_upsert: FunctionVersionRecord | None + function_after_upsert: FunctionVersionRecord | None + function_after_second_upsert: FunctionVersionRecord | None + relation_v1: PhysicalRelationRecord | None + relation_missing: PhysicalRelationRecord | None + relations: tuple[PhysicalRelationRecord, ...] + environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] + active_lock_keys: tuple[str, ...] + expired_lock_keys: tuple[str, ...] + + +def exercise_state_read_contract( + *, backend: StateBackend, connection: Any, schema: str +) -> StateReadContractObservation: + function_before_upsert: FunctionVersionRecord | None = backend.get_function_version( + connection=connection, + schema=schema, + function_name=STATE_READ_CONTRACT_FUNCTION.function_name, + version_hash=STATE_READ_CONTRACT_FUNCTION.version_hash, + ) + backend.upsert_function_version( + connection=connection, schema=schema, record=STATE_READ_CONTRACT_FUNCTION + ) + function_after_upsert: FunctionVersionRecord | None = backend.get_function_version( + connection=connection, + schema=schema, + function_name=STATE_READ_CONTRACT_FUNCTION.function_name, + version_hash=STATE_READ_CONTRACT_FUNCTION.version_hash, + ) + backend.upsert_function_version( + connection=connection, + schema=schema, + record=replace(STATE_READ_CONTRACT_FUNCTION, status=ModelVersionStatus.FAILED), + ) + function_after_second_upsert: FunctionVersionRecord | None = backend.get_function_version( + connection=connection, + schema=schema, + function_name=STATE_READ_CONTRACT_FUNCTION.function_name, + version_hash=STATE_READ_CONTRACT_FUNCTION.version_hash, + ) + backend.upsert_physical_relation( + connection=connection, schema=schema, record=STATE_READ_CONTRACT_RELATION_V1 + ) + backend.upsert_physical_relation( + connection=connection, schema=schema, record=STATE_READ_CONTRACT_RELATION_V2 + ) + relation_v1: PhysicalRelationRecord | None = backend.get_physical_relation_for_artifact( + connection=connection, + schema=schema, + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + version_hash="orders-v1", + ) + relation_missing: PhysicalRelationRecord | None = backend.get_physical_relation_for_artifact( + connection=connection, + schema=schema, + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + version_hash="orders-v3", + ) + relations: tuple[PhysicalRelationRecord, ...] = backend.list_physical_relations_for_artifact( + connection=connection, + schema=schema, + artifact_type=PhysicalArtifactType.MODEL, + artifact_name="orders", + ) + backend.upsert_virtual_environment( + connection=connection, + schema=schema, + record=VirtualEnvironmentRecord("dev", VirtualEnvironmentStatus.ACTIVE), + ) + backend.upsert_virtual_environment( + connection=connection, + schema=schema, + record=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.FINALIZED), + ) + environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] = tuple( + (record.virtual_environment_name, record.status) + for record in backend.list_virtual_environments(connection=connection, schema=schema) + ) + backend.acquire_lock( + connection=connection, + schema=schema, + lock_key="lock_expired", + owner_id="run-1", + expires_at=datetime.now() - timedelta(days=2), + ) + backend.acquire_lock( + connection=connection, + schema=schema, + lock_key="lock_active", + owner_id="run-1", + expires_at=datetime.now() + timedelta(days=2), + ) + active_lock_keys: tuple[str, ...] = tuple( + lock.lock_key for lock in backend.list_active_locks(connection=connection, schema=schema) + ) + expired_lock_keys: tuple[str, ...] = tuple( + lock.lock_key for lock in backend.list_expired_locks(connection=connection, schema=schema) + ) + return StateReadContractObservation( + function_before_upsert=function_before_upsert, + function_after_upsert=function_after_upsert, + function_after_second_upsert=function_after_second_upsert, + relation_v1=relation_v1, + relation_missing=relation_missing, + relations=relations, + environments=environments, + active_lock_keys=active_lock_keys, + expired_lock_keys=expired_lock_keys, + ) + + +class StateRefContractObservation(NamedTuple): + """Values read back after writing the shared state ref contract fixture.""" + + node_refs_after_first_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] + node_refs_after_second_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] + untouched_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + function_refs: tuple[VirtualEnvironmentFunctionRefRecord, ...] + grouped_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + grouped_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + published_environment: VirtualEnvironmentRecord | None + published_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + deleted_environment: VirtualEnvironmentRecord | None + deleted_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + surviving_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] + + +def exercise_state_ref_contract( + *, backend: StateBackend, connection: Any, schema: str +) -> StateRefContractObservation: + backend.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + node_type="seed", + refs=(VirtualEnvironmentNodeRefRecord("dev", "seed", "countries", "seed-v1"),), + ) + backend.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + node_type="model", + refs=( + VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v1"), + VirtualEnvironmentNodeRefRecord("dev", "model", "customers", "customers-v1"), + ), + ) + node_refs_after_first_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="dev", node_type="model" + ) + ) + backend.replace_virtual_environment_node_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + node_type="model", + refs=(VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v2"),), + ) + node_refs_after_second_replace: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="dev", node_type="model" + ) + ) + untouched_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="dev", node_type="seed" + ) + ) + backend.replace_virtual_environment_function_refs( + connection=connection, + schema=schema, + virtual_environment_name="dev", + refs=( + VirtualEnvironmentFunctionRefRecord("dev", "table_fn", "customer_orders", "fn-v1"), + VirtualEnvironmentFunctionRefRecord("dev", "udf", "normalize_email", "fn-v2"), + ), + ) + function_refs: tuple[VirtualEnvironmentFunctionRefRecord, ...] = ( + backend.get_virtual_environment_function_refs( + connection=connection, schema=schema, virtual_environment_name="dev" + ) + ) + backend.replace_virtual_environment_node_ref_groups( + connection=connection, + schema=schema, + virtual_environment_name="qa", + refs_by_node_type={ + "model": (VirtualEnvironmentNodeRefRecord("qa", "model", "orders", "orders-v3"),), + "seed": (VirtualEnvironmentNodeRefRecord("qa", "seed", "countries", "seed-v2"),), + }, + ) + grouped_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="qa", node_type="model" + ) + ) + grouped_seed_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="qa", node_type="seed" + ) + ) + backend.upsert_virtual_environment_and_replace_node_ref_groups( + connection=connection, + schema=schema, + record=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.ACTIVE), + refs_by_node_type={ + "model": (VirtualEnvironmentNodeRefRecord("prod", "model", "orders", "orders-v1"),) + }, + ) + published_environment: VirtualEnvironmentRecord | None = backend.get_virtual_environment( + connection=connection, schema=schema, virtual_environment_name="prod" + ) + published_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="prod", node_type="model" + ) + ) + backend.delete_virtual_environment( + connection=connection, schema=schema, virtual_environment_name="prod" + ) + deleted_environment: VirtualEnvironmentRecord | None = backend.get_virtual_environment( + connection=connection, schema=schema, virtual_environment_name="prod" + ) + deleted_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="prod", node_type="model" + ) + ) + surviving_model_refs: tuple[VirtualEnvironmentNodeRefRecord, ...] = ( + backend.get_virtual_environment_node_refs( + connection=connection, schema=schema, virtual_environment_name="qa", node_type="model" + ) + ) + return StateRefContractObservation( + node_refs_after_first_replace=node_refs_after_first_replace, + node_refs_after_second_replace=node_refs_after_second_replace, + untouched_seed_refs=untouched_seed_refs, + function_refs=function_refs, + grouped_model_refs=grouped_model_refs, + grouped_seed_refs=grouped_seed_refs, + published_environment=published_environment, + published_model_refs=published_model_refs, + deleted_environment=deleted_environment, + deleted_model_refs=deleted_model_refs, + surviving_model_refs=surviving_model_refs, + ) + + +EXPECTED_STATE_REF_CONTRACT_OBSERVATION: StateRefContractObservation = StateRefContractObservation( + node_refs_after_first_replace=( + VirtualEnvironmentNodeRefRecord("dev", "model", "customers", "customers-v1"), + VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v1"), + ), + node_refs_after_second_replace=( + VirtualEnvironmentNodeRefRecord("dev", "model", "orders", "orders-v2"), + ), + untouched_seed_refs=(VirtualEnvironmentNodeRefRecord("dev", "seed", "countries", "seed-v1"),), + function_refs=( + VirtualEnvironmentFunctionRefRecord("dev", "table_fn", "customer_orders", "fn-v1"), + VirtualEnvironmentFunctionRefRecord("dev", "udf", "normalize_email", "fn-v2"), + ), + grouped_model_refs=(VirtualEnvironmentNodeRefRecord("qa", "model", "orders", "orders-v3"),), + grouped_seed_refs=(VirtualEnvironmentNodeRefRecord("qa", "seed", "countries", "seed-v2"),), + published_environment=VirtualEnvironmentRecord("prod", VirtualEnvironmentStatus.ACTIVE), + published_model_refs=(VirtualEnvironmentNodeRefRecord("prod", "model", "orders", "orders-v1"),), + deleted_environment=None, + deleted_model_refs=(), + surviving_model_refs=(VirtualEnvironmentNodeRefRecord("qa", "model", "orders", "orders-v3"),), +) diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py index d93e797c00..f993dd7f92 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/_test_types.py @@ -2,8 +2,10 @@ from dataclasses import dataclass +from sqlbuild.virtual.state.types import VirtualEnvironmentStatus from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ConditionalPublicationPayload, + StateRefContractObservation, ) @@ -209,3 +211,19 @@ class PostgresRetiredMicrobatchRecordTestCase: description: str retired_record_type: str expected_event_count: int + + +@dataclass(frozen=True) +class PostgresStateReadContractTestCase: + description: str + sqlbuild_version: str + expected_environments: tuple[tuple[str, VirtualEnvironmentStatus], ...] + expected_active_lock_keys: tuple[str, ...] + expected_expired_lock_keys: tuple[str, ...] + + +@dataclass(frozen=True) +class PostgresStateBackendRefContractTestCase: + description: str + expected_observation: StateRefContractObservation + sqlbuild_version: str diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py index e0fd36e5bf..2c97970d33 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/postgres/test_postgres.py @@ -2,6 +2,7 @@ from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import replace from datetime import UTC, datetime, timedelta from typing import Any @@ -66,6 +67,7 @@ CHECKPOINT_ENVIRONMENT_PAYLOAD, CHECKPOINT_ID_PAYLOAD, DETACHED_CHECKPOINT_PAYLOAD, + EXPECTED_STATE_REF_CONTRACT_OBSERVATION, FAILED_CHECKPOINT_PAYLOAD, FINALIZING_CHECKPOINT_PAYLOAD, FUNCTION_OMISSION_PAYLOAD, @@ -75,7 +77,14 @@ REF_ENVIRONMENT_PAYLOAD, REF_NODE_TYPE_PAYLOAD, SEED_EXTRA_PAYLOAD, + STATE_READ_CONTRACT_FUNCTION, + STATE_READ_CONTRACT_RELATION_V1, + STATE_READ_CONTRACT_RELATION_V2, VALID_PAYLOAD, + StateReadContractObservation, + StateRefContractObservation, + exercise_state_read_contract, + exercise_state_ref_contract, ) from tests.integration.src.sqlbuild.virtual.state.classes.postgres._test_types import ( PostgresAtomicFinalizedVirtualPublishTestCase, @@ -94,11 +103,13 @@ PostgresStateBackendNodeResultTestCase, PostgresStateBackendOperationEventTestCase, PostgresStateBackendPythonNodeIdentityTestCase, + PostgresStateBackendRefContractTestCase, PostgresStateBackendSeedRefTestCase, PostgresStateBackendSourceFreshnessTestCase, PostgresStateBackendTableCreationTestCase, PostgresStateBackendTransactionRollbackTestCase, PostgresStateBackendValidationTestCase, + PostgresStateReadContractTestCase, ) from tests.integration.src.sqlbuild.virtual.state.classes.postgres.helpers import ( fetch_all, @@ -2499,3 +2510,83 @@ def _insert_then_fail(**kwargs: Any) -> None: ) == checkpoint_refs ) + + +@pytest.mark.parametrize( + "test_case", + [ + PostgresStateReadContractTestCase( + description="round trips function versions relations environments and locks", + sqlbuild_version="0.0.test", + expected_environments=( + ("prod", VirtualEnvironmentStatus.FINALIZED), + ("dev", VirtualEnvironmentStatus.ACTIVE), + ), + expected_active_lock_keys=("lock_active",), + expected_expired_lock_keys=("lock_expired",), + ) + ], + ids=lambda case: case.description, +) +def test_given_postgres_state_backend_when_reading_state_then_matches_read_contract( + test_case: PostgresStateReadContractTestCase, + postgres_state_backend: PostgresStateBackend, + postgres_state_connection: Any, + postgres_state_schema: str, +) -> None: + postgres_state_backend.initialize( + connection=postgres_state_connection, + schema=postgres_state_schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateReadContractObservation = exercise_state_read_contract( + backend=postgres_state_backend, + connection=postgres_state_connection, + schema=postgres_state_schema, + ) + + assert observation.function_before_upsert is None + assert observation.function_after_upsert == STATE_READ_CONTRACT_FUNCTION + assert observation.function_after_second_upsert == replace( + STATE_READ_CONTRACT_FUNCTION, status=ModelVersionStatus.FAILED + ) + assert observation.relation_v1 == STATE_READ_CONTRACT_RELATION_V1 + assert observation.relation_missing is None + assert observation.relations == ( + STATE_READ_CONTRACT_RELATION_V2, + STATE_READ_CONTRACT_RELATION_V1, + ) + assert observation.environments == test_case.expected_environments + assert observation.active_lock_keys == test_case.expected_active_lock_keys + assert observation.expired_lock_keys == test_case.expected_expired_lock_keys + + +@pytest.mark.parametrize( + "test_case", + [ + PostgresStateBackendRefContractTestCase( + description="replaces node function and grouped refs per environment and type", + expected_observation=EXPECTED_STATE_REF_CONTRACT_OBSERVATION, + sqlbuild_version="0.0.test", + ) + ], + ids=lambda case: case.description, +) +def test_given_postgres_state_backend_when_replacing_refs_then_matches_ref_contract( + test_case: PostgresStateBackendRefContractTestCase, + postgres_state_backend: PostgresStateBackend, + postgres_state_connection: Any, + postgres_state_schema: str, +) -> None: + postgres_state_backend.initialize( + connection=postgres_state_connection, + schema=postgres_state_schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateRefContractObservation = exercise_state_ref_contract( + backend=postgres_state_backend, + connection=postgres_state_connection, + schema=postgres_state_schema, + ) + + assert observation == test_case.expected_observation diff --git a/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py b/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py index 177c6687e9..66e92b8089 100644 --- a/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py +++ b/tests/integration/src/sqlbuild/virtual/state/classes/test_duckdb.py @@ -90,6 +90,7 @@ DuckDbStateBackendNodeResultTestCase, DuckDbStateBackendOperationEventTestCase, DuckDbStateBackendPythonNodeIdentityTestCase, + DuckDbStateBackendRefContractTestCase, DuckDbStateBackendRollbackTestCase, DuckDbStateBackendSeedRefTestCase, DuckDbStateBackendSourceFreshnessTestCase, @@ -97,6 +98,7 @@ DuckDbStateBackendTransactionRollbackTestCase, DuckDbStateBackendValidationTestCase, MicrobatchStateRoundTripTestCase, + StateReadContractTestCase, ) from tests.integration.src.sqlbuild.virtual.state.classes.helpers import ( ACTIVE_CHECKPOINT_PAYLOAD, @@ -105,6 +107,7 @@ CHECKPOINT_ENVIRONMENT_PAYLOAD, CHECKPOINT_ID_PAYLOAD, DETACHED_CHECKPOINT_PAYLOAD, + EXPECTED_STATE_REF_CONTRACT_OBSERVATION, FAILED_CHECKPOINT_PAYLOAD, FINALIZING_CHECKPOINT_PAYLOAD, FUNCTION_OMISSION_PAYLOAD, @@ -114,7 +117,14 @@ REF_ENVIRONMENT_PAYLOAD, REF_NODE_TYPE_PAYLOAD, SEED_EXTRA_PAYLOAD, + STATE_READ_CONTRACT_FUNCTION, + STATE_READ_CONTRACT_RELATION_V1, + STATE_READ_CONTRACT_RELATION_V2, VALID_PAYLOAD, + StateReadContractObservation, + StateRefContractObservation, + exercise_state_read_contract, + exercise_state_ref_contract, fetch_all, open_duckdb_state_backend, ) @@ -2802,3 +2812,85 @@ def test_given_duckdb_state_backend_when_two_connections_acquire_same_lock_then_ finally: backend.close(first_connection) second_backend.close(second_connection) + + +@pytest.mark.parametrize( + "test_case", + [ + StateReadContractTestCase( + description="round trips function versions relations environments and locks", + schema="sqlbuild_state", + sqlbuild_version="0.0.test", + expected_environments=( + ("prod", VirtualEnvironmentStatus.FINALIZED), + ("dev", VirtualEnvironmentStatus.ACTIVE), + ), + expected_active_lock_keys=("lock_active",), + expected_expired_lock_keys=("lock_expired",), + ) + ], + ids=lambda case: case.description, +) +def test_given_duckdb_state_backend_when_reading_state_then_matches_read_contract( + test_case: StateReadContractTestCase, + tmp_path: Path, +) -> None: + backend, connection = open_duckdb_state_backend(db_path=tmp_path / "state.duckdb") + try: + backend.initialize( + connection=connection, + schema=test_case.schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateReadContractObservation = exercise_state_read_contract( + backend=backend, connection=connection, schema=test_case.schema + ) + finally: + backend.close(connection) + + assert observation.function_before_upsert is None + assert observation.function_after_upsert == STATE_READ_CONTRACT_FUNCTION + assert observation.function_after_second_upsert == replace( + STATE_READ_CONTRACT_FUNCTION, status=ModelVersionStatus.FAILED + ) + assert observation.relation_v1 == STATE_READ_CONTRACT_RELATION_V1 + assert observation.relation_missing is None + assert observation.relations == ( + STATE_READ_CONTRACT_RELATION_V2, + STATE_READ_CONTRACT_RELATION_V1, + ) + assert observation.environments == test_case.expected_environments + assert observation.active_lock_keys == test_case.expected_active_lock_keys + assert observation.expired_lock_keys == test_case.expected_expired_lock_keys + + +@pytest.mark.parametrize( + "test_case", + [ + DuckDbStateBackendRefContractTestCase( + description="replaces node function and grouped refs per environment and type", + expected_observation=EXPECTED_STATE_REF_CONTRACT_OBSERVATION, + schema="sqlbuild_state", + sqlbuild_version="0.0.test", + ) + ], + ids=lambda case: case.description, +) +def test_given_duckdb_state_backend_when_replacing_refs_then_matches_ref_contract( + test_case: DuckDbStateBackendRefContractTestCase, + tmp_path: Path, +) -> None: + backend, connection = open_duckdb_state_backend(db_path=tmp_path / "state.duckdb") + try: + backend.initialize( + connection=connection, + schema=test_case.schema, + sqlbuild_version=test_case.sqlbuild_version, + ) + observation: StateRefContractObservation = exercise_state_ref_contract( + backend=backend, connection=connection, schema=test_case.schema + ) + finally: + backend.close(connection) + + assert observation == test_case.expected_observation diff --git a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py index 00b008da24..6d016fad87 100644 --- a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py +++ b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/_test_types.py @@ -83,3 +83,18 @@ class LineageFingerprintAvailabilityTestCase: relative_path: str config: str expected_available: bool + + +@dataclass(frozen=True) +class LineagePathSelectorTestCase: + description: str + select: tuple[str, ...] + expected_node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class LineagePathSelectorErrorTestCase: + description: str + select: tuple[str, ...] + expected_error_code: str + expected_error_fragment: str diff --git a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py index 78c628baeb..f23331dd46 100644 --- a/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py +++ b/tests/unit/src/sqlbuild/cli/commands/_helpers/lineage/test_selection.py @@ -21,6 +21,8 @@ from sqlbuild.compiler.pipeline.models import ProjectGraph from tests.unit.src.sqlbuild.cli.commands._helpers.lineage._test_types import ( ColumnLineageSelectionTestCase, + LineagePathSelectorErrorTestCase, + LineagePathSelectorTestCase, LineageSelectionTestCase, LineageSelectorDepthErrorTestCase, ) @@ -250,3 +252,85 @@ def test_given_unclear_selector_anchor_with_depth_when_selecting_then_raises_use ) assert test_case.expected_error_fragment in str(error.value) + + +@pytest.mark.parametrize( + "test_case", + [ + LineagePathSelectorTestCase( + description="keeps only nodes on paths between endpoints", + select=("raw_orders~daily_rollup",), + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + LineagePathSelectorTestCase( + description="adds upstreams of the path start", + select=("+stg_orders~daily_rollup",), + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + LineagePathSelectorTestCase( + description="adds downstreams of the path end", + select=("raw_orders~stg_orders+",), + expected_node_ids=( + "model:daily_rollup", + "model:fact_orders", + "model:stg_orders", + "source:raw_orders", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_path_selector_when_selecting_lineage_then_returns_path_nodes( + test_case: LineagePathSelectorTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + result: LineageGraph = select_selector_lineage( + graph=graph, + select=test_case.select, + exclude=(), + depth=None, + ) + + assert node_ids(result.nodes) == test_case.expected_node_ids + + +@pytest.mark.parametrize( + "test_case", + [ + LineagePathSelectorErrorTestCase( + description="rejects an end that is not downstream of the start", + select=("daily_rollup~raw_orders",), + expected_error_code="C319", + expected_error_fragment=( + "'source:raw_orders' is not downstream of 'model:daily_rollup'" + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_unreachable_path_selector_when_selecting_lineage_then_raises_user_error( + test_case: LineagePathSelectorErrorTestCase, +) -> None: + graph: ProjectGraph = build_lineage_test_graph() + + with pytest.raises(CliUserError) as error: + select_selector_lineage( + graph=graph, + select=test_case.select, + exclude=(), + depth=None, + ) + + assert error.value.code == test_case.expected_error_code + assert test_case.expected_error_fragment in str(error.value) diff --git a/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py index 6c6bb9abaf..ae1f3dd4ca 100644 --- a/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/auditing/main/_test_types.py @@ -109,3 +109,19 @@ class InvalidMeasuredValueTestCase: description: str measured_value: float expected_error_message: str + + +@dataclass(frozen=True) +class ParseAuditInstancesTestCase: + description: str + raw_audits: object | None + null_as_empty: bool + expected_definition_names: tuple[str, ...] + + +@dataclass(frozen=True) +class ParseAuditInstancesErrorTestCase: + description: str + raw_audits: object | None + null_as_empty: bool + expected_message: str diff --git a/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py b/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py new file mode 100644 index 0000000000..d1c526fa55 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/auditing/main/test_parse_audit_instances.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from sqlbuild.compiler.auditing.main._parse_audit_instances import parse_audit_instances +from sqlbuild.spec.contracts.models import SchemaAuditInstance +from tests.unit.src.sqlbuild.compiler.auditing.main._test_types import ( + ParseAuditInstancesErrorTestCase, + ParseAuditInstancesTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + [ + ParseAuditInstancesTestCase( + description="parses string and mapping audits in order", + raw_audits=["not_null", {"accepted_values": {"values": ["open", "closed"]}}], + null_as_empty=False, + expected_definition_names=("not_null", "accepted_values"), + ), + ParseAuditInstancesTestCase( + description="treats null as empty when requested", + raw_audits=None, + null_as_empty=True, + expected_definition_names=(), + ), + ParseAuditInstancesTestCase( + description="accepts an empty list", + raw_audits=[], + null_as_empty=False, + expected_definition_names=(), + ), + ], + ids=lambda case: case.description, +) +def test_given_audit_list_when_parsing_then_returns_audit_instances( + test_case: ParseAuditInstancesTestCase, +) -> None: + result: tuple[SchemaAuditInstance, ...] = parse_audit_instances( + raw_audits=test_case.raw_audits, + file_path=Path("models/orders.yml"), + label="model orders", + error_class=ValueError, + null_as_empty=test_case.null_as_empty, + ) + + assert tuple(audit.definition_name for audit in result) == (test_case.expected_definition_names) + + +@pytest.mark.parametrize( + "test_case", + [ + ParseAuditInstancesErrorTestCase( + description="rejects null when null is not empty", + raw_audits=None, + null_as_empty=False, + expected_message="models/orders.yml model orders audits must be a list", + ), + ParseAuditInstancesErrorTestCase( + description="rejects a mapping", + raw_audits={"not_null": {}}, + null_as_empty=True, + expected_message="models/orders.yml model orders audits must be a list", + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_audit_list_when_parsing_then_raises_injected_error( + test_case: ParseAuditInstancesErrorTestCase, +) -> None: + with pytest.raises(ValueError) as error: + parse_audit_instances( + raw_audits=test_case.raw_audits, + file_path=Path("models/orders.yml"), + label="model orders", + error_class=ValueError, + null_as_empty=test_case.null_as_empty, + ) + + assert str(error.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py new file mode 100644 index 0000000000..0f86900ab7 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/authored_values/main/_test_types.py @@ -0,0 +1,18 @@ +from collections.abc import Callable +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AuthoredValueTestCase: + description: str + parse: Callable[[object | None], object] + raw_value: object | None + expected_value: object + + +@dataclass(frozen=True) +class AuthoredValueErrorTestCase: + description: str + parse: Callable[[object | None], object] + raw_value: object | None + expected_message: str diff --git a/tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py b/tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py new file mode 100644 index 0000000000..2b3a1b94bc --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/authored_values/main/helpers.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from pathlib import Path + +from sqlbuild.compiler.authored_values.main._optional_bool import optional_bool +from sqlbuild.compiler.authored_values.main._optional_mapping import optional_mapping +from sqlbuild.compiler.authored_values.main._optional_named_bool import optional_named_bool +from sqlbuild.compiler.authored_values.main._optional_named_string import ( + optional_named_string, +) +from sqlbuild.compiler.authored_values.main._optional_non_empty_string import ( + optional_non_empty_string, +) +from sqlbuild.compiler.authored_values.main._optional_string_tuple import ( + optional_string_tuple, +) +from sqlbuild.compiler.authored_values.main._require_non_empty_string import ( + require_non_empty_string, +) + +AUTHORED_VALUES_FILE_PATH: Path = Path("models/orders.yml") +AUTHORED_VALUES_LABEL: str = "model 'orders'" + + +class AuthoredValueTestError(Exception): + """Error class injected into authored value parsers under test.""" + + +def _entry(raw_value: object | None) -> dict[str, object]: + return {"setting": raw_value} + + +def parse_required_string(raw_value: object | None) -> object: + return require_non_empty_string( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_optional_string(raw_value: object | None) -> object: + return optional_non_empty_string( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_named_string(raw_value: object | None) -> object: + return optional_named_string( + raw_value=raw_value, + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + key="setting", + error_class=AuthoredValueTestError, + ) + + +def parse_optional_bool(raw_value: object | None) -> object: + return optional_bool( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_named_bool_default_false(raw_value: object | None) -> object: + return optional_named_bool( + raw_value=raw_value, + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + key="setting", + error_class=AuthoredValueTestError, + default=False, + ) + + +def parse_named_bool_default_none(raw_value: object | None) -> object: + return optional_named_bool( + raw_value=raw_value, + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + key="setting", + error_class=AuthoredValueTestError, + default=None, + ) + + +def parse_optional_mapping(raw_value: object | None) -> object: + return optional_mapping( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) + + +def parse_optional_string_tuple(raw_value: object | None) -> object: + return optional_string_tuple( + entry=_entry(raw_value), + key="setting", + file_path=AUTHORED_VALUES_FILE_PATH, + label=AUTHORED_VALUES_LABEL, + error_class=AuthoredValueTestError, + ) diff --git a/tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py b/tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py new file mode 100644 index 0000000000..518c6ae372 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/authored_values/main/test_authored_values.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import pytest + +from tests.unit.src.sqlbuild.compiler.authored_values.main._test_types import ( + AuthoredValueErrorTestCase, + AuthoredValueTestCase, +) +from tests.unit.src.sqlbuild.compiler.authored_values.main.helpers import ( + AuthoredValueTestError, + parse_named_bool_default_false, + parse_named_bool_default_none, + parse_named_string, + parse_optional_bool, + parse_optional_mapping, + parse_optional_string, + parse_optional_string_tuple, + parse_required_string, +) + + +@pytest.mark.parametrize( + "test_case", + [ + AuthoredValueTestCase( + description="required string keeps surrounding whitespace", + parse=parse_required_string, + raw_value=" orders ", + expected_value=" orders ", + ), + AuthoredValueTestCase( + description="optional string defaults to none", + parse=parse_optional_string, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="optional string returns value", + parse=parse_optional_string, + raw_value="orders", + expected_value="orders", + ), + AuthoredValueTestCase( + description="named string defaults to none", + parse=parse_named_string, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="optional bool defaults to none", + parse=parse_optional_bool, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="optional bool returns false", + parse=parse_optional_bool, + raw_value=False, + expected_value=False, + ), + AuthoredValueTestCase( + description="named bool uses false default", + parse=parse_named_bool_default_false, + raw_value=None, + expected_value=False, + ), + AuthoredValueTestCase( + description="named bool uses none default", + parse=parse_named_bool_default_none, + raw_value=None, + expected_value=None, + ), + AuthoredValueTestCase( + description="named bool returns true over default", + parse=parse_named_bool_default_false, + raw_value=True, + expected_value=True, + ), + AuthoredValueTestCase( + description="optional mapping defaults to empty", + parse=parse_optional_mapping, + raw_value=None, + expected_value={}, + ), + AuthoredValueTestCase( + description="optional mapping returns value", + parse=parse_optional_mapping, + raw_value={"region": "west"}, + expected_value={"region": "west"}, + ), + AuthoredValueTestCase( + description="optional string list defaults to empty", + parse=parse_optional_string_tuple, + raw_value=None, + expected_value=(), + ), + AuthoredValueTestCase( + description="optional string list returns tuple", + parse=parse_optional_string_tuple, + raw_value=["orders", "customers"], + expected_value=("orders", "customers"), + ), + ], + ids=lambda case: case.description, +) +def test_given_authored_value_when_parsing_then_returns_expected_value( + test_case: AuthoredValueTestCase, +) -> None: + assert test_case.parse(test_case.raw_value) == test_case.expected_value + + +@pytest.mark.parametrize( + "test_case", + [ + AuthoredValueErrorTestCase( + description="required string rejects missing value", + parse=parse_required_string, + raw_value=None, + expected_message=( + "models/orders.yml model 'orders' must define non-empty string 'setting'" + ), + ), + AuthoredValueErrorTestCase( + description="optional string rejects blank value", + parse=parse_optional_string, + raw_value=" ", + expected_message="models/orders.yml model 'orders' 'setting' must be a non-empty string", + ), + AuthoredValueErrorTestCase( + description="named string rejects non-string value", + parse=parse_named_string, + raw_value=3, + expected_message="models/orders.yml model 'orders' 'setting' must be a non-empty string", + ), + AuthoredValueErrorTestCase( + description="optional bool rejects string value", + parse=parse_optional_bool, + raw_value="yes", + expected_message="models/orders.yml model 'orders' 'setting' must be a boolean", + ), + AuthoredValueErrorTestCase( + description="named bool rejects integer value", + parse=parse_named_bool_default_false, + raw_value=1, + expected_message="models/orders.yml model 'orders' 'setting' must be a boolean", + ), + AuthoredValueErrorTestCase( + description="optional mapping rejects list value", + parse=parse_optional_mapping, + raw_value=["region"], + expected_message="models/orders.yml model 'orders' 'setting' must be a mapping", + ), + AuthoredValueErrorTestCase( + description="optional string list rejects mapping value", + parse=parse_optional_string_tuple, + raw_value={"region": "west"}, + expected_message="models/orders.yml model 'orders' 'setting' must be a list", + ), + AuthoredValueErrorTestCase( + description="optional string list rejects non-string entries", + parse=parse_optional_string_tuple, + raw_value=["orders", 3], + expected_message="models/orders.yml model 'orders' 'setting' entries must be strings", + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_authored_value_when_parsing_then_raises_injected_error( + test_case: AuthoredValueErrorTestCase, +) -> None: + with pytest.raises(AuthoredValueTestError) as error: + test_case.parse(test_case.raw_value) + + assert str(error.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py index 1f5c1c2482..2981975de0 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/_test_types.py @@ -884,3 +884,29 @@ class CachedLineageRoundTripTestCase: transform_kind: ColumnTransformKind confidence: ColumnLineageConfidence expected_compact: bool + + +@dataclass(frozen=True) +class ClassifyDirectLogicSqlTestCtesErrorTestCase: + description: str + mode: SqlTestMode + ctes: tuple[tuple[str, str], ...] + expected_message: str + + +@dataclass(frozen=True) +class CteScannerMessageTestCase: + description: str + sql: str + expected_message: str + + +@dataclass(frozen=True) +class ExpectedProjectionScanTestCase: + description: str + sql: str + expected_branches: tuple[str, ...] + expected_select_list_end: int + expected_commas: tuple[str, ...] + expected_alias: str | None + expected_contains_select_star: bool diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py new file mode 100644 index 0000000000..78559b4c50 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_direct_logic_sql_test_classification.py @@ -0,0 +1,286 @@ +"""Characterisation tests for macro, UDF, and table_fn SQL-test CTE classification errors.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile._helpers.sql_tests.core import classify_sql_test_ctes +from sqlbuild.compiler.compile.exceptions import CompileInputError +from sqlbuild.compiler.compile.models import CompileSqlTestCte +from sqlbuild.compiler.compile.types import SqlTestMode +from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + ClassifyDirectLogicSqlTestCtesErrorTestCase, +) + +_FILE: str = "tests/unit/orders.sql" + + +@pytest.mark.parametrize( + "test_case", + [ + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro duplicate actual", + mode=SqlTestMode.MACRO, + ctes=(("__macro_actual__", "SELECT 1"), ("__macro_actual__", "SELECT 2")), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' must define exactly one __macro_actual__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro duplicate expected", + mode=SqlTestMode.MACRO, + ctes=( + ("__macro_expected__", "SELECT 1 AS value"), + ("__macro_expected__", "SELECT 2 AS value"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' must define exactly one __macro_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro model-test cte", + mode=SqlTestMode.MACRO, + ctes=(("__ref__orders", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'macro' but defines model-test CTE '__ref__orders'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro udf-test cte", + mode=SqlTestMode.MACRO, + ctes=(("__udf_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'macro' but defines UDF-test CTE '__udf_actual__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro table_fn-test cte", + mode=SqlTestMode.MACRO, + ctes=(("__table_fn_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'macro' but defines table_fn-test CTE " + "'__table_fn_expected__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro reserved helper name", + mode=SqlTestMode.MACRO, + ctes=(("__missing__", "SELECT 1"),), + expected_message=f"SQL test '{_FILE}' uses reserved helper CTE name '__missing__'", + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro missing expected", + mode=SqlTestMode.MACRO, + ctes=(("__macro_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' must define exactly one __macro_actual__ CTE " + "and exactly one __macro_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro helper calls macro", + mode=SqlTestMode.MACRO, + ctes=( + ("helper", "SELECT @normalize_status('paid') AS status"), + ("__macro_actual__", "SELECT 1 AS status"), + ("__macro_expected__", "SELECT 1 AS status"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' helper CTE 'helper' must not call macros; " + "call macros only in __macro_actual__" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro expected calls macro", + mode=SqlTestMode.MACRO, + ctes=( + ("__macro_actual__", "SELECT 1 AS status"), + ("__macro_expected__", "SELECT @normalize_status('paid') AS status"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'macro' CTE __macro_expected__ must not call macros" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="macro expected unaliased projection names the expected cte", + mode=SqlTestMode.MACRO, + ctes=( + ("__macro_actual__", "SELECT 1 AS status"), + ("__macro_expected__", "SELECT 1 + 1"), + ), + expected_message=( + f"SQL test '{_FILE}' must alias every non-trivial __macro_expected__ projection" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn expected branch mismatch names the expected cte", + mode=SqlTestMode.TABLE_FN, + ctes=( + ("__table_fn_actual__", "SELECT 1 AS value"), + ("__table_fn_expected__", "SELECT 1 AS value UNION ALL SELECT 2 AS other"), + ), + expected_message=( + f"SQL test '{_FILE}' must use the same __table_fn_expected__ projection names " + "and order in every set-operation branch; branch 2 does not match branch 1" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf duplicate actual", + mode=SqlTestMode.UDF, + ctes=(("__udf_actual__", "SELECT 1"), ("__udf_actual__", "SELECT 2")), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' must define exactly one __udf_actual__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf duplicate expected", + mode=SqlTestMode.UDF, + ctes=( + ("__udf_expected__", "SELECT 1 AS value"), + ("__udf_expected__", "SELECT 2 AS value"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' must define exactly one __udf_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf model-test cte", + mode=SqlTestMode.UDF, + ctes=(("__expected__orders", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'udf' but defines model-test CTE '__expected__orders'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf macro-test cte", + mode=SqlTestMode.UDF, + ctes=(("__macro_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'udf' but defines macro-test CTE '__macro_expected__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf table_fn-test cte", + mode=SqlTestMode.UDF, + ctes=(("__table_fn_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'udf' but defines table_fn-test CTE " + "'__table_fn_actual__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf reserved helper name", + mode=SqlTestMode.UDF, + ctes=(("__actual", "SELECT 1"),), + expected_message=f"SQL test '{_FILE}' uses reserved helper CTE name '__actual'", + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf missing actual", + mode=SqlTestMode.UDF, + ctes=(("__udf_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' must define exactly one __udf_actual__ CTE " + "and exactly one __udf_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf helper calls macro", + mode=SqlTestMode.UDF, + ctes=( + ("helper", "SELECT @normalize_status('paid') AS status"), + ("__udf_actual__", "SELECT 1 AS formatted"), + ("__udf_expected__", "SELECT 1 AS formatted"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' helper CTE 'helper' must not call macros" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="udf expected calls udf", + mode=SqlTestMode.UDF, + ctes=( + ("__udf_actual__", "SELECT 1 AS formatted"), + ("__udf_expected__", 'SELECT __udf("format_cents")(1250) AS formatted'), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'udf' CTE __udf_expected__ must not call udf; " + "call reusable logic only in __udf_actual__" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn duplicate actual", + mode=SqlTestMode.TABLE_FN, + ctes=(("__table_fn_actual__", "SELECT 1"), ("__table_fn_actual__", "SELECT 2")), + expected_message=( + f"SQL test '{_FILE}' mode 'table_fn' must define exactly one " + "__table_fn_actual__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn duplicate expected", + mode=SqlTestMode.TABLE_FN, + ctes=( + ("__table_fn_expected__", "SELECT 1 AS value"), + ("__table_fn_expected__", "SELECT 2 AS value"), + ), + expected_message=( + f"SQL test '{_FILE}' mode 'table_fn' must define exactly one " + "__table_fn_expected__ CTE" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn model-test cte", + mode=SqlTestMode.TABLE_FN, + ctes=(("__source__raw_orders", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'table_fn' but defines model-test CTE " + "'__source__raw_orders'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn macro cte uses generic direct-logic wording", + mode=SqlTestMode.TABLE_FN, + ctes=(("__macro_actual__", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'table_fn' but defines another direct-logic CTE " + "'__macro_actual__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn udf cte uses generic direct-logic wording", + mode=SqlTestMode.TABLE_FN, + ctes=(("__udf_expected__", "SELECT 1 AS value"),), + expected_message=( + f"SQL test '{_FILE}' is mode 'table_fn' but defines another direct-logic CTE " + "'__udf_expected__'" + ), + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn reserved helper name", + mode=SqlTestMode.TABLE_FN, + ctes=(("__unexpected__", "SELECT 1"),), + expected_message=f"SQL test '{_FILE}' uses reserved helper CTE name '__unexpected__'", + ), + ClassifyDirectLogicSqlTestCtesErrorTestCase( + description="table_fn missing both", + mode=SqlTestMode.TABLE_FN, + ctes=(("helper", "SELECT 1"),), + expected_message=( + f"SQL test '{_FILE}' mode 'table_fn' must define exactly one " + "__table_fn_actual__ CTE and exactly one __table_fn_expected__ CTE" + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_direct_logic_ctes_when_classifying_then_raises_mode_specific_message( + test_case: ClassifyDirectLogicSqlTestCtesErrorTestCase, +) -> None: + ctes: tuple[CompileSqlTestCte, ...] = tuple( + CompileSqlTestCte(name=name, sql_body=body) for name, body in test_case.ctes + ) + + with pytest.raises(CompileInputError) as error_info: + _ = classify_sql_test_ctes(ctes=ctes, file_label=_FILE, mode=test_case.mode) + + assert str(error_info.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py new file mode 100644 index 0000000000..af71f97607 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_expected_projection_scanning.py @@ -0,0 +1,94 @@ +"""Characterisation tests for the top-level scanners behind SQL-test expected projections.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile._helpers.sql_tests import core as sql_test_core +from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + ExpectedProjectionScanTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + ( + ExpectedProjectionScanTestCase( + description="plain select splits on top-level commas", + sql="SELECT a, b AS c FROM orders", + expected_branches=("SELECT a, b AS c FROM orders",), + expected_select_list_end=len("SELECT a, b AS c "), + expected_commas=("SELECT a", "b AS c FROM orders"), + expected_alias=None, + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="quotes and comments hide keywords, commas and parentheses", + sql="SELECT 'x, FROM (' AS a, /* UNION , */ \"b)\" -- FROM\nFROM t", + expected_branches=("SELECT 'x, FROM (' AS a, /* UNION , */ \"b)\" -- FROM\nFROM t",), + expected_select_list_end=len("SELECT 'x, FROM (' AS a, /* UNION , */ \"b)\" -- FROM\n"), + expected_commas=("SELECT 'x, FROM (' AS a", '/* UNION , */ "b)" -- FROM\nFROM t'), + expected_alias=None, + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="nested parentheses hide commas, FROM and AS", + sql="SELECT f(a, (SELECT 1 AS x FROM y)) AS total", + expected_branches=("SELECT f(a, (SELECT 1 AS x FROM y)) AS total",), + expected_select_list_end=len("SELECT f(a, (SELECT 1 AS x FROM y)) AS total"), + expected_commas=("SELECT f(a, (SELECT 1 AS x FROM y)) AS total",), + expected_alias="total", + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="union all splits branches", + sql="SELECT 1 AS a UNION ALL SELECT 2 AS a", + expected_branches=("SELECT 1 AS a", "SELECT 2 AS a"), + expected_select_list_end=len("SELECT 1 AS a UNION ALL SELECT 2 AS a"), + expected_commas=("SELECT 1 AS a UNION ALL SELECT 2 AS a",), + expected_alias="a", + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="union distinct splits branches", + sql="SELECT 1 AS a UNION DISTINCT SELECT 2 AS a", + expected_branches=("SELECT 1 AS a", "SELECT 2 AS a"), + expected_select_list_end=len("SELECT 1 AS a UNION DISTINCT SELECT 2 AS a"), + expected_commas=("SELECT 1 AS a UNION DISTINCT SELECT 2 AS a",), + expected_alias="a", + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="backtick identifiers are quoted", + sql="SELECT `a, b` AS `c`, * FROM t", + expected_branches=("SELECT `a, b` AS `c`, * FROM t",), + expected_select_list_end=len("SELECT `a, b` AS `c`, * "), + expected_commas=("SELECT `a, b` AS `c`", "* FROM t"), + expected_alias=None, + expected_contains_select_star=False, + ), + ExpectedProjectionScanTestCase( + description="select star inside a subquery is detected", + sql="SELECT a FROM (SELECT * FROM t)", + expected_branches=("SELECT a FROM (SELECT * FROM t)",), + expected_select_list_end=len("SELECT a "), + expected_commas=("SELECT a FROM (SELECT * FROM t)",), + expected_alias=None, + expected_contains_select_star=True, + ), + ), + ids=lambda case: case.description, +) +def test_given_expected_cte_sql_when_scanning_top_level_then_results_are_characterised( + test_case: ExpectedProjectionScanTestCase, +) -> None: + assert sql_test_core._split_set_operation_branches(test_case.sql) == test_case.expected_branches + assert ( + sql_test_core._find_select_list_end(sql=test_case.sql, start=0) + == test_case.expected_select_list_end + ) + assert sql_test_core._split_top_level_commas(test_case.sql) == test_case.expected_commas + assert sql_test_core._extract_as_alias(test_case.sql) == test_case.expected_alias + assert sql_test_core._contains_select_star(test_case.sql) is ( + test_case.expected_contains_select_star + ) diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py index bb84908477..3b191b7ba2 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_scenarios.py @@ -5,7 +5,11 @@ import pytest from sqlbuild.compiler.compile._helpers.attachment.sql_tests import build_scenario_inputs -from sqlbuild.compiler.compile._helpers.scenarios.core import extract_sql_scenario_ctes +from sqlbuild.compiler.compile._helpers.scenarios.core import ( + extract_sql_scenario_ctes, + extract_sql_scenario_expected_model_names, +) +from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.models import ( CompileSqlScenarioCtes, CompileSqlScenarioInput, @@ -20,6 +24,7 @@ from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( BuildScenarioInputsErrorTestCase, BuildScenarioInputsTestCase, + CteScannerMessageTestCase, ExtractSqlScenarioCtesErrorTestCase, ExtractSqlScenarioCtesTestCase, ) @@ -439,3 +444,107 @@ def test_given_invalid_scenario_source_refs_when_building_inputs_then_it_raises_ loaded_macros={}, declaration_expansion=DUCKDB_DECLARATION_EXPANSION_CONTEXT, ) + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="scenario without a with clause uses scenario wording", + sql="SELECT * FROM orders", + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must declare fixture CTEs and at least " + "one __expected__ or __assert__ CTE before `SELECT 1`" + ), + ), + CteScannerMessageTestCase( + description="unclosed comment before a fixture uses scenario wording", + sql="WITH /* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), + CteScannerMessageTestCase( + description="unclosed trailing comment uses scenario wording", + sql="WITH __expected__orders AS (SELECT 1 AS order_id) SELECT 1; /* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), + CteScannerMessageTestCase( + description="scenario with an invalid cte name uses scenario wording", + sql="WITH 1orders AS (SELECT 1) SELECT 1", + expected_message="SQL scenario 'tests/scenarios/orders.sql' expected a CTE name", + ), + CteScannerMessageTestCase( + description="scenario without a ceremonial select uses scenario wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__orders AS (SELECT 1 AS order_id), " + "result AS (SELECT 1) SELECT * FROM result" + ), + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must end with a ceremonial top-level " + "`SELECT 1` after its CTEs" + ), + ), + CteScannerMessageTestCase( + description="scenario fixture without a target uses scenario wording", + sql=( + "WITH __source__ AS (SELECT 1 AS order_id), " + "__expected__orders AS (SELECT 1 AS order_id) SELECT 1" + ), + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must use __source__ to " + "identify a target" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_scenario_ctes_when_extracting_then_messages_name_the_scenario( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_scenario_ctes(sql=test_case.sql, file_label="tests/scenarios/orders.sql") + + assert str(error_info.value) == test_case.expected_message + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="expected model scan without AS uses scenario wording", + sql="WITH __source__raw_orders (SELECT 1 AS order_id) SELECT 1", + expected_message="SQL scenario 'tests/scenarios/orders.sql' expected keyword AS", + ), + CteScannerMessageTestCase( + description="expected model scan unclosed comment uses scenario wording", + sql="WITH /* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), + CteScannerMessageTestCase( + description="expected model scan leading unclosed comment uses scenario wording", + sql="/* unfinished", + expected_message="SQL scenario contains an unclosed block comment", + ), + CteScannerMessageTestCase( + description="expected model scan without a target uses scenario wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__ AS (SELECT 1 AS order_id) SELECT 1" + ), + expected_message=( + "SQL scenario 'tests/scenarios/orders.sql' must use __expected__ to " + "identify a target" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_scenario_ctes_when_scanning_expected_models_then_names_the_scenario( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_scenario_expected_model_names( + sql=test_case.sql, file_label="tests/scenarios/orders.sql" + ) + + assert str(error_info.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py index 6d9b6a2586..0476e95425 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/test_tests.py @@ -5,13 +5,16 @@ from sqlbuild.compiler.compile._helpers.sql_tests.core import ( CompileSqlTestCtes, extract_sql_test_ctes, + extract_sql_test_expected_model_names, ) +from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.models import ( CompileDirectLogicSqlTestCtes, CompileModelSqlTestCtes, ) from sqlbuild.compiler.compile.types import SqlTestMode from tests.unit.src.sqlbuild.compiler.compile._helpers._test_types import ( + CteScannerMessageTestCase, ExtractSqlTestCtesErrorTestCase, ExtractSqlTestCtesTestCase, ) @@ -833,3 +836,75 @@ def test_given_invalid_sql_test_cte_variants_when_extracting_then_it_raises_clea file_label="tests/unit/orders.sql", mode=test_case.mode, ) + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="sql test without a with clause keeps sql test wording", + sql="SELECT * FROM orders", + expected_message=( + "SQL test 'tests/unit/orders.sql' must declare mock CTEs and one " + "__expected__ CTE before `SELECT 1`" + ), + ), + CteScannerMessageTestCase( + description="sql test with an invalid cte name keeps sql test wording", + sql="WITH 1orders AS (SELECT 1) SELECT 1", + expected_message="SQL test 'tests/unit/orders.sql' expected a CTE name", + ), + CteScannerMessageTestCase( + description="sql test without a ceremonial select keeps sql test wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__orders AS (SELECT 1 AS order_id), " + "result AS (SELECT 1) SELECT * FROM result" + ), + expected_message=( + "SQL test 'tests/unit/orders.sql' must end with a ceremonial top-level " + "`SELECT 1` after its CTEs" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_sql_test_ctes_when_extracting_then_messages_name_the_sql_test( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_test_ctes(sql=test_case.sql, file_label="tests/unit/orders.sql") + + assert str(error_info.value) == test_case.expected_message + + +@pytest.mark.parametrize( + "test_case", + ( + CteScannerMessageTestCase( + description="expected model scan without AS keeps sql test wording", + sql="WITH __source__raw_orders (SELECT 1 AS order_id) SELECT 1", + expected_message="SQL test 'tests/unit/orders.sql' expected keyword AS", + ), + CteScannerMessageTestCase( + description="expected model scan without a target keeps sql test wording", + sql=( + "WITH __source__raw_orders AS (SELECT 1 AS order_id), " + "__expected__ AS (SELECT 1 AS order_id) SELECT 1" + ), + expected_message=( + "SQL test 'tests/unit/orders.sql' must use __expected__ to identify a target" + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_malformed_sql_test_ctes_when_scanning_expected_models_then_names_the_sql_test( + test_case: CteScannerMessageTestCase, +) -> None: + with pytest.raises(CompileInputError) as error_info: + _ = extract_sql_test_expected_model_names( + sql=test_case.sql, file_label="tests/unit/orders.sql" + ) + + assert str(error_info.value) == test_case.expected_message diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py index ef91b1d8cd..b57da39078 100644 --- a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/_test_types.py @@ -828,3 +828,21 @@ class LoadTargetRetentionPoliciesErrorTestCase: description: str target_lines: tuple[str, ...] expected_error_fragment: str + + +@dataclass(frozen=True) +class DiscoverFileFaultsTestCase: + description: str + discovery_name: str + unreadable_files: tuple[str, ...] + expected_fault_paths: tuple[Path, ...] + expected_error_type: type[Exception] = UnicodeDecodeError + + +@dataclass(frozen=True) +class DiscoverModuleImportFailureTestCase: + description: str + discovery_name: str + broken_file: str + expected_error_type: type[Exception] + expected_error_fragment: str diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py index ac65e0bb96..55ce12db8c 100644 --- a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/helpers.py @@ -231,3 +231,13 @@ def discovered_test_case_values( for test_case in cases: case_values.append(tuple(value.value for _name, value in test_case.values)) return tuple(case_values) + + +def write_unreadable_files(*, project_dir: Path, relative_paths: tuple[str, ...]) -> None: + """Write files whose bytes are not valid UTF-8 under project_dir.""" + + relative_path: str + for relative_path in relative_paths: + file_path: Path = project_dir / relative_path + file_path.parent.mkdir(parents=True, exist_ok=True) + _ = file_path.write_bytes(b"\xff\xfe\xfa") diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py new file mode 100644 index 0000000000..b82e41149b --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_filesystem_file_faults.py @@ -0,0 +1,178 @@ +"""Tests for per-file fault isolation and module import failures in filesystem discovery.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +import pytest + +from sqlbuild.compiler.discovery._helpers.filesystem.command_output_sinks import ( + discover_command_output_sink_declarations, +) +from sqlbuild.compiler.discovery._helpers.filesystem.core import ( + discover_audit_files, + discover_constant_files, + discover_enum_files, + discover_event_exporter_declarations, + discover_materialization_files, + discover_provider_classes, + discover_scenario_files, + discover_source_files, + discover_sql_function_files, + discover_sql_hook_files, +) +from sqlbuild.compiler.discovery.exceptions import ( + EventExporterDiscoveryError, + ProviderDiscoveryError, + PythonNodeDiscoveryError, +) +from sqlbuild.compiler.discovery.models import DiscoveryFileFault +from tests.unit.src.sqlbuild.compiler.discovery._helpers._test_types import ( + DiscoverFileFaultsTestCase, + DiscoverModuleImportFailureTestCase, +) +from tests.unit.src.sqlbuild.compiler.discovery._helpers.helpers import write_unreadable_files + +_FILE_DISCOVERIES: dict[str, Callable[..., tuple[object, ...]]] = { + "audits": discover_audit_files, + "constants": discover_constant_files, + "enums": discover_enum_files, + "scenarios": discover_scenario_files, + "sources": discover_source_files, + "sql_functions": discover_sql_function_files, + "sql_hooks": discover_sql_hook_files, +} +_MODULE_DISCOVERIES: dict[str, Callable[..., object]] = { + "command_output_sinks": discover_command_output_sink_declarations, + "event_exporters": discover_event_exporter_declarations, + "materializations": discover_materialization_files, + "providers": discover_provider_classes, +} + + +@pytest.mark.parametrize( + "test_case", + ( + DiscoverFileFaultsTestCase( + description="audits fault every nested SQL file", + discovery_name="audits", + unreadable_files=("audits/orders.sql", "audits/nested/_customers.sql"), + expected_fault_paths=(Path("audits/nested/_customers.sql"), Path("audits/orders.sql")), + ), + DiscoverFileFaultsTestCase( + description="constants fault declaration files", + discovery_name="constants", + unreadable_files=("constants/order_limits.sql",), + expected_fault_paths=(Path("constants/order_limits.sql"),), + ), + DiscoverFileFaultsTestCase( + description="enums fault declaration files", + discovery_name="enums", + unreadable_files=("enums/order_status.sql",), + expected_fault_paths=(Path("enums/order_status.sql"),), + ), + DiscoverFileFaultsTestCase( + description="scenarios fault nested SQL files", + discovery_name="scenarios", + unreadable_files=("tests/scenarios/orders/_refunds.sql",), + expected_fault_paths=(Path("tests/scenarios/orders/_refunds.sql"),), + ), + DiscoverFileFaultsTestCase( + description="sources fault only top-level YAML files", + discovery_name="sources", + unreadable_files=( + "sources/orders.yml", + "sources/customers.yaml", + "sources/notes.txt", + "sources/nested/products.yml", + ), + expected_fault_paths=(Path("sources/customers.yaml"), Path("sources/orders.yml")), + ), + DiscoverFileFaultsTestCase( + description="sql functions fault underscore-prefixed files", + discovery_name="sql_functions", + unreadable_files=("functions/sql/_order_total.sql", "functions/sql/tax.sql"), + expected_fault_paths=( + Path("functions/sql/_order_total.sql"), + Path("functions/sql/tax.sql"), + ), + ), + DiscoverFileFaultsTestCase( + description="sql hooks skip underscore-prefixed files", + discovery_name="sql_hooks", + unreadable_files=("hooks/sql/_shared.sql", "hooks/sql/grants/grant_access.sql"), + expected_fault_paths=(Path("hooks/sql/grants/grant_access.sql"),), + ), + ), + ids=lambda case: case.description, +) +def test_given_unreadable_files_when_discovering_then_faults_each_file_or_raises( + tmp_path: Path, + test_case: DiscoverFileFaultsTestCase, +) -> None: + write_unreadable_files(project_dir=tmp_path, relative_paths=test_case.unreadable_files) + discover: Callable[..., tuple[object, ...]] = _FILE_DISCOVERIES[test_case.discovery_name] + faults: list[DiscoveryFileFault] = [] + + discovered: tuple[object, ...] = discover(project_dir=tmp_path, on_fault=faults.append) + + assert discovered == () + assert tuple(fault.path for fault in faults) == test_case.expected_fault_paths + assert all("utf-8" in fault.message for fault in faults) + with pytest.raises(test_case.expected_error_type): + _ = discover(project_dir=tmp_path) + + +@pytest.mark.parametrize( + "test_case", + ( + DiscoverModuleImportFailureTestCase( + description="command output sink import failure identifies the sink file", + discovery_name="command_output_sinks", + broken_file="sinks/broken_orders.py", + expected_error_type=EventExporterDiscoveryError, + expected_error_fragment="Failed to import sink file sinks/broken_orders.py: ", + ), + DiscoverModuleImportFailureTestCase( + description="lifecycle sink import failure identifies the sink file", + discovery_name="event_exporters", + broken_file="sinks/broken_orders.py", + expected_error_type=EventExporterDiscoveryError, + expected_error_fragment="Failed to import sink file sinks/broken_orders.py: ", + ), + DiscoverModuleImportFailureTestCase( + description="materialization import failure identifies the materialization file", + discovery_name="materializations", + broken_file="materializations/broken_orders.py", + expected_error_type=PythonNodeDiscoveryError, + expected_error_fragment=( + "Failed to import materialization file materializations/broken_orders.py: " + ), + ), + DiscoverModuleImportFailureTestCase( + description="provider import failure identifies the provider file", + discovery_name="providers", + broken_file="providers/broken_orders.py", + expected_error_type=ProviderDiscoveryError, + expected_error_fragment="Failed to import provider file providers/broken_orders.py: ", + ), + ), + ids=lambda case: case.description, +) +def test_given_module_import_failure_when_discovering_then_raises_kind_specific_error( + tmp_path: Path, + test_case: DiscoverModuleImportFailureTestCase, +) -> None: + file_path: Path = tmp_path / test_case.broken_file + file_path.parent.mkdir(parents=True) + _ = file_path.write_text("raise RuntimeError('orders module failed')\n", encoding="utf-8") + + with pytest.raises(test_case.expected_error_type) as error_info: + _ = _MODULE_DISCOVERIES[test_case.discovery_name](project_dir=tmp_path) + + assert str(error_info.value) == f"{test_case.expected_error_fragment}orders module failed" + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py index 2adc2a2267..fec22682e2 100644 --- a/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py +++ b/tests/unit/src/sqlbuild/compiler/discovery/_helpers/test_yml_sources.py @@ -1119,6 +1119,15 @@ def test_given_dlt_sources_yaml_when_parsing_then_expands_managed_sources( """, expected_error_fragment="source audits must be a list", ), + ParseSourcesYamlErrorTestCase( + description="raises when source audits is null", + contents=""" + sources: + - name: raw_orders + audits: + """, + expected_error_fragment="source audits must be a list", + ), ParseSourcesYamlErrorTestCase( description="raises when source columns is not a list", contents=""" diff --git a/tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py new file mode 100644 index 0000000000..cd0f8f07c7 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/discovery/main/_test_types.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SerializeHookEntriesTestCase: + description: str + value: object + sql_fields: tuple[str, ...] + python_hook_fields: dict[str, dict[str, object]] + expected_hooks: list[dict[str, object]] diff --git a/tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py b/tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py new file mode 100644 index 0000000000..2df0c8979e --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/discovery/main/test_serialize_hook_entries.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from sqlbuild.compiler.discovery.constants import SQL_HOOK_OUTPUT_FIELDS +from sqlbuild.compiler.discovery.main.serialize_hook_entries import serialize_hook_entries +from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from tests.unit.src.sqlbuild.compiler.discovery.main._test_types import ( + SerializeHookEntriesTestCase, +) + +_NAMED_SQL_HOOK: SqlHookEntry = SqlHookEntry( + statement="INSERT INTO order_audit SELECT 'orders'", + name="record_order_access", + relative_path=Path("hooks/sql/record_order_access.sql"), + definition_sql="INSERT INTO order_audit SELECT {{ table_name }}", + kwargs={"table_name": "'orders'"}, + description="Record order table access", +) + + +@pytest.mark.parametrize( + "test_case", + [ + SerializeHookEntriesTestCase( + description="serializes every output field of a named SQL hook in order", + value=[_NAMED_SQL_HOOK], + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={}, + expected_hooks=[ + { + "type": "sql", + "statement": "INSERT INTO order_audit SELECT 'orders'", + "name": "record_order_access", + "relative_path": "hooks/sql/record_order_access.sql", + "definition_sql": "INSERT INTO order_audit SELECT {{ table_name }}", + "kwargs": {"table_name": "'orders'"}, + "description": "Record order table access", + } + ], + ), + SerializeHookEntriesTestCase( + description="omits unset fields of an inline SQL hook", + value=(SqlHookEntry(statement="ANALYZE orders"),), + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={}, + expected_hooks=[{"type": "sql", "statement": "ANALYZE orders"}], + ), + SerializeHookEntriesTestCase( + description="keeps only allowlisted SQL hook fields", + value=[_NAMED_SQL_HOOK], + sql_fields=("statement", "kwargs"), + python_hook_fields={}, + expected_hooks=[ + { + "type": "sql", + "statement": "INSERT INTO order_audit SELECT 'orders'", + "kwargs": {"table_name": "'orders'"}, + } + ], + ), + SerializeHookEntriesTestCase( + description="adds per-name fields to Python hooks and keeps hook order", + value=[ + PythonHookEntry(name="notify_orders", kwargs={"channel": "orders"}), + SqlHookEntry(statement="ANALYZE orders"), + PythonHookEntry(name="refresh_customers", kwargs={}), + ], + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={"notify_orders": {"version_hash": "notify-hash"}}, + expected_hooks=[ + { + "type": "python", + "name": "notify_orders", + "kwargs": {"channel": "orders"}, + "version_hash": "notify-hash", + }, + {"type": "sql", "statement": "ANALYZE orders"}, + {"type": "python", "name": "refresh_customers", "kwargs": {}}, + ], + ), + SerializeHookEntriesTestCase( + description="returns no hooks for a non-sequence value", + value="ANALYZE orders", + sql_fields=SQL_HOOK_OUTPUT_FIELDS, + python_hook_fields={}, + expected_hooks=[], + ), + ], + ids=lambda case: case.description, +) +def test_given_hook_entries_when_serializing_then_returns_allowlisted_payloads( + test_case: SerializeHookEntriesTestCase, +) -> None: + hooks: list[dict[str, object]] = serialize_hook_entries( + value=test_case.value, + sql_fields=test_case.sql_fields, + python_hook_fields=test_case.python_hook_fields, + ) + + assert hooks == test_case.expected_hooks + assert [list(hook) for hook in hooks] == [list(hook) for hook in test_case.expected_hooks] diff --git a/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py index 1f38e264b7..b613351a59 100644 --- a/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/graph/main/_test_types.py @@ -17,6 +17,16 @@ class TransitiveClosureTestCase: expected_nodes: frozenset[str] +@dataclass(frozen=True) +class TransitiveClosureManyTestCase: + description: str + edges: dict[str, tuple[str, ...]] + starts: tuple[str, ...] + include_starts: bool + expected_nodes: frozenset[str] + max_depth: int | None = None + + @dataclass(frozen=True) class PathNodesTestCase: description: str diff --git a/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py b/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py index f2b2db9a2c..232e5c32f5 100644 --- a/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py +++ b/tests/unit/src/sqlbuild/compiler/graph/main/test_graph_algorithms.py @@ -5,12 +5,32 @@ from sqlbuild.compiler.graph.main.invert_edges import invert_edges from sqlbuild.compiler.graph.main.path_nodes import path_nodes from sqlbuild.compiler.graph.main.transitive_closure import transitive_closure +from sqlbuild.compiler.graph.main.transitive_closure_many import transitive_closure_many from tests.unit.src.sqlbuild.compiler.graph.main._test_types import ( InvertEdgesTestCase, PathNodesTestCase, + TransitiveClosureManyTestCase, TransitiveClosureTestCase, ) +_DIAMOND_UPSTREAM: dict[str, tuple[str, ...]] = { + "orders_report": ("orders_left", "orders_right"), + "orders_left": ("stg_orders",), + "orders_right": ("stg_orders",), + "stg_orders": (), + "inventory_report": ("stg_inventory",), + "stg_inventory": (), +} + +_SHORTCUT_DOWNSTREAM: dict[str, tuple[str, ...]] = { + "orders": ("orders_left", "orders_right"), + "orders_right": ("orders_mid",), + "orders_mid": ("orders_joined",), + "orders_left": ("orders_joined",), + "orders_joined": ("orders_report",), + "orders_report": (), +} + @pytest.mark.parametrize( "test_case", @@ -51,6 +71,15 @@ def test_given_edges_when_inverting_then_returns_expected_edges( max_depth=1, expected_nodes=frozenset({"b"}), ), + TransitiveClosureTestCase( + description="bounds depth by the shortest route", + edges=_SHORTCUT_DOWNSTREAM, + start="orders", + max_depth=3, + expected_nodes=frozenset( + {"orders_left", "orders_right", "orders_mid", "orders_joined", "orders_report"} + ), + ), ], ids=lambda case: case.description, ) @@ -66,6 +95,110 @@ def test_given_edges_when_finding_transitive_closure_then_returns_expected_nodes assert result == test_case.expected_nodes +@pytest.mark.parametrize( + "test_case", + [ + TransitiveClosureManyTestCase( + description="excludes a single start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=False, + expected_nodes=frozenset({"orders_left", "orders_right", "stg_orders"}), + ), + TransitiveClosureManyTestCase( + description="includes a single start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=True, + expected_nodes=frozenset( + {"orders_report", "orders_left", "orders_right", "stg_orders"} + ), + ), + TransitiveClosureManyTestCase( + description="unions disconnected starts", + edges=_DIAMOND_UPSTREAM, + starts=("orders_left", "inventory_report"), + include_starts=False, + expected_nodes=frozenset({"stg_orders", "stg_inventory"}), + ), + TransitiveClosureManyTestCase( + description="keeps an excluded start reachable from another start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report", "orders_left"), + include_starts=False, + expected_nodes=frozenset({"orders_left", "orders_right", "stg_orders"}), + ), + TransitiveClosureManyTestCase( + description="includes starts missing from the edge mapping", + edges=_DIAMOND_UPSTREAM, + starts=("customers",), + include_starts=True, + expected_nodes=frozenset({"customers"}), + ), + TransitiveClosureManyTestCase( + description="returns nothing without starts", + edges=_DIAMOND_UPSTREAM, + starts=(), + include_starts=True, + expected_nodes=frozenset(), + ), + TransitiveClosureManyTestCase( + description="terminates on cycles", + edges={"orders": ("customers",), "customers": ("orders",)}, + starts=("orders",), + include_starts=False, + expected_nodes=frozenset({"orders", "customers"}), + ), + TransitiveClosureManyTestCase( + description="returns nothing at zero depth without starts", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=False, + max_depth=0, + expected_nodes=frozenset(), + ), + TransitiveClosureManyTestCase( + description="returns only starts at zero depth with starts", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report",), + include_starts=True, + max_depth=0, + expected_nodes=frozenset({"orders_report"}), + ), + TransitiveClosureManyTestCase( + description="bounds depth from every start", + edges=_DIAMOND_UPSTREAM, + starts=("orders_report", "inventory_report"), + include_starts=False, + max_depth=1, + expected_nodes=frozenset({"orders_left", "orders_right", "stg_inventory"}), + ), + TransitiveClosureManyTestCase( + description="bounds multi-start depth by the shortest route", + edges=_SHORTCUT_DOWNSTREAM, + starts=("orders", "orders_right"), + include_starts=False, + max_depth=3, + expected_nodes=frozenset( + {"orders_left", "orders_right", "orders_mid", "orders_joined", "orders_report"} + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_edges_when_finding_multi_start_closure_then_returns_expected_nodes( + test_case: TransitiveClosureManyTestCase, +) -> None: + result: frozenset[str] = transitive_closure_many( + starts=test_case.starts, + edges=test_case.edges, + include_starts=test_case.include_starts, + max_depth=test_case.max_depth, + ) + + assert result == test_case.expected_nodes + + @pytest.mark.parametrize( "test_case", [ diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py index 64a01b056a..cea5f69ec4 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/_test_types.py @@ -65,3 +65,12 @@ class ForeignSnapshotRelationTestCase: desired_days: int expected_direction: str expected_inspections: int + + +@dataclass(frozen=True) +class PlannerScopesTestCase: + description: str + model_deps: dict[str, tuple[str, ...]] + select: tuple[str, ...] + expected_selected_names: frozenset[str] + expected_stale_warning_names: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py new file mode 100644 index 0000000000..e6f37aefd9 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/planning/test_scopes.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.compiler.planner._helpers.planning.scopes import resolve_planner_scopes +from sqlbuild.compiler.planner.models import ( + PlannerPolicies, + PlannerScopeResolution, + PlannerSelection, +) +from tests.unit.src.sqlbuild.compiler.planner._helpers.helpers import build_test_project +from tests.unit.src.sqlbuild.compiler.planner._helpers.planning._test_types import ( + PlannerScopesTestCase, +) + +_MODEL_DEPS: dict[str, tuple[str, ...]] = { + "stg_orders": (), + "stg_customers": (), + "orders_enriched": ("stg_orders", "stg_customers"), + "orders_summary": ("orders_enriched",), + "orders_report": ("orders_summary",), + "inventory": (), +} + + +@pytest.mark.parametrize( + "test_case", + [ + PlannerScopesTestCase( + description="stale warning scope is the selection plus all upstreams", + model_deps=_MODEL_DEPS, + select=("orders_summary",), + expected_selected_names=frozenset({"orders_summary"}), + expected_stale_warning_names=( + "orders_enriched", + "orders_summary", + "stg_customers", + "stg_orders", + ), + ), + PlannerScopesTestCase( + description="stale warning scope unions upstreams of every selected model", + model_deps=_MODEL_DEPS, + select=("orders_enriched", "inventory"), + expected_selected_names=frozenset({"orders_enriched", "inventory"}), + expected_stale_warning_names=( + "inventory", + "orders_enriched", + "stg_customers", + "stg_orders", + ), + ), + PlannerScopesTestCase( + description="stale warning scope of a root model is the model itself", + model_deps=_MODEL_DEPS, + select=("stg_orders",), + expected_selected_names=frozenset({"stg_orders"}), + expected_stale_warning_names=("stg_orders",), + ), + ], + ids=lambda case: case.description, +) +def test_given_selection_when_resolving_scopes_then_stale_warning_scope_is_upstream_closure( + test_case: PlannerScopesTestCase, +) -> None: + project: CompiledProject = build_test_project(model_deps=test_case.model_deps) + + resolution: PlannerScopeResolution = resolve_planner_scopes( + project=project, + selection=PlannerSelection(select=test_case.select), + policies=PlannerPolicies(), + ) + + selected_names: frozenset[str] = frozenset( + key.name for key in resolution.selected_scope.selected_keys + ) + stale_scope_names: tuple[str, ...] = tuple(sorted(resolution.stale_warning_scope.all_keys)) + stale_order_names: tuple[str, ...] = tuple( + key.name for key in resolution.stale_warning_scope.execution_order + ) + assert selected_names == test_case.expected_selected_names + assert resolution.stale_warning_scope.selected_keys == resolution.selected_scope.selected_keys + assert stale_scope_names == test_case.expected_stale_warning_names + assert tuple(sorted(stale_order_names)) == test_case.expected_stale_warning_names + assert resolution.inspection_scope is resolution.selected_scope diff --git a/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py index a7acd351aa..53a0099164 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/main/identity/_test_types.py @@ -30,3 +30,19 @@ class CursorRoleIdentityTestCase: original_config: dict[str, object] changed_config: dict[str, object] expected_equal: bool + + +@dataclass(frozen=True) +class HookIdentityTestCase: + description: str + original_hooks: dict[str, object] + changed_hooks: dict[str, object] + expected_changed: bool + + +@dataclass(frozen=True) +class HookIdentityPayloadTestCase: + description: str + hooks: dict[str, object] + expected_fragments: tuple[str, ...] + forbidden_fragments: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py b/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py index 69a18b815c..65fecd6cca 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/planner/main/identity/helpers.py @@ -14,6 +14,9 @@ from sqlbuild.compiler.planner._helpers.identity.direct import ( build_direct_model_version_identities, ) +from sqlbuild.compiler.planner.main.identity.version_identity_model_metadata import ( + build_model_version_identity_metadata_json, +) from sqlbuild.compiler.planner.models import DirectModelVersionIdentities, PlannerScope from tests.unit.src.sqlbuild.compiler.planner._helpers.helpers import build_compiled_function @@ -88,3 +91,18 @@ def _build_model( qualified_name=key.name, ), ) + + +def build_hooked_model_identity_json(*, hook_config: dict[str, object]) -> str: + model: CompiledModel = replace( + _build_model( + key=CompiledObjectKey(resource_type=CompiledResourceType.MODEL, name="orders"), + query_sql="SELECT 1 AS order_id", + deps=(), + ), + config=CompileModelConfig(values=hook_config), + ) + return build_model_version_identity_metadata_json( + model=model, + hook_version_hashes={"notify_orders": "notify-hash-v1"}, + ) diff --git a/tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py b/tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py new file mode 100644 index 0000000000..b84f7ff474 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/main/identity/test_version_identity_hooks.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +import pytest + +from sqlbuild.compiler.discovery.models import PythonHookEntry, SqlHookEntry +from tests.unit.src.sqlbuild.compiler.planner.main.identity._test_types import ( + HookIdentityPayloadTestCase, + HookIdentityTestCase, +) +from tests.unit.src.sqlbuild.compiler.planner.main.identity.helpers import ( + build_hooked_model_identity_json, +) + +_RECORD_ACCESS: SqlHookEntry = SqlHookEntry( + statement="INSERT INTO order_audit SELECT 'orders'", + name="record_order_access", + relative_path=Path("hooks/sql/record_order_access.sql"), + definition_sql="INSERT INTO order_audit SELECT {{ table_name }}", + kwargs={"table_name": "'orders'"}, + description="Record order table access", +) +_ANALYZE: SqlHookEntry = SqlHookEntry(statement="ANALYZE orders") +_NOTIFY: PythonHookEntry = PythonHookEntry(name="notify_orders", kwargs={"channel": "orders"}) + + +@pytest.mark.parametrize( + "test_case", + [ + HookIdentityTestCase( + description="ignores a hook description change", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={"pre_hooks": [replace(_RECORD_ACCESS, description="Audit order reads")]}, + expected_changed=False, + ), + HookIdentityTestCase( + description="ignores a hook file location change", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [ + replace( + _RECORD_ACCESS, + relative_path=Path("hooks/sql/audit/record_order_access.sql"), + ) + ] + }, + expected_changed=False, + ), + HookIdentityTestCase( + description="changes when hook SQL changes", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [ + replace(_RECORD_ACCESS, statement="INSERT INTO order_audit SELECT 'all'") + ] + }, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook definition changes", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [ + replace( + _RECORD_ACCESS, + definition_sql="INSERT INTO order_audit SELECT {{ table_name }} -- v2", + ) + ] + }, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook arguments change", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={ + "pre_hooks": [replace(_RECORD_ACCESS, kwargs={"table_name": "'customers'"})] + }, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook name changes", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={"pre_hooks": [replace(_RECORD_ACCESS, name="record_access")]}, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when a hook moves between phases", + original_hooks={"pre_hooks": [_RECORD_ACCESS]}, + changed_hooks={"post_hooks": [_RECORD_ACCESS]}, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when hook order changes", + original_hooks={"post_hooks": [_RECORD_ACCESS, _ANALYZE]}, + changed_hooks={"post_hooks": [_ANALYZE, _RECORD_ACCESS]}, + expected_changed=True, + ), + HookIdentityTestCase( + description="changes when Python hook arguments change", + original_hooks={"post_hooks": [_NOTIFY]}, + changed_hooks={"post_hooks": [replace(_NOTIFY, kwargs={"channel": "customers"})]}, + expected_changed=True, + ), + ], + ids=lambda case: case.description, +) +def test_given_hook_change_when_building_model_identity_then_tracks_only_execution_fields( + test_case: HookIdentityTestCase, +) -> None: + original_json: str = build_hooked_model_identity_json(hook_config=test_case.original_hooks) + changed_json: str = build_hooked_model_identity_json(hook_config=test_case.changed_hooks) + + assert (original_json != changed_json) is test_case.expected_changed + + +@pytest.mark.parametrize( + "test_case", + [ + HookIdentityPayloadTestCase( + description="keeps execution fields and omits documentation fields", + hooks={"pre_hooks": [_RECORD_ACCESS], "post_hooks": [_NOTIFY]}, + expected_fragments=( + '"definition_sql":"INSERT INTO order_audit SELECT {{ table_name }}"', + '"name":"record_order_access"', + '"version_hash":"notify-hash-v1"', + ), + forbidden_fragments=( + "Record order table access", + "hooks/sql/record_order_access.sql", + '"description"', + '"relative_path"', + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_documented_hooks_when_building_model_identity_then_omits_doc_fields( + test_case: HookIdentityPayloadTestCase, +) -> None: + identity_json: str = build_hooked_model_identity_json(hook_config=test_case.hooks) + + assert all(fragment in identity_json for fragment in test_case.expected_fragments) + assert not any(fragment in identity_json for fragment in test_case.forbidden_fragments) diff --git a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py index 128296fa5b..968741bc70 100644 --- a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/_test_types.py @@ -126,3 +126,9 @@ class PythonNodeIdentityChangeTestCase: expected_version_hash_changed: bool before_decorator_config: dict[str, object] | None = None after_decorator_config: dict[str, object] | None = None + + +@dataclass(frozen=True) +class PythonNodeIdentityConfigTestCase: + description: str + expected_decorator_config_keys_by_selector: dict[str, tuple[str, ...]] diff --git a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py index 12056c1e06..b6eea07a48 100644 --- a/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py +++ b/tests/unit/src/sqlbuild/compiler/python_nodes/_helpers/test_inventory.py @@ -2,8 +2,10 @@ from __future__ import annotations +import json from collections import defaultdict from pathlib import Path +from typing import cast import pytest @@ -15,14 +17,23 @@ DiscoveredProvider, DiscoveredTaskFunction, ) -from sqlbuild.compiler.python_nodes._helpers.inventory import build_python_node_graph -from sqlbuild.compiler.python_nodes.models import DiscoveredPythonNode, PythonNodeGraph +from sqlbuild.compiler.discovery.types import LoaderConnectionMode +from sqlbuild.compiler.python_nodes._helpers.inventory import ( + build_python_node_graph, + build_python_nodes, +) +from sqlbuild.compiler.python_nodes.models import ( + DiscoveredPythonNode, + PythonNodeGraph, + PythonNodeIdentity, +) from sqlbuild.compiler.python_nodes.types import PythonNodeKind from sqlbuild.python_nodes.models import RetryPolicy from sqlbuild.python_nodes.types import PythonCheckSeverity from sqlbuild.spec.contracts.models import LocalConfig, ProjectConfig, SourceColumnEntry from tests.unit.src.sqlbuild.compiler.python_nodes._helpers._test_types import ( PythonNodeGraphInventoryTestCase, + PythonNodeIdentityConfigTestCase, ) from tests.unit.src.sqlbuild.compiler.python_nodes._helpers.helpers import ( SlackProvider, @@ -217,3 +228,113 @@ def test_given_discovered_python_functions_when_building_graph_then_indexes_node for edge in graph.dependency_edges: edges_by_downstream[edge.downstream_name].append((edge.upstream_name, edge.downstream_name)) assert tuple(edges_by_downstream["notify_orders"]) == () + + +@pytest.mark.parametrize( + "test_case", + [ + PythonNodeIdentityConfigTestCase( + description="identity decorator config keeps each kind's authored field set", + expected_decorator_config_keys_by_selector={ + "loader:load_events": ( + "contract", + "cursor_column", + "destination", + "unique_key", + "write_strategy", + ), + "task:prepare_orders": ("description", "group", "meta", "tags"), + "asset:export_orders": ( + "column_lineage", + "columns", + "description", + "group", + "meta", + "tags", + ), + "check:check_orders_export": ( + "description", + "group", + "meta", + "severity", + "tags", + ), + }, + ) + ], + ids=lambda case: case.description, +) +def test_given_fully_configured_python_nodes_when_building_then_identity_config_is_per_kind( + test_case: PythonNodeIdentityConfigTestCase, +) -> None: + retry: RetryPolicy = RetryPolicy(max_attempts=2, retry_on=RuntimeError, jitter=False) + columns: tuple[SourceColumnEntry, ...] = (SourceColumnEntry(name="order_id", type="INTEGER"),) + + nodes: tuple[DiscoveredPythonNode, ...] = build_python_nodes( + loader_functions=( + DiscoveredLoaderFunction( + file_path=Path("/project/loaders/events.py"), + relative_path=Path("loaders/events.py"), + name="load_events", + function=load_events, + destination="raw.events", + cursor_column="updated_at", + unique_key=("event_id",), + columns=columns, + contract="enforced", + connection_mode=LoaderConnectionMode.EXTERNAL, + ), + ), + task_functions=( + DiscoveredTaskFunction( + file_path=Path("/project/tasks/orders.py"), + relative_path=Path("tasks/orders.py"), + name="prepare_orders", + function=prepare_orders, + tags=("orders",), + group="ingestion", + description="Prepare order export inputs.", + meta={"owner": "data-eng"}, + retry=retry, + ), + ), + asset_functions=( + DiscoveredAssetFunction( + file_path=Path("/project/assets/orders.py"), + relative_path=Path("assets/orders.py"), + name="export_orders", + function=export_orders, + tags=("orders",), + group="exports", + description="Export orders.", + meta={"owner": "data-eng"}, + columns=columns, + retry=retry, + ), + ), + check_functions=( + DiscoveredCheckFunction( + file_path=Path("/project/checks/orders.py"), + relative_path=Path("checks/orders.py"), + name="check_orders_export", + function=check_orders_export, + depends_on=(export_orders,), + severity=PythonCheckSeverity.WARN, + tags=("orders",), + group="exports", + description="Check exported orders.", + meta={"owner": "data-eng"}, + ), + ), + ) + + assert { + f"{node.kind.value}:{node.name}": tuple( + sorted( + json.loads(cast(PythonNodeIdentity, node.identity).definition_json)[ + "decorator_config" + ] + ) + ) + for node in nodes + } == test_case.expected_decorator_config_keys_by_selector diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py index fcc46e88d0..4fa311faf2 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/_test_types.py @@ -20,3 +20,12 @@ class SkipQuotedTextErrorTestCase: sql: str context: str expected_error: str + + +@dataclass(frozen=True) +class IterCodePositionsTestCase: + """Code-position scan case.""" + + description: str + sql: str + expected_positions: tuple[tuple[int, int], ...] diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py index cdd6c6c355..3245e9b7aa 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/_helpers/test_scanning.py @@ -3,8 +3,12 @@ import pytest from sqlbuild.compiler.compile.exceptions import CompileInputError -from sqlbuild.compiler.sql_analysis._helpers.scanning import skip_quoted_text_impl +from sqlbuild.compiler.sql_analysis._helpers.scanning import ( + iter_code_positions_impl, + skip_quoted_text_impl, +) from tests.unit.src.sqlbuild.compiler.sql_analysis._helpers._test_types import ( + IterCodePositionsTestCase, SkipQuotedTextErrorTestCase, SkipQuotedTextSuccessTestCase, ) @@ -68,3 +72,30 @@ def test_given_unclosed_quote_when_skipping_then_raises_contextual_error( ) -> None: with pytest.raises(CompileInputError, match=test_case.expected_error): skip_quoted_text_impl(sql=test_case.sql, start=0, context=test_case.context) + + +@pytest.mark.parametrize( + "test_case", + ( + IterCodePositionsTestCase( + description="parentheses change depth and are not yielded", + sql="a(b)c", + expected_positions=((0, 0), (2, 1), (4, 0)), + ), + IterCodePositionsTestCase( + description="quotes and comments are skipped", + sql="a'(b'`c`--d\n/*e*/f", + expected_positions=((0, 0), (17, 0)), + ), + IterCodePositionsTestCase( + description="unbalanced close parenthesis goes below zero depth", + sql=")a", + expected_positions=((1, -1),), + ), + ), + ids=lambda case: case.description, +) +def test_given_sql_when_iterating_code_positions_then_yields_code_offsets_with_depth( + test_case: IterCodePositionsTestCase, +) -> None: + assert tuple(iter_code_positions_impl(sql=test_case.sql)) == test_case.expected_positions diff --git a/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py index efbd3a3413..aaedf2c0f3 100644 --- a/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/build/_helpers/_test_types.py @@ -191,3 +191,12 @@ class BatchedRetentionReconciliationTestCase: expected_batch_request_ids: tuple[tuple[str, ...], ...] expected_single_inspections: tuple[str, ...] expected_statements: tuple[str, ...] + + +@dataclass(frozen=True) +class DownstreamBlockedKeysTestCase: + description: str + failed_name: str + downstream_deps: dict[str, tuple[str, ...]] + selected_names: frozenset[str] + expected_names: frozenset[str] diff --git a/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py b/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py index 13286a522e..da4a54d744 100644 --- a/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/executor/build/_helpers/helpers.py @@ -262,3 +262,18 @@ def batched_retention_request_ids(*, adapter: Mock) -> tuple[tuple[str, ...], .. for batch in adapter.inspect_retentions.call_args_list: batches.append(tuple(request.request_id for request in batch.kwargs["requests"])) return tuple(batches) + + +def model_key(name: str) -> CompiledObjectKey: + return CompiledObjectKey(resource_type=CompiledResourceType.MODEL, name=name) + + +def build_model_key_edges( + *, edges: dict[str, tuple[str, ...]] +) -> dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]]: + key_edges: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = {} + name: str + children: tuple[str, ...] + for name, children in edges.items(): + key_edges[model_key(name)] = tuple(model_key(child) for child in children) + return key_edges diff --git a/tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py b/tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py new file mode 100644 index 0000000000..e8bf977cb7 --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/build/_helpers/test_blocking.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.compile.models import CompiledObjectKey +from sqlbuild.executor.build._helpers.blocking import downstream_blocked_keys +from tests.unit.src.sqlbuild.executor.build._helpers._test_types import ( + DownstreamBlockedKeysTestCase, +) +from tests.unit.src.sqlbuild.executor.build._helpers.helpers import ( + build_model_key_edges, + model_key, +) + + +@pytest.mark.parametrize( + "test_case", + [ + DownstreamBlockedKeysTestCase( + description="blocks selected transitive downstreams but not the failed key", + failed_name="stg_orders", + downstream_deps={ + "stg_orders": ("orders_summary", "orders_audit"), + "orders_summary": ("orders_report",), + "orders_audit": (), + "orders_report": (), + "customers": ("orders_report",), + }, + selected_names=frozenset( + {"stg_orders", "orders_summary", "orders_audit", "orders_report", "customers"} + ), + expected_names=frozenset({"orders_summary", "orders_audit", "orders_report"}), + ), + DownstreamBlockedKeysTestCase( + description="walks through unselected keys to selected descendants", + failed_name="stg_orders", + downstream_deps={ + "stg_orders": ("orders_summary",), + "orders_summary": ("orders_report",), + "orders_report": (), + }, + selected_names=frozenset({"stg_orders", "orders_report"}), + expected_names=frozenset({"orders_report"}), + ), + DownstreamBlockedKeysTestCase( + description="blocks a diamond descendant once", + failed_name="stg_orders", + downstream_deps={ + "stg_orders": ("orders_left", "orders_right"), + "orders_left": ("orders_joined",), + "orders_right": ("orders_joined",), + "orders_joined": (), + }, + selected_names=frozenset({"orders_left", "orders_right", "orders_joined"}), + expected_names=frozenset({"orders_left", "orders_right", "orders_joined"}), + ), + DownstreamBlockedKeysTestCase( + description="returns nothing for a failed leaf", + failed_name="orders_report", + downstream_deps={"stg_orders": ("orders_report",), "orders_report": ()}, + selected_names=frozenset({"stg_orders", "orders_report"}), + expected_names=frozenset(), + ), + ], + ids=lambda case: case.description, +) +def test_given_failed_key_when_computing_blocked_keys_then_returns_selected_descendants( + test_case: DownstreamBlockedKeysTestCase, +) -> None: + downstream_deps: dict[CompiledObjectKey, tuple[CompiledObjectKey, ...]] = build_model_key_edges( + edges=test_case.downstream_deps + ) + + blocked: frozenset[CompiledObjectKey] = downstream_blocked_keys( + failed_key=model_key(test_case.failed_name), + downstream_deps=downstream_deps, + selected_keys=frozenset(model_key(name) for name in test_case.selected_names), + ) + + assert blocked == frozenset(model_key(name) for name in test_case.expected_names) diff --git a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py index a56870ac9a..6222f977da 100644 --- a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/_test_types.py @@ -91,6 +91,16 @@ class ScenarioFailureHelpTestCase: expected_help: str | None +@dataclass(frozen=True) +class ScenarioConnectionLifecycleTestCase: + """One warehouse scenario pipeline target-connection lifecycle case.""" + + description: str + expected_event_types: tuple[str, ...] + expected_operation_names: tuple[str, ...] + expected_connection_events: tuple[str, ...] + + @dataclass(frozen=True) class ScenarioTestPipelineTestCase: """One scenario test pipeline orchestration case.""" 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 81c1caf395..795c0b8b82 100644 --- a/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py +++ b/tests/unit/src/sqlbuild/executor/pipeline/_helpers/test_scenario.py @@ -17,11 +17,16 @@ SCENARIO_LOCAL_SNAPSHOT_MISSING, SCENARIO_LOCAL_SNAPSHOT_STALE, ) -from sqlbuild.executor.scenario.models import ScenarioRunResult, ScenarioSnapshotStateResult +from sqlbuild.executor.scenario.models import ( + ScenarioCaptureSettings, + ScenarioRunResult, + ScenarioSnapshotStateResult, +) from sqlbuild.executor.scenario.types import ScenarioLocalRunStatus, ScenarioSnapshotState from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.observability import EventDispatcher, LifecycleEvent, dispatcher_scope from tests.unit.src.sqlbuild.executor.pipeline._helpers._test_types import ( + ScenarioConnectionLifecycleTestCase, ScenarioFailureHelpTestCase, ScenarioLocalPipelineTestCase, ScenarioPlanningLifecycleTestCase, @@ -119,19 +124,108 @@ def complete_scenario( assert tuple(adapter.events) == test_case.expected_connection_events assert results[-1].error_message is not None assert test_case.expected_error_fragment in results[-1].error_message + resource_events: list[LifecycleEvent] = lifecycle_events[2:] assert tuple(event.event_type for event in lifecycle_events) == ( + "operation_started", + "operation_completed", "resource_attempt_started", "resource_attempt_completed", "resource_attempt_started", "resource_attempt_failed", ) - assert tuple(event.resource_id for event in lifecycle_events) == ( + assert tuple(event.resource_id for event in resource_events) == ( "sql_scenario:passing_scenario", "sql_scenario:passing_scenario", "sql_scenario:planning_failure", "sql_scenario:planning_failure", ) - assert tuple(event.run_id for event in lifecycle_events) == ("scenario-test-run",) * 4 + assert tuple(event.run_id for event in resource_events) == ("scenario-test-run",) * 4 + + +@pytest.mark.parametrize( + "test_case", + ( + ScenarioConnectionLifecycleTestCase( + description="warehouse test pipeline observes the target connection", + expected_event_types=("operation_started", "operation_completed"), + expected_operation_names=("scenario_target_connection",) * 2, + expected_connection_events=("connect:pipeline.duckdb", "close"), + ), + ), + ids=lambda case: case.description, +) +def test_given_warehouse_test_pipeline_when_connecting_then_publishes_connection_operation( + test_case: ScenarioConnectionLifecycleTestCase, +) -> None: + pipeline_result: CompilePipelineResult = build_scenario_pipeline_result(scenario_names=()) + adapter: ScenarioPipelineTestAdapter = ScenarioPipelineTestAdapter() + lifecycle_events: list[LifecycleEvent] = [] + dispatcher: EventDispatcher = EventDispatcher() + dispatcher.subscribe_lifecycle(subscriber=lifecycle_events.append, accepts_opaque=False) + + with dispatcher_scope(dispatcher): + _ = scenario_pipeline.run_scenario_test_pipeline( + pipeline_result=pipeline_result, + scenarios=(), + connection_config={"database": "pipeline.duckdb"}, + adapter=adapter, + project_name="waffle_shop", + retain=False, + ) + + assert tuple(event.event_type for event in lifecycle_events) == test_case.expected_event_types + assert ( + tuple(event.payload.get("operation_name") for event in lifecycle_events) + == test_case.expected_operation_names + ) + assert tuple(adapter.events) == test_case.expected_connection_events + + +@pytest.mark.parametrize( + "test_case", + ( + ScenarioConnectionLifecycleTestCase( + description="capture pipeline observes the target connection", + expected_event_types=("operation_started", "operation_completed"), + expected_operation_names=("scenario_target_connection",) * 2, + expected_connection_events=("connect:pipeline.duckdb", "close"), + ), + ), + ids=lambda case: case.description, +) +def test_given_capture_pipeline_when_connecting_then_publishes_connection_operation( + tmp_path: Path, + test_case: ScenarioConnectionLifecycleTestCase, +) -> None: + pipeline_result: CompilePipelineResult = build_scenario_pipeline_result(scenario_names=()) + adapter: ScenarioPipelineTestAdapter = ScenarioPipelineTestAdapter() + lifecycle_events: list[LifecycleEvent] = [] + dispatcher: EventDispatcher = EventDispatcher() + dispatcher.subscribe_lifecycle(subscriber=lifecycle_events.append, accepts_opaque=False) + + with dispatcher_scope(dispatcher): + _ = scenario_pipeline.run_scenario_capture_pipeline( + project_dir=tmp_path, + pipeline_result=pipeline_result, + scenarios=(), + connection_config={"database": "pipeline.duckdb"}, + adapter=adapter, + project_name="waffle_shop", + settings=ScenarioCaptureSettings( + captured_at="2026-01-01T00:00:00Z", + capture_adapter="duckdb", + capture_dialect="duckdb", + sqlbuild_version="0.0.0", + retain=False, + ), + ) + + assert tuple(event.event_type for event in lifecycle_events) == test_case.expected_event_types + assert ( + tuple(event.payload.get("operation_name") for event in lifecycle_events) + == test_case.expected_operation_names + ) + assert tuple(adapter.events) == test_case.expected_connection_events @pytest.mark.parametrize( diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py index 3da48e4ec9..be60a39820 100644 --- a/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/_test_types.py @@ -168,3 +168,20 @@ class ScenarioCaptureRunIdentityTestCase: expected_run_id: str expected_seed_resource_ids: tuple[str, ...] expected_operation_name: str + + +@dataclass(frozen=True) +class ScenarioCaptureMissingFixtureMessageTestCase: + description: str + fixture_error_message: str | None + expected_error_message: str + + +@dataclass(frozen=True) +class FirstFailureDetailsTestCase: + description: str + step_outcomes: tuple[tuple[ExecutionStatus, str | None, str | None, str | None], ...] + fallback_code: str | None + expected_error_code: str | None + expected_error_help: str | None + expected_error_message: str | None diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py index 8ec24e5a4f..352c0adcb5 100644 --- a/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_capture.py @@ -6,15 +6,19 @@ import pytest from sqlbuild.compiler.planner.models import ScenarioExecutionPlan, SeedPlanEntry +from sqlbuild.compiler.planner.types import ScenarioArtifactKind +from sqlbuild.executor.scenario._helpers.capture import core as capture_core from sqlbuild.executor.scenario._helpers.capture.core import execute_scenario_snapshot_capture_steps from sqlbuild.executor.scenario.models import ( ScenarioCaptureSettings, + ScenarioFixtureExecutionResult, ScenarioSnapshotCaptureRunResult, ) from sqlbuild.executor.scheduling.types import ExecutionStatus from sqlbuild.observability import EventDispatcher, LifecycleEvent, dispatcher_scope from tests.unit.src.sqlbuild.executor.scenario._helpers._test_types import ( ExecuteScenarioSnapshotCaptureStepsTestCase, + ScenarioCaptureMissingFixtureMessageTestCase, ScenarioCaptureRunIdentityTestCase, ) from tests.unit.src.sqlbuild.executor.scenario._helpers.helpers import ( @@ -205,3 +209,60 @@ def test_given_multiple_capture_seeds_when_executing_then_operation_and_resource test_case.expected_seed_resource_ids[1], test_case.expected_seed_resource_ids[1], ) + + +@pytest.mark.parametrize( + "test_case", + ( + ScenarioCaptureMissingFixtureMessageTestCase( + description="failed fixture without a message uses the capture fallback message", + fixture_error_message=None, + expected_error_message="scenario snapshot capture failed", + ), + ScenarioCaptureMissingFixtureMessageTestCase( + description="failed fixture with an empty message uses the capture fallback message", + fixture_error_message="", + expected_error_message="scenario snapshot capture failed", + ), + ScenarioCaptureMissingFixtureMessageTestCase( + description="failed fixture message is reported unchanged", + fixture_error_message="orders fixture failed", + expected_error_message="orders fixture failed", + ), + ), + ids=lambda case: case.description, +) +def test_given_failed_fixture_message_when_capturing_then_reports_message_or_fallback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + test_case: ScenarioCaptureMissingFixtureMessageTestCase, +) -> None: + failed_fixture: ScenarioFixtureExecutionResult = ScenarioFixtureExecutionResult( + scenario_name=SCENARIO_PLAN.name, + kind=ScenarioArtifactKind.SOURCE, + logical_name="raw__orders", + target_relation="scenario_schema.raw__orders", + status=ExecutionStatus.FAILED, + error_message=test_case.fixture_error_message, + ) + monkeypatch.setattr( + capture_core, "execute_scenario_fixtures", lambda **_kwargs: (failed_fixture,) + ) + + result: ScenarioSnapshotCaptureRunResult = execute_scenario_snapshot_capture_steps( + project_dir=tmp_path, + scenario_plan=SCENARIO_PLAN, + adapter=ScenarioSnapshotCaptureStepsTestAdapter(), + connection=object(), + run_id="capture-test-run", + settings=ScenarioCaptureSettings( + captured_at="2026-05-09T00:00:00Z", + capture_adapter="duckdb", + capture_dialect="duckdb", + sqlbuild_version="0.1.0", + retain=False, + ), + ) + + assert result.status == ExecutionStatus.FAILED + assert result.error_message == test_case.expected_error_message diff --git a/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py new file mode 100644 index 0000000000..540d94ae32 --- /dev/null +++ b/tests/unit/src/sqlbuild/executor/scenario/_helpers/test_failures.py @@ -0,0 +1,92 @@ +"""Tests for projecting failure details from scenario step results.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.planner.types import ScenarioArtifactKind +from sqlbuild.executor.scenario._helpers.lifecycle.failures import first_failure_details +from sqlbuild.executor.scenario.models import ( + ScenarioFailureDetails, + ScenarioFixtureExecutionResult, +) +from sqlbuild.executor.scheduling.types import ExecutionStatus +from tests.unit.src.sqlbuild.executor.scenario._helpers._test_types import ( + FirstFailureDetailsTestCase, +) + +_FAILED: ExecutionStatus = ExecutionStatus.FAILED +_SUCCESS: ExecutionStatus = ExecutionStatus.SUCCESS + + +@pytest.mark.parametrize( + "test_case", + ( + FirstFailureDetailsTestCase( + description="no failed step yields no message and the fallback code", + step_outcomes=((_SUCCESS, "S100", "help", "ignored"),), + fallback_code="S999", + expected_error_code="S999", + expected_error_help=None, + expected_error_message=None, + ), + FirstFailureDetailsTestCase( + description="first failed step supplies every detail", + step_outcomes=( + (_SUCCESS, "S100", "success help", "success message"), + (_FAILED, "S200", "orders help", "orders failed"), + (_FAILED, "S300", "customers help", "customers failed"), + ), + fallback_code=None, + expected_error_code="S200", + expected_error_help="orders help", + expected_error_message="orders failed", + ), + FirstFailureDetailsTestCase( + description="missing message uses the fallback message", + step_outcomes=((_FAILED, "S200", None, None),), + fallback_code=None, + expected_error_code="S200", + expected_error_help=None, + expected_error_message="step failed", + ), + FirstFailureDetailsTestCase( + description="code and help come from the first failed step that has them", + step_outcomes=( + (_FAILED, None, None, "orders failed"), + (_FAILED, "S300", "customers help", "customers failed"), + ), + fallback_code="S999", + expected_error_code="S300", + expected_error_help="customers help", + expected_error_message="orders failed", + ), + ), + ids=lambda case: case.description, +) +def test_given_step_results_when_projecting_failure_then_returns_expected_details( + test_case: FirstFailureDetailsTestCase, +) -> None: + results: tuple[ScenarioFixtureExecutionResult, ...] = tuple( + ScenarioFixtureExecutionResult( + scenario_name="orders_refund", + kind=ScenarioArtifactKind.SOURCE, + logical_name=f"fixture_{index}", + target_relation=f"scenario_schema.fixture_{index}", + status=status, + error_code=code, + error_help=help_text, + error_message=message, + ) + for index, (status, code, help_text, message) in enumerate(test_case.step_outcomes) + ) + + details: ScenarioFailureDetails = first_failure_details( + results=results, fallback_message="step failed", fallback_code=test_case.fallback_code + ) + + assert details == ScenarioFailureDetails( + error_code=test_case.expected_error_code, + error_help=test_case.expected_error_help, + error_message=test_case.expected_error_message, + ) diff --git a/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py b/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py index 11b90ee991..043c02b7af 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/_test_types.py @@ -289,3 +289,10 @@ class DagsterProjectPrepareFailureTestCase: class DagsterProjectDecoratorTestCase: description: str expected_asset_keys: tuple[tuple[str, ...], ...] + + +@dataclass(frozen=True) +class DagsterAssetSelectionErrorTestCase: + description: str + select: str + expected_error_fragment: str diff --git a/tests/unit/src/sqlbuild/integrations/dagster/helpers.py b/tests/unit/src/sqlbuild/integrations/dagster/helpers.py index d66bc16284..c89331807f 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/helpers.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/helpers.py @@ -568,3 +568,47 @@ def assert_json_output_file_behavior(*, command: tuple[str, ...]) -> None: assert "--json" not in command json_output_index: int = command.index("--json-output") + 1 assert not Path(command[json_output_index]).exists() + + +def _selection_model_node(name: str) -> dict[str, Any]: + return { + "id": f"model:{name}", + "kind": "model", + "name": name, + "asset_key": ["analytics", name], + "path": f"models/{name}.sql", + "materialization_type": "table", + } + + +def _selection_edge(from_name: str, to_name: str) -> dict[str, str]: + return {"from_id": f"model:{from_name}", "to_id": f"model:{to_name}"} + + +def build_dagster_selection_graph_dag() -> Mapping[str, Any]: + return { + "version": 1, + "project_name": "dagster_selection_project", + "nodes": [ + _selection_model_node("stg_orders"), + _selection_model_node("orders"), + _selection_model_node("customers"), + _selection_model_node("raw_events"), + _selection_model_node("events_left"), + _selection_model_node("events_right"), + _selection_model_node("events_joined"), + _selection_model_node("events_report"), + _selection_model_node("events_side"), + ], + "edges": [ + _selection_edge("stg_orders", "orders"), + _selection_edge("stg_orders", "customers"), + _selection_edge("raw_events", "events_left"), + _selection_edge("raw_events", "events_right"), + _selection_edge("raw_events", "events_side"), + _selection_edge("events_left", "events_joined"), + _selection_edge("events_right", "events_joined"), + _selection_edge("events_joined", "events_report"), + ], + "checks": [], + } diff --git a/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py b/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py index e3379076b0..a6afd656e0 100644 --- a/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py +++ b/tests/unit/src/sqlbuild/integrations/dagster/test_assets.py @@ -21,6 +21,7 @@ from sqlbuild.integrations.dagster.models import SqlBuildProject from tests.unit.src.sqlbuild.integrations.dagster._test_types import ( DagsterAssetCheckFilterTestCase, + DagsterAssetSelectionErrorTestCase, DagsterAssetSelectionTestCase, DagsterAssetSpecTestCase, DagsterConflictingInputTestCase, @@ -32,6 +33,7 @@ DagsterSourceAssetKeyTestCase, ) from tests.unit.src.sqlbuild.integrations.dagster.helpers import ( + build_dagster_selection_graph_dag, build_dagster_test_dag, build_python_augmented_dagster_test_dag, ) @@ -226,6 +228,94 @@ def assets_def() -> dg.MaterializeResult: ) +@pytest.mark.parametrize( + "test_case", + ( + DagsterAssetSelectionTestCase( + description="both-direction expansion keeps only ancestors and descendants", + select="+orders+", + exclude=None, + expected_asset_keys=(("analytics", "orders"), ("analytics", "stg_orders")), + ), + DagsterAssetSelectionTestCase( + description="downstream expansion includes every child", + select="stg_orders+", + exclude=None, + expected_asset_keys=( + ("analytics", "customers"), + ("analytics", "orders"), + ("analytics", "stg_orders"), + ), + ), + DagsterAssetSelectionTestCase( + description="path selector keeps every node on every path", + select="raw_events~events_report", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_report"), + ("analytics", "events_right"), + ("analytics", "raw_events"), + ), + ), + DagsterAssetSelectionTestCase( + description="path selector keeps both equal-length diamond branches", + select="raw_events~events_joined", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_right"), + ("analytics", "raw_events"), + ), + ), + DagsterAssetSelectionTestCase( + description="path selector expands upstream of start and downstream of end", + select="+events_left~events_joined+", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_report"), + ("analytics", "raw_events"), + ), + ), + DagsterAssetSelectionTestCase( + description="both-direction expansion of a middle node skips siblings", + select="+events_left+", + exclude=None, + expected_asset_keys=( + ("analytics", "events_joined"), + ("analytics", "events_left"), + ("analytics", "events_report"), + ("analytics", "raw_events"), + ), + ), + ), + ids=lambda case: case.description, +) +def test_given_branching_dag_when_building_asset_selection_then_expands_from_matched_nodes( + test_case: DagsterAssetSelectionTestCase, +) -> None: + dag: Mapping[str, Any] = build_dagster_selection_graph_dag() + + @sqlbuild_assets(dag=dag) + def assets_def() -> dg.MaterializeResult: + return dg.MaterializeResult() + + selection: Any = build_sqlbuild_asset_selection( + sqlbuild_assets=[assets_def], + dag=dag, + sqlbuild_select=test_case.select, + sqlbuild_exclude=test_case.exclude, + ) + + assert tuple(sorted(tuple(key.path) for key in selection.resolve([assets_def]))) == tuple( + sorted(test_case.expected_asset_keys) + ) + + @pytest.mark.parametrize( "test_case", [ @@ -529,3 +619,36 @@ def test_given_sqlbuild_dag_when_building_check_specs_then_can_exclude_scenarios assert check_names == test_case.expected_check_names assert not set(test_case.unexpected_check_names).intersection(check_names) + + +@pytest.mark.parametrize( + "test_case", + [ + DagsterAssetSelectionErrorTestCase( + description="rejects a path whose end is not downstream of its start", + select="events_report~raw_events", + expected_error_fragment="no SQLBuild DAG path exists between", + ), + DagsterAssetSelectionErrorTestCase( + description="rejects a path from a node to itself like the planner", + select="orders~orders", + expected_error_fragment="no SQLBuild DAG path exists between", + ), + ], + ids=lambda case: case.description, +) +def test_given_unreachable_path_selector_when_building_asset_selection_then_raises_input_error( + test_case: DagsterAssetSelectionErrorTestCase, +) -> None: + dag: Mapping[str, Any] = build_dagster_selection_graph_dag() + + @sqlbuild_assets(dag=dag) + def assets_def() -> dg.MaterializeResult: + return dg.MaterializeResult() + + with pytest.raises(DagsterDagInputError, match=test_case.expected_error_fragment): + build_sqlbuild_asset_selection( + sqlbuild_assets=[assets_def], + dag=dag, + sqlbuild_select=test_case.select, + )