From 9820d9ee287b0abd68b01fb5ac14a8ed9a62511c Mon Sep 17 00:00:00 2001 From: bmoore813 Date: Fri, 18 Sep 2026 01:08:09 -0400 Subject: [PATCH 1/3] feat(snowflake): add opt-in table write-audit-publish --- .../unreleased/Features-20260918-002011.yaml | 7 + crates/dbt-adapter-sql/src/ident.rs | 22 +- .../dbt-adapter/src/metadata/snowflake/mod.rs | 78 ++ .../dbt-adapter/src/relation/relation_impl.rs | 82 +- crates/dbt-compilation/src/schedule.rs | 4 +- .../src/phases/compile_and_run_context.rs | 362 ++++++- crates/dbt-jinja-utils/src/phases/mod.rs | 6 +- crates/dbt-main/src/compilation.rs | 31 +- crates/dbt-main/src/dbt_lib.rs | 8 +- crates/dbt-main/src/partial_parse.rs | 56 +- crates/dbt-main/src/retry.rs | 108 +- .../dbt-parser/src/resolve/resolve_models.rs | 5 +- .../dbt-parser/src/resolve/validate_models.rs | 84 +- .../tests/data/snowflake_wap/README.md | 60 ++ .../tests/data/snowflake_wap/dbt_project.yml | 9 + .../data/snowflake_wap/macros/acceptance.sql | 41 + .../data/snowflake_wap/models/downstream.sql | 3 + .../data/snowflake_wap/models/orders.sql | 10 + .../data/snowflake_wap/models/schema.yml | 8 + .../data/snowflake_wap/tests/nonnegative.sql | 3 + .../tests/snowflake_wap_acceptance.py | 294 ++++++ .../tests/test_snowflake_wap_acceptance.py | 70 ++ .../src/schemas/manifest/manifest_nodes.rs | 8 + crates/dbt-schemas/src/schemas/nodes.rs | 77 ++ .../dbt-schemas/src/schemas/prev_state/mod.rs | 5 + .../schemas/project/configs/model_config.rs | 62 +- crates/dbt-schemas/src/schemas/serde.rs | 15 + crates/dbt-tasks-core/src/context.rs | 74 +- crates/dbt-tasks-core/src/context_factory.rs | 2 +- crates/dbt-tasks-core/src/lib.rs | 1 + .../src/run_cache/run_cache_service.rs | 3 +- crates/dbt-tasks-core/src/wap.rs | 774 ++++++++++++++ crates/dbt-tasks-sa/src/graph.rs | 459 +++++++- crates/dbt-tasks-sa/src/lib.rs | 1 + crates/dbt-tasks-sa/src/materialize.rs | 4 +- .../src/renderable/renderable/default.rs | 112 +- crates/dbt-tasks-sa/src/runnable/model.rs | 20 +- .../dbt-tasks-sa/src/runnable/runnable/mod.rs | 40 +- crates/dbt-tasks-sa/src/runnable/test.rs | 1 + crates/dbt-tasks-sa/src/visitor.rs | 41 + crates/dbt-tasks-sa/src/wap.rs | 976 ++++++++++++++++++ crates/dbt-tasks-sa/src/wap/clone_tests.rs | 181 ++++ crates/dbt-tasks-sa/src/wap/header_tests.rs | 91 ++ .../dbt-tasks-sa/src/wap/publication_tests.rs | 201 ++++ docs/snowflake-table-wap-plan.md | 183 ++++ 45 files changed, 4619 insertions(+), 63 deletions(-) create mode 100644 .changes/unreleased/Features-20260918-002011.yaml create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/README.md create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/dbt_project.yml create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/models/downstream.sql create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/models/schema.yml create mode 100644 crates/dbt-sa-cli/tests/data/snowflake_wap/tests/nonnegative.sql create mode 100644 crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py create mode 100644 crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py create mode 100644 crates/dbt-tasks-core/src/wap.rs create mode 100644 crates/dbt-tasks-sa/src/wap.rs create mode 100644 crates/dbt-tasks-sa/src/wap/clone_tests.rs create mode 100644 crates/dbt-tasks-sa/src/wap/header_tests.rs create mode 100644 crates/dbt-tasks-sa/src/wap/publication_tests.rs create mode 100644 docs/snowflake-table-wap-plan.md diff --git a/.changes/unreleased/Features-20260918-002011.yaml b/.changes/unreleased/Features-20260918-002011.yaml new file mode 100644 index 00000000000..5665132b5a6 --- /dev/null +++ b/.changes/unreleased/Features-20260918-002011.yaml @@ -0,0 +1,7 @@ +kind: Features +body: Add opt-in write-audit-publish builds for native Snowflake SQL table models using same-schema working tables and a required passing data-test gate. +time: 2026-09-18T00:20:11.745413-04:00 +custom: + author: "" + issue: "" + project: dbt-core diff --git a/crates/dbt-adapter-sql/src/ident.rs b/crates/dbt-adapter-sql/src/ident.rs index 4d38e9ed240..66b033a3251 100644 --- a/crates/dbt-adapter-sql/src/ident.rs +++ b/crates/dbt-adapter-sql/src/ident.rs @@ -263,11 +263,13 @@ pub fn quote_identifier(id: &str, backend: AdapterType) -> String { /// value can never terminate the literal early. The literal-value member of /// the family next to [`sanitize_identifier`] (strip) and [`quote_identifier`] /// (quote), which handle identifiers. -pub fn escape_string_literal(s: &str, _backend: AdapterType) -> String { - // ANSI '' doubling — correct for every currently supported backend. - // Dialects with different literal-escape rules (e.g. backslash-escaped - // strings) grow a match arm on `_backend` when they arrive. - s.replace('\'', "''") +pub fn escape_string_literal(s: &str, backend: AdapterType) -> String { + match backend { + // Snowflake interprets backslash escapes inside single-quoted literals. + // https://docs.snowflake.com/en/sql-reference/data-types-text#escape-sequences-in-single-quoted-string-constants + AdapterType::Snowflake => s.replace('\\', "\\\\").replace('\'', "''"), + _ => s.replace('\'', "''"), + } } #[cfg(test)] @@ -301,6 +303,16 @@ mod sanitize_tests { assert_eq!(e("plain"), "plain"); } + #[test] + fn snowflake_string_literals_preserve_backslashes_and_quotes() { + let escape = |s: &str| escape_string_literal(s, AdapterType::Snowflake); + assert_eq!(escape(r"path\name"), r"path\\name"); + assert_eq!(escape(r"\'; DROP TABLE x; --"), r"\\''; DROP TABLE x; --"); + assert_eq!(escape("trailing\\"), "trailing\\\\"); + assert_eq!(escape("o'brien"), "o''brien"); + assert_eq!(escape("plain"), "plain"); + } + #[test] fn duckdb_sanitize_strips_sql_metacharacters() { let s = |n: &str| sanitize_identifier(n, AdapterType::DuckDB); diff --git a/crates/dbt-adapter/src/metadata/snowflake/mod.rs b/crates/dbt-adapter/src/metadata/snowflake/mod.rs index d9a8c0d63fb..1fbc2224e79 100644 --- a/crates/dbt-adapter/src/metadata/snowflake/mod.rs +++ b/crates/dbt-adapter/src/metadata/snowflake/mod.rs @@ -25,6 +25,7 @@ use dbt_adapter_engine::{ConnectionFactory, MapReduce}; use dbt_adbc::{Connection, QueryCtx}; use dbt_common::AsyncAdapterResult; use dbt_common::ErrorCode; +use dbt_common::FsResult; use dbt_common::cancellation::Cancellable; use dbt_common::cancellation::CancellationToken; use dbt_common::tracing::dbt_emit::{emit_debug_log_message, emit_warn_log_message}; @@ -44,6 +45,45 @@ use std::sync::Arc; const SNOWFLAKE_METADATA_NODE_ID: &str = "snowflake-metadata"; +/// Inspect an exact relation without depending on schema-list pagination limits. +/// +/// Callers must filter the returned names for exact equality: `STARTS WITH` may +/// also return longer names. Two rows let callers detect temporary-table +/// shadowing. SHOW orders names lexicographically, so an exact match comes first. +/// https://docs.snowflake.com/en/sql-reference/sql/show-objects +pub fn relation_lookup_sql(relation: &dyn BaseRelation) -> FsResult { + lookup_sql(relation, "OBJECTS") +} + +/// Inspect the lifecycle and native-table flags not exposed by SHOW OBJECTS. +/// https://docs.snowflake.com/en/sql-reference/sql/show-tables +pub fn table_lookup_sql(relation: &dyn BaseRelation) -> FsResult { + lookup_sql(relation, "TABLES") +} + +fn lookup_sql(relation: &dyn BaseRelation, object_type: &str) -> FsResult { + let metadata_error = |error| dbt_common::FsError::from_jinja_err(error, "WAP relation lookup"); + let database = quote_identifier( + &relation + .database_as_resolved_str() + .map_err(metadata_error)?, + AdapterType::Snowflake, + ); + let schema = quote_identifier( + &relation.schema_as_resolved_str().map_err(metadata_error)?, + AdapterType::Snowflake, + ); + let identifier = escape_string_literal( + &relation + .identifier_as_resolved_str() + .map_err(metadata_error)?, + AdapterType::Snowflake, + ); + Ok(format!( + "SHOW {object_type} IN {database}.{schema} STARTS WITH '{identifier}' LIMIT 2" + )) +} + fn metadata_warehouse_error(err: impl Display) -> AdapterError { AdapterError::new(AdapterErrorKind::Configuration, err.to_string()) } @@ -2100,6 +2140,44 @@ mod tests { use arrow_schema::{DataType, Field}; use std::sync::Mutex; + #[test] + fn relation_lookup_resolves_unquoted_names_and_bounds_results() { + let relation = Relation::new( + AdapterType::Snowflake, + "database".to_string(), + "schema".to_string(), + "orders".to_string(), + ) + .with_quoting(ResolvedQuoting::falses()); + assert_eq!( + relation_lookup_sql(&relation).unwrap(), + "SHOW OBJECTS IN \"DATABASE\".\"SCHEMA\" STARTS WITH 'ORDERS' LIMIT 2" + ); + assert_eq!( + table_lookup_sql(&relation).unwrap(), + "SHOW TABLES IN \"DATABASE\".\"SCHEMA\" STARTS WITH 'ORDERS' LIMIT 2" + ); + } + + #[test] + fn relation_lookup_preserves_quoted_case_and_escapes_literal_characters() { + let relation = Relation::new( + AdapterType::Snowflake, + "My\"Database".to_string(), + "My.Schema".to_string(), + r"O\'Brien".to_string(), + ) + .with_quoting(ResolvedQuoting::trues()); + assert_eq!( + relation_lookup_sql(&relation).unwrap(), + r#"SHOW OBJECTS IN "My""Database"."My.Schema" STARTS WITH 'O\\''Brien' LIMIT 2"# + ); + assert_eq!( + table_lookup_sql(&relation).unwrap(), + r#"SHOW TABLES IN "My""Database"."My.Schema" STARTS WITH 'O\\''Brien' LIMIT 2"# + ); + } + struct FakeConnectionFactory { recycled: Arc>, } diff --git a/crates/dbt-adapter/src/relation/relation_impl.rs b/crates/dbt-adapter/src/relation/relation_impl.rs index 794a482c582..7fb620d9c5b 100644 --- a/crates/dbt-adapter/src/relation/relation_impl.rs +++ b/crates/dbt-adapter/src/relation/relation_impl.rs @@ -14,7 +14,7 @@ use crate::relation::{RelationObject, StaticBaseRelation}; use crate::value::none_value; use dbt_adapter_core::AdapterType; -use dbt_adapter_sql::ident::max_identifier_length; +use dbt_adapter_sql::ident::{max_identifier_length, quote_identifier}; use dbt_common::{ErrorCode, FsResult, constants::DBT_CTE_PREFIX, fs_err}; use dbt_frontend_common::ident::Identifier; use dbt_schema_store::CanonicalFqn; @@ -285,7 +285,7 @@ impl BaseRelationProperties for Relation { } }; - let schema = if self.quote_policy().database { + let schema = if self.quote_policy().schema { schema_str } else { match self.adapter_type { @@ -295,7 +295,7 @@ impl BaseRelationProperties for Relation { } }; - let ident = if self.quote_policy().database { + let ident = if self.quote_policy().identifier { ident_str } else { match self.adapter_type { @@ -1176,6 +1176,18 @@ impl BaseRelation for Relation { } } + fn quoted(&self, component: &str) -> String { + match self.adapter_type { + AdapterType::Snowflake => quote_identifier(component, self.adapter_type), + _ => format!( + "{}{}{}", + self.quote_character(), + component, + self.quote_character() + ), + } + } + fn render_self_as_str(&self) -> String { if self.adapter_type == AdapterType::DuckDB && let Some(external) = &self.external @@ -1582,6 +1594,70 @@ mod tests { assert_eq!(relation.semantic_fqn(), "\"MyDB\".\"myschema\".\"MyTable\""); } + #[test] + fn relation_rendering_escapes_embedded_quotes_in_each_component() { + let relation = Relation::new( + AdapterType::Snowflake, + "db\"name".to_owned(), + "schema.with.dot".to_owned(), + "table\"name".to_owned(), + ) + .with_quoting(Policy::enabled()); + let expected = "\"db\"\"name\".\"schema.with.dot\".\"table\"\"name\""; + assert_eq!(relation.render_self_as_str(), expected); + assert_eq!(relation.semantic_fqn(), expected); + + let mixed = relation.with_quoting(Policy { + database: true, + schema: false, + identifier: true, + }); + assert_eq!( + mixed.render_self_as_str(), + "\"db\"\"name\".schema.with.dot.\"table\"\"name\"" + ); + } + + #[test] + fn test_canonical_fqn_uses_each_component_quote_policy() { + for (adapter, normalized) in [ + (AdapterType::Snowflake, ["MYDB", "MYSCHEMA", "MYTABLE"]), + (AdapterType::Postgres, ["mydb", "myschema", "mytable"]), + ] { + for (quoting, expected) in [ + ( + Policy { + database: true, + schema: false, + identifier: false, + }, + ["MyDB", normalized[1], normalized[2]], + ), + ( + Policy { + database: false, + schema: true, + identifier: true, + }, + [normalized[0], "MySchema", "MyTable"], + ), + ] { + let fqn = Relation::new( + adapter, + "MyDB".to_string(), + "MySchema".to_string(), + "MyTable".to_string(), + ) + .with_quoting(quoting) + .get_canonical_fqn() + .unwrap(); + assert_eq!(fqn.catalog().as_str(), expected[0]); + assert_eq!(fqn.schema().as_str(), expected[1]); + assert_eq!(fqn.table().as_str(), expected[2]); + } + } + } + fn filter_relation() -> Relation { Relation::new( AdapterType::Postgres, diff --git a/crates/dbt-compilation/src/schedule.rs b/crates/dbt-compilation/src/schedule.rs index b30504fa69a..a27231dae46 100644 --- a/crates/dbt-compilation/src/schedule.rs +++ b/crates/dbt-compilation/src/schedule.rs @@ -6,7 +6,7 @@ use dbt_common::node_selector::IndirectSelection; use dbt_common::path::DbtPath; use dbt_schemas::state::CacheState; -#[derive(Debug)] +#[derive(Debug, Default)] pub struct DbtCustomScheduleDescription { pub unique_ids: Vec, pub include_parents: bool, @@ -17,6 +17,8 @@ pub struct DbtCustomScheduleDescription { /// to reflect the semantics they want (e.g. `Empty` when the originating /// command does not execute tests). pub indirect_selection: IndirectSelection, + /// Retry rebuilds WAP models and reruns every audit of their fresh candidates. + pub is_retry: bool, } pub enum DbtScheduleDescription<'a> { diff --git a/crates/dbt-jinja-utils/src/phases/compile_and_run_context.rs b/crates/dbt-jinja-utils/src/phases/compile_and_run_context.rs index 9f6856bfa40..2b55e4e9347 100644 --- a/crates/dbt-jinja-utils/src/phases/compile_and_run_context.rs +++ b/crates/dbt-jinja-utils/src/phases/compile_and_run_context.rs @@ -258,7 +258,7 @@ impl MicrobatchRefContext { /// - Package-qualified refs: `ref('package_name', 'model_name')` /// - Versioned refs: `ref('model_name', version=1)` /// - Microbatch-aware filtering when `microbatch_context` is set -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct RefFunction { node_resolver: Arc, package_name: String, @@ -270,6 +270,8 @@ pub struct RefFunction { /// The unique_id of the node that owns this ref context. /// Used for O(1) defer decisions via `NodeResolver::prefers_deferred`. current_node_unique_id: String, + /// Physical relations visible only to this execution context, keyed by node identity. + relation_overrides: Option>>, } impl RefFunction { @@ -287,6 +289,7 @@ impl RefFunction { validation_config: DependencyValidationConfig::default(), microbatch_context: None, current_node_unique_id: String::new(), + relation_overrides: None, } } @@ -305,6 +308,7 @@ impl RefFunction { validation_config, microbatch_context: None, current_node_unique_id, + relation_overrides: None, } } @@ -328,6 +332,7 @@ impl RefFunction { validation_config, microbatch_context: Some(microbatch_context), current_node_unique_id, + relation_overrides: None, } } @@ -341,6 +346,18 @@ impl RefFunction { } } + /// Bind selected node identities to execution-local physical relations. + /// Dependency validation still runs before an override is considered. + pub fn with_relation_overrides( + self, + relation_overrides: Arc>, + ) -> Self { + Self { + relation_overrides: Some(relation_overrides), + ..self + } + } + fn resolve_args( &self, args: &[MinijinjaValue], @@ -432,6 +449,49 @@ To fix this, add the following hint to the top of the model \"{model_name}\": } } +/// Overlay both spellings of `ref` without changing the shared resolver or base context. +/// Invalid context objects are rejected before either binding is changed. +pub fn apply_ref_relation_overrides( + context: &mut BTreeMap, + relation_overrides: Arc>, +) -> Result<(), MinijinjaError> { + let ref_function = context + .get("ref") + .and_then(MinijinjaValue::as_object) + .and_then(|object| object.downcast_ref::()) + .ok_or_else(|| { + MinijinjaError::new( + MinijinjaErrorKind::InvalidOperation, + "execution relation overrides require the built-in ref function", + ) + })?; + let ref_value = MinijinjaValue::from_object( + ref_function + .clone() + .with_relation_overrides(relation_overrides), + ); + let mut builtins = match context.get("builtins") { + Some(value) => value + .as_object() + .and_then(|object| object.downcast_ref::>()) + .cloned() + .ok_or_else(|| { + MinijinjaError::new( + MinijinjaErrorKind::InvalidOperation, + "execution relation overrides require the built-in context map", + ) + })?, + None => BTreeMap::new(), + }; + builtins.insert("ref".to_string(), ref_value.clone()); + context.insert("ref".to_string(), ref_value); + context.insert( + "builtins".to_string(), + MinijinjaValue::from_object(builtins), + ); + Ok(()) +} + impl Object for RefFunction { fn call( self: &Arc, @@ -450,15 +510,23 @@ impl Object for RefFunction { Ok((unique_id, relation, _, deferred_relation)) => { // Validate that this ref is allowed (only if validation is configured) self.validate_dependency(&unique_id, &package_name, &model_name)?; - // Use phase-aware defer logic: check if we should use the deferred - // (production) relation for this specific upstream node. - let prefers = self - .node_resolver - .prefers_deferred(&self.current_node_unique_id, &unique_id); - let resolved_relation = match (prefers, deferred_relation) { - (true, Some(deferred)) => deferred, - _ => relation, - }; + let resolved_relation = self + .relation_overrides + .as_deref() + .and_then(|overrides| overrides.get(&unique_id)) + .cloned() + // Only ordinary refs consult phase-aware deferral. A scoped + // execution relation must take precedence over production. + .unwrap_or_else(|| match deferred_relation { + Some(deferred) + if self + .node_resolver + .prefers_deferred(&self.current_node_unique_id, &unique_id) => + { + deferred + } + _ => relation, + }); for listener in listeners { listener.on_ref_or_source_resolved(&unique_id); @@ -1011,10 +1079,284 @@ impl Object for LazyFlatGraph { #[cfg(test)] mod tests { use super::*; + use crate::node_resolver::NodeResolver; + use dbt_adapter::AdapterType; + use dbt_adapter::relation::create_relation_from_node; + use dbt_schemas::schemas::common::{DbtMaterialization, ResolvedQuoting}; + use dbt_schemas::schemas::nodes::{CommonAttributes, DbtModel, DbtTest, NodeBaseAttributes}; + use dbt_schemas::schemas::serde::StringOrInteger; use dbt_schemas::state::DummyNodeResolverTracker; + use dbt_schemas::state::ModelStatus; use dbt_test_utils::TestEnvGuard; use std::env; + fn ref_model(name: &str, alias: &str) -> DbtModel { + DbtModel { + __common_attr__: CommonAttributes { + unique_id: format!("model.pkg.{name}"), + name: name.to_string(), + package_name: "pkg".to_string(), + ..Default::default() + }, + __base_attr__: NodeBaseAttributes { + database: "DB".to_string(), + schema: "SCHEMA".to_string(), + alias: alias.to_string(), + materialized: DbtMaterialization::Table, + quoting: ResolvedQuoting::enabled(), + ..Default::default() + }, + ..Default::default() + } + } + + fn relation_value(model: &DbtModel) -> MinijinjaValue { + RelationObject::new( + create_relation_from_node(AdapterType::Snowflake, model, None) + .unwrap() + .into(), + ) + .into_value() + } + + fn ref_context(resolver: NodeResolver) -> BTreeMap { + let ref_value = MinijinjaValue::from_object(RefFunction::new_unvalidated( + Arc::new(resolver), + "pkg".to_string(), + Arc::new(DbtRuntimeConfig::default()), + )); + BTreeMap::from([ + ("ref".to_string(), ref_value.clone()), + ( + "builtins".to_string(), + MinijinjaValue::from_object(BTreeMap::from([ + ("ref".to_string(), ref_value), + ( + "config".to_string(), + MinijinjaValue::from("original config"), + ), + ])), + ), + ]) + } + + #[test] + fn execution_refs_override_defer_for_both_spellings_without_leaking() { + let mut model = ref_model("orders", "Public Orders"); + model.__model_attr__.version = Some(StringOrInteger::Integer(1)); + model.__model_attr__.latest_version = Some(StringOrInteger::Integer(1)); + let upstream = ref_model("upstream", "Upstream"); + let mut deferred = model.clone(); + deferred.__base_attr__.database = "DEFERRED".to_string(); + let mut candidate = model.clone(); + candidate.__base_attr__.alias = "__DBT_WAP_TEST".to_string(); + + let mut resolver = NodeResolver::default(); + resolver + .insert_ref(&model, AdapterType::Snowflake, ModelStatus::Enabled, false) + .unwrap(); + resolver + .insert_ref( + &upstream, + AdapterType::Snowflake, + ModelStatus::Enabled, + false, + ) + .unwrap(); + resolver + .update_ref_with_deferral(&deferred, AdapterType::Snowflake, false) + .unwrap(); + assert!(resolver.prefers_deferred("test.pkg.audit", &model.__common_attr__.unique_id)); + + let mut original = ref_context(resolver); + original.extend([ + ( + "this".to_string(), + MinijinjaValue::from("audit_failure_table"), + ), + ( + "source".to_string(), + MinijinjaValue::from_function(|_: String, _: String| "public_source"), + ), + ]); + let mut audit = original.clone(); + apply_ref_relation_overrides( + &mut audit, + Arc::new(BTreeMap::from([( + model.__common_attr__.unique_id.clone(), + relation_value(&candidate), + )])), + ) + .unwrap(); + + let env = minijinja::Environment::new(); + let template = "{{ ref('pkg', 'orders', version=1) }}|{{ builtins.ref('orders', v=1) }}|{{ ref('upstream') }}|{{ this }}|{{ source('raw', 'orders') }}"; + assert_eq!( + env.render_str(template, &audit, &[]).unwrap(), + format!( + "{0}|{0}|{1}|audit_failure_table|public_source", + relation_value(&candidate), + relation_value(&upstream) + ), + ); + assert_eq!( + env.render_str( + "{{ ref('orders', version=1) }}|{{ builtins.ref('orders', version=1) }}", + &original, + &[], + ) + .unwrap(), + format!("{0}|{0}", relation_value(&deferred)), + ); + + // The run overlay replaces `builtins.config` and `this`, while both + // audit ref bindings must survive into the test materialization. + let test = DbtTest { + __common_attr__: CommonAttributes { + unique_id: "test.pkg.audit".to_string(), + name: "audit".to_string(), + package_name: "pkg".to_string(), + ..Default::default() + }, + __base_attr__: NodeBaseAttributes { + alias: "audit_failures".to_string(), + ..model.__base_attr__.clone() + }, + ..Default::default() + }; + let (run_context, _) = crate::phases::run::build_run_node_context( + &test, + &test.deprecated_config, + AdapterType::Snowflake, + None, + &audit, + &dbt_common::io_args::IoArgs::default(), + dbt_telemetry::ExecutionPhase::Run, + None, + BTreeSet::new(), + ); + assert_eq!( + env.render_str( + "{{ ref('orders', version=1) }}|{{ builtins.ref('orders', version=1) }}|{{ this.identifier }}", + &run_context, + &[], + ).unwrap(), + format!("{0}|{0}|audit_failures", relation_value(&candidate)), + ); + } + + #[test] + fn execution_refs_preserve_wrappers_identity_and_other_execution_contexts() { + let model = ref_model("orders", "Public Orders"); + let mut other_package = ref_model("orders", "Other Package Orders"); + other_package.__common_attr__.package_name = "other".to_string(); + other_package.__common_attr__.unique_id = "model.other.orders".to_string(); + let mut resolver = NodeResolver::default(); + for node in [&model, &other_package] { + resolver + .insert_ref(node, AdapterType::Snowflake, ModelStatus::Enabled, false) + .unwrap(); + } + let original = ref_context(resolver); + let mut candidates = Vec::new(); + // Two independent overlays coexist on one resolver, as concurrent builds + // or separate audited models do. Neither may mutate the original binding. + for alias in ["__DBT_WAP_FIRST", "__DBT_WAP_SECOND"] { + let mut candidate = model.clone(); + candidate.__base_attr__.alias = alias.to_string(); + let mut context = original.clone(); + apply_ref_relation_overrides( + &mut context, + Arc::new(BTreeMap::from([( + model.__common_attr__.unique_id.clone(), + relation_value(&candidate), + )])), + ) + .unwrap(); + candidates.push((candidate, context)); + } + let env = minijinja::Environment::new(); + let template = "{% macro ref(name) %}{{ builtins.ref(name) }}{% endmacro %}{{ ref('orders') }}|{{ builtins.ref.id('orders') }}|{{ builtins.ref('other', 'orders') }}|{{ builtins.config }}"; + for (candidate, context) in &candidates { + assert_eq!( + env.render_str(template, context, &[]).unwrap(), + format!( + "{}|{}|{}|original config", + relation_value(candidate), + model.__common_attr__.unique_id, + relation_value(&other_package), + ), + ); + } + assert_eq!( + env.render_str("{{ ref('orders') }}", &original, &[]) + .unwrap(), + relation_value(&model).to_string(), + ); + } + + #[test] + fn execution_ref_overlays_reject_invalid_contexts_without_partial_changes() { + let model = ref_model("orders", "Public Orders"); + let mut candidate = model.clone(); + candidate.__base_attr__.alias = "__DBT_WAP_TEST".to_string(); + let mut resolver = NodeResolver::default(); + resolver + .insert_ref(&model, AdapterType::Snowflake, ModelStatus::Enabled, false) + .unwrap(); + let original = ref_context(resolver); + let overrides = Arc::new(BTreeMap::from([( + model.__common_attr__.unique_id.clone(), + relation_value(&candidate), + )])); + let env = minijinja::Environment::new(); + + let mut invalid_builtins = original.clone(); + invalid_builtins.insert("builtins".to_string(), MinijinjaValue::from("invalid")); + assert!( + apply_ref_relation_overrides(&mut invalid_builtins, Arc::clone(&overrides)).is_err() + ); + assert_eq!( + env.render_str("{{ ref('orders') }}", &invalid_builtins, &[]) + .unwrap(), + relation_value(&model).to_string(), + ); + + let mut invalid_ref = original; + invalid_ref.insert("ref".to_string(), MinijinjaValue::from("custom binding")); + assert!(apply_ref_relation_overrides(&mut invalid_ref, overrides).is_err()); + assert_eq!( + env.render_str("{{ ref }}|{{ builtins.ref('orders') }}", &invalid_ref, &[]) + .unwrap(), + format!("custom binding|{}", relation_value(&model)), + ); + } + + #[test] + fn execution_refs_do_not_bypass_dependency_validation() { + let model = ref_model("orders", "Public Orders"); + let mut resolver = NodeResolver::default(); + resolver + .insert_ref(&model, AdapterType::Snowflake, ModelStatus::Enabled, false) + .unwrap(); + let ref_function = RefFunction::new_with_validation( + Arc::new(resolver), + "pkg".to_string(), + Arc::new(DbtRuntimeConfig::default()), + DependencyValidationConfig::new_validated(), + "test.pkg.audit".to_string(), + ) + .with_relation_overrides(Arc::new(BTreeMap::from([( + model.__common_attr__.unique_id.clone(), + relation_value(&model), + )]))); + let ctx = BTreeMap::from([("ref", MinijinjaValue::from_object(ref_function))]); + let err = minijinja::Environment::new() + .render_str("{{ ref('orders') }}", ctx, &[]) + .unwrap_err(); + assert!(err.to_string().contains("unable to infer all dependencies")); + } + #[test] fn test_dbt_metadata_envs_populated_from_env() { // Isolate env to avoid interference across tests diff --git a/crates/dbt-jinja-utils/src/phases/mod.rs b/crates/dbt-jinja-utils/src/phases/mod.rs index a34be767352..7489ae41447 100644 --- a/crates/dbt-jinja-utils/src/phases/mod.rs +++ b/crates/dbt-jinja-utils/src/phases/mod.rs @@ -6,7 +6,7 @@ pub mod run; mod utils; pub use compile_and_run_context::{ - MacroLookupContext, MicrobatchRefContext, RefFunction, SourceFunction, build_compile_base_ctx, - build_operation_context, build_operation_context_btreemap, - configure_compile_and_run_jinja_environment, + MacroLookupContext, MicrobatchRefContext, RefFunction, SourceFunction, + apply_ref_relation_overrides, build_compile_base_ctx, build_operation_context, + build_operation_context_btreemap, configure_compile_and_run_jinja_environment, }; diff --git a/crates/dbt-main/src/compilation.rs b/crates/dbt-main/src/compilation.rs index 84276917cff..32ed1facee6 100644 --- a/crates/dbt-main/src/compilation.rs +++ b/crates/dbt-main/src/compilation.rs @@ -804,6 +804,7 @@ pub type DbtRunTasksResult = ( use crate::partial_parse::{ PrevCompilationResult, try_lazy_load_fast_path, try_load_prev_compilation, + wap_full_parse_reason, }; use dbt_compilation::traits::{CompilationCache, CompiledProject}; use dbt_state::selector::RunCacheStateSelectorArgs; @@ -1270,6 +1271,25 @@ impl DbtProjectCompilation { .await } Ok((mut compilation, jinja_env, changes)) => { + if let Some(reason) = wap_full_parse_reason( + partial_load_filter_applied, + &compilation.resolved_state.nodes, + ) { + tracing::debug!("Partial parse: {reason}, falling back to full parse"); + return DbtProjectCompilation::initialize( + feature_stack, + arg, + cli, + config, + event_emitter, + jinja_type_checking_event_listener_factory, + None, + token, + version_check_handle, + artifacts_sink, + ) + .await; + } compilation.partial_load_filter_applied = partial_load_filter_applied; Ok((compilation, jinja_env, changes)) } @@ -1603,12 +1623,21 @@ impl DbtProjectCompilation { .await? } DbtScheduleDescription::Custom(custom_schedule_desc) => { + let unique_ids = + if custom_schedule_desc.is_retry && arg.command == FsCommand::Build { + crate::retry::expand_wap_retry_ids( + &custom_schedule_desc.unique_ids, + &self.resolved_state.nodes, + )? + } else { + custom_schedule_desc.unique_ids.clone() + }; schedule_with_unique_ids( &self.resolved_state, scheduler_args, maybe_previous_state.as_ref().map(|x| x.as_ref()), run_cache_state_selector_args, - &custom_schedule_desc.unique_ids, + &unique_ids, custom_schedule_desc.include_parents, custom_schedule_desc.include_children, custom_schedule_desc.indirect_selection, diff --git a/crates/dbt-main/src/dbt_lib.rs b/crates/dbt-main/src/dbt_lib.rs index 7012a1410a9..4a621e59854 100644 --- a/crates/dbt-main/src/dbt_lib.rs +++ b/crates/dbt-main/src/dbt_lib.rs @@ -755,15 +755,18 @@ impl<'a> AllPhasesExecutor<'a> { // This keeps retry bounded to nodes recorded in run_results.json instead // of broadening the selection by expanding descendants. // - // Retry executes exactly the node ids recorded in run_results.json, so + // Ordinary retry executes the node ids recorded in run_results.json, so // indirect selection must be `Empty`. The recorded set already contains // every node that has to re-run -- in particular, tests skipped behind a // failed model are recorded `skipped`, which is retryable -- so expanding - // is never needed. Any expansion is by construction a node the original + // is normally unnecessary. Any expansion is by construction a node the original // command did not run: an already-passed unit test, or a test excluded by // the original --indirect-selection / --exclude. This matches dbt-core, // whose retry path replaces the graph queue outright and never consults // indirect selection (dbt-labs/dbt-core#14536). + // WAP is the explicit exception: after resolving the current manifest, + // compilation expands a retried WAP owner to every required audit, + // because the rebuilt candidate has never been audited. // Check ids are **not schedulable**: checks run before the task graph exists, so they // are not nodes in it. Putting one in a custom schedule produced a false green — // `dbt retry` after a failing check selected only that id, row scoping then found no @@ -782,6 +785,7 @@ impl<'a> AllPhasesExecutor<'a> { None } else { Some(DbtCustomScheduleDescription { + is_retry: true, unique_ids: retry_node_ids, include_parents: false, include_children: false, diff --git a/crates/dbt-main/src/partial_parse.rs b/crates/dbt-main/src/partial_parse.rs index e2818ce5db7..46356e53894 100644 --- a/crates/dbt-main/src/partial_parse.rs +++ b/crates/dbt-main/src/partial_parse.rs @@ -15,7 +15,7 @@ use dbt_metadata::{ }, }; use dbt_schemas::{ - schemas::common::ResolvedQuoting, + schemas::{Nodes, common::ResolvedQuoting}, state::{DbtRuntimeConfig, DbtState, NodeResolverTracker, ResolverState}, }; @@ -32,6 +32,18 @@ pub enum PrevCompilationResult { None, } +/// A filtered manifest cannot establish the complete mandatory WAP audit set. +/// This is checked both on cached nodes and after incremental SQL re-resolution, +/// because an inline config can enable WAP during the incremental pass. +pub(crate) fn wap_full_parse_reason(filtered: bool, nodes: &Nodes) -> Option<&'static str> { + (filtered + && nodes + .models + .values() + .any(|model| model.deprecated_config.wap == Some(true))) + .then_some("WAP audit discovery requires an unfiltered manifest") +} + /// Attempt to load and reconstruct a previous compilation from the parquet parse cache. /// /// Returns: @@ -126,6 +138,11 @@ pub fn try_load_prev_compilation( return (PrevCompilationResult::None, use_lazy_filter); }; + if let Some(reason) = wap_full_parse_reason(use_lazy_filter, &state.nodes) { + tracing::debug!("Partial parse: {reason}, falling back to full parse"); + return (PrevCompilationResult::FullParse, false); + } + if let Some(reason) = state.validate(&cli.common_args.vars) { tracing::debug!("Partial parse: {reason}, invalidating cache"); return (PrevCompilationResult::None, use_lazy_filter); @@ -325,3 +342,40 @@ pub fn try_lazy_load_fast_path( } } } + +#[cfg(test)] +mod wap_tests { + use super::*; + use dbt_schemas::schemas::DbtModel; + + #[test] + fn filtered_wap_nodes_require_full_parse_even_when_audits_are_absent() { + let mut nodes = Nodes::default(); + let mut model = DbtModel::default(); + model.deprecated_config.wap = Some(true); + nodes + .models + .insert("model.pkg.orders".to_owned(), Arc::new(model)); + // Empty/Cautious indirect selection can omit these tests entirely at load time. + assert!(nodes.tests.is_empty()); + assert!(wap_full_parse_reason(true, &nodes).is_some()); + assert!(wap_full_parse_reason(false, &nodes).is_none()); + } + + #[test] + fn newly_enabled_wap_requires_reloading_the_filtered_manifest() { + let mut nodes = Nodes::default(); + nodes + .models + .insert("model.pkg.orders".to_owned(), Arc::new(DbtModel::default())); + assert!(wap_full_parse_reason(true, &nodes).is_none()); + Arc::make_mut(nodes.models.get_mut("model.pkg.orders").unwrap()) + .deprecated_config + .wap = Some(true); + assert!(wap_full_parse_reason(true, &nodes).is_some()); + Arc::make_mut(nodes.models.get_mut("model.pkg.orders").unwrap()) + .deprecated_config + .wap = Some(false); + assert!(wap_full_parse_reason(true, &nodes).is_none()); + } +} diff --git a/crates/dbt-main/src/retry.rs b/crates/dbt-main/src/retry.rs index b0aa6edda2a..21e332ea4d4 100644 --- a/crates/dbt-main/src/retry.rs +++ b/crates/dbt-main/src/retry.rs @@ -3,8 +3,9 @@ use dbt_clap_core::*; use dbt_common::io_args::StaticAnalysisKind; use dbt_common::{ErrorCode, FsResult, err}; -use dbt_schemas::schemas::{BatchResults, RunResultsArtifact}; -use std::collections::HashMap; +use dbt_schemas::schemas::{BatchResults, InternalDbtNode, Nodes, RunResultsArtifact}; +use dbt_tasks_core::wap::WapPlan; +use std::collections::{BTreeSet, HashMap}; use std::path::Path; use std::str::FromStr; @@ -21,6 +22,61 @@ pub const RETRIABLE_COMMANDS: &[&str] = &[ "run", "build", "test", "seed", "snapshot", "compile", "check", ]; +/// A retried WAP build creates a new candidate, so previously passed audits must rerun. +pub(crate) fn expand_wap_retry_ids(ids: &[String], nodes: &Nodes) -> FsResult> { + let mut selected: BTreeSet = ids.iter().cloned().collect(); + let mut owners = BTreeSet::new(); + for id in ids { + if nodes + .models + .get(id) + .is_some_and(|model| model.deprecated_config.wap.unwrap_or(false)) + { + owners.insert(id.clone()); + } + if let Some(test) = nodes.tests.get(id) { + let potential_owners = match test.__test_attr__.attached_node.as_ref() { + Some(owner) => vec![owner], + None => test.base().depends_on.nodes.iter().collect(), + }; + for owner in potential_owners { + if nodes + .models + .get(owner) + .is_some_and(|model| model.deprecated_config.wap.unwrap_or(false)) + { + owners.insert(owner.clone()); + } + } + } + if let Some(test) = nodes.unit_tests.get(id) + && let Some(owner) = test.base().depends_on.nodes.first() + && nodes + .models + .get(owner) + .is_some_and(|model| model.deprecated_config.wap.unwrap_or(false)) + { + owners.insert(owner.clone()); + } + } + for owner in owners { + selected.insert(owner.clone()); + selected.extend(WapPlan::required_audit_ids(&owner, nodes)?); + selected.extend( + nodes + .unit_tests + .iter() + .filter(|(_, test)| { + test.base().enabled + && test.deprecated_config.enabled != Some(false) + && test.base().depends_on.nodes.first() == Some(&owner) + }) + .map(|(id, _)| id.clone()), + ); + } + Ok(selected.into_iter().collect()) +} + /// Holds the state extracted from a previous run's run_results.json /// needed to execute a retry command. #[derive(Debug)] @@ -212,6 +268,54 @@ mod tests { use std::io::Write; use tempfile::NamedTempFile; + #[test] + fn wap_retry_rebuilds_owner_and_all_audits_without_unrelated_nodes() { + use dbt_schemas::schemas::{DbtModel, DbtTest, DbtUnitTest}; + use std::sync::Arc; + + let model_id = "model.project.orders"; + let mut nodes = Nodes::default(); + let mut model = DbtModel::default(); + model.deprecated_config.wap = Some(true); + nodes.models.insert(model_id.to_string(), Arc::new(model)); + for id in ["test.failed", "test.passed", "test.disabled"] { + let mut test = DbtTest::default(); + test.__base_attr__.enabled = id != "test.disabled"; + test.__base_attr__.depends_on.nodes = vec![model_id.to_string()]; + test.__test_attr__.attached_node = Some(model_id.to_string()); + nodes.tests.insert(id.to_string(), Arc::new(test)); + } + let mut unit = DbtUnitTest::default(); + unit.__base_attr__.enabled = true; + unit.__base_attr__.depends_on.nodes = vec![model_id.to_string()]; + nodes + .unit_tests + .insert("unit_test.passed".to_string(), Arc::new(unit)); + + let expected = vec![ + model_id.to_string(), + "test.failed".to_string(), + "test.passed".to_string(), + "unit_test.passed".to_string(), + ]; + assert_eq!( + expand_wap_retry_ids(&["test.failed".to_string()], &nodes).unwrap(), + expected + ); + assert_eq!( + expand_wap_retry_ids(&[model_id.to_string()], &nodes).unwrap(), + expected + ); + assert_eq!( + expand_wap_retry_ids(&["unit_test.passed".to_string()], &nodes).unwrap(), + expected + ); + assert_eq!( + expand_wap_retry_ids(&["model.other".to_string()], &nodes).unwrap(), + ["model.other"] + ); + } + #[test] fn test_retryable_statuses_contains_expected() { assert!(RETRYABLE_STATUSES.contains(&"error")); diff --git a/crates/dbt-parser/src/resolve/resolve_models.rs b/crates/dbt-parser/src/resolve/resolve_models.rs index b2cd4f8c0d7..0833aab6dd6 100644 --- a/crates/dbt-parser/src/resolve/resolve_models.rs +++ b/crates/dbt-parser/src/resolve/resolve_models.rs @@ -102,7 +102,7 @@ use super::resolve_tests::persist_generic_data_tests::{ TestUnrenderedConfigs, extract_test_unrendered_configs, }; use super::resolve_utils::{validate_compute, validate_node_adapter}; -use super::validate_models::validate_model; +use super::validate_models::{validate_model, validate_wap_model}; /// Parses `ref('name')`, `ref('pkg', 'name')`, `ref('name', version=N)`, or /// `ref('pkg', 'name', version=N)` from a constraint `to:` string (also accepts `v=` alias). @@ -1144,6 +1144,9 @@ async fn build_model_nodes( __other__: BTreeMap::new(), }; + validate_wap_model(&dbt_model, selected_adapter) + .map_err(|error| error.with_location(dbt_asset.path.clone()))?; + let components = RelationComponents { database: model_config.database.clone().into_inner().unwrap_or(None), schema: model_config.schema.clone().into_inner().unwrap_or(None), diff --git a/crates/dbt-parser/src/resolve/validate_models.rs b/crates/dbt-parser/src/resolve/validate_models.rs index 6f546f76b34..45dfa580c35 100644 --- a/crates/dbt-parser/src/resolve/validate_models.rs +++ b/crates/dbt-parser/src/resolve/validate_models.rs @@ -1,7 +1,13 @@ +use dbt_adapter_core::AdapterType; use dbt_common::{ErrorCode, FsError, FsResult, fs_err}; +use dbt_schemas::schemas::DbtModel; use dbt_schemas::schemas::properties::ModelProperties; use std::collections::HashSet; +pub(super) fn validate_wap_model(model: &DbtModel, adapter: AdapterType) -> FsResult<()> { + model.validate_wap_config(adapter) +} + /// Validates time spine configuration for semantic models according to the rules ported from Python dbt. /// This checks: /// - Standard granularity column exists in model columns @@ -163,13 +169,89 @@ pub fn validate_model(model_props: &ModelProperties) -> FsResult> { #[cfg(test)] mod tests { use super::*; - use dbt_schemas::schemas::common::Versions; + use dbt_schemas::schemas::common::{ConstraintType, DbtMaterialization, OnError, Versions}; use dbt_schemas::schemas::dbt_column::{ColumnProperties, Granularity}; use dbt_schemas::schemas::properties::model_properties::{ ModelPropertiesTimeSpine, TimeSpineCustomGranularity, }; use dbt_schemas::schemas::serde::FloatOrString; + fn wap_model() -> DbtModel { + let mut model = DbtModel::default(); + model.__common_attr__.unique_id = "model.project.orders".to_string(); + model.__common_attr__.language = Some("sql".to_string()); + model.__base_attr__.materialized = DbtMaterialization::Table; + model.deprecated_config.wap = Some(true); + model + } + + #[test] + fn wap_accepts_native_sql_tables_and_false_is_inert() { + let mut model = wap_model(); + assert!(validate_wap_model(&model, AdapterType::Snowflake).is_ok()); + model.deprecated_config.wap = Some(false); + model.__common_attr__.language = Some("python".to_string()); + assert!(validate_wap_model(&model, AdapterType::Bigquery).is_ok()); + } + + #[test] + fn wap_rejects_unsupported_model_kinds() { + let mut model = wap_model(); + assert!(validate_wap_model(&model, AdapterType::Bigquery).is_err()); + model.__common_attr__.language = Some("python".to_string()); + assert!(validate_wap_model(&model, AdapterType::Snowflake).is_err()); + model.__common_attr__.language = Some("sql".to_string()); + model.__base_attr__.materialized = DbtMaterialization::View; + assert!(validate_wap_model(&model, AdapterType::Snowflake).is_err()); + model.__base_attr__.materialized = DbtMaterialization::Table; + model.deprecated_config.table_format = Some("iceberg".to_string()); + assert!(validate_wap_model(&model, AdapterType::Snowflake).is_err()); + } + + #[test] + fn wap_rejects_hooks_headers_continue_and_governance_options() { + use dbt_schemas::schemas::common::Hooks; + + let cases: &[fn(&mut DbtModel)] = &[ + |m| { + m.deprecated_config.pre_hook = + dbt_yaml::Verbatim::from(Some(Hooks::String("select 1".to_string()))) + }, + |m| { + m.deprecated_config.post_hook = + dbt_yaml::Verbatim::from(Some(Hooks::String("select 1".to_string()))) + }, + |m| m.deprecated_config.sql_header = Some("set x=1;".to_string()), + |m| m.deprecated_config.on_error = Some(OnError::Continue), + |m| { + m.deprecated_config + .__warehouse_specific_config__ + .row_access_policy = Some("policy on (id)".to_string()) + }, + |m| { + m.deprecated_config.__warehouse_specific_config__.table_tag = + Some("tag='value'".to_string()) + }, + |m| m.deprecated_config.__warehouse_specific_config__.copy_tags = Some(true), + |m| { + m.__model_attr__.constraints.push( + dbt_schemas::schemas::properties::ModelConstraint { + type_: ConstraintType::Custom, + expression: Some("check (id > 0)".to_string()), + ..Default::default() + }, + ) + }, + ]; + for change in cases { + let mut model = wap_model(); + change(&mut model); + let error = validate_wap_model(&model, AdapterType::Snowflake).unwrap_err(); + assert_eq!(error.code, ErrorCode::InvalidConfig); + assert!(error.context.contains("wap=true")); + } + } + fn create_test_model_properties(name: &str) -> ModelProperties { ModelProperties { name: name.to_string(), diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md b/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md new file mode 100644 index 00000000000..a0883cf94c8 --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md @@ -0,0 +1,60 @@ +# Snowflake WAP live acceptance + +This fixture exercises a built dbt binary against an existing Snowflake profile, +database, and schema. It uses Python 3.9+ standard-library modules and the +project's own `run-operation` assertion macros; no Python dbt or Snowflake driver +is needed. + +From the repository root: + +```sh +python3 crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py \ + --dbt-bin target/debug/dbt \ + --profiles-dir ~/.dbt \ + --profile YOUR_PROFILE \ + --target YOUR_SNOWFLAKE_TARGET +``` + +The selected role needs its usual warehouse access plus permission to create, +read, replace, and drop tables in the existing target schema. The runner does not +create or change credentials, profiles, databases, or schemas. All warehouse +work is explicitly invoked by this command. Transformation, audit queries, and +retained tables consume ordinary Snowflake resources. + +The default run covers both permanent and transient WAP tables. Use +`--table-kind permanent` or `--table-kind transient` for one variant. Each variant +checks: + +- An error-severity audit fails; the public sentinel stays `[999]`, the candidate + retains `[-1, 2]`, and the downstream table is not created. +- A subsequent standalone `dbt test` passes against the public sentinel, despite + the retained failing candidate and compiled SQL from the preceding build. +- `dbt retry` creates a different candidate, reruns all three audits including + the two that previously passed, publishes `[2, 3]`, and builds the downstream + through its ordinary public `ref`. +- A warning-severity audit also prevents publication and retains its candidate. +- A successful first build publishes exactly `[1, 2]` and drops its candidate. +- A failing first build leaves the public table absent. + +Assertions check actual warehouse contents, `run_results.json`, and compiled +downstream SQL. Each model must have one canonical result. Audit SQL includes +both generic `not_null`/`unique` tests and a singular test. Builds use four threads +and do not use fail-fast. + +Every run uses a copied temporary project and UUID-based public aliases in the +target schema. The runner first verifies those public aliases are absent. It +records candidate ownership only from WAP's post-preflight staging message, so a +pre-existing collision reported as “retained if created” is never treated as a +table the runner owns. It never scans or deletes tables using a wildcard. + +Successful and failed runs clean up only their recorded public tables and exact +candidate identifiers. `--keep-objects` retains them for inspection. Logs, +per-command artifacts, saved retry state, and `owned_objects.json` stay in the +temporary directory printed at startup. On interruption or incomplete cleanup, +use that inventory with `wap_fixture_cleanup(identifiers=...)` and the recorded +variables/profile/target in the copied project. In particular, do not drop every +table whose name begins with `__DBT_WAP_`. + +The fixture has no CI credential assumptions and is not part of ordinary local +unit tests. A successful offline parse or unit-test run does not count as this +live acceptance test passing. diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/dbt_project.yml b/crates/dbt-sa-cli/tests/data/snowflake_wap/dbt_project.yml new file mode 100644 index 00000000000..f0258233149 --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/dbt_project.yml @@ -0,0 +1,9 @@ +name: wap_live_acceptance +version: '1.0' +config-version: 2 +profile: wap_live_acceptance +model-paths: [models] +test-paths: [tests] +macro-paths: [macros] +quoting: + identifier: true diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql b/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql new file mode 100644 index 00000000000..2271036a26d --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql @@ -0,0 +1,41 @@ +{% macro wap_fixture_relation(identifier) %} + {% if target.type != 'snowflake' %} + {{ exceptions.raise_compiler_error('This acceptance fixture requires a Snowflake target') }} + {% endif %} + {% set public = ref('orders') %} + {{ return(public.incorporate(path={'identifier': identifier})) }} +{% endmacro %} + +{% macro wap_fixture_assert(identifier, present=true, expected_ids=none) %} + {% set relation = wap_fixture_relation(identifier) %} + {% set actual = adapter.get_relation(database=relation.database, schema=relation.schema, identifier=relation.identifier) %} + {% if (actual is not none) != present %} + {{ exceptions.raise_compiler_error('Unexpected presence for ' ~ relation ~ ': expected present=' ~ present) }} + {% endif %} + {% if present and expected_ids is not none %} + {% set result = run_query('select id from ' ~ relation ~ ' order by id') %} + {% set actual_ids = [] %} + {% for row in result.rows %} + {% do actual_ids.append(row[0] | int) %} + {% endfor %} + {% if actual_ids != expected_ids %} + {{ exceptions.raise_compiler_error('Unexpected rows in ' ~ relation ~ ': ' ~ actual_ids ~ ', expected ' ~ expected_ids) }} + {% endif %} + {% endif %} +{% endmacro %} + +{% macro wap_fixture_prepare() %} + {% set relation = ref('orders') %} + {% do wap_fixture_assert(relation.identifier, present=false) %} + {% do run_query('create ' ~ ('transient ' if var('transient', false) else '') ~ 'table ' ~ relation ~ ' as select 999::integer as id') %} +{% endmacro %} + +{% macro wap_fixture_cleanup(identifiers) %} + {% for identifier in identifiers %} + {% if not identifier.startswith(var('prefix') ~ '_') and not identifier.startswith('__DBT_WAP_') %} + {{ exceptions.raise_compiler_error('Refusing unexpected cleanup identifier: ' ~ identifier) }} + {% endif %} + {% set relation = wap_fixture_relation(identifier) %} + {% do run_query('drop table if exists ' ~ relation) %} + {% endfor %} +{% endmacro %} diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/models/downstream.sql b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/downstream.sql new file mode 100644 index 00000000000..478d00fa61b --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/downstream.sql @@ -0,0 +1,3 @@ +{{ config(materialized='table', alias=var('prefix') ~ '_DOWNSTREAM') }} + +select id from {{ ref('orders') }} diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql new file mode 100644 index 00000000000..a80f8df62de --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql @@ -0,0 +1,10 @@ +{{ config( + materialized='table', + wap=true, + alias=var('prefix') ~ '_ORDERS', + transient=var('transient', false) +) }} + +select {{ var('audit_value', 1) }}::integer as id +union all +select 2::integer as id diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/models/schema.yml b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/schema.yml new file mode 100644 index 00000000000..cf5d1f326a6 --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/schema.yml @@ -0,0 +1,8 @@ +version: 2 +models: + - name: orders + columns: + - name: id + data_tests: + - not_null + - unique diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/tests/nonnegative.sql b/crates/dbt-sa-cli/tests/data/snowflake_wap/tests/nonnegative.sql new file mode 100644 index 00000000000..558f76a2d28 --- /dev/null +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/tests/nonnegative.sql @@ -0,0 +1,3 @@ +{{ config(severity=var('audit_severity', 'error')) }} + +select * from {{ ref('orders') }} where id < 0 diff --git a/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py b/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py new file mode 100644 index 00000000000..d6249b1f30a --- /dev/null +++ b/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +"""Run WAP acceptance checks against an explicitly selected Snowflake target.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import re +import shutil +import subprocess +import sys +import tempfile +from typing import Any +import uuid + + +PROJECT = "wap_live_acceptance" +MODEL_ID = f"model.{PROJECT}.orders" +DOWNSTREAM_ID = f"model.{PROJECT}.downstream" +STAGED_CANDIDATE = re.compile( + r"WAP: building[^\r\n]*?in working table[^\r\n]*?" + r"(__DBT_WAP_[0-9A-F]{32}_[0-9A-F]{32})" +) + + +def require(condition: bool, message: str) -> None: + if not condition: + raise RuntimeError(message) + + +class Acceptance: + def __init__(self, args: argparse.Namespace) -> None: + self.args = args + self.work = Path(tempfile.mkdtemp(prefix="dbt-snowflake-wap-")) + self.project = self.work / "project" + fixture = Path(__file__).parent / "data" / "snowflake_wap" + shutil.copytree(fixture, self.project) + self.artifacts = self.project / "target" + self.sequence = 0 + self.variables: dict[str, Any] = {} + self.owned: dict[str, dict[str, Any]] = {} + self.token = uuid.uuid4().hex.upper() + print(f"Acceptance logs and artifacts: {self.work}", flush=True) + + def persist_owned(self) -> None: + # An exact inventory also supports manual cleanup after interruption. + (self.work / "owned_objects.json").write_text( + json.dumps(list(self.owned.values()), indent=2) + "\n" + ) + + def invoke( + self, command: list[str], *, success: bool = True, results: bool = False, + staged: bool = False, + ) -> tuple[dict[str, Any], set[str], Path]: + self.sequence += 1 + capture = self.work / f"{self.sequence:02d}_{command[0]}" + capture.mkdir() + result_path = self.artifacts / "run_results.json" + result_path.unlink(missing_ok=True) + argv = [ + str(self.args.dbt_bin), + *command, + "--project-dir", str(self.project), + "--profiles-dir", str(self.args.profiles_dir), + "--profile", self.args.profile, + "--target", self.args.target, + "--target-path", str(self.artifacts), + "--vars", json.dumps(self.variables), + ] + environment = dict(os.environ) + environment.update(DBT_QUIET="false", DBT_USE_COLORS="false") + (capture / "command.json").write_text(json.dumps(argv, indent=2) + "\n") + completed = subprocess.run( + argv, cwd=self.project, env=environment, text=True, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + (capture / "stdout.log").write_text(completed.stdout) + (capture / "stderr.log").write_text(completed.stderr) + output = completed.stdout + "\n" + completed.stderr + candidates = set(STAGED_CANDIDATE.findall(output)) + # Only the post-preflight "building" message establishes ownership. + # A collision error's "retained if created" message does not. + prefix = self.variables.get("prefix") + if prefix in self.owned: + identifiers = self.owned[prefix]["identifiers"] + identifiers[:] = sorted(set(identifiers) | candidates) + self.persist_owned() + for filename in ("run_results.json", "manifest.json"): + artifact = self.artifacts / filename + if artifact.exists(): + shutil.copy2(artifact, capture / filename) + require( + (completed.returncode == 0) == success, + f"Unexpected exit {completed.returncode}: {' '.join(command)}; see {capture}", + ) + if not results: + return {}, candidates, capture + require(result_path.exists(), f"Missing run_results.json; see {capture}") + artifact = json.loads(result_path.read_text()) + invocation = uuid.UUID(artifact["metadata"]["invocation_id"]).hex.upper() + require( + all(name.startswith(f"__DBT_WAP_{invocation}_") for name in candidates), + f"Candidate does not belong to this invocation; see {capture}", + ) + require( + len(candidates) == int(staged), + f"Expected {int(staged)} staged candidates; see {capture}", + ) + return artifact, candidates, capture + + def operation(self, name: str, **kwargs: Any) -> None: + self.invoke(["run-operation", name, "--args", json.dumps(kwargs)]) + + def check(self, identifier: str, *, present: bool = True, ids: Any = None) -> None: + self.operation( + "wap_fixture_assert", identifier=identifier, + present=present, expected_ids=ids, + ) + + def begin(self, kind: str, scenario: str, *, sentinel: bool) -> None: + prefix = f"DBT_WAP_ACCEPT_{self.token}_{kind.upper()}_{scenario.upper()}" + self.variables = { + "prefix": prefix, "transient": kind == "transient", + "audit_value": -1, "audit_severity": "error", + } + identifiers = [f"{prefix}_ORDERS", f"{prefix}_DOWNSTREAM"] + for identifier in identifiers: + self.check(identifier, present=False) + # Register only names verified absent in the selected schema. + self.owned[prefix] = { + "profile": self.args.profile, "target": self.args.target, + "variables": dict(self.variables), "identifiers": identifiers, + } + self.persist_owned() + if sentinel: + self.operation("wap_fixture_prepare") + + def public(self, suffix: str = "ORDERS") -> str: + return f"{self.variables['prefix']}_{suffix}" + + def build(self, *, success: bool) -> tuple[dict[str, Any], set[str], Path]: + return self.invoke( + ["build", "--select", "orders+", "--threads", "4"], + success=success, results=True, staged=True, + ) + + @staticmethod + def verify_audits(artifact: dict[str, Any], *, verdict: str) -> set[str]: + test_rows = [row for row in artifact["results"] if row["unique_id"].startswith("test.")] + tests = {row["unique_id"]: row for row in test_rows} + require(len(tests) == len(test_rows), "Expected one result per audit") + require(len(tests) == 3, f"Expected all three audits, got {list(tests)}") + singular_id = f"test.{PROJECT}.nonnegative" + require(singular_id in tests, "Missing nonnegative singular audit") + for unique_id, row in tests.items(): + expected = verdict if unique_id == singular_id else "pass" + require(row["status"] == expected, f"Unexpected audit result: {row}") + return set(tests) + + @staticmethod + def verify_results(artifact: dict[str, Any], *, verdict: str) -> set[str]: + rows = artifact["results"] + model_rows = [row for row in rows if row["unique_id"] == MODEL_ID] + require(len(model_rows) == 1, "Expected exactly one canonical model result") + expected_model = {"pass": "success", "fail": "skipped", "warn": "error"}[verdict] + require(model_rows[0]["status"] == expected_model, f"Unexpected model result: {model_rows}") + audits = Acceptance.verify_audits(artifact, verdict=verdict) + downstream = [row for row in rows if row["unique_id"] == DOWNSTREAM_ID] + expected_downstream = "success" if verdict == "pass" else "skipped" + require( + len(downstream) == 1 and downstream[0]["status"] == expected_downstream, + f"Unexpected downstream result: {downstream}", + ) + return audits + + def verify_published(self, ids: list[int], candidates: set[str]) -> None: + self.check(self.public(), ids=ids) + self.check(self.public("DOWNSTREAM"), ids=ids) + for candidate in candidates: + self.check(candidate, present=False) + compiled = list((self.artifacts / "compiled").rglob("downstream.sql")) + require(len(compiled) == 1, "Expected compiled downstream SQL") + sql = compiled[0].read_text() + require(self.public() in sql, "Downstream ref did not use the public relation") + require("__DBT_WAP_" not in sql, "Candidate relation leaked into downstream SQL") + + def run_kind(self, kind: str) -> None: + print(f"Checking {kind}: failure, retry, warning, first publish, first failure", flush=True) + self.begin(kind, "fail_retry", sentinel=True) + failed, retained, state = self.build(success=False) + audits = self.verify_results(failed, verdict="fail") + self.check(self.public(), ids=[999]) + self.check(self.public("DOWNSTREAM"), present=False) + for candidate in retained: + self.check(candidate, ids=[-1, 2]) + + # A standalone test must use the public sentinel even with failing + # candidate SQL left in the same target directory by the prior build. + standalone, _, _ = self.invoke( + ["test", "--select", "orders", "--threads", "4"], results=True, + ) + require( + self.verify_audits(standalone, verdict="pass") == audits, + "Standalone test omitted an audit", + ) + require(len(standalone["results"]) == 3, "Standalone test unexpectedly rebuilt a model") + self.check(self.public(), ids=[999]) + for candidate in retained: + self.check(candidate, ids=[-1, 2]) + + # The two generic audits passed in the failed build. They must appear + # again when retry builds and certifies a new candidate. + self.variables["audit_value"] = 3 + retried, new_candidates, _ = self.invoke( + ["retry", "--state", str(state), "--threads", "4"], results=True, staged=True, + ) + require(self.verify_results(retried, verdict="pass") == audits, "Retry omitted audits") + require(retained.isdisjoint(new_candidates), "Retry reused a failed candidate") + self.verify_published([2, 3], new_candidates) + for candidate in retained: + self.check(candidate, ids=[-1, 2]) + + self.begin(kind, "warning", sentinel=True) + self.variables["audit_severity"] = "warn" + warned, candidates, _ = self.build(success=False) + self.verify_results(warned, verdict="warn") + self.check(self.public(), ids=[999]) + self.check(self.public("DOWNSTREAM"), present=False) + for candidate in candidates: + self.check(candidate, ids=[-1, 2]) + + self.begin(kind, "first_publish", sentinel=False) + self.variables["audit_value"] = 1 + passed, candidates, _ = self.build(success=True) + self.verify_results(passed, verdict="pass") + self.verify_published([1, 2], candidates) + + self.begin(kind, "first_failure", sentinel=False) + failed, candidates, _ = self.build(success=False) + self.verify_results(failed, verdict="fail") + self.check(self.public(), present=False) + self.check(self.public("DOWNSTREAM"), present=False) + for candidate in candidates: + self.check(candidate, ids=[-1, 2]) + + def cleanup(self) -> None: + failures = [] + for entry in self.owned.values(): + self.variables = entry["variables"] + try: + self.operation("wap_fixture_cleanup", identifiers=entry["identifiers"]) + except Exception as error: + failures.append(str(error)) + require(not failures, "Cleanup incomplete: " + "; ".join(failures)) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dbt-bin", type=Path, required=True, help="Built binary with WAP changes") + parser.add_argument("--profile", required=True) + parser.add_argument("--target", required=True) + parser.add_argument("--profiles-dir", type=Path, default=Path.home() / ".dbt") + parser.add_argument("--table-kind", choices=("both", "permanent", "transient"), default="both") + parser.add_argument("--keep-objects", action="store_true", help="Retain test-owned tables for inspection") + args = parser.parse_args() + args.dbt_bin = args.dbt_bin.expanduser().resolve() + args.profiles_dir = args.profiles_dir.expanduser().resolve() + require(args.dbt_bin.is_file(), f"No dbt binary: {args.dbt_bin}") + require((args.profiles_dir / "profiles.yml").is_file(), "Missing profiles.yml") + fixture = Acceptance(args) + passed = False + try: + kinds = ("permanent", "transient") if args.table_kind == "both" else (args.table_kind,) + for kind in kinds: + fixture.run_kind(kind) + passed = True + finally: + if args.keep_objects: + print(f"Objects retained; exact inventory: {fixture.work / 'owned_objects.json'}") + else: + fixture.cleanup() + if passed: + print(f"PASS: Snowflake WAP acceptance checks; evidence: {fixture.work}") + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except Exception as error: + print(f"FAIL: {error}", file=sys.stderr) + sys.exit(1) diff --git a/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py b/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py new file mode 100644 index 00000000000..176c33721fe --- /dev/null +++ b/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py @@ -0,0 +1,70 @@ +"""Offline checks for the live acceptance runner's evidence handling.""" + +import unittest + +from snowflake_wap_acceptance import ( + Acceptance, + DOWNSTREAM_ID, + MODEL_ID, + PROJECT, + STAGED_CANDIDATE, +) + + +class AcceptanceEvidenceTests(unittest.TestCase): + def artifact(self, verdict): + return {"results": [ + {"unique_id": MODEL_ID, + "status": {"pass": "success", "fail": "skipped", "warn": "error"}[verdict]}, + {"unique_id": DOWNSTREAM_ID, + "status": "success" if verdict == "pass" else "skipped"}, + {"unique_id": f"test.{PROJECT}.nonnegative", "status": verdict}, + {"unique_id": f"test.{PROJECT}.not_null_orders_id.abc", "status": "pass"}, + {"unique_id": f"test.{PROJECT}.unique_orders_id.def", "status": "pass"}, + ]} + + def test_expected_verdicts(self): + for verdict in ("pass", "fail", "warn"): + with self.subTest(verdict=verdict): + audits = Acceptance.verify_results(self.artifact(verdict), verdict=verdict) + self.assertEqual(len(audits), 3) + + def test_missing_previously_passed_audit_is_rejected(self): + artifact = self.artifact("pass") + artifact["results"].pop() + with self.assertRaisesRegex(RuntimeError, "all three audits"): + Acceptance.verify_results(artifact, verdict="pass") + + def test_duplicate_canonical_result_is_rejected(self): + artifact = self.artifact("pass") + artifact["results"].append(dict(artifact["results"][0])) + with self.assertRaisesRegex(RuntimeError, "one canonical model result"): + Acceptance.verify_results(artifact, verdict="pass") + + def test_duplicate_audit_cannot_hide_a_failure(self): + artifact = self.artifact("fail") + duplicate = dict(artifact["results"][2], status="pass") + artifact["results"].append(duplicate) + with self.assertRaisesRegex(RuntimeError, "one result per audit"): + Acceptance.verify_audits(artifact, verdict="pass") + + def test_standalone_tests_validate_only_audit_results(self): + artifact = self.artifact("pass") + artifact["results"] = artifact["results"][2:] + self.assertEqual(len(Acceptance.verify_audits(artifact, verdict="pass")), 3) + artifact["results"][0]["status"] = "fail" + with self.assertRaisesRegex(RuntimeError, "Unexpected audit result"): + Acceptance.verify_audits(artifact, verdict="pass") + + def test_collision_notice_does_not_establish_candidate_ownership(self): + candidate = "__DBT_WAP_" + "A" * 32 + "_" + "B" * 32 + self.assertEqual(STAGED_CANDIDATE.findall( + f"WAP working table retained if created: DB.SCHEMA.{candidate}" + ), []) + self.assertEqual(STAGED_CANDIDATE.findall( + f'WAP: building "DB"."SCHEMA"."PUBLIC" in working table "DB"."SCHEMA"."{candidate}"' + ), [candidate]) + + +if __name__ == "__main__": + unittest.main() diff --git a/crates/dbt-schemas/src/schemas/manifest/manifest_nodes.rs b/crates/dbt-schemas/src/schemas/manifest/manifest_nodes.rs index a0d64e33bfb..6a506bc5ba1 100644 --- a/crates/dbt-schemas/src/schemas/manifest/manifest_nodes.rs +++ b/crates/dbt-schemas/src/schemas/manifest/manifest_nodes.rs @@ -840,6 +840,12 @@ pub struct ManifestModel { } #[derive(Deserialize, Serialize, Debug, Default, Clone, PartialEq, DbtSchema)] pub struct ManifestModelConfig { + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::schemas::serde::strict_bool_or_string_bool" + )] + pub wap: Option, #[serde(default, deserialize_with = "bool_or_string_bool")] pub enabled: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -1155,6 +1161,7 @@ impl From for SeedConfig { impl From for ManifestModelConfig { fn from(config: ModelConfig) -> Self { Self { + wap: config.wap, enabled: config.enabled, compute: config.compute, alias: config.alias, @@ -1233,6 +1240,7 @@ impl From for ManifestModelConfig { impl From for ModelConfig { fn from(config: ManifestModelConfig) -> Self { Self { + wap: config.wap, enabled: config.enabled, alias: config.alias, database: config.database, diff --git a/crates/dbt-schemas/src/schemas/nodes.rs b/crates/dbt-schemas/src/schemas/nodes.rs index 440538f77ae..da4e1b200df 100644 --- a/crates/dbt-schemas/src/schemas/nodes.rs +++ b/crates/dbt-schemas/src/schemas/nodes.rs @@ -5552,6 +5552,83 @@ pub struct DbtModel { } impl DbtModel { + /// Validate the model options supported by table write-audit-publish. + /// Catalog and materialization macro resolution are checked again at execution. + pub fn validate_wap_config(&self, adapter: AdapterType) -> FsResult<()> { + use crate::schemas::common::{ConstraintType, OnError}; + + let config = &self.deprecated_config; + if !config.wap.unwrap_or(false) { + return Ok(()); + } + + let warehouse = &config.__warehouse_specific_config__; + let has_hooks = [&config.pre_hook, &config.post_hook].iter().any(|hooks| { + hooks.as_ref().as_ref().is_some_and(|hooks| { + hooks.to_hook_config_array().iter().any(|hook| { + hook.sql + .as_deref() + .is_some_and(|sql| !sql.trim().is_empty()) + }) + }) + }); + let has_custom_constraints = self + .__model_attr__ + .constraints + .iter() + .chain(config.constraints.iter().flatten()) + .any(|constraint| constraint.type_ == ConstraintType::Custom) + || self.__base_attr__.columns.iter().any(|column| { + column + .constraints + .iter() + .any(|constraint| constraint.type_ == ConstraintType::Custom) + }); + + let unsupported = if adapter != AdapterType::Snowflake + || self.__common_attr__.language.as_deref() != Some("sql") + || self.__base_attr__.materialized != DbtMaterialization::Table + { + Some("requires a Snowflake SQL model materialized as table") + } else if config + .table_format + .as_deref() + .is_some_and(|format| !format.eq_ignore_ascii_case("default")) + || warehouse.external_volume.is_some() + || warehouse.base_location_root.is_some() + || warehouse.base_location_subpath.is_some() + { + Some("supports only native Snowflake tables, not Iceberg or external catalogs") + } else if has_hooks + || config + .sql_header + .as_deref() + .is_some_and(|sql| !sql.trim().is_empty()) + { + Some("does not support model pre-hooks, post-hooks, or sql_header") + } else if config.on_error == Some(OnError::Continue) { + Some("does not support on_error: continue") + } else if warehouse.row_access_policy.is_some() + || warehouse.table_tag.is_some() + || warehouse.copy_tags.unwrap_or(false) + || has_custom_constraints + { + Some("does not support row_access_policy, table_tag, copy_tags, or custom constraints") + } else { + None + }; + + if let Some(reason) = unsupported { + return Err(dbt_common::fs_err!( + ErrorCode::InvalidConfig, + "Model '{}': wap=true {}", + self.__common_attr__.unique_id, + reason + )); + } + Ok(()) + } + /// Transcribes dbt-core's `ModelNode.same_ref_representation` (dbt-mantle /// `core/dbt/contracts/graph/nodes.py:684-691`), which is ANDed into `ModelNode.same_contents` /// at `:677-682`: diff --git a/crates/dbt-schemas/src/schemas/prev_state/mod.rs b/crates/dbt-schemas/src/schemas/prev_state/mod.rs index e6ea5ee0e0b..559fbe2770c 100644 --- a/crates/dbt-schemas/src/schemas/prev_state/mod.rs +++ b/crates/dbt-schemas/src/schemas/prev_state/mod.rs @@ -2628,6 +2628,11 @@ mod tests { use dbt_yaml::{Spanned, Verbatim}; let cases: Vec<(&str, ExcludeKind, Box)> = vec![ + ( + "wap", + ExcludeKind::Relevant, + Box::new(|n| n.deprecated_config.wap = Some(true)), + ), // --- fields `ModelConfig::same_config` actually compares --- ( "enabled", diff --git a/crates/dbt-schemas/src/schemas/project/configs/model_config.rs b/crates/dbt-schemas/src/schemas/project/configs/model_config.rs index 8038bb61d68..89052165d88 100644 --- a/crates/dbt-schemas/src/schemas/project/configs/model_config.rs +++ b/crates/dbt-schemas/src/schemas/project/configs/model_config.rs @@ -340,6 +340,12 @@ pub struct ProjectModelConfig { pub matched_condition: Option, #[serde(rename = "+materialized")] pub materialized: Option, + #[serde( + default, + rename = "+wap", + deserialize_with = "crate::schemas::serde::strict_bool_or_string_bool" + )] + pub wap: Option, #[serde(rename = "+max_staleness")] pub max_staleness: Option, #[serde( @@ -673,7 +679,8 @@ impl TypedRecursiveConfig for ProjectModelConfig { } fn has_set_fields(&self) -> bool { - self.access.is_some() + self.wap.is_some() + || self.access.is_some() || self.adapter_properties.is_some() || self.alias.is_some() || self.automatic_clustering.is_some() @@ -870,6 +877,13 @@ pub struct ModelConfig { pub group: Option, #[resolved(promote, default = DbtMaterialization::View)] pub materialized: Option, + /// Audit a Snowflake SQL table candidate before publishing it during `build`. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::schemas::serde::strict_bool_or_string_bool" + )] + pub wap: Option, pub incremental_strategy: Option, pub incremental_predicates: Option, // Hashed into the run-cache key only; the functional source of a model's constraints is the @@ -1007,6 +1021,7 @@ impl From for ModelConfig { location: config.location, lookback: config.lookback, materialized: config.materialized, + wap: config.wap, merge_exclude_columns: config.merge_exclude_columns, merge_update_columns: config.merge_update_columns, meta: config.meta.0, @@ -1210,6 +1225,7 @@ impl From for ProjectModelConfig { location: config.location, lookback: config.lookback, materialized: config.materialized, + wap: config.wap, merge_exclude_columns: config.merge_exclude_columns, merge_update_columns: config.merge_update_columns, meta: config.meta.into(), @@ -1514,6 +1530,7 @@ impl ModelConfig { let propagate_eq = self.propagate == other.propagate; let meta_eq_result = meta_eq(&self.meta, &other.meta); // Custom comparison for meta let materialized_eq_result = materialized_eq(&self.materialized, &other.materialized); + let wap_eq = self.wap.unwrap_or(false) == other.wap.unwrap_or(false); let incremental_strategy_eq = self.incremental_strategy == other.incremental_strategy; // incremental_predicates can differ because of environment, i.e. dev vs prod // so we don't compare them. To compare them we will need a SQL AST whose @@ -1575,6 +1592,7 @@ impl ModelConfig { && propagate_eq && meta_eq_result && materialized_eq_result + && wap_eq && incremental_strategy_eq && batch_size_eq && lookback_eq_result @@ -1874,6 +1892,7 @@ impl ModelConfig { format!("{:?}", &other.predicates), )), ), + ("wap", wap_eq, None), ("warehouse_config", warehouse_config_eq, None), ], ); @@ -1918,6 +1937,8 @@ impl ConfigKeys for ModelConfig { field_names.insert("dataset".to_string()); // alias for schema field_names.insert("post-hook".to_string()); // might be serialized as post_hook field_names.insert("pre-hook".to_string()); // might be serialized as pre_hook + // Optional extension omitted from default serialized configs. + field_names.insert("wap".to_string()); field_names } @@ -2062,6 +2083,45 @@ mod tests { use crate::schemas::serde::{AdapterTypeOrArray, RefreshableConfig, StringOrArrayOfStrings}; use dbt_adapter_core::AdapterType; + #[test] + fn wap_config_inherits_overrides_and_round_trips() { + let project: ProjectModelConfig = + dbt_yaml::from_str("+wap: true\n__additional_properties__: {}\n").unwrap(); + let parent: ModelConfig = project.into(); + let mut inherited = ModelConfig::default(); + inherited.default_to(&parent); + assert_eq!(inherited.wap, Some(true)); + + let mut child: ModelConfig = + dbt_yaml::from_str("wap: false\n__warehouse_specific_config__: {}\n").unwrap(); + child.default_to(&parent); + assert_eq!(child.wap, Some(false)); + let manifest: ManifestModelConfig = inherited.into(); + let serialized = serde_json::to_string(&manifest).unwrap(); + let manifest: ManifestModelConfig = serde_json::from_str(&serialized).unwrap(); + let restored: ModelConfig = manifest.into(); + assert_eq!(restored.wap, Some(true)); + let project: ProjectModelConfig = restored.into(); + assert_eq!(project.wap, Some(true)); + } + + #[test] + fn wap_is_a_known_model_config_and_changes_state_only_when_enabled() { + use super::ConfigKeys; + + assert!(ModelConfig::valid_field_names().contains("wap")); + let omitted = ModelConfig::default(); + let mut configured = omitted.clone(); + configured.wap = Some(false); + assert!(omitted.same_config(&configured)); + configured.wap = Some(true); + assert!(!omitted.same_config(&configured)); + assert!( + dbt_yaml::from_str::("wap: invalid\n__warehouse_specific_config__: {}\n") + .is_err() + ); + } + /// `+propagate` rides the same project -> node -> project path `+adapter` /// does, in both its single-value and list forms. #[test] diff --git a/crates/dbt-schemas/src/schemas/serde.rs b/crates/dbt-schemas/src/schemas/serde.rs index faa6838da4a..e754af61b75 100644 --- a/crates/dbt-schemas/src/schemas/serde.rs +++ b/crates/dbt-schemas/src/schemas/serde.rs @@ -252,6 +252,21 @@ where .or_else(|| value.as_str().map(|s| s.to_lowercase() == "true"))) } +/// Deserialize a boolean flag without silently treating unknown values as false. +pub fn strict_bool_or_string_bool<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let value = dbt_yaml::Value::deserialize(deserializer)?; + match value { + dbt_yaml::Value::Bool(value, _) => Ok(Some(value)), + dbt_yaml::Value::Null(_) => Ok(None), + dbt_yaml::Value::String(value, _) if value.eq_ignore_ascii_case("true") => Ok(Some(true)), + dbt_yaml::Value::String(value, _) if value.eq_ignore_ascii_case("false") => Ok(Some(false)), + _ => Err(de::Error::custom("expected true, false, or null")), + } +} + pub fn bool_or_string_bool_default<'de, D>(deserializer: D) -> Result where D: Deserializer<'de>, diff --git a/crates/dbt-tasks-core/src/context.rs b/crates/dbt-tasks-core/src/context.rs index ad5b405e172..e1b3a491437 100644 --- a/crates/dbt-tasks-core/src/context.rs +++ b/crates/dbt-tasks-core/src/context.rs @@ -12,17 +12,18 @@ use crate::run_cache::run_cache_service::{ use arrow::array::RecordBatch; use arrow_schema::SchemaRef; use dbt_adapter::AdapterStore; -use dbt_adapter::relation::create_relation_from_node; +use dbt_adapter::relation::{RelationObject, create_relation_from_node}; use dbt_adapter::response::AdapterResponse; use dbt_adapter_core::AdapterType; -use dbt_common::FsResult; use dbt_common::collections::{DashMap, SccHashMap}; use dbt_common::io_args::OptimizeTestsOptions; use dbt_common::path::DbtPath; use dbt_common::stats::{NodeStatus, Stat}; +use dbt_common::{ErrorCode, FsError, FsResult, fs_err}; use dbt_dag::schedule::Schedule; use dbt_frontend_common::sources_extractor::SourcesExtractor; use dbt_jinja_utils::jinja_environment::{JinjaEnv, adapter_api_value}; +use dbt_jinja_utils::phases::apply_ref_relation_overrides; use dbt_jinja_utils::phases::compile::{ DependencyValidationConfig, build_compile_node_context_inner, }; @@ -46,6 +47,7 @@ use crate::task::Task; use crate::test_aggregation::{GenericTestRelationships, is_data_test_static_analysis_eligible}; use crate::unit_test_schema::UnitTestSchemaState; use crate::visitor::SkipReason; +use crate::wap::WapPlan; use dbt_schemas::schemas::common::DbtMaterialization; @@ -162,6 +164,9 @@ pub struct TaskRunnerCtxInner { pub base_context: BTreeMap, pub analyze_stats: DashMap, pub run_stats: DashMap, + /// Candidate builds are provisional until their publication task succeeds. + pub wap_stage_stats: DashMap, + pub wap_plan: WapPlan, pub data_test_execution_results: DashMap, pub batch_results_map: DashMap, pub main_adapter_responses: DashMap, @@ -220,7 +225,22 @@ impl TaskRunnerCtxInner { adapter_store: Arc, sources_extractor: Arc, run_cache_ctx: RunCacheCtx, - ) -> Self { + ) -> FsResult { + let wap_plan = WapPlan::build(&arg, &schedule, &resolver_state.nodes)?; + if !wap_plan.models.is_empty() + && adapter_store.default_adapter_type() != AdapterType::Snowflake + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP requires a Snowflake profile target; a per-model adapter override is not supported" + )); + } + if !wap_plan.models.is_empty() { + wap_plan.validate_test_storage_relations( + &resolver_state.nodes, + adapter_store.default_adapter()?.engine().quoting(), + )?; + } let runnable_set = schedule .selected_nodes .iter() @@ -249,7 +269,7 @@ impl TaskRunnerCtxInner { let infer_schema_registry = Arc::new(dbt_common::infer_schema_registry::InferSchemaRegistry::new()); - TaskRunnerCtxInner { + Ok(TaskRunnerCtxInner { arg, worker_id, schedule, @@ -257,6 +277,8 @@ impl TaskRunnerCtxInner { base_context, analyze_stats: DashMap::default(), run_stats: DashMap::default(), + wap_stage_stats: DashMap::default(), + wap_plan, data_test_execution_results: DashMap::default(), batch_results_map, main_adapter_responses: DashMap::default(), @@ -282,7 +304,7 @@ impl TaskRunnerCtxInner { preview_error: parking_lot::Mutex::new(None), unit_test_schema: UnitTestSchemaState::default(), run_cache_ctx, - } + }) } pub fn span_manager(&self) -> Arc, SkipReason>> { @@ -409,6 +431,9 @@ impl TaskRunnerCtx { } pub async fn is_data_test_statically_skippable(&self, unique_id: &str) -> bool { + if self.inner.wap_plan.audit_owner(unique_id).is_some() { + return false; + } if !self .inner .arg @@ -556,6 +581,23 @@ impl TaskRunnerCtx { ref_validation_config, )?; + // A candidate is a scoped copy of the model, while the resolver still + // stores the public relation. Bind `this` from that copy explicitly. + // The alias check keeps unit-test contexts for the canonical model unchanged. + if let Some(entry) = self + .inner + .wap_plan + .model(&model.common().unique_id) + .filter(|entry| entry.candidate_identifier == model.base().alias) + { + let relation = entry.candidate_relation()?; + ctx.insert( + "this".to_string(), + RelationObject::new(relation.into()).into_value(), + ); + } + self.apply_wap_ref_overrides(&model.common().unique_id, &mut ctx)?; + // `adapter` and `api` are environment globals bound to a single adapter, // and the environment is shared by every render. Shadow them in this // node's own context -- the same context-then-globals lookup `DIALECT` @@ -574,6 +616,28 @@ impl TaskRunnerCtx { Ok((ctx, config_map)) } + /// Route an audit's refs to its candidate without changing other node contexts. + /// Runtime materialization contexts use this too, so macro-time refs agree + /// with the relations used while rendering the audit SQL. + pub fn apply_wap_ref_overrides( + &self, + node_unique_id: &str, + context: &mut BTreeMap, + ) -> FsResult<()> { + let Some(entry) = self.inner.wap_plan.audit_owner(node_unique_id) else { + return Ok(()); + }; + let relation = entry.candidate_relation()?; + let overrides = BTreeMap::from([( + entry.model.common().unique_id.clone(), + RelationObject::new(relation.into()).into_value(), + )]); + apply_ref_relation_overrides(context, Arc::new(overrides)).map_err(|error| { + FsError::from_jinja_err(error, "binding WAP audit references to the candidate") + })?; + Ok(()) + } + pub fn on_test_failure( &self, node: &Arc, diff --git a/crates/dbt-tasks-core/src/context_factory.rs b/crates/dbt-tasks-core/src/context_factory.rs index 1584fd646c1..3b906e06f7f 100644 --- a/crates/dbt-tasks-core/src/context_factory.rs +++ b/crates/dbt-tasks-core/src/context_factory.rs @@ -169,7 +169,7 @@ pub trait TaskRunnerCtxFactory: Send + Sync + 'static { adapter_store, sources_extractor, run_cache_ctx, - )), + )?), schema_cache, data_store, resolver_state, diff --git a/crates/dbt-tasks-core/src/lib.rs b/crates/dbt-tasks-core/src/lib.rs index f1ea949173f..043378044b8 100644 --- a/crates/dbt-tasks-core/src/lib.rs +++ b/crates/dbt-tasks-core/src/lib.rs @@ -26,6 +26,7 @@ pub mod test_aggregation; pub mod unit_test_schema; pub mod utils; pub mod visitor; +pub mod wap; use std::any::Any; use std::collections::{HashMap, HashSet}; diff --git a/crates/dbt-tasks-core/src/run_cache/run_cache_service.rs b/crates/dbt-tasks-core/src/run_cache/run_cache_service.rs index 70444cb96d9..ed96c63c373 100644 --- a/crates/dbt-tasks-core/src/run_cache/run_cache_service.rs +++ b/crates/dbt-tasks-core/src/run_cache/run_cache_service.rs @@ -8006,7 +8006,8 @@ mod tests { telemetry_session_ended: std::sync::atomic::AtomicBool::new(false), telemetry_dispatcher: std::sync::OnceLock::new(), }, - ); + ) + .expect("test context must have a valid execution plan"); TaskRunnerCtx { inner: Arc::new(inner), diff --git a/crates/dbt-tasks-core/src/wap.rs b/crates/dbt-tasks-core/src/wap.rs new file mode 100644 index 00000000000..264fb60958c --- /dev/null +++ b/crates/dbt-tasks-core/src/wap.rs @@ -0,0 +1,774 @@ +//! Invocation-local planning for audited publication of Snowflake tables. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use dbt_adapter::relation::{create_relation, create_relation_from_node}; +use dbt_adapter_core::AdapterType; +use dbt_common::hashing::code_hash; +use dbt_common::io_args::{ComputeArg, FsCommand, LocalExecutionBackendKind}; +use dbt_common::{ErrorCode, FsResult, fs_err}; +use dbt_dag::schedule::Schedule; +use dbt_schemas::dbt_types::RelationType; +use dbt_schemas::schemas::common::{DbtMaterialization, ResolvedQuoting, StoreFailuresAs}; +use dbt_schemas::schemas::relations::base::BaseRelation; +use dbt_schemas::schemas::telemetry::NodeType; +use dbt_schemas::schemas::{DbtModel, InternalDbtNode, InternalDbtNodeAttributes, Nodes}; + +use crate::RunTasksArgs; + +/// A canonical model and the private execution identity used by this invocation. +#[derive(Debug, Clone)] +pub struct WapModel { + pub model: Arc, + pub candidate_identifier: String, + pub audit_ids: BTreeSet, +} + +impl WapModel { + pub fn public_relation(&self) -> FsResult> { + create_relation_from_node(self.model.node_adapter(), self.model.as_ref(), None) + } + + pub fn candidate_relation(&self) -> FsResult> { + self.candidate_relation_with_quoting(self.model.quoting()) + } + + fn candidate_relation_with_quoting( + &self, + quoting: ResolvedQuoting, + ) -> FsResult> { + create_relation( + self.model.node_adapter(), + self.model.database(), + self.model.schema(), + Some(self.candidate_identifier.clone()), + Some(RelationType::Table), + quoting, + ) + } + + /// The built-in materialization creates its table using the adapter's quoting. + pub fn validate_materialization_quoting( + &self, + effective_quoting: ResolvedQuoting, + ) -> FsResult<()> { + let expected = relation_identity(self.candidate_relation()?.as_ref())?; + let actual = relation_identity( + self.candidate_relation_with_quoting(effective_quoting)? + .as_ref(), + )?; + if actual != expected { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP model '{}': the table materialization resolves a different database or schema under the profile's quoting policy; align the model and profile quoting settings", + self.model.common().unique_id + )); + } + Ok(()) + } + + /// Create the scoped model used for candidate rendering and materialization. + /// The manifest's model is never mutated. + pub fn execution_model(&self) -> FsResult { + let mut model = self.model.as_ref().clone(); + model.__base_attr__.alias = self.candidate_identifier.clone(); + model.deprecated_config.alias = Some(self.candidate_identifier.clone()); + model.__base_attr__.relation_name = Some(self.candidate_relation()?.render_self_as_str()); + // Access is applied to the public relation only after the audit passes. + model.deprecated_config.grants = Default::default(); + let warehouse = &mut model.deprecated_config.__warehouse_specific_config__; + warehouse.copy_grants = Some(false); + warehouse.copy_tags = Some(false); + if let Some(snowflake) = model.__adapter_attr__.snowflake_attr.as_mut() { + snowflake.copy_grants = Some(false); + snowflake.copy_tags = Some(false); + } + Ok(model) + } + + fn relation_names(&self) -> FsResult<([String; 3], [String; 3])> { + Ok(( + relation_identity(self.public_relation()?.as_ref())?, + relation_identity(self.candidate_relation()?.as_ref())?, + )) + } +} + +/// Keep canonical components separate and compare their normalized strings exactly: +/// CanonicalFqn's Ident equality ignores case even for quoted relation names. +fn relation_identity(relation: &dyn BaseRelation) -> FsResult<[String; 3]> { + let canonical = relation.get_canonical_fqn()?; + Ok([canonical.catalog(), canonical.schema(), canonical.table()] + .map(|part| part.as_str().to_owned())) +} + +/// Validated WAP work for the selected models in a single invocation. +#[derive(Debug, Clone, Default)] +pub struct WapPlan { + pub models: BTreeMap, +} + +impl WapPlan { + /// Discover all required audits before any warehouse writes are scheduled. + pub fn build( + args: &RunTasksArgs, + schedule: &Schedule, + nodes: &Nodes, + ) -> FsResult { + let mut plan = Self::default(); + for unique_id in &schedule.selected_nodes { + let Some(model) = nodes.models.get(unique_id) else { + continue; + }; + if !model.deprecated_config.wap.unwrap_or(false) { + continue; + } + model.validate_wap_config(model.node_adapter())?; + match args.command { + FsCommand::Run | FsCommand::Clone => { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "Model '{unique_id}' has wap=true; use dbt build to audit and publish it" + )); + } + FsCommand::Build => {} + // Read-only commands and standalone tests use canonical relations. + _ => continue, + } + if args.empty || args.sample.is_some() || !args.sample_renaming.is_empty() { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "Model '{unique_id}': wap=true does not support --empty or --sample" + )); + } + if args.local_execution_backend != LocalExecutionBackendKind::Remote + || model + .base() + .compute + .is_some_and(|compute| compute != ComputeArg::Remote) + || !model.node_propagate().is_empty() + || args.infer_schemas_and_typeless + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "Model '{unique_id}': wap=true requires execution on Snowflake without compute propagation or schema-only inference" + )); + } + let audit_ids = Self::required_audit_ids(unique_id, nodes)?; + if audit_ids.is_empty() { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "Model '{unique_id}': wap=true requires at least one enabled single-relation data test" + )); + } + let missing = audit_ids + .difference(&schedule.selected_nodes) + .map(String::as_str) + .collect::>(); + if !missing.is_empty() { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "Model '{unique_id}': all WAP audits must be selected; missing: {}", + missing.join(", ") + )); + } + for audit_id in &audit_ids { + let audit = &nodes.tests[audit_id]; + if audit.node_adapter() != AdapterType::Snowflake + || audit + .base() + .compute + .is_some_and(|compute| compute != ComputeArg::Remote) + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP audit '{audit_id}' must execute on Snowflake" + )); + } + } + let candidate_identifier = format!( + "__DBT_WAP_{}_{}", + args.io.invocation_id.simple(), + code_hash(unique_id) + ) + .to_ascii_uppercase(); + plan.models.insert( + unique_id.clone(), + WapModel { + model: Arc::clone(model), + candidate_identifier, + audit_ids, + }, + ); + } + plan.validate_relation_names(nodes)?; + Ok(plan) + } + + fn validate_relation_names(&self, nodes: &Nodes) -> FsResult<()> { + for (model_id, entry) in &self.models { + let (target_name, candidate_name) = entry.relation_names()?; + for (other_id, other) in nodes.iter() { + if !other.base().enabled + || other.node_adapter() != AdapterType::Snowflake + || !matches!( + other.resource_type(), + NodeType::Model | NodeType::Seed | NodeType::Snapshot | NodeType::Source + ) + || other.materialized() == DbtMaterialization::Ephemeral + { + continue; + } + let other_relation = + create_relation_from_node(AdapterType::Snowflake, other, None)?; + let other_name = relation_identity(other_relation.as_ref())?; + if other_name == candidate_name { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP candidate for '{model_id}' collides with relation of '{other_id}'" + )); + } + if other_id != model_id + && other.resource_type() != NodeType::Source + && other_name == target_name + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP model '{model_id}' shares its public relation with '{other_id}'" + )); + } + } + } + Ok(()) + } + + /// Failure storage must not overwrite a public or working WAP table. + /// Test materializations use `api.Relation.create` with the adapter's quoting, + /// which can differ from the test node's configured quoting. + pub fn validate_test_storage_relations( + &self, + nodes: &Nodes, + effective_quoting: ResolvedQuoting, + ) -> FsResult<()> { + for (model_id, entry) in &self.models { + let (target_name, candidate_name) = entry.relation_names()?; + for (test_id, test) in &nodes.tests { + let config = &test.deprecated_config; + if !test.base().enabled + || config.enabled == Some(false) + || test.node_adapter() != AdapterType::Snowflake + || !config.store_failures.unwrap_or(matches!( + config.store_failures_as, + Some(StoreFailuresAs::Table | StoreFailuresAs::View) + )) + { + continue; + } + let storage = create_relation( + AdapterType::Snowflake, + test.database(), + test.schema(), + Some(test.base().alias.clone()), + None, + effective_quoting, + )?; + let storage_name = relation_identity(storage.as_ref())?; + if storage_name == target_name || storage_name == candidate_name { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "Failure storage for test '{test_id}' collides with the public or candidate table of WAP model '{model_id}'" + )); + } + } + } + Ok(()) + } + + /// Required enabled audits, using the full manifest rather than pruned selection. + pub fn required_audit_ids(model_id: &str, nodes: &Nodes) -> FsResult> { + let mut audit_ids = BTreeSet::new(); + for (test_id, test) in &nodes.tests { + if !test.base().enabled || test.deprecated_config.enabled == Some(false) { + continue; + } + let dependencies = &test.base().depends_on.nodes; + let owned = match test.__test_attr__.attached_node.as_deref() { + Some(owner) => owner == model_id, + None => dependencies.iter().any(|dependency| dependency == model_id), + }; + if !owned { + continue; + } + if dependencies.is_empty() + || dependencies.iter().any(|dependency| dependency != model_id) + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP audit '{test_id}' for '{model_id}' must depend only on that model; multi-relation audits are not supported" + )); + } + audit_ids.insert(test_id.clone()); + } + Ok(audit_ids) + } + + pub fn model(&self, unique_id: &str) -> Option<&WapModel> { + self.models.get(unique_id) + } + + pub fn audit_owner(&self, unique_id: &str) -> Option<&WapModel> { + self.models + .values() + .find(|model| model.audit_ids.contains(unique_id)) + } + + pub fn contains_node(&self, unique_id: &str) -> bool { + self.models.contains_key(unique_id) || self.audit_owner(unique_id).is_some() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dbt_adapter::relation::{RelationObject, factory::create_static_relation}; + use dbt_schemas::schemas::DbtTest; + + fn fixtures() -> (RunTasksArgs, Schedule, Nodes) { + let args = RunTasksArgs { + command: FsCommand::Build, + ..Default::default() + }; + let mut nodes = Nodes::default(); + let mut model = DbtModel::default(); + model.__common_attr__.unique_id = "model.pkg.orders".to_owned(); + model.__common_attr__.language = Some("sql".to_owned()); + model.__base_attr__.enabled = true; + model.__base_attr__.materialized = DbtMaterialization::Table; + model.__base_attr__.database = "DB".to_owned(); + model.__base_attr__.schema = "PUBLIC".to_owned(); + model.__base_attr__.alias = "ORDERS".to_owned(); + model.deprecated_config.wap = Some(true); + nodes.models.insert(model.unique_id(), Arc::new(model)); + let mut test = DbtTest::default(); + test.__common_attr__.unique_id = "test.pkg.orders_not_null".to_owned(); + test.__base_attr__.enabled = true; + test.__base_attr__.depends_on.nodes = vec!["model.pkg.orders".to_owned()]; + test.__test_attr__.attached_node = Some("model.pkg.orders".to_owned()); + nodes.tests.insert(test.unique_id(), Arc::new(test)); + let schedule = Schedule { + selected_nodes: BTreeSet::from([ + "model.pkg.orders".to_owned(), + "test.pkg.orders_not_null".to_owned(), + ]), + ..Default::default() + }; + (args, schedule, nodes) + } + + #[test] + fn candidates_are_scoped_and_leave_manifest_relations_unchanged() { + let (mut args, schedule, nodes) = fixtures(); + let first = WapPlan::build(&args, &schedule, &nodes).unwrap(); + let entry = first.model("model.pkg.orders").unwrap(); + let candidate = entry.execution_model().unwrap(); + assert_ne!(candidate.base().alias, entry.model.base().alias); + assert_eq!(candidate.base().schema, entry.model.base().schema); + assert_eq!(candidate.unique_id(), entry.model.unique_id()); + assert_eq!(nodes.models["model.pkg.orders"].base().alias, "ORDERS"); + assert!(first.audit_owner("test.pkg.orders_not_null").is_some()); + args.io.invocation_id = "00000000-0000-0000-0000-000000000001".parse().unwrap(); + let second = WapPlan::build(&args, &schedule, &nodes).unwrap(); + assert_ne!( + entry.candidate_identifier, + second.models["model.pkg.orders"].candidate_identifier + ); + } + + #[test] + fn candidate_rendering_preserves_quoted_database_and_schema_names() { + let (args, schedule, mut nodes) = fixtures(); + let model = Arc::make_mut(nodes.models.get_mut("model.pkg.orders").unwrap()); + model.__base_attr__.database = "DB\"name".to_owned(); + model.__base_attr__.schema = "schema.with.dot".to_owned(); + model.__base_attr__.alias = "Ord\"ers".to_owned(); + let plan = WapPlan::build(&args, &schedule, &nodes).unwrap(); + let entry = &plan.models["model.pkg.orders"]; + + let candidate = entry.execution_model().unwrap(); + let candidate_sql = entry.candidate_relation().unwrap().render_self_as_str(); + assert_eq!( + candidate_sql, + format!( + "\"DB\"\"name\".\"schema.with.dot\".\"{}\"", + entry.candidate_identifier + ) + ); + assert_eq!( + candidate.base().relation_name.as_deref(), + Some(candidate_sql.as_str()) + ); + assert_eq!( + entry.public_relation().unwrap().render_self_as_str(), + "\"DB\"\"name\".\"schema.with.dot\".\"Ord\"\"ers\"" + ); + assert_eq!(nodes.models["model.pkg.orders"].base().alias, "Ord\"ers"); + } + + #[test] + fn materialization_quoting_guard_matches_jinja_relation_creation() { + let env = minijinja::Environment::new(); + let expression = env + .compile_expression( + "api.Relation.create(database=database, schema=schema, identifier=identifier, type='table')", + ) + .unwrap(); + for (database, schema, model_quoting, effective_quoting, compatible) in [ + ( + "DB", + "PUBLIC", + ResolvedQuoting::trues(), + ResolvedQuoting::falses(), + true, + ), + ( + "Db", + "PUBLIC", + ResolvedQuoting::trues(), + ResolvedQuoting::falses(), + false, + ), + ( + "DB", + "Public", + ResolvedQuoting::trues(), + ResolvedQuoting::falses(), + false, + ), + ( + "db", + "public", + ResolvedQuoting::falses(), + ResolvedQuoting::trues(), + false, + ), + ( + "db", + "public", + ResolvedQuoting::trues(), + ResolvedQuoting::trues(), + true, + ), + ] { + let (args, schedule, mut nodes) = fixtures(); + let model = Arc::make_mut(nodes.models.get_mut("model.pkg.orders").unwrap()); + model.__base_attr__.database = database.to_owned(); + model.__base_attr__.schema = schema.to_owned(); + model.__base_attr__.quoting = model_quoting; + let plan = WapPlan::build(&args, &schedule, &nodes).unwrap(); + let entry = &plan.models["model.pkg.orders"]; + let api = BTreeMap::from([( + "Relation", + create_static_relation(AdapterType::Snowflake, effective_quoting).unwrap(), + )]); + let value = expression + .eval( + minijinja::context! { + api => api, + database => database, + schema => schema, + identifier => entry.candidate_identifier.as_str(), + }, + &[], + ) + .unwrap(); + let actual = value + .downcast_object_ref::() + .unwrap() + .inner(); + let expected = entry.candidate_relation().unwrap(); + assert_eq!( + relation_identity(actual.as_ref()).unwrap() + == relation_identity(expected.as_ref()).unwrap(), + compatible, + "database={database}, schema={schema}", + ); + assert_eq!( + entry + .validate_materialization_quoting(effective_quoting) + .is_ok(), + compatible, + "database={database}, schema={schema}", + ); + } + } + + #[test] + fn missing_and_empty_audit_selections_fail_closed() { + let (args, mut schedule, mut nodes) = fixtures(); + schedule.selected_nodes.remove("test.pkg.orders_not_null"); + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap_err() + .to_string() + .contains("all WAP audits") + ); + nodes.tests.clear(); + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap_err() + .to_string() + .contains("at least one") + ); + } + + #[test] + fn generic_test_owner_is_attachment_not_every_dependency() { + let (args, schedule, mut nodes) = fixtures(); + let mut relationship = DbtTest::default(); + relationship.__common_attr__.unique_id = "test.pkg.downstream_relationship".to_owned(); + relationship.__base_attr__.enabled = true; + relationship.__base_attr__.depends_on.nodes = vec![ + "model.pkg.orders".to_owned(), + "model.pkg.downstream".to_owned(), + ]; + relationship.__test_attr__.attached_node = Some("model.pkg.downstream".to_owned()); + nodes + .tests + .insert(relationship.unique_id(), Arc::new(relationship.clone())); + let plan = WapPlan::build(&args, &schedule, &nodes).unwrap(); + assert!( + plan.audit_owner("test.pkg.downstream_relationship") + .is_none() + ); + relationship.__test_attr__.attached_node = Some("model.pkg.orders".to_owned()); + nodes + .tests + .insert(relationship.unique_id(), Arc::new(relationship)); + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap_err() + .to_string() + .contains("multi-relation") + ); + } + + #[test] + fn singular_audits_support_one_relation_and_reject_ambiguous_ownership() { + let (args, schedule, mut nodes) = fixtures(); + let audit = Arc::make_mut(nodes.tests.get_mut("test.pkg.orders_not_null").unwrap()); + audit.__test_attr__.attached_node = None; + assert!(WapPlan::build(&args, &schedule, &nodes).is_ok()); + Arc::make_mut(nodes.tests.get_mut("test.pkg.orders_not_null").unwrap()) + .__base_attr__ + .depends_on + .nodes + .push("source.pkg.raw".to_owned()); + assert!(WapPlan::build(&args, &schedule, &nodes).is_err()); + } + + #[test] + fn disabled_audits_are_not_required() { + let (args, schedule, mut nodes) = fixtures(); + let mut disabled = nodes.tests["test.pkg.orders_not_null"].as_ref().clone(); + disabled.__common_attr__.unique_id = "test.pkg.disabled".to_owned(); + disabled.deprecated_config.enabled = Some(false); + disabled + .__base_attr__ + .depends_on + .nodes + .push("model.pkg.other".to_owned()); + nodes.tests.insert(disabled.unique_id(), Arc::new(disabled)); + assert!(WapPlan::build(&args, &schedule, &nodes).is_ok()); + } + + #[test] + fn commands_cannot_bypass_audits() { + let (mut args, schedule, nodes) = fixtures(); + for command in [FsCommand::Run, FsCommand::Clone] { + args.command = command; + assert!(WapPlan::build(&args, &schedule, &nodes).is_err()); + } + for command in [FsCommand::Test, FsCommand::Compile] { + args.command = command; + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap() + .models + .is_empty() + ); + } + args.command = FsCommand::Build; + args.empty = true; + assert!(WapPlan::build(&args, &schedule, &nodes).is_err()); + args.empty = false; + args.sample = Some("10 rows".to_owned()); + assert!(WapPlan::build(&args, &schedule, &nodes).is_err()); + args.sample = None; + args.local_execution_backend = LocalExecutionBackendKind::Inline; + assert!(WapPlan::build(&args, &schedule, &nodes).is_err()); + } + + #[test] + fn test_only_builds_and_unselected_wap_parents_use_public_relations() { + let (args, mut schedule, mut nodes) = fixtures(); + schedule.selected_nodes.remove("model.pkg.orders"); + schedule + .frontier_nodes + .insert("model.pkg.orders".to_owned()); + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap() + .models + .is_empty() + ); + + schedule + .selected_nodes + .insert("model.pkg.orders".to_owned()); + schedule.frontier_nodes.clear(); + Arc::make_mut(nodes.models.get_mut("model.pkg.orders").unwrap()) + .deprecated_config + .wap = Some(false); + nodes.tests.clear(); + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap() + .models + .is_empty() + ); + } + + #[test] + fn public_relation_collisions_respect_snowflake_identifier_case() { + let (args, schedule, mut nodes) = fixtures(); + let mut other = nodes.models["model.pkg.orders"].as_ref().clone(); + other.__common_attr__.unique_id = "model.pkg.other".to_owned(); + other.__base_attr__.alias = "orders".to_owned(); + other.__base_attr__.quoting.identifier = false; + nodes.models.insert(other.unique_id(), Arc::new(other)); + assert!( + WapPlan::build(&args, &schedule, &nodes) + .unwrap_err() + .to_string() + .contains("shares its public relation") + ); + Arc::make_mut(nodes.models.get_mut("model.pkg.other").unwrap()) + .__base_attr__ + .quoting + .identifier = true; + let plan = WapPlan::build(&args, &schedule, &nodes); + assert!(plan.is_ok(), "{plan:?}"); + } + + #[test] + fn relation_identity_preserves_component_boundaries() { + let (args, schedule, mut nodes) = fixtures(); + let model = Arc::make_mut(nodes.models.get_mut("model.pkg.orders").unwrap()); + model.__base_attr__.database = "A.B".to_owned(); + model.__base_attr__.schema = "C".to_owned(); + let mut other = model.clone(); + other.__common_attr__.unique_id = "model.pkg.other".to_owned(); + other.__base_attr__.database = "A".to_owned(); + other.__base_attr__.schema = "B.C".to_owned(); + nodes.models.insert(other.unique_id(), Arc::new(other)); + assert!(WapPlan::build(&args, &schedule, &nodes).is_ok()); + } + + #[test] + fn audit_failure_storage_cannot_replace_public_or_candidate_tables() { + let (args, schedule, mut nodes) = fixtures(); + let plan = WapPlan::build(&args, &schedule, &nodes).unwrap(); + let candidate = plan.models["model.pkg.orders"].candidate_identifier.clone(); + for alias in ["ORDERS", candidate.as_str()] { + for (store_failures, store_failures_as) in [ + (Some(true), None), + (Some(true), Some(StoreFailuresAs::Table)), + (None, Some(StoreFailuresAs::View)), + ] { + let audit = Arc::make_mut(nodes.tests.get_mut("test.pkg.orders_not_null").unwrap()); + audit.__base_attr__.database = "DB".to_owned(); + audit.__base_attr__.schema = "PUBLIC".to_owned(); + audit.__base_attr__.alias = alias.to_ascii_lowercase(); + audit.__base_attr__.quoting.identifier = true; + audit.deprecated_config.store_failures = store_failures; + audit.deprecated_config.store_failures_as = store_failures_as; + let error = plan + .validate_test_storage_relations( + &nodes, + ResolvedQuoting { + database: true, + schema: true, + identifier: false, + }, + ) + .unwrap_err(); + assert!(error.to_string().contains("test.pkg.orders_not_null")); + } + } + } + + #[test] + fn distinct_failure_storage_and_non_storing_tests_remain_supported() { + let (args, schedule, mut nodes) = fixtures(); + let plan = WapPlan::build(&args, &schedule, &nodes).unwrap(); + for storage_type in [StoreFailuresAs::Table, StoreFailuresAs::View] { + let audit = Arc::make_mut(nodes.tests.get_mut("test.pkg.orders_not_null").unwrap()); + audit.__base_attr__.database = "DB".to_owned(); + audit.__base_attr__.schema = "PUBLIC".to_owned(); + audit.__base_attr__.alias = "ORDERS_FAILURES".to_owned(); + audit.deprecated_config.store_failures = Some(true); + audit.deprecated_config.store_failures_as = Some(storage_type); + assert!( + plan.validate_test_storage_relations(&nodes, ResolvedQuoting::trues()) + .is_ok() + ); + } + + let audit = Arc::make_mut(nodes.tests.get_mut("test.pkg.orders_not_null").unwrap()); + audit.__base_attr__.alias = "ORDERS".to_owned(); + audit.deprecated_config.store_failures = Some(false); + audit.deprecated_config.store_failures_as = Some(StoreFailuresAs::Ephemeral); + assert!( + plan.validate_test_storage_relations(&nodes, ResolvedQuoting::trues()) + .is_ok() + ); + + let mut disabled = nodes.tests["test.pkg.orders_not_null"].as_ref().clone(); + disabled.__common_attr__.unique_id = "test.pkg.disabled_storage".to_owned(); + disabled.deprecated_config.store_failures = Some(true); + disabled.deprecated_config.store_failures_as = Some(StoreFailuresAs::Table); + disabled.deprecated_config.enabled = Some(false); + nodes.tests.insert(disabled.unique_id(), Arc::new(disabled)); + assert!( + plan.validate_test_storage_relations(&nodes, ResolvedQuoting::trues()) + .is_ok() + ); + } + + #[test] + fn storage_collision_check_uses_runtime_quoting_and_all_enabled_tests() { + let (args, schedule, mut nodes) = fixtures(); + let plan = WapPlan::build(&args, &schedule, &nodes).unwrap(); + let mut other = nodes.tests["test.pkg.orders_not_null"].as_ref().clone(); + other.__common_attr__.unique_id = "test.pkg.other_storage".to_owned(); + other.__test_attr__.attached_node = Some("model.pkg.other".to_owned()); + other.__base_attr__.depends_on.nodes = vec!["model.pkg.other".to_owned()]; + other.__base_attr__.database = "DB".to_owned(); + other.__base_attr__.schema = "PUBLIC".to_owned(); + other.__base_attr__.alias = "orders".to_owned(); + other.__base_attr__.quoting.identifier = false; + other.deprecated_config.store_failures = Some(true); + nodes.tests.insert(other.unique_id(), Arc::new(other)); + assert!( + plan.validate_test_storage_relations(&nodes, ResolvedQuoting::trues()) + .is_ok() + ); + assert!( + plan.validate_test_storage_relations(&nodes, ResolvedQuoting::falses()) + .unwrap_err() + .to_string() + .contains("test.pkg.other_storage") + ); + } +} diff --git a/crates/dbt-tasks-sa/src/graph.rs b/crates/dbt-tasks-sa/src/graph.rs index 8c977d0d9f0..b122fc891b0 100644 --- a/crates/dbt-tasks-sa/src/graph.rs +++ b/crates/dbt-tasks-sa/src/graph.rs @@ -1,7 +1,7 @@ use dbt_common::io_args::{FsCommand, OptimizeTestsOptions}; use dbt_common::static_analysis::is_strict_static_analysis; use dbt_common::tracing::dbt_emit::emit_warn_log_message; -use dbt_common::{ErrorCode, FsResult}; +use dbt_common::{ErrorCode, FsResult, fs_err}; use dbt_dag::deps_mgmt::{find_all_upstream_deps, restrict_with_transitive}; use dbt_dag::schedule::Schedule; use dbt_schemas::schemas::InternalDbtNode; @@ -26,12 +26,14 @@ use dbt_tasks_core::task::{TP, TasksForNode}; use dbt_tasks_core::test_aggregation::{ GenericTestAggregation, GenericTestRelationships, create_generic_test_aggregation, }; +use dbt_tasks_core::wap::WapPlan; use crate::barrier::BarrierTask; use crate::cloneable::RunCloneTask; use crate::cloneable::cloneable_task; use crate::renderable::unit_test::build_unit_test_overrides_map; use crate::task::TasksForNodeFactory; +use crate::wap::PublishTask; const PHASES_RENDER_ANALYZE_RUN: &[TP] = &[TP::Render, TP::Analyze, TP::Run]; const PHASES_RENDER_ANALYZE: &[TP] = &[TP::Render, TP::Analyze]; @@ -104,6 +106,7 @@ impl GraphBuilder { resolver_state: &Arc, ) -> FsResult<(Graph, ()>, GenericTestRelationships)> { let nodes = &resolver_state.nodes; + let wap_plan = WapPlan::build(self.arg.as_ref(), schedule, nodes)?; // Test aggregation let generic_test_aggregation = if self @@ -112,7 +115,20 @@ impl GraphBuilder { .contains(&OptimizeTestsOptions::TestAggregation) && command_uses_generic_test_aggregation(self.arg.command) { - create_generic_test_aggregation(&self.arg.io, schedule, nodes, self.execute)? + // Audit nodes must keep their identities and candidate-scoped contexts. + let aggregation_schedule = (!wap_plan.models.is_empty()).then(|| { + let mut schedule = schedule.clone(); + schedule + .selected_nodes + .retain(|unique_id| wap_plan.audit_owner(unique_id).is_none()); + schedule + }); + create_generic_test_aggregation( + &self.arg.io, + aggregation_schedule.as_ref().unwrap_or(schedule), + nodes, + self.execute, + )? } else { None }; @@ -166,7 +182,8 @@ impl GraphBuilder { self.static_analysis_buckets.as_ref(), self.arg.infer_schemas_and_typeless, aggregation, - ) + &wap_plan, + )? } else { // Freshness and `jinja-check` legitimately produce an // empty task graph. @@ -205,7 +222,8 @@ impl GraphBuilder { buckets: &dyn StaticAnalysisBuckets, infer_schemas: bool, generic_test_aggregation: Option<&GenericTestAggregation>, - ) -> (DiGraph, ()>, BTreeSet) { + wap_plan: &WapPlan, + ) -> FsResult<(DiGraph, ()>, BTreeSet)> { // Build reverse dependencies map once for efficient propagation let reverse_deps = build_reverse_deps(&schedule.deps); @@ -232,6 +250,7 @@ impl GraphBuilder { &unit_test_overrides_map, generic_test_aggregation, &reverse_deps, + wap_plan, ); // Add phase transitions (e.g. render -> analyze -> run) for each node @@ -448,6 +467,8 @@ impl GraphBuilder { ); } + add_wap_publication_tasks(&mut graph, &node_indices, wap_plan)?; + // Insert barrier only for static nodes if buckets .global_static_analysis() @@ -488,10 +509,68 @@ impl GraphBuilder { assert_graph(&graph); - (graph, nodes_with_no_tasks) + if !wap_plan.models.is_empty() && petgraph::algo::is_cyclic_directed(&graph) { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "The execution graph contains a cycle; WAP audits must depend only on their owning model" + )); + } + + Ok((graph, nodes_with_no_tasks)) } } +/// Give consumers a publication dependency while audits consume the staged table. +fn add_wap_publication_tasks( + graph: &mut DiGraph, ()>, + node_indices: &BTreeMap<(TP, String), NodeIndex>, + wap_plan: &WapPlan, +) -> FsResult<()> { + for (model_id, model) in &wap_plan.models { + let stage = *node_indices + .get(&(TP::Run, model_id.clone())) + .ok_or_else(|| { + fs_err!( + ErrorCode::InvalidConfig, + "WAP model '{model_id}' has no execution task" + ) + })?; + let publish = graph.add_node(Arc::new(PublishTask::new(Arc::clone(&model.model)))); + let consumers: Vec<_> = graph + .neighbors_directed(stage, petgraph::Outgoing) + .filter(|index| !model.audit_ids.contains(graph[*index].work_node_id())) + .collect(); + for consumer in consumers { + if let Some(edge) = graph.find_edge(stage, consumer) { + graph.remove_edge(edge); + } + graph.update_edge(publish, consumer, ()); + } + graph.update_edge(stage, publish, ()); + for audit_id in &model.audit_ids { + let render = node_indices + .get(&(TP::Render, audit_id.clone())) + .ok_or_else(|| { + fs_err!( + ErrorCode::InvalidConfig, + "WAP audit '{audit_id}' has no render task" + ) + })?; + let run = node_indices + .get(&(TP::Run, audit_id.clone())) + .ok_or_else(|| { + fs_err!( + ErrorCode::InvalidConfig, + "WAP audit '{audit_id}' has no execution task" + ) + })?; + graph.update_edge(stage, *render, ()); + graph.update_edge(*run, publish, ()); + } + } + Ok(()) +} + /// Create a new schedule and nodes based on test aggregation. fn create_aggregated_schedule_and_nodes( schedule: &Schedule, @@ -610,6 +689,7 @@ fn initialize_graph( unit_test_overrides_map: &BTreeMap, generic_test_aggregation: Option<&GenericTestAggregation>, reverse_deps: &HashMap>, + wap_plan: &WapPlan, ) -> ( DiGraph, ()>, BTreeMap>>, @@ -716,6 +796,13 @@ fn initialize_graph( expected_node_phases.clear(); } + if wap_plan.audit_owner(unique_id).is_some() { + // Audits must read the newly staged table on Snowflake. Its invocation-local + // relation is not part of the schema snapshot used by static analysis. + // Select phases before constructing channels so Render feeds Run directly. + expected_node_phases.retain(|phase| *phase != TP::Analyze); + } + let TasksForNode { renderable, analyzeable, @@ -1142,3 +1229,365 @@ pub fn print_task_graph(graph: &DiGraph, ()>) -> Vec<(String, Stri } res } + +#[cfg(test)] +mod wap_tests { + use super::*; + use std::future::Future; + use std::pin::Pin; + + use dbt_common::stats::NodeStatus; + use dbt_schemas::schemas::telemetry::NodeType; + use dbt_schemas::schemas::{DbtModel, DbtTest, InternalDbtNodeAttributes}; + use dbt_tasks_core::context::TaskRunnerCtx; + use dbt_tasks_core::wap::WapModel; + + struct PhasedFactory; + + impl TasksForNodeFactory for PhasedFactory { + fn tasks_for_node( + &self, + unique_id: &str, + nodes: &Nodes, + _schedule: &Schedule, + _execute: Execute, + phases: &[TP], + _aggregation: Option<&GenericTestAggregation>, + _unit_test_overrides: Option<&UnitTestOverrides>, + _reverse_deps: &HashMap>, + ) -> TasksForNode { + let node = nodes.get_node_owned(unique_id).unwrap(); + let task = |phase| { + phases.contains(&phase).then(|| { + Arc::new(GraphTask { + node: Arc::clone(&node), + phase, + }) as Arc + }) + }; + TasksForNode { + renderable: task(TP::Render), + analyzeable: task(TP::Analyze), + runnable: task(TP::Run), + showable: None, + } + } + + fn tasks_for_generic_test_group( + &self, + _: &[TP], + _: &Arc, + ) -> TasksForNode { + unreachable!() + } + fn runnable_task( + &self, + _: &str, + _: &Nodes, + _: &[TP], + _: Execute, + ) -> Option> { + unreachable!() + } + fn analyze_task( + &self, + _: Arc, + _: Option>, + _: Option>, + ) -> Option> { + unreachable!() + } + fn analyzeable_task( + &self, + _: &Nodes, + _: &str, + _: &[TP], + ) -> Option> { + unreachable!() + } + fn create_show_task_hooks( + &self, + ) -> Arc { + unreachable!() + } + fn create_render_task_hooks( + &self, + ) -> Arc { + unreachable!() + } + fn create_run_task_hooks(&self) -> Arc { + unreachable!() + } + } + + #[derive(Default)] + struct StrictBuckets { + deferred: HashMap, + } + + impl StaticAnalysisBuckets for StrictBuckets { + fn global_static_analysis(&self) -> Option { + Some(dbt_common::io_args::StaticAnalysisKind::Strict) + } + fn deferred_unique_ids(&self) -> &HashMap { + &self.deferred + } + fn in_off_closure(&self, _: &str) -> bool { + false + } + fn in_baseline_closure(&self, _: &str) -> bool { + false + } + fn in_dynamic_closure(&self, _: &str) -> bool { + false + } + fn dynamic_node(&self, _: &str) -> Option { + None + } + fn has_dynamic_closure(&self) -> bool { + false + } + fn will_build_phased_task_graph(&self, _: &RunTasksArgs, _: &Nodes) {} + fn did_build_phased_task_graph(&self, _: &RunTasksArgs, _: &BTreeSet) {} + } + + struct GraphTask { + node: Arc, + phase: TP, + } + + impl Task for GraphTask { + fn run_task<'a>( + &'a self, + _ctx: &'a mut TaskRunnerCtx, + ) -> Pin> + Send + 'a>> { + Box::pin(async { Ok(NodeStatus::Succeeded) }) + } + + fn resource_type(&self) -> NodeType { + self.node.resource_type() + } + fn task_type(&self) -> &str { + match self.phase { + TP::Render => "render", + TP::Analyze => "analyze", + _ => "run", + } + } + fn work_node_id(&self) -> &str { + &self.node.common().unique_id + } + fn dbt_nodes(&self) -> Vec> { + vec![Arc::clone(&self.node)] + } + fn task_phase(&self) -> Option { + Some(self.phase) + } + } + + fn add_task( + graph: &mut DiGraph, ()>, + indices: &mut BTreeMap<(TP, String), NodeIndex>, + unique_id: &str, + phase: TP, + ) -> NodeIndex { + let node: Arc = if unique_id.starts_with("test.") { + let mut node = DbtTest::default(); + node.__common_attr__.unique_id = unique_id.to_owned(); + Arc::new(node) + } else { + let mut node = DbtModel::default(); + node.__common_attr__.unique_id = unique_id.to_owned(); + Arc::new(node) + }; + let index = graph.add_node(Arc::new(GraphTask { node, phase })); + indices.insert((phase, unique_id.to_owned()), index); + index + } + + fn entry(model_id: &str, audit_id: &str) -> WapModel { + let mut model = DbtModel::default(); + model.__common_attr__.unique_id = model_id.to_owned(); + WapModel { + model: Arc::new(model), + candidate_identifier: format!("candidate_{model_id}"), + audit_ids: BTreeSet::from([audit_id.to_owned()]), + } + } + + #[test] + fn audit_runs_before_publication_and_downstream_rendering() { + let mut graph = DiGraph::new(); + let mut indices = BTreeMap::new(); + let stage = add_task(&mut graph, &mut indices, "model.pkg.a", TP::Run); + let render = add_task(&mut graph, &mut indices, "test.pkg.a", TP::Render); + let analyze = add_task(&mut graph, &mut indices, "test.pkg.a", TP::Analyze); + let audit = add_task(&mut graph, &mut indices, "test.pkg.a", TP::Run); + let second_render = add_task(&mut graph, &mut indices, "test.pkg.a_second", TP::Render); + let second_audit = add_task(&mut graph, &mut indices, "test.pkg.a_second", TP::Run); + let child_render = add_task(&mut graph, &mut indices, "model.pkg.b", TP::Render); + let child_run = add_task(&mut graph, &mut indices, "model.pkg.b", TP::Run); + let unrelated = add_task(&mut graph, &mut indices, "model.pkg.unrelated", TP::Run); + for (from, to) in [ + (render, analyze), + (analyze, audit), + (stage, audit), + (second_render, second_audit), + (stage, second_audit), + (stage, child_render), + (stage, child_run), + ] { + graph.update_edge(from, to, ()); + } + let mut model = entry("model.pkg.a", "test.pkg.a"); + model.audit_ids.insert("test.pkg.a_second".to_owned()); + let plan = WapPlan { + models: BTreeMap::from([("model.pkg.a".to_owned(), model)]), + }; + add_wap_publication_tasks(&mut graph, &indices, &plan).unwrap(); + let publish = graph + .node_indices() + .find(|index| graph[*index].task_type() == "wap_publish_run") + .unwrap(); + assert!(graph.contains_edge(stage, render)); + assert!(graph.contains_edge(stage, second_render)); + assert!(graph.contains_edge(audit, publish)); + assert!(graph.contains_edge(second_audit, publish)); + assert!(graph.contains_edge(publish, child_render)); + assert!(graph.contains_edge(publish, child_run)); + assert!(!graph.contains_edge(stage, child_render)); + assert!(!graph.contains_edge(stage, child_run)); + assert_eq!( + graph + .neighbors_directed(unrelated, petgraph::Incoming) + .count(), + 0 + ); + let sorted = stable_toposort_task_graph(&graph).unwrap(); + let position = |index| { + sorted + .iter() + .position(|candidate| *candidate == index) + .unwrap() + }; + assert!(position(stage) < position(render)); + assert!(position(audit) < position(publish)); + assert!(position(second_audit) < position(publish)); + assert!(position(publish) < position(child_render)); + } + + #[test] + fn chained_wap_models_each_wait_for_upstream_publication() { + let mut graph = DiGraph::new(); + let mut indices = BTreeMap::new(); + let mut plan = WapPlan::default(); + let mut previous = None; + for suffix in ["a", "b"] { + let model_id = format!("model.pkg.{suffix}"); + let audit_id = format!("test.pkg.{suffix}"); + let stage = add_task(&mut graph, &mut indices, &model_id, TP::Run); + let render = add_task(&mut graph, &mut indices, &audit_id, TP::Render); + let audit = add_task(&mut graph, &mut indices, &audit_id, TP::Run); + graph.update_edge(render, audit, ()); + graph.update_edge(stage, audit, ()); + if let Some(upstream) = previous { + graph.update_edge(upstream, stage, ()); + } + previous = Some(stage); + plan.models + .insert(model_id.clone(), entry(&model_id, &audit_id)); + } + add_wap_publication_tasks(&mut graph, &indices, &plan).unwrap(); + let publish_a = graph + .node_indices() + .find(|index| { + graph[*index].work_node_id() == "model.pkg.a" + && graph[*index].task_type() == "wap_publish_run" + }) + .unwrap(); + let stage_b = indices[&(TP::Run, "model.pkg.b".to_owned())]; + assert!(graph.contains_edge(publish_a, stage_b)); + assert!(stable_toposort_task_graph(&graph).is_ok()); + assert_graph(&graph); + } + + #[test] + fn missing_audit_execution_task_rejects_publication_graph() { + let mut graph = DiGraph::new(); + let mut indices = BTreeMap::new(); + add_task(&mut graph, &mut indices, "model.pkg.a", TP::Run); + add_task(&mut graph, &mut indices, "test.pkg.a", TP::Render); + let plan = WapPlan { + models: BTreeMap::from([( + "model.pkg.a".to_owned(), + entry("model.pkg.a", "test.pkg.a"), + )]), + }; + assert!( + add_wap_publication_tasks(&mut graph, &indices, &plan) + .unwrap_err() + .to_string() + .contains("no execution task") + ); + } + + #[test] + fn strict_analysis_omits_audit_analysis_and_preserves_model_barrier() { + let model_id = "model.pkg.a".to_owned(); + let audit_id = "test.pkg.a".to_owned(); + let wap = entry(&model_id, &audit_id); + let mut nodes = Nodes::default(); + nodes + .models + .insert(model_id.clone(), Arc::clone(&wap.model)); + let mut audit = DbtTest::default(); + audit.__common_attr__.unique_id = audit_id.clone(); + nodes.tests.insert(audit_id.clone(), Arc::new(audit)); + let schedule = Schedule { + deps: BTreeMap::from([ + (model_id.clone(), BTreeSet::new()), + (audit_id.clone(), BTreeSet::from([model_id.clone()])), + ]), + selected_nodes: BTreeSet::from([model_id.clone(), audit_id.clone()]), + sorted_nodes: vec![model_id.clone(), audit_id.clone()], + ..Default::default() + }; + let plan = WapPlan { + models: BTreeMap::from([(model_id.clone(), wap)]), + }; + let (graph, _) = GraphBuilder::build_phased_task_graph( + &schedule, + &PhasedFactory, + &nodes, + Execute::Remote, + PHASES_RENDER_ANALYZE_RUN, + &StrictBuckets::default(), + false, + None, + &plan, + ) + .unwrap(); + assert!( + !graph + .node_weights() + .any(|task| task.work_node_id() == audit_id + && task.task_phase() == Some(TP::Analyze)) + ); + assert!( + graph + .node_weights() + .any(|task| task.work_node_id() == model_id + && task.task_phase() == Some(TP::Analyze)) + ); + assert!( + graph + .node_weights() + .any(|task| task.task_type() == "barrier") + ); + assert!(stable_toposort_task_graph(&graph).is_ok()); + let rows = print_task_graph(&graph); + let (_, dependencies) = rows.iter().find(|(id, _)| id == "test.pkg.a/run").unwrap(); + assert!(dependencies.contains("test.pkg.a/render")); + assert!(!dependencies.contains("test.pkg.a/analyze")); + } +} diff --git a/crates/dbt-tasks-sa/src/lib.rs b/crates/dbt-tasks-sa/src/lib.rs index d78d8a33c8b..aee2c886fb2 100644 --- a/crates/dbt-tasks-sa/src/lib.rs +++ b/crates/dbt-tasks-sa/src/lib.rs @@ -34,3 +34,4 @@ pub mod task_runner; pub mod task_runner_hooks; pub mod utils; pub mod visitor; +pub mod wap; diff --git a/crates/dbt-tasks-sa/src/materialize.rs b/crates/dbt-tasks-sa/src/materialize.rs index b07d9f25814..abbc9453e41 100644 --- a/crates/dbt-tasks-sa/src/materialize.rs +++ b/crates/dbt-tasks-sa/src/materialize.rs @@ -130,7 +130,7 @@ fn execute_materialization_macro( }) } -fn apply_node_overrides( +pub(crate) fn apply_node_overrides( adapter: &Adapter, adapter_type: AdapterType, custom_warehouse: Option, @@ -160,7 +160,7 @@ fn apply_node_overrides( /// The fingerprint check in `borrow_tlocal_connection_impl` does not help here: /// the connection's *configuration* is unchanged, only its session scope is /// wrong. -fn reset_node_overrides( +pub(crate) fn reset_node_overrides( adapter: &Adapter, unique_id: &str, targets: &[NodeOverride], diff --git a/crates/dbt-tasks-sa/src/renderable/renderable/default.rs b/crates/dbt-tasks-sa/src/renderable/renderable/default.rs index 348ff163a14..776bf663c8f 100644 --- a/crates/dbt-tasks-sa/src/renderable/renderable/default.rs +++ b/crates/dbt-tasks-sa/src/renderable/renderable/default.rs @@ -5,10 +5,11 @@ use std::sync::Arc; use dbt_common::collections::DashMap; use dbt_common::constants::DBT_EPHEMERAL_DIR_NAME; use dbt_common::constants::RENDERING; +use dbt_common::io_args::IoArgs; use dbt_common::serde_utils::convert_yml_to_dash_map; use dbt_common::stats::NodeStatus; use dbt_common::tracing::emit::emit_debug_event; -use dbt_common::{FsResult, MacroSpansOnly, stdfs}; +use dbt_common::{CompiledSpans, FsResult, MacroSpansOnly, stdfs}; use dbt_jinja_utils::phases::compile::DependencyValidationConfig; use dbt_jinja_utils::utils::{ add_task_context, inject_and_persist_ephemeral_models, macro_spans_to_macro_span_vec, @@ -17,7 +18,8 @@ use dbt_jinja_utils::utils::{ use dbt_scheduler::instructions::SqlInstruction; use dbt_schemas::schemas::common::DbtMaterialization; use dbt_schemas::schemas::properties::UnitTestOverrides; -use dbt_schemas::schemas::{InternalDbtNodeAttributes, NodePathKind}; +use dbt_schemas::schemas::{CommonAttributes, InternalDbtNodeAttributes, NodePathKind}; +use dbt_tasks_core::CompiledSqlCache; use dbt_tasks_core::context::TaskRunnerCtx; use dbt_tasks_core::task::TaskOp; use dbt_telemetry::{CompiledCode, NodeType}; @@ -53,12 +55,29 @@ fn render_default( ctx: &mut TaskRunnerCtx, local_exec_unit_test_overrides: &Option, ) -> FsResult<(SqlInstruction, Arc>)> { + // Analysis registers the model's output under its public identity so + // downstream refs can use that schema after publication. + let canonical_fqn = vec![node.database(), node.schema(), node.alias()]; + // Execution-local copies keep the canonical manifest and resolver intact. + let execution_node = ctx + .inner + .wap_plan + .model(&node.common().unique_id) + .map(|entry| { + entry + .execution_model() + .map(|model| Arc::new(model) as Arc) + }) + .transpose()?; + let node = execution_node.as_ref().unwrap_or(node); + let cacheable = !ctx.inner.wap_plan.contains_node(&node.common().unique_id); report_rendering_progress(node, ctx); - if let Some((rendered_sql_maybe_with_cte, macro_spans, reclassify_spans)) = ctx - .inner - .compiled_sql_cache - .try_get_compiled_sql(&ctx.inner.arg.io, node.common()) + if cacheable + && let Some((rendered_sql_maybe_with_cte, macro_spans, reclassify_spans)) = ctx + .inner + .compiled_sql_cache + .try_get_compiled_sql(&ctx.inner.arg.io, node.common()) { let config_map = Arc::new(convert_yml_to_dash_map(node.serialized_config())); emit_compiled_code(node, ctx, &rendered_sql_maybe_with_cte); @@ -150,18 +169,20 @@ fn render_default( .rendering_listener_factory .compiled_spans(macro_spans, &render_file_path); - ctx.inner.compiled_sql_cache.set_compiled_sql( + persist_rendered_sql( + ctx.inner.compiled_sql_cache.as_ref(), &ctx.inner.arg.io, node.common(), &rendered_sql_maybe_with_cte, spans.as_ref(), + cacheable, )?; emit_compiled_code(node, ctx, &rendered_sql_maybe_with_cte); Ok(( SqlInstruction { - fqn: vec![node.database(), node.schema(), node.alias()], + fqn: canonical_fqn, sql: rendered_sql_maybe_with_cte, original_path: node.common().original_file_path.to_path_buf(), spans, @@ -170,6 +191,30 @@ fn render_default( )) } +fn persist_rendered_sql( + cache: &dyn CompiledSqlCache, + io: &IoArgs, + common: &CommonAttributes, + sql: &str, + spans: &dyn CompiledSpans, + cacheable: bool, +) -> FsResult<()> { + if cacheable { + cache.set_compiled_sql(io, common, sql, spans) + } else { + // Materialization contexts lazily read `model.compiled_code` from this + // artifact. Persist the current SQL without making it reusable by a + // later invocation (or a standalone test using the same cache). + cache.clear(&common.unique_id); + let path = cache.get_compiled_sql_path(io, common); + if let Some(parent) = path.parent() { + stdfs::create_dir_all(parent)?; + } + stdfs::write(path, sql)?; + Ok(()) + } +} + fn report_rendering_progress(node: &Arc, ctx: &TaskRunnerCtx) { let io = &ctx.inner.arg.io; @@ -264,3 +309,54 @@ fn render_python_model( config_map, )) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::compiled_sql_cache::CompiledSqlCacheImpl; + + #[test] + fn candidate_sql_replaces_artifacts_without_becoming_reusable() { + let directory = tempfile::tempdir().unwrap(); + let io = IoArgs { + in_dir: directory.path().to_path_buf(), + out_dir: directory.path().join("target"), + ..Default::default() + }; + let common = CommonAttributes { + unique_id: "test.pkg.audit".to_string(), + name: "audit".to_string(), + package_name: "pkg".to_string(), + path: PathBuf::from("tests/audit.sql").into(), + original_file_path: PathBuf::from("tests/audit.sql").into(), + ..Default::default() + }; + let cache = CompiledSqlCacheImpl::default(); + let spans = MacroSpansOnly::default(); + let public_sql = "select * from PUBLIC_ORDERS"; + persist_rendered_sql(&cache, &io, &common, public_sql, &spans, true).unwrap(); + assert_eq!( + cache.try_get_compiled_sql(&io, &common).unwrap().0, + public_sql + ); + + // A previous cache entry must be invalidated, and every new candidate + // must replace the artifact that run-time model.compiled_code reads. + for candidate in ["__DBT_WAP_FIRST", "__DBT_WAP_SECOND"] { + let sql = format!("select * from {candidate}"); + persist_rendered_sql(&cache, &io, &common, &sql, &spans, false).unwrap(); + assert!(cache.try_get_compiled_sql(&io, &common).is_none()); + assert_eq!( + stdfs::read_to_string(cache.get_compiled_sql_path(&io, &common)).unwrap(), + sql, + ); + } + + // Standalone tests may cache their freshly rendered public SQL again. + persist_rendered_sql(&cache, &io, &common, public_sql, &spans, true).unwrap(); + assert_eq!( + cache.try_get_compiled_sql(&io, &common).unwrap().0, + public_sql + ); + } +} diff --git a/crates/dbt-tasks-sa/src/runnable/model.rs b/crates/dbt-tasks-sa/src/runnable/model.rs index e9f88edcc60..37ac35d2026 100644 --- a/crates/dbt-tasks-sa/src/runnable/model.rs +++ b/crates/dbt-tasks-sa/src/runnable/model.rs @@ -319,6 +319,19 @@ pub fn execute_model_remote( ctx: &TaskRunnerCtx, task_result: &TaskResult, ) -> FsResult { + let sql_header = task_result + .config_map + .get("sql_header") + .map(|v| v.value().clone()); + let wap = ctx.inner.wap_plan.model(&model.common().unique_id); + let execution_model = if let Some(wap) = wap { + crate::wap::validate_runtime_sql_header(sql_header.as_ref())?; + crate::wap::preflight_stage(wap, ctx)?; + Some(wap.execution_model()?) + } else { + None + }; + let model = execution_model.as_ref().unwrap_or(model); let mut base_context = ctx.inner.base_context.clone(); add_task_context(&mut base_context, model.common(), &ctx.thread_id); @@ -330,11 +343,6 @@ pub fn execute_model_remote( return Ok(NodeStatus::NoOp); } - let sql_header = task_result - .config_map - .get("sql_header") - .map(|v| v.value().clone()); - // Traditional warehouse execution via Jinja materialization macros match materialize_model( &task_result.sql_instruction.sql, @@ -361,7 +369,7 @@ pub fn execute_model_remote( } // After successful materialization, create the latest version pointer view if applicable - if should_create_latest_version_pointer(model, ctx.runtime_config()) { + if wap.is_none() && should_create_latest_version_pointer(model, ctx.runtime_config()) { let relations_map = materialize_latest_version_pointer( model, model.node_adapter(), diff --git a/crates/dbt-tasks-sa/src/runnable/runnable/mod.rs b/crates/dbt-tasks-sa/src/runnable/runnable/mod.rs index 87df62098c7..991d7084a51 100644 --- a/crates/dbt-tasks-sa/src/runnable/runnable/mod.rs +++ b/crates/dbt-tasks-sa/src/runnable/runnable/mod.rs @@ -109,9 +109,25 @@ impl Task for RunTask { ) -> Pin> + Send + 'a>> { Box::pin(async move { let unique_id = self.node.unique_id(); + let is_wap_node = ctx.inner.wap_plan.contains_node(&unique_id); + let is_wap_stage = ctx.inner.wap_plan.model(&unique_id).is_some(); let mut result_receiver = { self.result_receiver.lock().take() }; let task_result = receive_task_result(&unique_id, &mut result_receiver)?; let start_time = chrono::Utc::now(); + if is_wap_stage { + // Keep a retryable provisional outcome even if cancellation interrupts staging. + ctx.inner.wap_stage_stats.insert( + unique_id.clone(), + Stat::new( + unique_id.clone(), + start_time.into(), + None, + NodeStatus::SkippedUpstreamFailed, + Some("WAP: working table not published".to_string()), + ctx.thread_id, + ), + ); + } // Status lines keep source path for Models for readability let display_path_kind = if self.node.resource_type() == NodeType::Model { NodePathKind::Definition @@ -140,7 +156,8 @@ impl Task for RunTask { // the original recording when the service was active. Replay // must be driven solely by recorded events (see // `maybe_replay_remote_run` / `maybe_replay_run_cache_clone`). - let cache_enabled = self.execution_path == RunExecutionPath::Remote + let cache_enabled = !is_wap_node + && self.execution_path == RunExecutionPath::Remote && !ctx.inner.run_cache_ctx.run_cache_service_requested && !is_replaying() && ctx.inner.arg.run_cache_mode.write_cache() @@ -208,7 +225,9 @@ impl Task for RunTask { // `sao_guard`), or Disabled. let run_cache_service_requested = ctx.inner.run_cache_ctx.run_cache_service_requested; - let decision = if let Some(decision) = replayed_cache_decision { + let decision = if is_wap_node { + RunCacheServiceDecision::Disabled + } else if let Some(decision) = replayed_cache_decision { decision } else if is_replaying() && ctx @@ -551,7 +570,7 @@ impl Task for RunTask { } }; - if result.is_ok() { + if result.is_ok() && !is_wap_node { run_cache_after_success_action( ctx, self.node.as_ref(), @@ -567,6 +586,16 @@ impl Task for RunTask { span_rows_affected = attrs.rows_affected.map(|n| n as i64); }); + if is_wap_stage && let Ok(status) = &result { + if let Some(mut stat) = ctx.inner.wap_stage_stats.get_mut(&unique_id) { + stat.rows_affected = span_rows_affected; + stat.end_time = std::time::SystemTime::now(); + } + // Stage completion only releases audits. The publish task owns the + // model's terminal result, progress completion, and usage reporting. + return Ok(status.clone()); + } + // Get status and insert stats // Note: Inner visit_run implementations may insert their own stats on success, // but we need to ensure stats are inserted even when errors occur early. @@ -604,6 +633,9 @@ impl Task for RunTask { // TODO: At some point, these should log as part of the same event let node_status = NodeStatus::Errored; let error_message = e.to_string(); + if is_wap_stage { + crate::wap::report_retained_candidate(ctx, &unique_id); + } report_completed( &NodeStatus::Errored, self.node.defined_at().cloned(), @@ -1076,7 +1108,7 @@ fn sao_status_for_task_status(task_status: &NodeStatus) -> Option<(SaoStatus, St }) } -fn emit_run_usage_stats( +pub(crate) fn emit_run_usage_stats( node: &dyn InternalDbtNodeAttributes, ctx: &TaskRunnerCtx, execution_path: RunExecutionPath, diff --git a/crates/dbt-tasks-sa/src/runnable/test.rs b/crates/dbt-tasks-sa/src/runnable/test.rs index 48969942063..229695c282b 100644 --- a/crates/dbt-tasks-sa/src/runnable/test.rs +++ b/crates/dbt-tasks-sa/src/runnable/test.rs @@ -582,6 +582,7 @@ pub fn execute_test_remote( let mut base_context = ctx.inner.base_context.clone(); add_task_context(&mut base_context, test.common(), &ctx.thread_id); + ctx.apply_wap_ref_overrides(unique_id, &mut base_context)?; let sql_instruction = match &task_result.lp_instruction { Some(_) => { diff --git a/crates/dbt-tasks-sa/src/visitor.rs b/crates/dbt-tasks-sa/src/visitor.rs index 00ab514f5f6..c25596480ce 100644 --- a/crates/dbt-tasks-sa/src/visitor.rs +++ b/crates/dbt-tasks-sa/src/visitor.rs @@ -312,6 +312,24 @@ fn record_skipped_stats( let now = chrono::Utc::now(); for node in skipped_nodes { let unique_id = node.common().unique_id.clone(); + if ctx.inner.wap_plan.model(&unique_id).is_some() { + // The stage and publish tasks share a model identity. A skipped + // publication must not overwrite the stage's execution error. + if ctx.inner.run_stats.contains_key(&unique_id) { + continue; + } + if let Some((_, mut stat)) = ctx.inner.wap_stage_stats.remove(&unique_id) { + stat.end_time = std::time::SystemTime::now(); + stat.status = NodeStatus::SkippedUpstreamFailed; + stat.message = Some( + "WAP: table not published because an audit or upstream phase failed" + .to_string(), + ); + ctx.inner.run_stats.insert(unique_id.clone(), stat); + crate::wap::report_retained_candidate(ctx, &unique_id); + continue; + } + } let (status, message) = match skip_reason { SkipReason::Reused => ( NodeStatus::ReusedNoChanges("Model reused".to_string()), @@ -546,6 +564,20 @@ async fn visit( e: CancelledError, waiting: HashMap| -> FsResult<()> { + // Completed stages are not terminal model successes. Include them in + // cancellation artifacts even when their publish task never started. + for stage in ctx.inner.wap_stage_stats.iter() { + if !ctx.inner.run_stats.contains_key(stage.key()) { + let mut stat = stage.value().clone(); + stat.end_time = std::time::SystemTime::now(); + stat.status = NodeStatus::Errored; + let response = ctx.inner.main_adapter_responses.get(stage.key()); + stat.message = + Some(crate::wap::interruption_message(response.as_deref()).to_string()); + ctx.inner.run_stats.insert(stage.key().clone(), stat); + crate::wap::report_retained_candidate(ctx, stage.key()); + } + } // Any tasks that are running and waiting need to be reported as failed. for (node_idx, task_span) in waiting { let maybe_node = schedule.node_weight(node_idx); @@ -780,6 +812,15 @@ fn report_node_evaluation( node: &dyn Task, node_status: Option<&NodeStatus>, ) { + if node.task_type() == "run" + && ctx.inner.wap_plan.model(node.work_node_id()).is_some() + && matches!( + node_status, + Some(NodeStatus::Succeeded | NodeStatus::SucceededWithWarning) + ) + { + return; + } if let Some(reporter) = &ctx.inner.arg.io.status_reporter { // For successful status, emit a node evaluation event let node_outcome = match node_status { diff --git a/crates/dbt-tasks-sa/src/wap.rs b/crates/dbt-tasks-sa/src/wap.rs new file mode 100644 index 00000000000..592d0e43364 --- /dev/null +++ b/crates/dbt-tasks-sa/src/wap.rs @@ -0,0 +1,976 @@ +//! Publication of audited working tables. The manifest always names the public relation. + +use std::collections::BTreeMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::SystemTime; + +use arrow::array::{Array, RecordBatch, StringArray}; +use dbt_adapter::Adapter; +use dbt_adapter::catalog_relation::CatalogRelation; +use dbt_adapter::connection::drop_thread_local_connection; +use dbt_adapter::record_batch::RecordBatchExt; +use dbt_adapter::relation::RelationObject; +use dbt_adapter::response::AdapterResponse; +use dbt_adapter_core::AdapterType; +use dbt_agate::AgateTable; +use dbt_common::stats::{NodeStatus, Stat}; +use dbt_common::status_reporter::report_completed; +use dbt_common::tracing::dbt_emit::{ + emit_error_log_from_fs_error, emit_info_log_message, emit_warn_log_message, +}; +use dbt_common::{ErrorCode, FsError, FsResult, fs_err}; +use dbt_jinja_utils::phases::run::build_run_node_context; +use dbt_jinja_utils::utils::add_task_context; +use dbt_schemas::schemas::dbt_catalogs_v2::{CatalogType, TableFormat}; +use dbt_schemas::schemas::relations::base::BaseRelation; +use dbt_schemas::schemas::{DbtModel, InternalDbtNode, InternalDbtNodeAttributes, NodePathKind}; +use dbt_tasks_core::context::TaskRunnerCtx; +use dbt_tasks_core::run_cache::run_cache_service::evict_node_metadata_for_untracked_rebuild; +use dbt_tasks_core::task::{TP, Task, TaskOp}; +use dbt_tasks_core::wap::WapModel; +use dbt_telemetry::{ExecutionPhase, NodeType}; +use minijinja::Value; + +use crate::materialize::{ + apply_node_overrides, materialize_latest_version_pointer, reset_node_overrides, + should_create_latest_version_pointer, +}; +use crate::runnable::cache::cache_materialization_return_value; +use crate::runnable::runnable::{RunExecutionPath, emit_run_usage_stats}; + +/// The model's terminal task: no public success is recorded by the preceding stage task. +pub struct PublishTask { + model: Arc, +} + +impl PublishTask { + pub fn new(model: Arc) -> Self { + Self { model } + } +} + +impl Task for PublishTask { + fn run_task<'a>( + &'a self, + ctx: &'a mut TaskRunnerCtx, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let unique_id = &self.model.common().unique_id; + let model = Arc::clone(&self.model); + let publish_ctx = ctx.clone(); + let result = TaskOp::Blocking(Box::new(move || publish_model(&model, &publish_ctx))) + .run() + .await + .and_then(|result| result); + + let mut stat = ctx + .inner + .wap_stage_stats + .remove(unique_id) + .map(|(_, stat)| stat) + .unwrap_or_else(|| { + Stat::new( + unique_id.clone(), + SystemTime::now(), + None, + NodeStatus::Errored, + None, + ctx.thread_id, + ) + }); + stat.end_time = SystemTime::now(); + match result { + Ok(()) => { + stat.status = NodeStatus::Succeeded; + stat.message = Some("WAP: all audits passed; table published".to_string()); + } + Err(error) => { + stat.status = NodeStatus::Errored; + stat.message = Some(error.to_string()); + emit_error_log_from_fs_error(*error); + report_retained_candidate(ctx, unique_id); + } + } + let status = stat.status.clone(); + ctx.inner.run_stats.insert(unique_id.clone(), stat); + if ctx.inner.arg.io.send_anonymous_usage_stats { + emit_run_usage_stats(self.model.as_ref(), ctx, RunExecutionPath::Remote); + } + report_completed( + &status, + self.model.defined_at().cloned(), + &self + .model + .get_node_path( + NodePathKind::Definition, + &ctx.inner.arg.io.in_dir, + &ctx.inner.arg.io.out_dir, + ) + .display() + .to_string(), + false, + ctx.inner.arg.io.status_reporter.as_ref(), + ); + Ok(status) + }) + } + + fn resource_type(&self) -> NodeType { + NodeType::Model + } + + fn task_type(&self) -> &str { + "wap_publish_run" + } + + fn work_node_id(&self) -> &str { + &self.model.common().unique_id + } + + fn dbt_nodes(&self) -> Vec> { + vec![self.model.clone()] + } + + fn task_phase(&self) -> Option { + Some(TP::Run) + } +} + +fn require_passing_audits<'a>( + audits: impl IntoIterator)>, +) -> FsResult<()> { + let mut count = 0; + for (id, status) in audits { + count += 1; + if status != Some(NodeStatus::TestPassed) { + return Err(fs_err!( + ErrorCode::ExecutionError, + "WAP: table not published because audit '{id}' did not PASS (result: {}). \ + Warnings, missing results, and skipped audits also prevent publication", + status.map_or_else(|| "missing".to_string(), |s| s.default_message()) + )); + } + } + if count == 0 { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP requires at least one passing audit before publication" + )); + } + Ok(()) +} + +fn model_context(model: &DbtModel, ctx: &TaskRunnerCtx) -> BTreeMap { + let mut base_context = ctx.inner.base_context.clone(); + add_task_context(&mut base_context, model.common(), &ctx.thread_id); + build_run_node_context( + model, + &model.deprecated_config, + model.node_adapter(), + None, + &base_context, + &ctx.inner.arg.io, + ExecutionPhase::Run, + None, + ctx.runtime_config().dependencies.keys().cloned().collect(), + ) + .0 +} + +fn eval( + ctx: &TaskRunnerCtx, + context: &BTreeMap, + expression: &str, +) -> FsResult { + ctx.env.compile_expression(expression)?.eval(context, &[]) +} + +fn relations(wap: &WapModel) -> FsResult<(Arc, Arc)> { + Ok(( + Arc::from(wap.public_relation()?), + Arc::from(wap.candidate_relation()?), + )) +} + +/// Rendering can set a header after the parsed model configuration was validated. +pub(crate) fn validate_runtime_sql_header(sql_header: Option<&Value>) -> FsResult<()> { + if sql_header.is_some_and(|header| { + !header.is_none() + && !header.is_undefined() + && !header.as_str().is_some_and(|text| text.trim().is_empty()) + }) { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP does not support sql_header, including headers added while rendering" + )); + } + Ok(()) +} + +/// Check the live warehouse, not the adapter cache, before touching a working table. +pub(crate) fn preflight_stage(wap: &WapModel, ctx: &TaskRunnerCtx) -> FsResult<()> { + if ctx + .inner + .materialization_resolver + .is_custom_materialization("table", wap.model.node_adapter()) + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP requires the built-in Snowflake table materialization: {}", + wap.model.common().unique_id + )); + } + let adapter = ctx + .env + .get_base_adapter() + .ok_or_else(|| fs_err!(ErrorCode::Unexpected, "Missing adapter for WAP staging"))?; + wap.validate_materialization_quoting(adapter.engine().quoting())?; + let (target, candidate) = relations(wap)?; + let mut context = model_context(&wap.model, ctx); + resolved_transient(ctx, &context)?; + inspect_relations(ctx, &mut context, &target, &candidate, false)?; + emit_info_log_message(format!( + "WAP: building {} in working table {}", + target.render_self_as_str(), + candidate.render_self_as_str() + )); + Ok(()) +} + +fn inspect_relations( + ctx: &TaskRunnerCtx, + context: &mut BTreeMap, + target: &Arc, + candidate: &Arc, + candidate_must_exist: bool, +) -> FsResult { + let mut target_exists = false; + for (relation, expectation) in [ + (target, RelationExpectation::OptionalTable), + ( + candidate, + if candidate_must_exist { + RelationExpectation::ExistingTable + } else { + RelationExpectation::Absent + }, + ), + ] { + let sql = match relation.adapter_type() { + AdapterType::Snowflake => { + dbt_adapter::metadata::snowflake::relation_lookup_sql(relation.as_ref())? + } + _ => return Err(fs_err!(ErrorCode::InvalidConfig, "WAP requires Snowflake")), + }; + let objects = fetch_relation_metadata(ctx, context, sql)?; + let name = relation + .identifier_as_resolved_str() + .map_err(|e| FsError::from_jinja_err(e, "WAP relation identifier"))?; + let exists = inspect_relation_rows(&objects, &name, expectation)?; + if exists { + let sql = match relation.adapter_type() { + AdapterType::Snowflake => { + dbt_adapter::metadata::snowflake::table_lookup_sql(relation.as_ref())? + } + _ => return Err(fs_err!(ErrorCode::InvalidConfig, "WAP requires Snowflake")), + }; + let tables = fetch_relation_metadata(ctx, context, sql)?; + inspect_table_rows(&tables, &name)?; + } + if matches!(expectation, RelationExpectation::OptionalTable) { + target_exists = exists; + } + } + Ok(target_exists) +} + +fn fetch_relation_metadata( + ctx: &TaskRunnerCtx, + context: &mut BTreeMap, + sql: String, +) -> FsResult> { + context.insert("__wap_sql".to_string(), Value::from(sql)); + let result = eval( + ctx, + context, + "adapter.execute(__wap_sql, auto_begin=false, fetch=true)", + )?; + let value = result + .get_item_by_index(1) + .map_err(|e| FsError::from_jinja_err(e, "WAP relation lookup"))?; + let table = value.downcast_object::().ok_or_else(|| { + fs_err!( + ErrorCode::Unexpected, + "WAP relation inspection did not return a table" + ) + })?; + Ok(table.to_record_batch()) +} + +#[derive(Clone, Copy)] +enum RelationExpectation { + Absent, + OptionalTable, + ExistingTable, +} + +fn inspect_relation_rows( + batch: &RecordBatch, + identifier: &str, + expectation: RelationExpectation, +) -> FsResult { + let row = exact_relation_row(batch, identifier)?; + if let Some(row) = row { + if matches!(expectation, RelationExpectation::Absent) { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP working relation '{identifier}' already exists; it will not be overwritten" + )); + } + require_relation_kind(batch, row, identifier, &["TABLE"])?; + require_disabled_flag(batch, row, identifier, "is_dynamic")?; + // The adapter also tolerates this column being absent on accounts without + // interactive tables. A present but unknown value cannot establish safety. + if batch.column_by_name("is_interactive").is_some() { + require_disabled_flag(batch, row, identifier, "is_interactive")?; + } + } + if matches!(expectation, RelationExpectation::ExistingTable) && row.is_none() { + return Err(fs_err!( + ErrorCode::ExecutionError, + "WAP audited working table '{identifier}' is missing; table not published" + )); + } + Ok(row.is_some()) +} + +fn inspect_table_rows(batch: &RecordBatch, identifier: &str) -> FsResult<()> { + let row = exact_relation_row(batch, identifier)?.ok_or_else(|| { + fs_err!( + ErrorCode::ExecutionError, + "WAP table '{identifier}' disappeared during inspection; table not published" + ) + })?; + require_relation_kind(batch, row, identifier, &["TABLE", "TRANSIENT"])?; + for flag in ["is_external", "is_event", "is_hybrid", "is_iceberg"] { + require_disabled_flag(batch, row, identifier, flag)?; + } + if batch.column_by_name("is_immutable").is_some() { + require_disabled_flag(batch, row, identifier, "is_immutable")?; + } + Ok(()) +} + +fn exact_relation_row(batch: &RecordBatch, identifier: &str) -> FsResult> { + let names = batch.column_values::("name")?; + let mut matched = None; + for row in 0..batch.num_rows() { + if names.is_null(row) { + return Err(fs_err!( + ErrorCode::Unexpected, + "WAP cannot inspect '{identifier}': SHOW returned a NULL relation name" + )); + } + if names.value(row) == identifier && matched.replace(row).is_some() { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP found shadowing relations for '{identifier}'; remove temporary name collisions first" + )); + } + } + Ok(matched) +} + +fn require_relation_kind( + batch: &RecordBatch, + row: usize, + identifier: &str, + allowed: &[&str], +) -> FsResult<()> { + let kinds = batch.column_values::("kind")?; + if kinds.is_null(row) + || !allowed + .iter() + .any(|kind| kinds.value(row).eq_ignore_ascii_case(kind)) + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP requires an ordinary native table; '{identifier}' has unsupported or unknown kind: {}", + if kinds.is_null(row) { + "NULL" + } else { + kinds.value(row) + } + )); + } + Ok(()) +} + +fn require_disabled_flag( + batch: &RecordBatch, + row: usize, + identifier: &str, + flag: &str, +) -> FsResult<()> { + let values = batch.column_values::(flag)?; + if values.is_null(row) + || (!values.value(row).eq_ignore_ascii_case("n") + && !values.value(row).eq_ignore_ascii_case("false")) + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP does not support '{identifier}' with {flag}={}", + if values.is_null(row) { + "NULL" + } else { + values.value(row) + } + )); + } + Ok(()) +} + +fn resolved_transient(ctx: &TaskRunnerCtx, context: &BTreeMap) -> FsResult { + let value = eval(ctx, context, "adapter.build_catalog_relation(config.model)")?; + let catalog = value.downcast_object::().ok_or_else(|| { + fs_err!( + ErrorCode::Unexpected, + "WAP could not resolve the model catalog" + ) + })?; + if catalog.catalog_type != CatalogType::SnowflakeNative + || catalog.table_format != TableFormat::Default + { + return Err(fs_err!( + ErrorCode::InvalidConfig, + "WAP requires a native Snowflake catalog and table format" + )); + } + catalog.is_transient.ok_or_else(|| { + fs_err!( + ErrorCode::InvalidConfig, + "WAP could not resolve the table lifecycle" + ) + }) +} + +fn execute_sql( + ctx: &TaskRunnerCtx, + context: &mut BTreeMap, + sql: String, +) -> FsResult { + context.insert("__wap_sql".to_string(), Value::from(sql)); + let result = eval( + ctx, + context, + "adapter.execute(__wap_sql, auto_begin=false, fetch=false)", + )?; + let response = result + .get_item_by_index(0) + .map_err(|e| FsError::from_jinja_err(e, "WAP statement response"))?; + Ok(AdapterResponse::try_from(response) + .map_err(|e| FsError::from_jinja_err(e, "WAP adapter response"))?) +} + +/// Only these operations may surround the public clone. Keeping the ordering in one +/// place lets failure tests exercise the same cleanup decisions as warehouse runs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PublicationStep { + Prepare, + SetQueryTag, + Clone, + Finalize, + ResetQueryTag, + ResetOverrides, + DropCandidate, +} + +fn run_publication_steps<'a>( + audits: impl IntoIterator)>, + target_name: &str, + candidate_name: &str, + mut execute: impl FnMut(PublicationStep) -> FsResult<()>, +) -> FsResult<()> { + use PublicationStep::*; + + require_passing_audits(audits)?; + let mut published = false; + let result = (|| { + execute(Prepare)?; + execute(SetQueryTag)?; + let result = execute(Clone).and_then(|()| { + published = true; + execute(Finalize) + }); + // Reset even when cloning or finalization failed; preserve the original error. + let reset_tag = execute(ResetQueryTag); + result.and(reset_tag) + })(); + let reset_overrides = execute(ResetOverrides); + result.and(reset_overrides).map_err(|error| { + if published { + let message = format!( + "WAP: {target_name} was published; finalization failed: {error}. \ + Working table retained: {candidate_name}" + ); + Box::new(error.with_context(message)) + } else { + error + } + })?; + + // A cleanup error cannot undo a successful publication. + if let Err(error) = execute(DropCandidate) { + emit_warn_log_message( + ErrorCode::ExecutionError, + format!( + "WAP published {target_name}; could not remove working table {candidate_name}: {error}" + ), + ); + } + Ok(()) +} + +fn publication_model(model: &DbtModel, transient: bool, target_exists: bool) -> DbtModel { + let mut publication = model.clone(); + let config = &mut publication.deprecated_config.__warehouse_specific_config__; + // CTAS resolves the catalog lifecycle, whereas the clone macro reads config.transient. + config.transient = Some(transient); + config.copy_grants = Some(target_exists && config.copy_grants == Some(true)); + publication +} + +fn publication_context( + model: &DbtModel, + ctx: &TaskRunnerCtx, + target: &Arc, + candidate: &Arc, +) -> FsResult> { + let mut context = model_context(model, ctx); + let target_exists = inspect_relations(ctx, &mut context, target, candidate, true)?; + let transient = resolved_transient(ctx, &context)?; + let publication = publication_model(model, transient, target_exists); + context = model_context(&publication, ctx); + context.insert( + "__wap_candidate".to_string(), + RelationObject::new(Arc::clone(candidate)).into_value(), + ); + context.insert( + "__wap_target_exists".to_string(), + Value::from(target_exists), + ); + Ok(context) +} + +fn finalize_publication( + model: &DbtModel, + ctx: &TaskRunnerCtx, + adapter: &Adapter, + context: &mut BTreeMap, + target: &Arc, +) -> FsResult<()> { + if ctx.inner.run_cache_ctx.run_cache_service_requested { + evict_node_metadata_for_untracked_rebuild(ctx, model); + } + adapter + .cache_added(&ctx.env.empty_state(), Arc::clone(target)) + .map_err(|e| FsError::from_jinja_err(e, "WAP published relation cache"))?; + eval( + ctx, + context, + "apply_grants(this, config.get('grants'), should_revoke=__wap_target_exists and config.get('copy_grants', false))", + )?; + if model + .deprecated_config + .__warehouse_specific_config__ + .automatic_clustering + == Some(true) + && eval(ctx, context, "config.get('cluster_by')")?.is_true() + { + execute_sql( + ctx, + context, + format!( + "alter table {} resume recluster", + target.render_self_as_str() + ), + )?; + } + if should_create_latest_version_pointer(model, ctx.runtime_config()) { + let mut base = ctx.inner.base_context.clone(); + add_task_context(&mut base, model.common(), &ctx.thread_id); + let result = materialize_latest_version_pointer( + model, + model.node_adapter(), + ctx.runtime_config(), + &ctx.inner.materialization_resolver, + ctx.env.clone(), + &base, + &ctx.inner.arg.io, + )?; + cache_materialization_return_value(ctx.env.clone(), &result) + .map_err(|e| FsError::from_jinja_err(e, "WAP version pointer cache"))?; + } + Ok(()) +} + +fn publish_model(model: &DbtModel, ctx: &TaskRunnerCtx) -> FsResult<()> { + let unique_id = &model.common().unique_id; + let wap = ctx + .inner + .wap_plan + .model(unique_id) + .ok_or_else(|| fs_err!(ErrorCode::Unexpected, "Missing WAP plan for '{unique_id}'"))?; + let (target, candidate) = relations(wap)?; + let adapter = ctx + .env + .get_base_adapter() + .ok_or_else(|| fs_err!(ErrorCode::Unexpected, "Missing adapter for WAP publication"))?; + let mut context = BTreeMap::new(); + let mut overrides = Vec::new(); + let audits = wap.audit_ids.iter().map(|id| { + let status = ctx.inner.run_stats.get(id).map(|stat| stat.status.clone()); + (id.as_str(), status) + }); + + run_publication_steps( + audits, + &target.render_self_as_str(), + &candidate.render_self_as_str(), + |step| { + match step { + PublicationStep::Prepare => { + context = publication_context(model, ctx, &target, &candidate)?; + overrides = apply_node_overrides( + &adapter, + model.node_adapter(), + model + .__adapter_attr__ + .snowflake_attr + .as_ref() + .and_then(|a| a.snowflake_warehouse.clone()), + &model.__base_attr__.database, + unique_id, + )?; + } + PublicationStep::SetQueryTag => { + let query_tag = eval(ctx, &context, "set_query_tag()")?; + context.insert("__wap_original_query_tag".to_string(), query_tag); + } + PublicationStep::Clone => { + let sql = eval( + ctx, + &context, + "dbt_snowflake.snowflake__create_or_replace_clone(this, __wap_candidate)", + )? + .to_string(); + adapter.cancellation_token().check_cancellation()?; + // A failed submission can have committed without returning a response. + let response = execute_sql(ctx, &mut context, sql).map_err(|error| { + let message = format!( + "WAP publication failed for {}: {error}. Publication may have committed; \ + inspect Snowflake query history before retrying. Working table: {}", + target.render_self_as_str(), + candidate.render_self_as_str() + ); + Box::new(error.with_context(message)) + })?; + ctx.inner.main_adapter_responses.insert( + unique_id.clone(), + response + .with("wap_published", true) + .with("wap_candidate", candidate.render_self_as_str()), + ); + } + PublicationStep::Finalize => { + finalize_publication(model, ctx, &adapter, &mut context, &target)?; + } + PublicationStep::ResetQueryTag => { + eval(ctx, &context, "unset_query_tag(__wap_original_query_tag)")?; + } + PublicationStep::ResetOverrides => { + reset_node_overrides(&adapter, unique_id, &overrides)?; + } + PublicationStep::DropCandidate => { + eval(ctx, &context, "adapter.drop_relation(__wap_candidate)")?; + } + } + Ok(()) + }, + ) + .inspect_err(|_| { + // Session changes can have committed even if setting/restoring the query + // tag failed. Never hand that uncertain connection to the next model. + drop_thread_local_connection(); + }) +} + +pub(crate) fn interruption_message(response: Option<&AdapterResponse>) -> &'static str { + if response.is_some_and(|response| { + Value::from_serialize(response) + .get_attr("wap_published") + .is_ok_and(|value| value.is_true()) + }) { + "WAP table was published; finalization was interrupted. Working table retained if present" + } else { + "WAP interrupted before completion; working table retained if created. \ + If publication was in flight, verify its outcome in Snowflake before retrying" + } +} + +pub(crate) fn report_retained_candidate(ctx: &TaskRunnerCtx, unique_id: &str) { + if let Some(wap) = ctx.inner.wap_plan.model(unique_id) + && let Ok((_, candidate)) = relations(wap) + { + emit_info_log_message(format!( + "WAP working table retained if created: {} (model {unique_id})", + candidate.render_self_as_str() + )); + } +} + +#[cfg(test)] +mod clone_tests; +#[cfg(test)] +mod header_tests; +#[cfg(test)] +mod publication_tests; + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{DataType, Field, Schema}; + + #[test] + fn wap_cancellation_preserves_confirmed_publication_outcome() { + let published = AdapterResponse::new().with("wap_published", true); + assert!(interruption_message(Some(&published)).contains("was published")); + assert!(interruption_message(Some(&published)).contains("finalization was interrupted")); + assert!(interruption_message(None).contains("verify its outcome")); + assert!(interruption_message(Some(&AdapterResponse::new())).contains("verify its outcome")); + } + + #[test] + fn wap_requires_actual_pass_for_every_audit() { + assert!(require_passing_audits([("a", Some(NodeStatus::TestPassed))]).is_ok()); + for status in [ + None, + Some(NodeStatus::TestWarned), + Some(NodeStatus::Errored), + Some(NodeStatus::SkippedUpstreamFailed), + Some(NodeStatus::StaticallyCheckedDataTest), + Some(NodeStatus::ReusedNoChanges("cached".to_string())), + ] { + assert!( + require_passing_audits([ + ("passed", Some(NodeStatus::TestPassed)), + ("blocking", status), + ]) + .is_err() + ); + } + assert!(require_passing_audits([]).is_err()); + } + + fn objects(names: &[&str], kinds: &[&str], dynamic: &[&str]) -> RecordBatch { + let columns = ["name", "kind", "is_dynamic"]; + let schema = Schema::new( + columns + .map(|name| Field::new(name, DataType::Utf8, false)) + .to_vec(), + ); + RecordBatch::try_new( + Arc::new(schema), + [names, kinds, dynamic] + .into_iter() + .map(|values| Arc::new(StringArray::from(values.to_vec())) as _) + .collect(), + ) + .unwrap() + } + + fn metadata_row(columns: &[(&str, Option<&str>)]) -> RecordBatch { + let schema = Schema::new( + columns + .iter() + .map(|(name, _)| Field::new(*name, DataType::Utf8, true)) + .collect::>(), + ); + RecordBatch::try_new( + Arc::new(schema), + columns + .iter() + .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as _) + .collect(), + ) + .unwrap() + } + + fn table(kind: &str, overrides: &[(&str, Option<&str>)]) -> RecordBatch { + let mut columns = vec![ + ("name", Some("PUBLIC")), + ("kind", Some(kind)), + ("is_external", Some("N")), + ("is_event", Some("N")), + ("is_hybrid", Some("N")), + ("is_iceberg", Some("N")), + ("is_immutable", Some("N")), + ]; + for (name, value) in &mut columns { + if let Some((_, replacement)) = overrides.iter().find(|(flag, _)| *flag == *name) { + *value = *replacement; + } + } + metadata_row(&columns) + } + + #[test] + fn wap_preflight_preserves_collisions_and_unsupported_targets() { + use RelationExpectation::{Absent, ExistingTable, OptionalTable}; + let empty = objects(&[], &[], &[]); + assert!(!inspect_relation_rows(&empty, "PUBLIC", OptionalTable).unwrap()); + assert!(inspect_relation_rows(&empty, "WORK", ExistingTable).is_err()); + let existing = objects(&["PUBLIC"], &["TABLE"], &["N"]); + assert!(inspect_relation_rows(&existing, "PUBLIC", OptionalTable).unwrap()); + let ready = objects(&["PUBLIC", "WORK"], &["TABLE", "TABLE"], &["N", "N"]); + assert!(inspect_relation_rows(&ready, "WORK", Absent).is_err()); + assert!(inspect_relation_rows(&ready, "WORK", ExistingTable).unwrap()); + for kind in ["VIEW", "TEMPORARY", "EXTERNAL TABLE"] { + assert!( + inspect_relation_rows( + &objects(&["PUBLIC"], &[kind], &["N"]), + "PUBLIC", + OptionalTable + ) + .is_err() + ); + } + assert!( + inspect_relation_rows( + &objects(&["PUBLIC"], &["TABLE"], &["Y"]), + "PUBLIC", + OptionalTable + ) + .is_err() + ); + assert!( + inspect_relation_rows( + &objects(&["PUBLIC", "PUBLIC"], &["TABLE", "TABLE"], &["N", "N"]), + "PUBLIC", + OptionalTable + ) + .is_err() + ); + } + + #[test] + fn wap_requires_native_permanent_or_transient_table_details() { + for kind in ["TABLE", "TRANSIENT"] { + assert!(inspect_table_rows(&table(kind, &[]), "PUBLIC").is_ok()); + } + for kind in [ + "TEMPORARY", + "EXTERNAL TABLE", + "VIEW", + "INTERACTIVE TABLE", + "", + ] { + assert!(inspect_table_rows(&table(kind, &[]), "PUBLIC").is_err()); + } + for flag in [ + "is_external", + "is_event", + "is_hybrid", + "is_iceberg", + "is_immutable", + ] { + assert!(inspect_table_rows(&table("TABLE", &[(flag, Some("Y"))]), "PUBLIC").is_err()); + } + assert!( + inspect_table_rows( + &table("TABLE", &[("name", Some("PUBLIC_SUFFIX"))]), + "PUBLIC" + ) + .is_err() + ); + assert!( + inspect_table_rows( + &objects(&["PUBLIC", "PUBLIC"], &["TABLE", "TABLE"], &["N", "N"]), + "PUBLIC" + ) + .is_err() + ); + } + + #[test] + fn wap_rejects_missing_or_unknown_type_metadata() { + for flag in ["is_external", "is_event", "is_hybrid", "is_iceberg"] { + let batch = table("TABLE", &[]); + let indices = batch + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| (field.name() != flag).then_some(index)) + .collect::>(); + assert!(inspect_table_rows(&batch.project(&indices).unwrap(), "PUBLIC").is_err()); + for value in [None, Some(""), Some("unknown")] { + assert!(inspect_table_rows(&table("TABLE", &[(flag, value)]), "PUBLIC").is_err()); + } + assert!( + inspect_table_rows(&table("TABLE", &[(flag, Some("false"))]), "PUBLIC").is_ok() + ); + } + for field in ["name", "kind"] { + assert!(inspect_table_rows(&table("TABLE", &[(field, None)]), "PUBLIC").is_err()); + } + for value in [None, Some(""), Some("unknown"), Some("Y")] { + let batch = metadata_row(&[ + ("name", Some("PUBLIC")), + ("kind", Some("TABLE")), + ("is_dynamic", value), + ]); + assert!( + inspect_relation_rows(&batch, "PUBLIC", RelationExpectation::OptionalTable) + .is_err() + ); + } + let missing_dynamic = metadata_row(&[("name", Some("PUBLIC")), ("kind", Some("TABLE"))]); + assert!( + inspect_relation_rows( + &missing_dynamic, + "PUBLIC", + RelationExpectation::OptionalTable + ) + .is_err() + ); + } + + #[test] + fn wap_rejects_interactive_tables_and_unknown_interactive_flags() { + for value in [None, Some(""), Some("unknown"), Some("Y")] { + let batch = metadata_row(&[ + ("name", Some("PUBLIC")), + ("kind", Some("TABLE")), + ("is_dynamic", Some("N")), + ("is_interactive", value), + ]); + assert!( + inspect_relation_rows(&batch, "PUBLIC", RelationExpectation::OptionalTable) + .is_err() + ); + } + let supported = metadata_row(&[ + ("name", Some("PUBLIC")), + ("kind", Some("TABLE")), + ("is_dynamic", Some("N")), + ("is_interactive", Some("N")), + ]); + assert!( + inspect_relation_rows(&supported, "PUBLIC", RelationExpectation::OptionalTable) + .unwrap() + ); + } +} diff --git a/crates/dbt-tasks-sa/src/wap/clone_tests.rs b/crates/dbt-tasks-sa/src/wap/clone_tests.rs new file mode 100644 index 00000000000..e47d6bdded3 --- /dev/null +++ b/crates/dbt-tasks-sa/src/wap/clone_tests.rs @@ -0,0 +1,181 @@ +//! Offline rendering checks for the actual Snowflake macro and runtime config. +//! Warehouse atomicity, privileges, and storage behavior require live acceptance tests. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use dbt_adapter::relation::RelationObject; +use dbt_adapter_core::AdapterType; +use dbt_common::io_args::IoArgs; +use dbt_jinja_utils::jinja_environment::JinjaEnv; +use dbt_jinja_utils::phases::run::{RunConfig, build_run_node_context}; +use dbt_schemas::schemas::DbtModel; +use dbt_schemas::schemas::common::{DbtMaterialization, ResolvedQuoting}; +use dbt_tasks_core::wap::WapModel; +use dbt_telemetry::ExecutionPhase; +use minijinja::{Environment, Value}; + +const CLONE_MACRO: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-snowflake/macros/materializations/clone.sql" +); + +fn fixture(copy_grants: Option) -> (tempfile::TempDir, IoArgs, WapModel) { + let directory = tempfile::tempdir().unwrap(); + std::fs::create_dir(directory.path().join("models")).unwrap(); + std::fs::write(directory.path().join("models/orders.sql"), "select 1 as id").unwrap(); + let io = IoArgs { + in_dir: directory.path().to_path_buf(), + out_dir: directory.path().join("target"), + ..Default::default() + }; + let mut model = DbtModel::default(); + model.__common_attr__.unique_id = "model.wap_test.orders".to_owned(); + model.__common_attr__.name = "orders".to_owned(); + model.__common_attr__.package_name = "wap_test".to_owned(); + model.__common_attr__.language = Some("sql".to_owned()); + model.__common_attr__.path = "orders.sql".into(); + model.__common_attr__.original_file_path = "models/orders.sql".into(); + model.__base_attr__.adapter = AdapterType::Snowflake; + model.__base_attr__.materialized = DbtMaterialization::Table; + model.__base_attr__.database = "My\"Database".to_owned(); + model.__base_attr__.schema = "My.Schema".to_owned(); + model.__base_attr__.alias = "Orders".to_owned(); + model.__base_attr__.quoting = ResolvedQuoting::trues(); + model.deprecated_config.wap = Some(true); + model.deprecated_config.materialized = Some(DbtMaterialization::Table); + model.deprecated_config.alias = Some("Orders".to_owned()); + model + .deprecated_config + .__warehouse_specific_config__ + .copy_grants = copy_grants; + model.deprecated_config.grants = + serde_json::from_value(serde_json::json!({"select": ["READER"]})).unwrap(); + let wap = WapModel { + model: Arc::new(model), + candidate_identifier: "__DBT_WAP_TEST".to_owned(), + audit_ids: BTreeSet::new(), + }; + (directory, io, wap) +} + +fn runtime_context(model: &DbtModel, io: &IoArgs) -> BTreeMap { + let (context, _) = build_run_node_context( + model, + &model.deprecated_config, + AdapterType::Snowflake, + None, + &BTreeMap::new(), + io, + ExecutionPhase::Run, + None, + BTreeSet::new(), + ); + assert!( + context["config"] + .downcast_object_ref::() + .is_some() + ); + context +} + +fn render_clone(wap: &WapModel, io: &IoArgs, transient: bool, target_exists: bool) -> String { + let model = super::publication_model(&wap.model, transient, target_exists); + let mut context = runtime_context(&model, io); + context.insert( + "__wap_candidate".to_owned(), + RelationObject::new(wap.candidate_relation().unwrap().into()).into_value(), + ); + let template = format!( + "{CLONE_MACRO}\n{{{{ snowflake__create_or_replace_clone(this, __wap_candidate) }}}}" + ); + let rendered = JinjaEnv::new(Environment::new()) + .render_str(&template, &context, &[]) + .unwrap(); + rendered.split_whitespace().collect::>().join(" ") +} + +#[test] +fn actual_clone_macro_uses_resolved_lifecycle_and_public_destination() { + let (_directory, io, mut wap) = fixture(Some(false)); + for transient in [false, true] { + // The catalog's resolved lifecycle must win over the original model option. + Arc::make_mut(&mut wap.model) + .deprecated_config + .__warehouse_specific_config__ + .transient = Some(!transient); + let rendered = render_clone(&wap, &io, transient, true); + let lifecycle = if transient { "transient " } else { "" }; + assert_eq!( + rendered, + format!( + "create or replace {lifecycle}table \"My\"\"Database\".\"My.Schema\".\"Orders\" \ + clone \"My\"\"Database\".\"My.Schema\".\"__DBT_WAP_TEST\"" + ) + ); + assert_eq!(wap.model.__base_attr__.alias, "Orders"); + assert_eq!( + wap.model + .deprecated_config + .__warehouse_specific_config__ + .transient, + Some(!transient), + "publication must leave the canonical model unchanged" + ); + } +} + +#[test] +fn actual_clone_macro_copies_grants_only_for_configured_replacements() { + for configured in [None, Some(false), Some(true)] { + let (_directory, io, wap) = fixture(configured); + for target_exists in [false, true] { + let rendered = render_clone(&wap, &io, true, target_exists); + assert_eq!( + rendered.ends_with(" copy grants"), + target_exists && configured == Some(true), + "copy_grants={configured:?}, target_exists={target_exists}: {rendered}" + ); + assert!(!rendered.contains("copy tags")); + assert_eq!( + wap.model + .deprecated_config + .__warehouse_specific_config__ + .copy_grants, + configured, + "publication must leave the canonical model unchanged" + ); + } + } +} + +#[test] +fn candidate_runtime_config_suppresses_public_grants_and_keeps_candidate_identity() { + let (_directory, io, wap) = fixture(Some(true)); + let candidate = wap.execution_model().unwrap(); + let context = runtime_context(&candidate, &io); + let canonical_context = runtime_context(&wap.model, &io); + let environment = Environment::new(); + for expression in ["config.get('grants')", "config.get('copy_grants')"] { + let expression = environment.compile_expression(expression).unwrap(); + assert!(!expression.eval(&context, &[]).unwrap().is_true()); + assert!(expression.eval(&canonical_context, &[]).unwrap().is_true()); + } + assert_eq!( + environment + .compile_expression("this") + .unwrap() + .eval(&context, &[]) + .unwrap() + .to_string(), + "\"My\"\"Database\".\"My.Schema\".\"__DBT_WAP_TEST\"" + ); + assert_eq!( + environment + .compile_expression("model.alias") + .unwrap() + .eval(&context, &[]) + .unwrap() + .as_str(), + Some("__DBT_WAP_TEST") + ); +} diff --git a/crates/dbt-tasks-sa/src/wap/header_tests.rs b/crates/dbt-tasks-sa/src/wap/header_tests.rs new file mode 100644 index 00000000000..7b3d15f596d --- /dev/null +++ b/crates/dbt-tasks-sa/src/wap/header_tests.rs @@ -0,0 +1,91 @@ +use std::collections::BTreeMap; +use std::sync::Arc; + +use dbt_adapter_core::AdapterType; +use dbt_common::ErrorCode; +use dbt_jinja_utils::node_resolver::NodeResolver; +use dbt_jinja_utils::phases::compile::{ + DependencyValidationConfig, build_compile_node_context_inner, +}; +use dbt_schemas::schemas::common::DbtMaterialization; +use dbt_schemas::schemas::{CommonAttributes, DbtModel, NodeBaseAttributes}; +use dbt_schemas::state::{DbtRuntimeConfig, ModelStatus, NodeResolverTracker}; +use minijinja::{Environment, Value}; + +use super::validate_runtime_sql_header; + +#[test] +fn wap_rejects_sql_header_added_only_when_execute_is_true() { + let mut model = DbtModel { + __common_attr__: CommonAttributes { + unique_id: "model.pkg.orders".to_string(), + name: "orders".to_string(), + package_name: "pkg".to_string(), + language: Some("sql".to_string()), + ..Default::default() + }, + __base_attr__: NodeBaseAttributes { + database: "DB".to_string(), + schema: "SCHEMA".to_string(), + alias: "orders".to_string(), + materialized: DbtMaterialization::Table, + ..Default::default() + }, + ..Default::default() + }; + model.deprecated_config.wap = Some(true); + assert!(model.validate_wap_config(AdapterType::Snowflake).is_ok()); + let mut resolver = NodeResolver::default(); + resolver + .insert_ref(&model, AdapterType::Snowflake, ModelStatus::Enabled, false) + .unwrap(); + let resolver = Arc::new(resolver); + let mut env = Environment::new(); + // Exercise the shipped macro and real CompileConfig::set mutation instead + // of inserting a sql_header entry directly into the result map. + env.add_template( + "configs.sql", + include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-adapters/macros/materializations/configs.sql" + ), + ) + .unwrap(); + let template = "{% from 'configs.sql' import set_sql_header %}\ + {% if execute %}\ + {% call set_sql_header(config) %}set wap_bypass = true;{% endcall %}\ + {% endif %}select 1 as id"; + + assert!(model.deprecated_config.sql_header.is_none()); + for execute in [false, true] { + let base = BTreeMap::from([("execute".to_string(), Value::from(execute))]); + let (context, config) = build_compile_node_context_inner( + &model, + AdapterType::Snowflake, + &base, + "pkg", + resolver.clone(), + Arc::new(DbtRuntimeConfig::default()), + DependencyValidationConfig::new_validated(), + ) + .unwrap(); + let sql = env.render_str(template, &context, &[]).unwrap(); + assert_eq!(sql, "select 1 as id"); + let header = config.get("sql_header").map(|entry| entry.value().clone()); + let result = validate_runtime_sql_header(header.as_ref()); + if execute { + assert_eq!(header.unwrap().as_str(), Some("set wap_bypass = true;")); + assert_eq!(result.unwrap_err().code, ErrorCode::InvalidConfig); + } else { + assert!(result.is_ok()); + } + } +} + +#[test] +fn wap_runtime_sql_header_allows_absent_or_empty_values() { + assert!(validate_runtime_sql_header(None).is_ok()); + for header in [Value::from(()), Value::from(""), Value::from(" \n\t")] { + assert!(validate_runtime_sql_header(Some(&header)).is_ok()); + } + assert!(validate_runtime_sql_header(Some(&Value::from("select 1"))).is_err()); +} diff --git a/crates/dbt-tasks-sa/src/wap/publication_tests.rs b/crates/dbt-tasks-sa/src/wap/publication_tests.rs new file mode 100644 index 00000000000..8b62a137fa3 --- /dev/null +++ b/crates/dbt-tasks-sa/src/wap/publication_tests.rs @@ -0,0 +1,201 @@ +use super::{PublicationStep, run_publication_steps}; +use dbt_common::stats::NodeStatus; +use dbt_common::{ErrorCode, FsResult, fs_err}; + +use PublicationStep::*; + +/// Simulate warehouse effects at the production publication boundary. A failed +/// statement leaves its previous state intact; tests below check what can run next. +struct Warehouse { + steps: Vec, + failures: Vec, + public_revision: u8, + candidate_exists: bool, + query_tag_set: bool, + overrides_set: bool, +} + +impl Warehouse { + fn new(failures: Vec) -> Self { + Self { + steps: Vec::new(), + failures, + public_revision: 1, + candidate_exists: true, + query_tag_set: false, + overrides_set: false, + } + } + + fn execute(&mut self, step: PublicationStep) -> FsResult<()> { + self.steps.push(step); + if self.failures.contains(&step) { + return Err(fs_err!( + ErrorCode::ExecutionError, + "injected {step:?} failure" + )); + } + match step { + Prepare => self.overrides_set = true, + SetQueryTag => self.query_tag_set = true, + Clone => self.public_revision = 2, + Finalize => {} + ResetQueryTag => self.query_tag_set = false, + ResetOverrides => self.overrides_set = false, + DropCandidate => self.candidate_exists = false, + } + Ok(()) + } + + fn publish(&mut self) -> FsResult<()> { + run_publication_steps( + [("audit", Some(NodeStatus::TestPassed))], + "DB.SCHEMA.ORDERS", + "DB.SCHEMA.__DBT_WAP_TEST", + |step| self.execute(step), + ) + } +} + +#[test] +fn publication_does_not_touch_warehouse_without_every_audit_passing() { + for status in [ + None, + Some(NodeStatus::Errored), + Some(NodeStatus::TestWarned), + Some(NodeStatus::SkippedUpstreamFailed), + Some(NodeStatus::StaticallyCheckedDataTest), + Some(NodeStatus::ReusedNoChanges("cached".to_owned())), + ] { + let mut warehouse = Warehouse::new(vec![]); + let result = run_publication_steps( + [ + ("passed", Some(NodeStatus::TestPassed)), + ("blocking", status), + ], + "ORDERS", + "WORK", + |step| warehouse.execute(step), + ); + assert!(result.is_err()); + assert!(warehouse.steps.is_empty()); + assert_eq!(warehouse.public_revision, 1); + assert!(warehouse.candidate_exists); + } + assert!( + run_publication_steps([], "ORDERS", "WORK", |_| { + panic!("no publication step may run without an audit") + }) + .is_err() + ); +} + +#[test] +fn publication_cleans_up_only_after_finalization_and_session_restoration() { + let mut warehouse = Warehouse::new(vec![]); + warehouse.publish().unwrap(); + assert_eq!(warehouse.public_revision, 2); + assert!(!warehouse.candidate_exists); + assert!(!warehouse.query_tag_set); + assert!(!warehouse.overrides_set); + assert_eq!( + warehouse.steps, + [ + Prepare, + SetQueryTag, + Clone, + Finalize, + ResetQueryTag, + ResetOverrides, + DropCandidate + ] + ); +} + +#[test] +fn publication_failures_before_clone_confirmation_retain_candidate_and_restore_session() { + for failure in [Prepare, SetQueryTag, Clone] { + let mut warehouse = Warehouse::new(vec![failure]); + let error = warehouse.publish().unwrap_err(); + assert!(error.to_string().contains(&format!("injected {failure:?}"))); + assert!(!error.to_string().contains("was published")); + assert_eq!(warehouse.public_revision, 1); + assert!(warehouse.candidate_exists); + assert!(!warehouse.query_tag_set); + assert!(!warehouse.overrides_set); + assert!(!warehouse.steps.contains(&Finalize)); + assert!(!warehouse.steps.contains(&DropCandidate)); + assert_eq!(warehouse.steps.last(), Some(&ResetOverrides)); + assert_eq!(warehouse.steps.contains(&ResetQueryTag), failure == Clone); + } +} + +#[test] +fn publication_retains_candidate_when_clone_outcome_is_unknown() { + let mut warehouse = Warehouse::new(vec![]); + let error = run_publication_steps( + [("audit", Some(NodeStatus::TestPassed))], + "ORDERS", + "WORK", + |step| { + warehouse.execute(step)?; + if step == Clone { + // Snowflake committed, but the client lost its response. The + // orchestration must not claim either confirmed success or rollback. + return Err(fs_err!( + ErrorCode::ExecutionError, + "Publication may have committed; inspect Snowflake query history" + )); + } + Ok(()) + }, + ) + .unwrap_err(); + assert!(error.to_string().contains("may have committed")); + assert!(!error.to_string().contains("was published")); + assert_eq!(warehouse.public_revision, 2); + assert!(warehouse.candidate_exists); + assert!(!warehouse.steps.contains(&Finalize)); + assert!(!warehouse.steps.contains(&DropCandidate)); + assert!(!warehouse.query_tag_set); + assert!(!warehouse.overrides_set); +} + +#[test] +fn publication_finalization_failures_report_published_and_retain_candidate() { + for failure in [Finalize, ResetQueryTag, ResetOverrides] { + let mut warehouse = Warehouse::new(vec![failure]); + let message = warehouse.publish().unwrap_err().to_string(); + assert!(message.contains("DB.SCHEMA.ORDERS was published")); + assert!(message.contains(&format!("injected {failure:?}"))); + assert!(message.contains("Working table retained: DB.SCHEMA.__DBT_WAP_TEST")); + assert_eq!(warehouse.public_revision, 2); + assert!(warehouse.candidate_exists); + assert!(warehouse.steps.contains(&ResetQueryTag)); + assert!(warehouse.steps.contains(&ResetOverrides)); + assert!(!warehouse.steps.contains(&DropCandidate)); + } +} + +#[test] +fn publication_preserves_primary_error_when_session_restoration_also_fails() { + for failure in [Clone, Finalize] { + let mut warehouse = Warehouse::new(vec![failure, ResetQueryTag, ResetOverrides]); + let message = warehouse.publish().unwrap_err().to_string(); + assert!(message.contains(&format!("injected {failure:?}"))); + assert!(!message.contains("injected Reset")); + assert!(warehouse.steps.contains(&ResetQueryTag)); + assert!(warehouse.steps.contains(&ResetOverrides)); + assert!(warehouse.candidate_exists); + } +} + +#[test] +fn publication_remains_successful_when_candidate_cleanup_fails() { + let mut warehouse = Warehouse::new(vec![DropCandidate]); + warehouse.publish().unwrap(); + assert_eq!(warehouse.public_revision, 2); + assert!(warehouse.candidate_exists); + assert!(!warehouse.query_tag_set); + assert!(!warehouse.overrides_set); +} diff --git a/docs/snowflake-table-wap-plan.md b/docs/snowflake-table-wap-plan.md new file mode 100644 index 00000000000..5cb2105bd08 --- /dev/null +++ b/docs/snowflake-table-wap-plan.md @@ -0,0 +1,183 @@ +# Snowflake table write-audit-publish + +## Summary + +Add an opt-in `wap` model configuration for native Snowflake SQL models using +`materialized='table'`. During `dbt build`, compute a uniquely named working +table, run the model's data tests against it, and publish it only when every +required audit reports `PASS`. The working table lives in the model's existing +database and schema. No extra schema or schema-creation permission is needed. + +```sql +{{ config(materialized='table', wap=true) }} + +select id, amount from {{ ref('stg_orders') }} +``` + +The flag can also be inherited through `dbt_project.yml`: + +```yaml +models: + my_project: + audited_tables: + +materialized: table + +wap: true +``` + +The default is false. A child model can set `wap=false`. Materialization remains +`table`, and the public model database, schema, alias, manifest identity, and +ordinary downstream `ref()` calls remain unchanged. + +## Build and publication behavior + +1. Validate the selected WAP models and their complete set of enabled audits + before scheduling warehouse writes. Require at least one enabled data test. +2. Build a working table with an invocation-specific identifier in the model's + canonical database and schema. Preserve ordinary SQL table construction, + contracts, documentation, and configured permanent/transient lifecycle. +3. Bind this model's audit references to the working relation using an + execution-local relation override. Do not rewrite SQL text or globally change + the manifest's model relation. Ordinary model dependencies continue to read + their published relations. +4. Wait for every required audit. `FAIL`, `WARN`, execution error, cancellation, + and skipped audits all prevent publication; only `PASS` permits it. +5. Publish the audited contents with one Snowflake clone statement: + + ```sql + CREATE OR REPLACE [TRANSIENT] TABLE + CLONE [COPY GRANTS]; + ``` + + The destination must be absent or a supported native table. Never pre-drop + the published relation to accommodate an incompatible destination type. +6. Reconcile configured production grants, restore configured automatic + clustering, refresh relation metadata, and release downstream tasks. Drop + the successful working table after publication completes. + +On audit failure, the previous published table remains available and downstream +models are skipped. On a first build, failure leaves the published relation +absent. Retain failed working tables for inspection and report their exact +qualified names. Operators remove them when they are no longer needed. + +Publication is atomic for one table: concurrent readers see the old or new +version. Each model publishes independently; a later failure does not undo +earlier successful publications. Snowflake DDL is not part of a transaction +spanning the tests and finalization. An error after a successful clone must +report that publication occurred, retain the working table, and stop downstream +execution rather than claim the original table is still published. +If submission fails without a definitive Snowflake response, report that the +publication outcome is unknown and direct the operator to query history. +[Snowflake CREATE TABLE](https://docs.snowflake.com/en/sql-reference/sql/create-table) + +## Supported commands and audits + +| Command or option | Behavior | +| --- | --- | +| `dbt build` | Builds, audits, and publishes selected WAP models. | +| `dbt run` / `dbt clone` selecting a WAP model | Rejects the operation; use `dbt build`. | +| `dbt test` | Tests published relations; never publishes a retained working table. | +| Parse, compile, listing, documentation | Keep canonical model identities. | +| `dbt retry` of a WAP build | Rebuilds a fresh working table and reruns all required audits, including audits that passed previously. | +| `wap=false` or omitted | Preserves existing behavior. | + +All enabled data tests owned by the WAP model must be selected. Excluding an +audit or using indirect selection that omits it produces an actionable error; +selection cannot silently weaken the publication gate. Disabled tests do not +participate. Generic tests use their attached model as owner; singular tests +must depend on only the WAP model. Multi-relation audits owned by the WAP model, +including its relationships tests, are excluded from the first version. +Relationships tests owned by downstream models continue to test the published +relation. Audits must use `ref()` or `builtins.ref()`; hardcoded public table names +cannot be redirected to the working table and are outside this contract. +Existing unit-test execution remains +a prerequisite according to normal build selection; unit tests do not replace +the required data audits. Retry also reruns enabled unit tests for rebuilt WAP +models. + +WAP audits execute directly on Snowflake after the working table exists. This +version skips static analysis and test aggregation for those audits, and never +reuses a cached test result. Model analysis retains the canonical relation +identity for downstream schema inference. Filtered partial parsing falls back +to a full parse for WAP so audit discovery includes every enabled test. +Tests that store failure rows may use their normal failure tables, provided those +relations do not collide with a WAP public table or working table. + +Reject `--empty`, sampled builds, local execution, and model +`on_error: continue` for WAP publication. The gate must validate the actual +complete Snowflake table and prevent downstream execution after failure. + +## Scope and compatibility + +The first version requires a Snowflake target and supports SQL, the built-in +Snowflake table materialization, and native permanent/transient tables. It +excludes incremental, view, dynamic, interactive, external, event, hybrid, and +Iceberg/catalog-linked tables; Python models; +custom materializations; effective model pre/post hooks; and `sql_header`, +including headers added during rendering. The model and profile quoting settings +must resolve the working table to the same database and schema. Reject a mismatch +before staging, because the built-in table macro creates its relation with the +profile's quoting policy. +Reject configured `row_access_policy`, `table_tag`, `copy_tags=true`, and custom +table/column constraints in this first version. Ordinary dbt selection tags and +query tags are unaffected. This feature does not preserve arbitrary properties +manually attached to the previous published table. + +Suppress explicit model grants on the working table. At replacement, +`copy_grants=true` preserves existing destination grants. On first publication, +omit `COPY GRANTS` so production future grants apply instead of copying working +table grants. Apply configured grants to the published relation afterward. +Because working tables occupy the same schema, existing future grants may make +them visible to other roles. Their names separate them from the public model; +they are **not private or hidden tables**. + +Cloning is inexpensive because it initially shares existing storage, but the +transformation, tests, retained failed working tables, and historical data still +incur costs. Cloning suspends automatic clustering, so configured clustering +must be restored after publication. Streams on replaced tables become stale; +stream continuity and object-identity preservation are outside this feature. +[Snowflake cloning reference](https://docs.snowflake.com/en/sql-reference/sql/create-clone) + +Working and published table lifecycles must agree: a transient working table +cannot be cloned into a permanent published table. +[Snowflake table design considerations](https://docs.snowflake.com/en/user-guide/table-considerations) + +Separate invocations use different working identifiers. This does not provide a +cross-invocation publication lock: users must serialize builds that target the +same model if publication order matters. + +## Implementation and validation + +- Add model-only typed `wap` configuration, project inheritance, manifest + round-tripping, known-key handling, state comparison, and parser validation. +- Build an invocation-local WAP plan from the full resolved manifest. Preserve + ordinary selection rules and add a publication dependency after the model's + audits. Do not report the model successful or submit published state before + publication completes. +- Materialize a copied runtime model targeting the working relation, and use + scoped audit relation overrides. Reuse the existing Snowflake clone SQL + helper and grant machinery. Keep vendored macros and dependency manifests + unchanged. +- Bypass shortcuts that could skip a new candidate or reuse an audit of a + different relation. On retry, expand only the failed/skipped WAP models or + owners of failed audits to include all required audits and unit prerequisites. + Never resume or publish an earlier retained working table. +- Test configuration inheritance and manifest/state compatibility; invalid + scopes; incomplete audit selection; passing/failing/warning/skipped audits; + dependency ordering; first publication; preservation of previous data after + failure; permanent/transient lifecycle; quoting; grants; clustering; retained + failure tables; retry; and ordinary non-WAP behavior. +- Run focused crate tests and a live Snowflake scenario proving that a failing + audit leaves sentinel production data unchanged, then that passing audits + replace it and release downstream models. Real Snowflake validation requires + credentials and should be reported separately from local test results. + +Offline regression tests exercise the shipped Snowflake clone macro with the real +runtime config, including lifecycle, quoted names, and first-build/replacement +grant behavior. Publication failure tests cover audit gating, clone and +finalization errors, session restoration, retained candidates, and cleanup-only +warnings. Ref and compiled-SQL cache tests verify that working-table identities +stay scoped to the current build; a later standalone test reads the public table. +These checks do not establish Snowflake's live DDL or permission behavior. + +The live acceptance runner and its invocation instructions are in +[`crates/dbt-sa-cli/tests/data/snowflake_wap/README.md`](../crates/dbt-sa-cli/tests/data/snowflake_wap/README.md). From 3e208ce72229c9daae38a3c3a5a51faab57a6f2c Mon Sep 17 00:00:00 2001 From: bmoore813 Date: Fri, 18 Sep 2026 01:27:58 -0400 Subject: [PATCH 2/3] test(snowflake): exercise WAP staging and failure isolation --- .../tests/data/snowflake_wap/README.md | 35 +- .../data/snowflake_wap/macros/acceptance.sql | 8 +- .../data/snowflake_wap/models/orders.sql | 4 + .../tests/snowflake_wap_acceptance.py | 96 ++++- .../tests/test_snowflake_wap_acceptance.py | 195 +++++++++- .../tests/test_snowflake_wap_cli.py | 116 ++++++ crates/dbt-tasks-sa/src/graph.rs | 7 +- .../src/graph/wap_tests/execution.rs | 334 ++++++++++++++++++ crates/dbt-tasks-sa/src/visitor.rs | 3 + .../src/visitor/wap_test_support.rs | 81 +++++ crates/dbt-tasks-sa/src/wap.rs | 4 +- crates/dbt-tasks-sa/src/wap/clone_tests.rs | 4 +- crates/dbt-tasks-sa/src/wap/stage_tests.rs | 244 +++++++++++++ docs/snowflake-table-wap-plan.md | 16 + 14 files changed, 1114 insertions(+), 33 deletions(-) create mode 100644 crates/dbt-sa-cli/tests/test_snowflake_wap_cli.py create mode 100644 crates/dbt-tasks-sa/src/graph/wap_tests/execution.rs create mode 100644 crates/dbt-tasks-sa/src/visitor/wap_test_support.rs create mode 100644 crates/dbt-tasks-sa/src/wap/stage_tests.rs diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md b/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md index a0883cf94c8..156d4ff67a4 100644 --- a/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/README.md @@ -25,14 +25,20 @@ The default run covers both permanent and transient WAP tables. Use `--table-kind permanent` or `--table-kind transient` for one variant. Each variant checks: -- An error-severity audit fails; the public sentinel stays `[999]`, the candidate - retains `[-1, 2]`, and the downstream table is not created. +- An error-severity audit fails; the public sentinel stays `[999]`, the existing + downstream sentinel stays `[777]`, and the candidate retains `[-1, 2]`. - A subsequent standalone `dbt test` passes against the public sentinel, despite the retained failing candidate and compiled SQL from the preceding build. - `dbt retry` creates a different candidate, reruns all three audits including the two that previously passed, publishes `[2, 3]`, and builds the downstream through its ordinary public `ref`. -- A warning-severity audit also prevents publication and retains its candidate. +- A deliberate SQL conversion error runs on Snowflake with static analysis off. + The model errors, all audits and downstream execution are skipped, and both + existing sentinel tables retain their original rows. The model result must + include the deliberate conversion-error marker, so an earlier configuration + or permissions error cannot satisfy this scenario. +- A warning-severity audit also prevents publication, preserves both existing + sentinel tables, and retains its candidate. - A successful first build publishes exactly `[1, 2]` and drops its candidate. - A failing first build leaves the public table absent. @@ -41,6 +47,12 @@ downstream SQL. Each model must have one canonical result. Audit SQL includes both generic `not_null`/`unique` tests and a singular test. Builds use four threads and do not use fail-fast. +Every staging message must identify its candidate in the same database and +schema as the public model. The runner checks fully qualified names, including +quoted names containing dots or embedded quotes, before recording ownership. +A candidate in another database or a scratch schema fails the check and is not +added to the cleanup inventory. The fixture never creates or drops a schema. + Every run uses a copied temporary project and UUID-based public aliases in the target schema. The runner first verifies those public aliases are absent. It records candidate ownership only from WAP's post-preflight staging message, so a @@ -58,3 +70,20 @@ table whose name begins with `__DBT_WAP_`. The fixture has no CI credential assumptions and is not part of ordinary local unit tests. A successful offline parse or unit-test run does not count as this live acceptance test passing. + +The runner's offline unit tests use mocked subprocesses and never load profiles +or make warehouse calls: + +```sh +python3 -m unittest discover -s crates/dbt-sa-cli/tests -p test_snowflake_wap_acceptance.py +``` + +Separate CLI integration checks parse temporary projects with an unusable +Snowflake profile. They verify actual config inheritance and validation, +canonical manifest identities and refs, and repeated parsing when WAP changes. +They make no warehouse calls and require a built dbt binary: + +```sh +DBT_WAP_TEST_BIN="$PWD/target/debug/dbt" \ + python3 -m unittest discover -s crates/dbt-sa-cli/tests -p test_snowflake_wap_cli.py -v +``` diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql b/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql index 2271036a26d..a6193c34bb2 100644 --- a/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/macros/acceptance.sql @@ -25,9 +25,11 @@ {% endmacro %} {% macro wap_fixture_prepare() %} - {% set relation = ref('orders') %} - {% do wap_fixture_assert(relation.identifier, present=false) %} - {% do run_query('create ' ~ ('transient ' if var('transient', false) else '') ~ 'table ' ~ relation ~ ' as select 999::integer as id') %} + {% for model_name, sentinel_id in [('orders', 999), ('downstream', 777)] %} + {% set relation = ref(model_name) %} + {% do wap_fixture_assert(relation.identifier, present=false) %} + {% do run_query('create ' ~ ('transient ' if var('transient', false) else '') ~ 'table ' ~ relation ~ ' as select ' ~ sentinel_id ~ '::integer as id') %} + {% endfor %} {% endmacro %} {% macro wap_fixture_cleanup(identifiers) %} diff --git a/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql index a80f8df62de..2194de5dbb6 100644 --- a/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql +++ b/crates/dbt-sa-cli/tests/data/snowflake_wap/models/orders.sql @@ -5,6 +5,10 @@ transient=var('transient', false) ) }} +{% if var('transformation_error', false) %} +select to_number('WAP_TRANSFORM_FAILURE')::integer as id +{% else %} select {{ var('audit_value', 1) }}::integer as id union all select 2::integer as id +{% endif %} diff --git a/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py b/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py index d6249b1f30a..d0a78581318 100644 --- a/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py +++ b/crates/dbt-sa-cli/tests/snowflake_wap_acceptance.py @@ -19,9 +19,17 @@ PROJECT = "wap_live_acceptance" MODEL_ID = f"model.{PROJECT}.orders" DOWNSTREAM_ID = f"model.{PROJECT}.downstream" +TRANSFORMATION_ERROR = "WAP_TRANSFORM_FAILURE" +CANDIDATE_IDENTIFIER = r"__DBT_WAP_[0-9A-F]{32}_[0-9A-F]{32}" STAGED_CANDIDATE = re.compile( r"WAP: building[^\r\n]*?in working table[^\r\n]*?" - r"(__DBT_WAP_[0-9A-F]{32}_[0-9A-F]{32})" + rf"({CANDIDATE_IDENTIFIER})" +) +RELATION_COMPONENT = r'(?:"(?:[^"\r\n]|"")*"|[A-Za-z_][A-Za-z0-9_$]*)' +RELATION_NAME = rf"{RELATION_COMPONENT}\.{RELATION_COMPONENT}\.{RELATION_COMPONENT}" +STAGED_RELATIONS = re.compile( + rf"WAP: building (?P{RELATION_NAME}) in working table " + rf"(?P{RELATION_NAME})" ) @@ -30,6 +38,35 @@ def require(condition: bool, message: str) -> None: raise RuntimeError(message) +def relation_components(relation: str) -> tuple[str, ...]: + require(re.fullmatch(RELATION_NAME, relation) is not None, f"Invalid relation: {relation}") + return tuple( + part[1:-1].replace('""', '"') if part.startswith('"') else part.upper() + for part in re.findall(RELATION_COMPONENT, relation) + ) + + +def staged_candidates(output: str, public_identifier: str) -> set[str]: + identifiers = set(STAGED_CANDIDATE.findall(output)) + located = set() + for match in STAGED_RELATIONS.finditer(output): + public = relation_components(match["public"]) + candidate = relation_components(match["candidate"]) + require(public[2] == public_identifier, f"Unexpected public relation: {match['public']}") + require( + public[:2] == candidate[:2], + f"Candidate must share the public database and schema: {match['candidate']}; " + f"public relation: {match['public']}", + ) + require( + re.fullmatch(CANDIDATE_IDENTIFIER, candidate[2]) is not None, + f"Unexpected candidate identifier: {candidate[2]}", + ) + located.add(candidate[2]) + require(located == identifiers, "Could not verify every staged candidate's database and schema") + return located + + class Acceptance: def __init__(self, args: argparse.Namespace) -> None: self.args = args @@ -79,7 +116,11 @@ def invoke( (capture / "stdout.log").write_text(completed.stdout) (capture / "stderr.log").write_text(completed.stderr) output = completed.stdout + "\n" + completed.stderr - candidates = set(STAGED_CANDIDATE.findall(output)) + for filename in ("run_results.json", "manifest.json"): + artifact = self.artifacts / filename + if artifact.exists(): + shutil.copy2(artifact, capture / filename) + candidates = staged_candidates(output, self.public()) # Only the post-preflight "building" message establishes ownership. # A collision error's "retained if created" message does not. prefix = self.variables.get("prefix") @@ -87,10 +128,6 @@ def invoke( identifiers = self.owned[prefix]["identifiers"] identifiers[:] = sorted(set(identifiers) | candidates) self.persist_owned() - for filename in ("run_results.json", "manifest.json"): - artifact = self.artifacts / filename - if artifact.exists(): - shutil.copy2(artifact, capture / filename) require( (completed.returncode == 0) == success, f"Unexpected exit {completed.returncode}: {' '.join(command)}; see {capture}", @@ -124,6 +161,7 @@ def begin(self, kind: str, scenario: str, *, sentinel: bool) -> None: self.variables = { "prefix": prefix, "transient": kind == "transient", "audit_value": -1, "audit_severity": "error", + "transformation_error": False, } identifiers = [f"{prefix}_ORDERS", f"{prefix}_DOWNSTREAM"] for identifier in identifiers: @@ -140,9 +178,14 @@ def begin(self, kind: str, scenario: str, *, sentinel: bool) -> None: def public(self, suffix: str = "ORDERS") -> str: return f"{self.variables['prefix']}_{suffix}" - def build(self, *, success: bool) -> tuple[dict[str, Any], set[str], Path]: + def build( + self, *, success: bool, static_analysis: str | None = None, + ) -> tuple[dict[str, Any], set[str], Path]: + command = ["build", "--select", "orders+", "--threads", "4"] + if static_analysis is not None: + command.extend(["--static-analysis", static_analysis]) return self.invoke( - ["build", "--select", "orders+", "--threads", "4"], + command, success=success, results=True, staged=True, ) @@ -155,7 +198,9 @@ def verify_audits(artifact: dict[str, Any], *, verdict: str) -> set[str]: singular_id = f"test.{PROJECT}.nonnegative" require(singular_id in tests, "Missing nonnegative singular audit") for unique_id, row in tests.items(): - expected = verdict if unique_id == singular_id else "pass" + expected = "skipped" if verdict == "transform_error" else ( + verdict if unique_id == singular_id else "pass" + ) require(row["status"] == expected, f"Unexpected audit result: {row}") return set(tests) @@ -164,8 +209,15 @@ def verify_results(artifact: dict[str, Any], *, verdict: str) -> set[str]: rows = artifact["results"] model_rows = [row for row in rows if row["unique_id"] == MODEL_ID] require(len(model_rows) == 1, "Expected exactly one canonical model result") - expected_model = {"pass": "success", "fail": "skipped", "warn": "error"}[verdict] + expected_model = { + "pass": "success", "fail": "skipped", "warn": "error", "transform_error": "error", + }[verdict] require(model_rows[0]["status"] == expected_model, f"Unexpected model result: {model_rows}") + if verdict == "transform_error": + require( + TRANSFORMATION_ERROR in str(model_rows[0].get("message", "")), + f"Expected the deliberate SQL transformation error: {model_rows}", + ) audits = Acceptance.verify_audits(artifact, verdict=verdict) downstream = [row for row in rows if row["unique_id"] == DOWNSTREAM_ID] expected_downstream = "success" if verdict == "pass" else "skipped" @@ -175,6 +227,10 @@ def verify_results(artifact: dict[str, Any], *, verdict: str) -> set[str]: ) return audits + def verify_existing_data_preserved(self) -> None: + self.check(self.public(), ids=[999]) + self.check(self.public("DOWNSTREAM"), ids=[777]) + def verify_published(self, ids: list[int], candidates: set[str]) -> None: self.check(self.public(), ids=ids) self.check(self.public("DOWNSTREAM"), ids=ids) @@ -187,12 +243,14 @@ def verify_published(self, ids: list[int], candidates: set[str]) -> None: require("__DBT_WAP_" not in sql, "Candidate relation leaked into downstream SQL") def run_kind(self, kind: str) -> None: - print(f"Checking {kind}: failure, retry, warning, first publish, first failure", flush=True) + print( + f"Checking {kind}: audit failure, retry, transformation error, warning, " + "first publish, first failure", flush=True, + ) self.begin(kind, "fail_retry", sentinel=True) failed, retained, state = self.build(success=False) audits = self.verify_results(failed, verdict="fail") - self.check(self.public(), ids=[999]) - self.check(self.public("DOWNSTREAM"), present=False) + self.verify_existing_data_preserved() for candidate in retained: self.check(candidate, ids=[-1, 2]) @@ -206,7 +264,7 @@ def run_kind(self, kind: str) -> None: "Standalone test omitted an audit", ) require(len(standalone["results"]) == 3, "Standalone test unexpectedly rebuilt a model") - self.check(self.public(), ids=[999]) + self.verify_existing_data_preserved() for candidate in retained: self.check(candidate, ids=[-1, 2]) @@ -222,12 +280,18 @@ def run_kind(self, kind: str) -> None: for candidate in retained: self.check(candidate, ids=[-1, 2]) + self.begin(kind, "transformation_error", sentinel=True) + self.variables["transformation_error"] = True + # Force warehouse execution so local analysis cannot satisfy this case. + errored, _, _ = self.build(success=False, static_analysis="off") + self.verify_results(errored, verdict="transform_error") + self.verify_existing_data_preserved() + self.begin(kind, "warning", sentinel=True) self.variables["audit_severity"] = "warn" warned, candidates, _ = self.build(success=False) self.verify_results(warned, verdict="warn") - self.check(self.public(), ids=[999]) - self.check(self.public("DOWNSTREAM"), present=False) + self.verify_existing_data_preserved() for candidate in candidates: self.check(candidate, ids=[-1, 2]) diff --git a/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py b/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py index 176c33721fe..d92db10c40b 100644 --- a/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py +++ b/crates/dbt-sa-cli/tests/test_snowflake_wap_acceptance.py @@ -1,6 +1,12 @@ -"""Offline checks for the live acceptance runner's evidence handling.""" +"""Offline checks for the live acceptance runner; no warehouse calls are made.""" +import argparse +import json +from pathlib import Path +import shutil +import subprocess import unittest +from unittest import mock from snowflake_wap_acceptance import ( Acceptance, @@ -8,6 +14,9 @@ MODEL_ID, PROJECT, STAGED_CANDIDATE, + TRANSFORMATION_ERROR, + relation_components, + staged_candidates, ) @@ -15,20 +24,37 @@ class AcceptanceEvidenceTests(unittest.TestCase): def artifact(self, verdict): return {"results": [ {"unique_id": MODEL_ID, - "status": {"pass": "success", "fail": "skipped", "warn": "error"}[verdict]}, + "status": {"pass": "success", "fail": "skipped", "warn": "error", + "transform_error": "error"}[verdict], + "message": f"Numeric value '{TRANSFORMATION_ERROR}' is not recognized"}, {"unique_id": DOWNSTREAM_ID, "status": "success" if verdict == "pass" else "skipped"}, - {"unique_id": f"test.{PROJECT}.nonnegative", "status": verdict}, - {"unique_id": f"test.{PROJECT}.not_null_orders_id.abc", "status": "pass"}, - {"unique_id": f"test.{PROJECT}.unique_orders_id.def", "status": "pass"}, + {"unique_id": f"test.{PROJECT}.nonnegative", + "status": "skipped" if verdict == "transform_error" else verdict}, + {"unique_id": f"test.{PROJECT}.not_null_orders_id.abc", + "status": "skipped" if verdict == "transform_error" else "pass"}, + {"unique_id": f"test.{PROJECT}.unique_orders_id.def", + "status": "skipped" if verdict == "transform_error" else "pass"}, ]} def test_expected_verdicts(self): - for verdict in ("pass", "fail", "warn"): + for verdict in ("pass", "fail", "warn", "transform_error"): with self.subTest(verdict=verdict): audits = Acceptance.verify_results(self.artifact(verdict), verdict=verdict) self.assertEqual(len(audits), 3) + def test_transformation_failure_requires_the_execution_error_marker(self): + artifact = self.artifact("transform_error") + artifact["results"][0]["message"] = "Missing permission during preflight" + with self.assertRaisesRegex(RuntimeError, "deliberate SQL transformation error"): + Acceptance.verify_results(artifact, verdict="transform_error") + + def test_transformation_failure_cannot_report_a_passed_audit(self): + artifact = self.artifact("transform_error") + artifact["results"][3]["status"] = "pass" + with self.assertRaisesRegex(RuntimeError, "Unexpected audit result"): + Acceptance.verify_results(artifact, verdict="transform_error") + def test_missing_previously_passed_audit_is_rejected(self): artifact = self.artifact("pass") artifact["results"].pop() @@ -66,5 +92,162 @@ def test_collision_notice_does_not_establish_candidate_ownership(self): ), [candidate]) +class CandidateLocationTests(unittest.TestCase): + candidate = "__DBT_WAP_" + "A" * 32 + "_" + "B" * 32 + + def test_quoted_components_preserve_dots_and_embedded_quotes(self): + self.assertEqual( + relation_components('"Db""Name"."schema.with.dot".orders'), + ('Db"Name', 'schema.with.dot', 'ORDERS'), + ) + output = ( + 'WAP: building "Db""Name"."schema.with.dot"."ORDERS" in working table ' + f'"Db""Name"."schema.with.dot"."{self.candidate}"' + ) + self.assertEqual(staged_candidates(output, "ORDERS"), {self.candidate}) + + def test_equivalent_unquoted_and_quoted_names_share_a_namespace(self): + output = ( + 'WAP: building db.public.orders in working table ' + f'"DB"."PUBLIC"."{self.candidate}"' + ) + self.assertEqual(staged_candidates(output, "ORDERS"), {self.candidate}) + + def test_unparseable_staging_relation_does_not_establish_ownership(self): + output = f"WAP: building DB.PUBLIC.ORDERS in working table {self.candidate}" + with self.assertRaisesRegex(RuntimeError, "Could not verify every staged candidate"): + staged_candidates(output, "ORDERS") + + +class AcceptanceInvocationTests(unittest.TestCase): + def setUp(self): + # The runner never sees real environment variables, credentials, or a profile. + environment = mock.patch.dict("snowflake_wap_acceptance.os.environ", {}, clear=True) + environment.start() + self.addCleanup(environment.stop) + args = argparse.Namespace( + dbt_bin=Path("/unused/dbt"), profiles_dir=Path("/unused/profiles"), + profile="offline", target="offline", + ) + with mock.patch("builtins.print"): + self.fixture = Acceptance(args) + self.addCleanup(shutil.rmtree, self.fixture.work) + self.prefix = "DBT_WAP_ACCEPT_OFFLINE" + self.public_names = [f"{self.prefix}_ORDERS", f"{self.prefix}_DOWNSTREAM"] + self.fixture.variables = {"prefix": self.prefix} + self.fixture.owned[self.prefix] = { + "profile": "offline", "target": "offline", + "variables": dict(self.fixture.variables), "identifiers": list(self.public_names), + } + self.fixture.persist_owned() + self.candidate = "__DBT_WAP_" + "A" * 32 + "_" + "B" * 32 + self.invocation = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa" + self.stdout = self.stage_message() + self.returncode = 0 + self.write_results = True + + def stage_message(self, database="DB", schema="PUBLIC"): + return ( + f'WAP: building "DB"."PUBLIC"."{self.public_names[0]}" in working table ' + f'"{database}"."{schema}"."{self.candidate}"' + ) + + def completed(self, argv, **kwargs): + result_path = self.fixture.artifacts / "run_results.json" + self.assertFalse(result_path.exists(), "A previous artifact was not removed") + self.fixture.artifacts.mkdir(parents=True, exist_ok=True) + if self.write_results: + result_path.write_text(json.dumps({ + "metadata": {"invocation_id": self.invocation}, "results": [], + })) + (self.fixture.artifacts / "manifest.json").write_text('{"nodes": {}}') + return subprocess.CompletedProcess( + argv, self.returncode, stdout=self.stdout, stderr="offline stderr", + ) + + def test_invoke_replaces_stale_results_captures_evidence_and_records_exact_ownership(self): + self.fixture.artifacts.mkdir() + (self.fixture.artifacts / "run_results.json").write_text('{"stale": true}') + with mock.patch("snowflake_wap_acceptance.subprocess.run", side_effect=self.completed) as run: + artifact, candidates, capture = self.fixture.invoke( + ["build"], results=True, staged=True, + ) + self.assertEqual(candidates, {self.candidate}) + self.assertEqual(artifact["metadata"]["invocation_id"], self.invocation) + self.assertEqual(json.loads((capture / "run_results.json").read_text()), artifact) + self.assertTrue((capture / "manifest.json").is_file()) + self.assertEqual((capture / "stdout.log").read_text(), self.stdout) + self.assertEqual((capture / "stderr.log").read_text(), "offline stderr") + inventory = json.loads((self.fixture.work / "owned_objects.json").read_text()) + self.assertEqual(set(inventory[0]["identifiers"]), set(self.public_names) | {self.candidate}) + self.assertEqual(run.call_args.kwargs["env"], {"DBT_QUIET": "false", "DBT_USE_COLORS": "false"}) + self.assertEqual(run.call_args.kwargs["cwd"], self.fixture.project) + self.assertFalse(run.call_args.kwargs["check"]) + + def test_unexpected_command_failure_still_records_owned_candidate_for_cleanup(self): + self.returncode = 1 + with mock.patch("snowflake_wap_acceptance.subprocess.run", side_effect=self.completed): + with self.assertRaisesRegex(RuntimeError, "Unexpected exit 1"): + self.fixture.invoke(["build"], results=True, staged=True) + self.assertIn(self.candidate, self.fixture.owned[self.prefix]["identifiers"]) + + def test_candidate_in_another_database_or_scratch_schema_is_rejected_before_ownership(self): + for database, schema in [("OTHER_DB", "PUBLIC"), ("DB", "DBT_WAP")]: + with self.subTest(database=database, schema=schema): + self.stdout = self.stage_message(database, schema) + with mock.patch("snowflake_wap_acceptance.subprocess.run", side_effect=self.completed): + with self.assertRaisesRegex(RuntimeError, "share the public database and schema"): + self.fixture.invoke(["build"], results=True, staged=True) + self.assertEqual(self.fixture.owned[self.prefix]["identifiers"], self.public_names) + + def test_candidate_from_another_invocation_is_rejected(self): + self.invocation = "cccccccc-cccc-cccc-cccc-cccccccccccc" + with mock.patch("snowflake_wap_acceptance.subprocess.run", side_effect=self.completed): + with self.assertRaisesRegex(RuntimeError, "does not belong to this invocation"): + self.fixture.invoke(["build"], results=True, staged=True) + + def test_missing_results_cannot_reuse_previous_artifacts(self): + self.write_results = False + self.fixture.artifacts.mkdir() + (self.fixture.artifacts / "run_results.json").write_text('{"stale": true}') + with mock.patch("snowflake_wap_acceptance.subprocess.run", side_effect=self.completed): + with self.assertRaisesRegex(RuntimeError, "Missing run_results.json"): + self.fixture.invoke(["build"], results=True, staged=True) + + def test_collision_notice_never_enters_cleanup_inventory(self): + self.returncode = 1 + self.stdout = f"WAP working table retained if created: DB.PUBLIC.{self.candidate}" + with mock.patch("snowflake_wap_acceptance.subprocess.run", side_effect=self.completed): + _, candidates, _ = self.fixture.invoke(["build"], success=False, results=True) + self.assertEqual(candidates, set()) + self.assertCountEqual(self.fixture.owned[self.prefix]["identifiers"], self.public_names) + + def test_cleanup_passes_only_the_exact_owned_identifiers(self): + identifiers = self.fixture.owned[self.prefix]["identifiers"] + identifiers.append(self.candidate) + with mock.patch.object(self.fixture, "invoke") as invoke: + self.fixture.cleanup() + command = invoke.call_args.args[0] + self.assertEqual(command[:3], ["run-operation", "wap_fixture_cleanup", "--args"]) + self.assertEqual(json.loads(command[3]), {"identifiers": identifiers}) + self.assertEqual(invoke.call_count, 1) + + def test_failure_preservation_checks_both_existing_tables(self): + with mock.patch.object(self.fixture, "check") as check: + self.fixture.verify_existing_data_preserved() + self.assertEqual(check.call_args_list, [ + mock.call(self.public_names[0], ids=[999]), + mock.call(self.public_names[1], ids=[777]), + ]) + + def test_transformation_error_build_disables_static_analysis(self): + with mock.patch.object(self.fixture, "invoke") as invoke: + self.fixture.build(success=False, static_analysis="off") + invoke.assert_called_once_with( + ["build", "--select", "orders+", "--threads", "4", "--static-analysis", "off"], + success=False, results=True, staged=True, + ) + + if __name__ == "__main__": unittest.main() diff --git a/crates/dbt-sa-cli/tests/test_snowflake_wap_cli.py b/crates/dbt-sa-cli/tests/test_snowflake_wap_cli.py new file mode 100644 index 00000000000..35b92038b3f --- /dev/null +++ b/crates/dbt-sa-cli/tests/test_snowflake_wap_cli.py @@ -0,0 +1,116 @@ +"""Offline CLI integration checks; set DBT_WAP_TEST_BIN to the built dbt binary.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import tempfile +import unittest + + +@unittest.skipUnless(os.environ.get("DBT_WAP_TEST_BIN"), "set DBT_WAP_TEST_BIN to run CLI checks") +class WapCliTests(unittest.TestCase): + def setUp(self) -> None: + temporary = tempfile.TemporaryDirectory(prefix="dbt-wap-cli-") + self.addCleanup(temporary.cleanup) + self.project = Path(temporary.name) + (self.project / "models").mkdir() + (self.project / "dbt_project.yml").write_text( + "name: wap_cli\nversion: '1.0'\nconfig-version: 2\nprofile: wap_cli\n" + "models:\n wap_cli:\n +materialized: table\n +wap: true\n" + ) + # Parse must not contact this deliberately unusable Snowflake target. + (self.project / "profiles.yml").write_text( + "wap_cli:\n target: offline\n outputs:\n offline:\n" + " type: snowflake\n account: snowflake.local\n" + " user: offline\n password: unused\n" + " database: WAP_DB\n schema: EXISTING_SCHEMA\n" + " warehouse: UNUSED\n threads: 2\n connect_timeout: 1\n" + ) + self.write_model("orders", "{{ config(alias='PUBLIC_ORDERS') }}\nselect 1 as id") + self.write_model( + "downstream", "{{ config(wap=false) }}\nselect * from {{ ref('orders') }}" + ) + (self.project / "models/schema.yml").write_text( + "version: 2\nmodels:\n - name: orders\n columns:\n" + " - name: id\n data_tests: [not_null, unique]\n" + ) + + def write_model(self, name: str, sql: str) -> None: + (self.project / f"models/{name}.sql").write_text(sql + "\n") + + def parse(self, *, success: bool = True, error: str = "wap") -> dict: + environment = dict(os.environ) + environment.update(DBT_USE_COLORS="false", DBT_QUIET="false") + completed = subprocess.run( + [ + str(Path(os.environ["DBT_WAP_TEST_BIN"]).resolve()), "parse", + "--project-dir", str(self.project), + "--profiles-dir", str(self.project), + "--target-path", str(self.project / "target"), + "--no-send-anonymous-usage-stats", "--no-version-check", + ], + cwd=self.project, env=environment, text=True, capture_output=True, timeout=60, + ) + output = completed.stdout + "\n" + completed.stderr + self.assertEqual(completed.returncode == 0, success, output[-6000:]) + if not success: + self.assertIn(error, output.lower(), output[-6000:]) + return {} + return json.loads((self.project / "target/manifest.json").read_text()) + + def test_parse_preserves_public_identity_and_ref_dependencies(self) -> None: + manifest = self.parse() + orders = manifest["nodes"]["model.wap_cli.orders"] + downstream = manifest["nodes"]["model.wap_cli.downstream"] + self.assertTrue(orders["config"]["wap"]) + self.assertFalse(downstream["config"]["wap"]) + self.assertEqual(orders["alias"], "PUBLIC_ORDERS") + self.assertEqual(orders["database"], "WAP_DB") + self.assertEqual(orders["schema"], "EXISTING_SCHEMA") + self.assertIn("PUBLIC_ORDERS", orders["relation_name"]) + self.assertNotIn("__DBT_WAP_", json.dumps(manifest)) + self.assertEqual(downstream["depends_on"]["nodes"], ["model.wap_cli.orders"]) + audits = [node for node in manifest["nodes"].values() if node["resource_type"] == "test"] + self.assertEqual(len(audits), 2) + for audit in audits: + self.assertEqual(audit["attached_node"], "model.wap_cli.orders") + + def test_parse_wap_false_allows_view(self) -> None: + self.write_model("orders", "{{ config(wap=false, materialized='view') }}\nselect 1 as id") + manifest = self.parse() + config = manifest["nodes"]["model.wap_cli.orders"]["config"] + self.assertFalse(config["wap"]) + self.assertEqual(config["materialized"], "view") + + def test_parse_rejects_unsupported_materializations(self) -> None: + for materialized in ("view", "incremental"): + with self.subTest(materialized=materialized): + self.write_model( + "orders", "{{ config(materialized='" + materialized + "') }}\nselect 1 as id" + ) + self.parse(success=False) + + def test_parse_rejects_invalid_wap_value(self) -> None: + self.write_model("orders", "{{ config(wap='typo') }}\nselect 1 as id") + self.parse(success=False, error="expected true, false, or null") + + def test_reparse_updates_wap_without_changing_public_identity(self) -> None: + original = self.parse()["nodes"]["model.wap_cli.orders"] + self.write_model( + "orders", "{{ config(wap=false, alias='PUBLIC_ORDERS') }}\nselect 1 as id" + ) + disabled = self.parse()["nodes"]["model.wap_cli.orders"] + self.write_model("orders", "{{ config(alias='PUBLIC_ORDERS') }}\nselect 1 as id") + enabled = self.parse()["nodes"]["model.wap_cli.orders"] + self.assertFalse(disabled["config"]["wap"]) + self.assertTrue(enabled["config"]["wap"]) + for node in (disabled, enabled): + for field in ("unique_id", "database", "schema", "alias", "relation_name"): + self.assertEqual(node[field], original[field], field) + + +if __name__ == "__main__": + unittest.main() diff --git a/crates/dbt-tasks-sa/src/graph.rs b/crates/dbt-tasks-sa/src/graph.rs index b122fc891b0..c69141782f3 100644 --- a/crates/dbt-tasks-sa/src/graph.rs +++ b/crates/dbt-tasks-sa/src/graph.rs @@ -1242,6 +1242,8 @@ mod wap_tests { use dbt_tasks_core::context::TaskRunnerCtx; use dbt_tasks_core::wap::WapModel; + mod execution; + struct PhasedFactory; impl TasksForNodeFactory for PhasedFactory { @@ -1323,17 +1325,18 @@ mod wap_tests { #[derive(Default)] struct StrictBuckets { deferred: HashMap, + off: bool, } impl StaticAnalysisBuckets for StrictBuckets { fn global_static_analysis(&self) -> Option { - Some(dbt_common::io_args::StaticAnalysisKind::Strict) + (!self.off).then_some(dbt_common::io_args::StaticAnalysisKind::Strict) } fn deferred_unique_ids(&self) -> &HashMap { &self.deferred } fn in_off_closure(&self, _: &str) -> bool { - false + self.off } fn in_baseline_closure(&self, _: &str) -> bool { false diff --git a/crates/dbt-tasks-sa/src/graph/wap_tests/execution.rs b/crates/dbt-tasks-sa/src/graph/wap_tests/execution.rs new file mode 100644 index 00000000000..1370dda9204 --- /dev/null +++ b/crates/dbt-tasks-sa/src/graph/wap_tests/execution.rs @@ -0,0 +1,334 @@ +use super::*; +use crate::visitor::wap_test_support::ScriptedSchedule; +use crate::wap::require_passing_audits; +use dbt_tasks_core::visitor::SkipReason; + +const MODEL: &str = "model.pkg.orders"; +const AUDIT: &str = "test.pkg.orders_not_null"; +const SECOND_AUDIT: &str = "test.pkg.orders_nonnegative"; +const DOWNSTREAM: &str = "model.pkg.downstream"; +const DOWNSTREAM_AUDIT: &str = "test.pkg.downstream_not_null"; +const INDEPENDENT: &str = "model.pkg.independent"; + +struct Fixture { + graph: DiGraph, ()>, +} + +impl Fixture { + fn new(wap: bool, chained: bool) -> Self { + let mut nodes = Nodes::default(); + let mut deps = BTreeMap::new(); + for (id, upstream) in [ + (MODEL, None), + (DOWNSTREAM, Some(MODEL)), + (INDEPENDENT, None), + ] { + let mut model = DbtModel::default(); + model.__common_attr__.unique_id = id.to_owned(); + nodes.models.insert(id.to_owned(), Arc::new(model)); + deps.insert( + id.to_owned(), + upstream.into_iter().map(str::to_owned).collect(), + ); + } + for (id, owner) in [(AUDIT, MODEL), (SECOND_AUDIT, MODEL)] + .into_iter() + .chain(chained.then_some((DOWNSTREAM_AUDIT, DOWNSTREAM))) + { + let mut test = DbtTest::default(); + test.__common_attr__.unique_id = id.to_owned(); + nodes.tests.insert(id.to_owned(), Arc::new(test)); + deps.insert(id.to_owned(), BTreeSet::from([owner.to_owned()])); + } + let schedule = Schedule { + selected_nodes: deps.keys().cloned().collect(), + all_selected_nodes: deps.keys().cloned().collect(), + sorted_nodes: deps.keys().cloned().collect(), + deps, + ..Default::default() + }; + let mut plan = WapPlan::default(); + if wap { + let mut model = entry(MODEL, AUDIT); + model.audit_ids.insert(SECOND_AUDIT.to_owned()); + plan.models.insert(MODEL.to_owned(), model); + } + if chained { + plan.models + .insert(DOWNSTREAM.to_owned(), entry(DOWNSTREAM, DOWNSTREAM_AUDIT)); + } + let (graph, _) = GraphBuilder::build_phased_task_graph( + &schedule, + &PhasedFactory, + &nodes, + Execute::Remote, + PHASES_RENDER_ANALYZE_RUN, + &StrictBuckets { + off: true, + ..Default::default() + }, + false, + None, + &plan, + ) + .unwrap(); + Self { graph } + } + + fn index(&self, id: &str, task_type: &str) -> NodeIndex { + self.graph + .node_indices() + .find(|index| { + let task = &self.graph[*index]; + task.work_node_id() == id && task.task_type() == task_type + }) + .unwrap_or_else(|| panic!("missing task {id}/{task_type}")) + } + + /// Complete compilation work while leaving every warehouse task in flight. + fn start_runs(&self, schedule: &mut ScriptedSchedule<'_>) -> BTreeSet { + let mut running = BTreeSet::new(); + loop { + let ready = schedule.start_ready(); + if ready.is_empty() { + return running; + } + for index in ready { + if self.graph[index].task_phase() == Some(TP::Run) { + running.insert(index); + } else { + schedule.complete(index, Ok(NodeStatus::Succeeded)); + } + } + } + } + + fn start_stage(&self) -> ScriptedSchedule<'_> { + let mut schedule = ScriptedSchedule::new(&self.graph); + assert_eq!( + self.start_runs(&mut schedule), + BTreeSet::from([self.index(MODEL, "run"), self.index(INDEPENDENT, "run")]) + ); + schedule.complete(self.index(INDEPENDENT, "run"), Ok(NodeStatus::Succeeded)); + schedule + } + + fn stage_succeeded(&self, schedule: &mut ScriptedSchedule<'_>) { + schedule.complete(self.index(MODEL, "run"), Ok(NodeStatus::Succeeded)); + assert_eq!( + self.start_runs(schedule), + BTreeSet::from([self.index(AUDIT, "run"), self.index(SECOND_AUDIT, "run")]) + ); + } +} + +#[test] +fn wap_publication_waits_for_every_audit_completion() { + let fixture = Fixture::new(true, false); + let mut schedule = fixture.start_stage(); + fixture.stage_succeeded(&mut schedule); + + schedule.complete(fixture.index(AUDIT, "run"), Ok(NodeStatus::TestPassed)); + assert!( + schedule.start_ready().is_empty(), + "second audit is still running" + ); + schedule.complete( + fixture.index(SECOND_AUDIT, "run"), + Ok(NodeStatus::TestPassed), + ); + let publication = fixture.index(MODEL, "wap_publish_run"); + assert_eq!(schedule.start_ready(), vec![publication]); + assert!( + schedule.start_ready().is_empty(), + "publication is still running" + ); + + let result = require_passing_audits([ + (AUDIT, Some(NodeStatus::TestPassed)), + (SECOND_AUDIT, Some(NodeStatus::TestPassed)), + ]) + .map(|()| NodeStatus::Succeeded); + schedule.complete(publication, result); + assert_eq!( + fixture.start_runs(&mut schedule), + BTreeSet::from([fixture.index(DOWNSTREAM, "run")]) + ); + schedule.complete(fixture.index(DOWNSTREAM, "run"), Ok(NodeStatus::Succeeded)); + assert!(schedule.start_ready().is_empty()); + schedule.assert_finished(); +} + +#[test] +fn wap_failed_transformation_skips_audits_publication_and_downstream() { + for result in [ + Ok(NodeStatus::Errored), + Err(fs_err!(ErrorCode::ExecutionError, "transformation failed")), + ] { + let fixture = Fixture::new(true, false); + let mut schedule = fixture.start_stage(); + schedule.complete(fixture.index(MODEL, "run"), result); + + assert!(schedule.start_ready().is_empty()); + for id in [AUDIT, SECOND_AUDIT, DOWNSTREAM] { + for phase in ["render", "run"] { + assert_eq!( + schedule.skip_reason(fixture.index(id, phase)), + Some(&SkipReason::FailedUpstream(MODEL.to_owned())) + ); + } + } + assert_eq!( + schedule.skip_reason(fixture.index(MODEL, "wap_publish_run")), + Some(&SkipReason::FailedPhase) + ); + assert!( + schedule + .skip_reason(fixture.index(INDEPENDENT, "run")) + .is_none() + ); + schedule.assert_finished(); + } +} + +#[test] +fn wap_audit_failure_keeps_publication_blocked_after_other_audits_finish() { + let fixture = Fixture::new(true, false); + let mut schedule = fixture.start_stage(); + fixture.stage_succeeded(&mut schedule); + + schedule.complete(fixture.index(AUDIT, "run"), Ok(NodeStatus::Errored)); + assert!(schedule.start_ready().is_empty()); + schedule.complete( + fixture.index(SECOND_AUDIT, "run"), + Ok(NodeStatus::TestPassed), + ); + + assert!(schedule.start_ready().is_empty()); + for index in [ + fixture.index(MODEL, "wap_publish_run"), + fixture.index(DOWNSTREAM, "render"), + fixture.index(DOWNSTREAM, "run"), + ] { + assert_eq!( + schedule.skip_reason(index), + Some(&SkipReason::FailedUpstream(AUDIT.to_owned())) + ); + } + schedule.assert_finished(); +} + +#[test] +fn wap_warned_audit_fails_publication_gate_and_blocks_downstream() { + let fixture = Fixture::new(true, false); + let mut schedule = fixture.start_stage(); + fixture.stage_succeeded(&mut schedule); + schedule.complete(fixture.index(AUDIT, "run"), Ok(NodeStatus::TestWarned)); + schedule.complete( + fixture.index(SECOND_AUDIT, "run"), + Ok(NodeStatus::TestPassed), + ); + let publication = fixture.index(MODEL, "wap_publish_run"); + assert_eq!(schedule.start_ready(), vec![publication]); + + // WARN is intentionally nonfatal to the visitor. The real WAP gate must + // still turn it into a publication error before consumers can run. + let result = require_passing_audits([ + (AUDIT, Some(NodeStatus::TestWarned)), + (SECOND_AUDIT, Some(NodeStatus::TestPassed)), + ]) + .map(|()| NodeStatus::Succeeded); + assert!(result.is_err()); + schedule.complete(publication, result); + + assert!(schedule.start_ready().is_empty()); + assert_eq!( + schedule.skip_reason(fixture.index(DOWNSTREAM, "render")), + Some(&SkipReason::FailedUpstream(MODEL.to_owned())) + ); + schedule.assert_finished(); +} + +#[test] +fn wap_chain_starts_next_candidate_only_after_upstream_publication() { + let fixture = Fixture::new(true, true); + let mut schedule = fixture.start_stage(); + fixture.stage_succeeded(&mut schedule); + for id in [AUDIT, SECOND_AUDIT] { + schedule.complete(fixture.index(id, "run"), Ok(NodeStatus::TestPassed)); + } + let first_publication = fixture.index(MODEL, "wap_publish_run"); + assert_eq!(schedule.start_ready(), vec![first_publication]); + schedule.complete(first_publication, Ok(NodeStatus::Succeeded)); + assert_eq!( + fixture.start_runs(&mut schedule), + BTreeSet::from([fixture.index(DOWNSTREAM, "run")]) + ); + schedule.complete(fixture.index(DOWNSTREAM, "run"), Ok(NodeStatus::Succeeded)); + assert_eq!( + fixture.start_runs(&mut schedule), + BTreeSet::from([fixture.index(DOWNSTREAM_AUDIT, "run")]) + ); + schedule.complete( + fixture.index(DOWNSTREAM_AUDIT, "run"), + Ok(NodeStatus::TestPassed), + ); + let second_publication = fixture.index(DOWNSTREAM, "wap_publish_run"); + assert_eq!(schedule.start_ready(), vec![second_publication]); + schedule.complete(second_publication, Ok(NodeStatus::Succeeded)); + assert!(schedule.start_ready().is_empty()); + schedule.assert_finished(); +} + +#[test] +fn wap_failure_skips_entire_downstream_wap_cycle() { + let fixture = Fixture::new(true, true); + let mut schedule = fixture.start_stage(); + fixture.stage_succeeded(&mut schedule); + schedule.complete(fixture.index(AUDIT, "run"), Ok(NodeStatus::Errored)); + schedule.complete( + fixture.index(SECOND_AUDIT, "run"), + Ok(NodeStatus::TestPassed), + ); + + assert!(schedule.start_ready().is_empty()); + for (id, phase) in [ + (DOWNSTREAM, "render"), + (DOWNSTREAM, "run"), + (DOWNSTREAM_AUDIT, "render"), + (DOWNSTREAM_AUDIT, "run"), + (DOWNSTREAM, "wap_publish_run"), + ] { + assert_eq!( + schedule.skip_reason(fixture.index(id, phase)), + Some(&SkipReason::FailedUpstream(AUDIT.to_owned())) + ); + } + schedule.assert_finished(); +} + +#[test] +fn wap_disabled_preserves_normal_warning_behavior_without_publication_task() { + let fixture = Fixture::new(false, false); + assert!( + !fixture + .graph + .node_weights() + .any(|task| task.task_type() == "wap_publish_run") + ); + let mut schedule = fixture.start_stage(); + fixture.stage_succeeded(&mut schedule); + schedule.complete(fixture.index(AUDIT, "run"), Ok(NodeStatus::TestWarned)); + schedule.complete( + fixture.index(SECOND_AUDIT, "run"), + Ok(NodeStatus::TestPassed), + ); + + assert_eq!( + fixture.start_runs(&mut schedule), + BTreeSet::from([fixture.index(DOWNSTREAM, "run")]) + ); + schedule.complete(fixture.index(DOWNSTREAM, "run"), Ok(NodeStatus::Succeeded)); + assert!(schedule.start_ready().is_empty()); + schedule.assert_finished(); +} diff --git a/crates/dbt-tasks-sa/src/visitor.rs b/crates/dbt-tasks-sa/src/visitor.rs index c25596480ce..bc6ec7b8c76 100644 --- a/crates/dbt-tasks-sa/src/visitor.rs +++ b/crates/dbt-tasks-sa/src/visitor.rs @@ -27,6 +27,9 @@ use dbt_tasks_core::visitor::SkipReason; use dbt_telemetry::{ExecutionPhase, NodeOutcome, NodeType}; use dbt_yaml::Span; +#[cfg(test)] +pub(crate) mod wap_test_support; + /// Pool of reusable logical worker slot IDs (1, 2, 3, …). /// /// Tasks are assigned a slot before execution and release it after. diff --git a/crates/dbt-tasks-sa/src/visitor/wap_test_support.rs b/crates/dbt-tasks-sa/src/visitor/wap_test_support.rs new file mode 100644 index 00000000000..41468401149 --- /dev/null +++ b/crates/dbt-tasks-sa/src/visitor/wap_test_support.rs @@ -0,0 +1,81 @@ +//! Drive production scheduling decisions with explicitly ordered task completions. +//! +//! These tests replace warehouse work, not dependency readiness or failure propagation. + +use super::*; + +pub(crate) struct ScriptedSchedule<'a> { + graph: &'a DiGraph, ()>, + strategy: VisitStrategy, + indegree: Vec, + dependents: Vec>, + pending: Vec, + waiting: HashMap, + skip_set: SkipSet, +} + +impl<'a> ScriptedSchedule<'a> { + pub(crate) fn new(graph: &'a DiGraph, ()>) -> Self { + let strategy = VisitStrategy::Parallel; + let (indegree, dependents) = get_indegree_and_dependents(graph); + Self { + graph, + strategy, + indegree, + dependents, + pending: strategy.new_pending_nodes(graph).unwrap(), + waiting: HashMap::new(), + skip_set: SkipSet::new(), + } + } + + pub(crate) fn start_ready(&mut self) -> Vec { + let mut started = Vec::new(); + while let Some(index) = self + .strategy + .try_pop_pending_node(&self.waiting, &mut self.pending) + { + if self.skip_set.skip.contains_key(&index) { + self.strategy.release_dependents( + index, + &self.dependents, + &mut self.indegree, + &mut self.pending, + ); + } else { + assert!(self.waiting.insert(index, tracing::Span::none()).is_none()); + started.push(index); + } + } + started + } + + pub(crate) fn complete(&mut self, index: NodeIndex, result: FsResult) { + assert!( + self.waiting.remove(&index).is_some(), + "task was not running" + ); + self.skip_set + .handle_task_result(result, index, &self.dependents, self.graph, false, false); + self.strategy.release_dependents( + index, + &self.dependents, + &mut self.indegree, + &mut self.pending, + ); + } + + pub(crate) fn skip_reason(&self, index: NodeIndex) -> Option<&SkipReason> { + self.skip_set.skip.get(&index) + } + + pub(crate) fn assert_finished(&self) { + assert!(self.pending.is_empty()); + assert!(self.waiting.is_empty()); + assert!( + self.strategy + .get_incomplete_tasks(self.graph, &self.indegree) + .is_empty() + ); + } +} diff --git a/crates/dbt-tasks-sa/src/wap.rs b/crates/dbt-tasks-sa/src/wap.rs index 592d0e43364..17c69e9ae2d 100644 --- a/crates/dbt-tasks-sa/src/wap.rs +++ b/crates/dbt-tasks-sa/src/wap.rs @@ -138,7 +138,7 @@ impl Task for PublishTask { } } -fn require_passing_audits<'a>( +pub(crate) fn require_passing_audits<'a>( audits: impl IntoIterator)>, ) -> FsResult<()> { let mut count = 0; @@ -737,6 +737,8 @@ mod clone_tests; mod header_tests; #[cfg(test)] mod publication_tests; +#[cfg(test)] +mod stage_tests; #[cfg(test)] mod tests { diff --git a/crates/dbt-tasks-sa/src/wap/clone_tests.rs b/crates/dbt-tasks-sa/src/wap/clone_tests.rs index e47d6bdded3..0848e1b7d23 100644 --- a/crates/dbt-tasks-sa/src/wap/clone_tests.rs +++ b/crates/dbt-tasks-sa/src/wap/clone_tests.rs @@ -19,7 +19,7 @@ const CLONE_MACRO: &str = include_str!( "../../../dbt-loader/src/dbt_macro_assets/dbt-snowflake/macros/materializations/clone.sql" ); -fn fixture(copy_grants: Option) -> (tempfile::TempDir, IoArgs, WapModel) { +pub(super) fn fixture(copy_grants: Option) -> (tempfile::TempDir, IoArgs, WapModel) { let directory = tempfile::tempdir().unwrap(); std::fs::create_dir(directory.path().join("models")).unwrap(); std::fs::write(directory.path().join("models/orders.sql"), "select 1 as id").unwrap(); @@ -58,7 +58,7 @@ fn fixture(copy_grants: Option) -> (tempfile::TempDir, IoArgs, WapModel) { (directory, io, wap) } -fn runtime_context(model: &DbtModel, io: &IoArgs) -> BTreeMap { +pub(super) fn runtime_context(model: &DbtModel, io: &IoArgs) -> BTreeMap { let (context, _) = build_run_node_context( model, &model.deprecated_config, diff --git a/crates/dbt-tasks-sa/src/wap/stage_tests.rs b/crates/dbt-tasks-sa/src/wap/stage_tests.rs new file mode 100644 index 00000000000..f517793f2b5 --- /dev/null +++ b/crates/dbt-tasks-sa/src/wap/stage_tests.rs @@ -0,0 +1,244 @@ +//! Execute the shipped materialization against a recording warehouse boundary. +//! These checks validate the submitted SQL, not Snowflake execution or atomicity. + +use std::collections::BTreeMap; +use std::rc::Rc; +use std::sync::Arc; + +use dbt_adapter::relation::{RelationObject, factory::create_static_relation}; +use dbt_adapter_core::AdapterType; +use dbt_common::FsResult; +use dbt_jinja_utils::JinjaEnvBuilder; +use minijinja::listener::RenderingEventListener; +use minijinja::value::{Kwargs, Object, from_args}; +use minijinja::{Error, ErrorKind, State, Value}; +use parking_lot::Mutex; + +use super::clone_tests::{fixture, runtime_context}; + +const TABLE_MACRO: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-snowflake/macros/materializations/table.sql" +); +const CREATE_MACROS: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-snowflake/macros/relations/table/create.sql" +); +const CREATE_DISPATCH_MACROS: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-adapters/macros/relations/table/create.sql" +); +const STATEMENT_MACROS: &str = + include_str!("../../../dbt-loader/src/dbt_macro_assets/dbt-adapters/macros/etc/statement.sql"); +const HOOK_MACROS: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-adapters/macros/materializations/hooks.sql" +); +const GRANT_MACROS: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-adapters/macros/adapters/apply_grants.sql" +); +const DOC_MACROS: &str = include_str!( + "../../../dbt-loader/src/dbt_macro_assets/dbt-adapters/macros/adapters/persist_docs.sql" +); +const ADAPTER_MACROS: &str = + include_str!("../../../dbt-loader/src/dbt_macro_assets/dbt-snowflake/macros/adapters.sql"); + +/// Resolve the limited dispatch surface used by this native table fixture to +/// the shipped macro definitions. Other adapter calls fail unless registered. +fn dispatch(args: &[Value]) -> Result { + let name = match args.first().and_then(Value::as_str) { + Some("create_table_as") => "snowflake__create_table_as", + Some("set_query_tag") => "snowflake__set_query_tag", + Some("unset_query_tag") => "snowflake__unset_query_tag", + Some("apply_grants") => "default__apply_grants", + Some("persist_docs") => "default__persist_docs", + unexpected => panic!("unexpected native table dispatch: {unexpected:?}"), + }; + Ok(Value::from_function( + move |state: &State, args: &[Value]| { + state.lookup(name, &[]).unwrap().call(state, args, &[]) + }, + )) +} + +#[derive(Debug)] +struct RecordedCall { + method: String, + args: Vec, +} + +#[derive(Debug)] +struct RecordingAdapter { + public_relation: Option, + transient: bool, + fail_execution: bool, + calls: Mutex>, +} + +impl Object for RecordingAdapter { + fn call_method( + self: &Arc, + _state: &State, + name: &str, + args: &[Value], + _listeners: &[Rc], + ) -> Result { + self.calls.lock().push(RecordedCall { + method: name.to_owned(), + args: args.to_vec(), + }); + match name { + "dispatch" => dispatch(args), + "get_relation" => { + let (_, kwargs) = from_args::<(&[Value], Kwargs)>(args)?; + let identifier: String = kwargs.get("identifier")?; + Ok(if identifier == "Orders" { + self.public_relation + .clone() + .unwrap_or_else(|| Value::from(())) + } else { + Value::from(()) + }) + } + "build_catalog_relation" => Ok(Value::from_serialize(serde_json::json!({ + "catalog_type": "INFO_SCHEMA", + "table_format": null, + "is_transient": self.transient + }))), + "execute" => { + if self.fail_execution { + Err(Error::new( + ErrorKind::InvalidOperation, + "injected CTAS failure", + )) + } else { + Ok(Value::from(vec![Value::from("SUCCESS"), Value::from(())])) + } + } + unexpected => panic!("unexpected native table adapter call: {unexpected}"), + } + } +} + +fn stage( + transient: bool, + public_exists: bool, + transformation: &str, + fail_execution: bool, +) -> (FsResult, Arc) { + let (_directory, io, mut wap) = fixture(Some(true)); + Arc::make_mut(&mut wap.model) + .deprecated_config + .__warehouse_specific_config__ + .transient = Some(transient); + let execution_model = wap.execution_model().unwrap(); + let mut context = runtime_context(&execution_model, &io); + context.insert("execute".to_owned(), Value::from(true)); + context.insert("compiled_code".to_owned(), Value::from(transformation)); + context.insert( + "api".to_owned(), + Value::from_serialize(BTreeMap::from([( + "Relation", + create_static_relation(AdapterType::Snowflake, wap.model.__base_attr__.quoting) + .unwrap(), + )])), + ); + let adapter = Arc::new(RecordingAdapter { + public_relation: public_exists + .then(|| RelationObject::new(wap.public_relation().unwrap().into()).into_value()), + transient, + fail_execution, + calls: Mutex::new(Vec::new()), + }); + context.insert( + "adapter".to_owned(), + Value::from_dyn_object(adapter.clone()), + ); + let environment = JinjaEnvBuilder::new().build(); + let template = [ + TABLE_MACRO, + CREATE_MACROS, + CREATE_DISPATCH_MACROS, + STATEMENT_MACROS, + HOOK_MACROS, + GRANT_MACROS, + DOC_MACROS, + ADAPTER_MACROS, + "{{ materialization_table_snowflake() }}", + ] + .join("\n"); + let result = environment.render_str(&template, &context, &[]); + assert_eq!(wap.model.__base_attr__.alias, "Orders"); + (result, adapter) +} + +fn submitted_sql(adapter: &RecordingAdapter) -> Vec { + adapter + .calls + .lock() + .iter() + .filter(|call| call.method == "execute") + .map(|call| { + call.args[0] + .as_str() + .unwrap() + .split_whitespace() + .collect::>() + .join(" ") + }) + .collect() +} + +#[test] +fn wap_actual_table_materialization_only_rebuilds_the_same_schema_candidate() { + for transient in [false, true] { + for public_exists in [false, true] { + for transformation in [ + "select 1 as id, 10 as amount", + "select 2 as id, -10 as amount", + ] { + let (result, adapter) = stage(transient, public_exists, transformation, false); + result.unwrap(); + let lifecycle = if transient { "transient " } else { "" }; + assert_eq!( + submitted_sql(&adapter), + [format!( + "create or replace {lifecycle}table \ + \"My\"\"Database\".\"My.Schema\".\"__DBT_WAP_TEST\" \ + as ({transformation} ) ;" + )], + "only the candidate CTAS may be submitted before auditing; \ + public_exists={public_exists}, transient={transient}" + ); + let calls = adapter.calls.lock(); + let lookups: Vec<_> = calls + .iter() + .filter(|call| call.method == "get_relation") + .collect(); + assert_eq!(lookups.len(), 1); + let (_, kwargs) = from_args::<(&[Value], Kwargs)>(&lookups[0].args).unwrap(); + assert_eq!( + kwargs.get::("identifier").unwrap(), + "__DBT_WAP_TEST" + ); + assert_eq!(kwargs.get::("database").unwrap(), "My\"Database"); + assert_eq!(kwargs.get::("schema").unwrap(), "My.Schema"); + } + } + } +} + +#[test] +fn wap_actual_table_materialization_propagates_transformation_failure() { + let (result, adapter) = stage(true, true, "select missing_column from upstream", true); + assert!( + result + .unwrap_err() + .to_string() + .contains("injected CTAS failure") + ); + assert_eq!(submitted_sql(&adapter).len(), 1); + let calls = adapter.calls.lock(); + let dispatched: Vec<_> = calls + .iter() + .filter(|call| call.method == "dispatch") + .map(|call| call.args[0].as_str().unwrap()) + .collect(); + assert_eq!(dispatched, ["set_query_tag", "create_table_as"]); +} diff --git a/docs/snowflake-table-wap-plan.md b/docs/snowflake-table-wap-plan.md index 5cb2105bd08..b0e78ad5d8e 100644 --- a/docs/snowflake-table-wap-plan.md +++ b/docs/snowflake-table-wap-plan.md @@ -54,6 +54,12 @@ ordinary downstream `ref()` calls remain unchanged. clustering, refresh relation metadata, and release downstream tasks. Drop the successful working table after publication completes. +Staging copies the model's execution definition to a new table name. The SQL +transformation builds that working table with `CREATE OR REPLACE TABLE ... AS`. +Snowflake `CLONE` is used at publication; the existing public table is not +cloned before the transformation. This also supports a first build when there +is no public table yet. + On audit failure, the previous published table remains available and downstream models are skipped. On a first build, failure leaves the published relation absent. Retain failed working tables for inspection and report their exact @@ -179,5 +185,15 @@ warnings. Ref and compiled-SQL cache tests verify that working-table identities stay scoped to the current build; a later standalone test reads the public table. These checks do not establish Snowflake's live DDL or permission behavior. +Further regression coverage executes the shipped table materialization and CTAS +macros against a recording adapter, requiring exactly one working-table CTAS in +the existing schema and checking transformation-error propagation. Scheduler +tests drive the production graph, readiness, and failure propagation with +scripted task completions, including delayed audits and chained WAP models. +Offline CLI tests parse actual projects and inspect manifests using a profile +that cannot connect to Snowflake. The live runner additionally checks that both +existing public and downstream sentinel data survive transformation errors, +audit failures, and warnings. + The live acceptance runner and its invocation instructions are in [`crates/dbt-sa-cli/tests/data/snowflake_wap/README.md`](../crates/dbt-sa-cli/tests/data/snowflake_wap/README.md). From 9aea5ec3e84df3860b88056cc27ac3bed12a2f21 Mon Sep 17 00:00:00 2001 From: bmoore813 Date: Fri, 18 Sep 2026 02:09:52 -0400 Subject: [PATCH 3/3] ci: run offline Snowflake WAP checks --- .github/workflows/ci.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a081d4fea97..4e86dd41f51 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -122,6 +122,17 @@ jobs: - name: Run tests run: cargo nextest run --workspace + - name: Build dbt for offline WAP checks + run: cargo build -p dbt-sa-cli --bin dbt + + - name: Run offline WAP Python checks + env: + DBT_WAP_TEST_BIN: ${{ github.workspace }}/target/debug/dbt + PYTHONDONTWRITEBYTECODE: "1" + run: | + test -x "$DBT_WAP_TEST_BIN" + python3 -m unittest discover -s crates/dbt-sa-cli/tests -p 'test_snowflake_wap_*.py' -v + - name: sccache stats if: always() run: sccache --show-stats || true