feat(pipeline): write-audit-publish for materialization data tests (#9911)

* feat(pipeline): write-audit-publish for materialization data tests (EE) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs: EE worktree E0583 troubleshooting + duckdb feature check row Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs: clarify EE symlink example (absolute target, EE repo layout) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(pipeline): move bootstrap DDL inside guarded WAP transaction Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(pipeline): move WAP guard SQL builder into EE, OSS keeps placement only Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * chore: bump ee-repo-ref to EE branch rebased on EE main Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * style: reword test comment as current invariant per AGENTS.md Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * chore: bump ee-repo-ref (EE module doc update) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(pipeline): OSS emits typed materialize plan, EE owns WAP transform Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: make rewrite assertion build-aware; refresh oss module doc Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * chore: update ee-repo-ref to 7be0bad1a6d6b5c3a107c0a2cd4bf003c36ec34c This commit updates the EE repository reference after PR #644 was merged in windmill-ee-private. Previous ee-repo-ref: 63cabae75329429f647e01083936d70f8197dc9e New ee-repo-ref: 7be0bad1a6d6b5c3a107c0a2cd4bf003c36ec34c Automated by sync-ee-ref workflow. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>
2026-08-18 16:02:10 +00:00 · 2026-07-04 15:19:49 +00:00
parent a368d49bd8
commit dce247c6d2
8 changed files with 289 additions and 39 deletions
@@ -1 +1 @@
-6f5fe0f7f56696fbef5a8349da38496c32e71666
+7be0bad1a6d6b5c3a107c0a2cd4bf003c36ec34c
@@ -440,7 +440,6 @@ impl<'a> MaterializeCodegen<'a> {
                "CREATE TABLE IF NOT EXISTS {t} AS SELECT * FROM ({sel}) WHERE false;"
            ));
        }
-
        out.push("BEGIN TRANSACTION;".to_string());
        // The rows to write, with the partition column appended when partitioned.
        let source = if self.partitioned {
@@ -578,7 +577,7 @@ impl<'a> MaterializeCodegen<'a> {
        ];
        // Hard-delete-close (`deletes=close`): the keys that vanished from the
        // snapshot — present-and-current in the table, absent from the SELECT.
-        // Captured before the transaction (like `changed`) and disjoint from it (a
+        // Captured before the close (like `changed`) and disjoint from it (a
        // key is either in the snapshot or not), so the two closes never overlap.
        if close_deleted {
            out.push(format!(
@@ -646,6 +645,44 @@ pub fn snapshot_capture_sql(alias: &str) -> String {
 /// from the target ducklake's config and passes it in as `target_attach`.
 pub const TARGET_ALIAS: &str = "_wm_target";

+/// Structural role of one statement in a [`MaterializePlan`]. The public build
+/// executes the plan verbatim, so the kinds are pure metadata there; they exist
+/// so a downstream assembler (`pipeline_advanced::finalize_materialize_query`)
+/// can reason about the plan without parsing SQL.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MaterializeStmtKind {
+    /// Pre-write statement: user setup, the target ATTACH, referenced-asset
+    /// ATTACHes.
+    Setup,
+    /// Write work against the target: bootstrap DDL, SCD2 temp-table captures,
+    /// the mutation itself, the `_current` view.
+    Write,
+    /// The `BEGIN TRANSACTION;` marker emitted by the strategy codegen.
+    TxnBegin,
+    /// The `COMMIT;` marker emitted by the strategy codegen.
+    TxnCommit,
+    /// The trailing one-row summary read (asset / rows / snapshot_id /
+    /// data_tests breakdown).
+    Summary,
+}
+
+/// One planned statement: its structural role and the SQL text.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MaterializeStmt {
+    pub kind: MaterializeStmtKind,
+    pub sql: String,
+}
+
+/// The full ordered materialization plan [`build_wrap_blocks`] produces:
+/// statements in execution order plus the compiled data-test checks (also
+/// embedded in the summary statement's breakdown). Assembled into the final
+/// statement list by `pipeline_advanced::finalize_materialize_query`.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MaterializePlan {
+    pub stmts: Vec<MaterializeStmt>,
+    pub checks: Vec<DataTestCheck>,
+}
+
 /// Assemble the full ordered statement list the DuckDB executor runs for a
 /// managed `// materialize` script. This is the single entry point the worker
 /// calls; it composes the already-tested pieces (classifier split → target
@@ -659,6 +696,14 @@ pub const TARGET_ALIAS: &str = "_wm_target";
 /// the full `<name>/<table>` for the result summary. The trailing statement is
 /// a one-row summary read (asset / rows / snapshot_id) that is both the job's
 /// result (a useful preview) and what the worker records.
+///
+/// Returns a [`MaterializePlan`] — the statements plus their structural role
+/// and the compiled data-test checks — rather than raw SQL: the executor hands
+/// the plan to `windmill_common::pipeline_advanced::finalize_materialize_query`
+/// (which this crate cannot depend on), whose public-build implementation
+/// assembles the statements verbatim. Everything this function produces runs
+/// as-is on the public build; the plan's structure is metadata about it, not a
+/// second mode.
 pub fn build_wrap_blocks(
    plan: &WrapPlan,
    target_attach: &str,
@@ -669,17 +714,9 @@ pub fn build_wrap_blocks(
    partitioned: bool,
    strategy: MaterializeStrategy,
    tests: &[DataTestResolved],
-) -> Result<Vec<String>, String> {
+) -> Result<MaterializePlan, String> {
    let target_qualified = format!("{TARGET_ALIAS}.{target_table}");
    let scd2 = matches!(strategy, MaterializeStrategy::Scd2 { .. });
-    let cg = MaterializeCodegen {
-        target_qualified: &target_qualified,
-        select_sql: &plan.output,
-        partition_col,
-        partition_value_sql,
-        partitioned,
-        strategy,
-    };
    let ctx = DataTestCtx {
        target_qualified: &target_qualified,
        asset_path,
@@ -689,27 +726,49 @@ pub fn build_wrap_blocks(
        scd2,
    };
    let test_sql = build_data_test_checks(tests, &ctx)?;
-    let mut blocks: Vec<String> = Vec::new();
+    let cg = MaterializeCodegen {
+        target_qualified: &target_qualified,
+        select_sql: &plan.output,
+        partition_col,
+        partition_value_sql,
+        partitioned,
+        strategy,
+    };
+    let mut stmts: Vec<MaterializeStmt> = Vec::new();
+    let setup = |sql: String| MaterializeStmt { kind: MaterializeStmtKind::Setup, sql };
    // Setup blocks come from the splitter with their `;` stripped — re-terminate
    // each so that when the executor re-joins and re-splits the assembled query,
    // adjacent statements (e.g. the user ATTACH and the synthetic target ATTACH)
    // don't merge into one malformed statement.
-    blocks.extend(plan.setup.iter().map(|s| terminate(s)));
-    blocks.push(target_attach.to_string());
+    stmts.extend(plan.setup.iter().map(|s| setup(terminate(s))));
+    stmts.push(setup(target_attach.to_string()));
    // Referenced-asset ATTACHes (relationships tests) — read-only, before the
    // write and the summary that probes them.
-    blocks.extend(test_sql.attaches);
-    blocks.extend(cg.statements());
+    stmts.extend(test_sql.attaches.into_iter().map(setup));
+    // Classify the codegen statements by matching the exact transaction-marker
+    // literals this module emits (`BEGIN TRANSACTION;` / `COMMIT;`); everything
+    // else the codegen produces is write work.
+    stmts.extend(cg.statements().into_iter().map(|sql| {
+        let kind = match sql.as_str() {
+            "BEGIN TRANSACTION;" => MaterializeStmtKind::TxnBegin,
+            "COMMIT;" => MaterializeStmtKind::TxnCommit,
+            _ => MaterializeStmtKind::Write,
+        };
+        MaterializeStmt { kind, sql }
+    }));
    // The summary read carries the per-test breakdown (when any tests apply).
-    blocks.push(materialize_result_sql(
-        &target_qualified,
-        asset_path,
-        partition_col,
-        partition_value_sql,
-        partitioned,
-        &test_sql.checks,
-    ));
-    Ok(blocks)
+    stmts.push(MaterializeStmt {
+        kind: MaterializeStmtKind::Summary,
+        sql: materialize_result_sql(
+            &target_qualified,
+            asset_path,
+            partition_col,
+            partition_value_sql,
+            partitioned,
+            &test_sql.checks,
+        ),
+    });
+    Ok(MaterializePlan { stmts, checks: test_sql.checks })
 }

 /// The trailing one-row summary the materialize run returns: the asset it
@@ -1454,7 +1513,7 @@ mod tests {
    #[test]
    fn build_wrap_blocks_orders_setup_attach_codegen_snapshot() {
        let plan = ok("ATTACH 'ducklake://main' AS dl;\n SELECT a FROM dl.orders WHERE d = '{p}'");
-        let blocks = build_wrap_blocks(
+        let blocks: Vec<String> = build_wrap_blocks(
            &plan,
            "ATTACH 'ducklake:postgres:…' AS _wm_target (DATA_PATH 's3://b/p');",
            "orders_daily",
@@ -1465,7 +1524,11 @@ mod tests {
            MaterializeStrategy::Replace,
            &[],
        )
-        .unwrap();
+        .unwrap()
+        .stmts
+        .into_iter()
+        .map(|s| s.sql)
+        .collect();
        // setup block first, then the target ATTACH, then codegen, then result.
        assert!(blocks[0].starts_with("ATTACH 'ducklake://main' AS dl"));
        // every setup block must be `;`-terminated so re-splitting can't merge it
@@ -1488,6 +1551,102 @@ mod tests {
        assert!(last.contains("ducklake_snapshots('_wm_target')"));
    }

+    // -- materialize plan structure ------------------------------------------
+
+    fn plan_for(strategy: MaterializeStrategy, partitioned: bool) -> MaterializePlan {
+        let plan = ok("SELECT a, b FROM src");
+        build_wrap_blocks(
+            &plan,
+            "ATTACH 'ducklake:…' AS _wm_target;",
+            "orders",
+            "main/orders",
+            "_wm_partition",
+            "'2026-06-19'",
+            partitioned,
+            strategy,
+            &[
+                DataTestResolved::BuiltIn(DataTest::NotNull { column: "a".into() }),
+                DataTestResolved::BuiltIn(DataTest::Unique { column: "b".into() }),
+            ],
+        )
+        .unwrap()
+    }
+
+    fn kidx(plan: &MaterializePlan, pred: impl Fn(&MaterializeStmt) -> bool) -> usize {
+        plan.stmts
+            .iter()
+            .position(|s| pred(s))
+            .expect("stmt present")
+    }
+
+    #[test]
+    fn plan_tags_structure_and_carries_checks() {
+        use MaterializeStmtKind::*;
+        let plan = plan_for(MaterializeStrategy::Replace, true);
+        // leading statements are Setup, ending with the target ATTACH
+        assert!(plan.stmts[0].kind == Setup);
+        assert!(plan
+            .stmts
+            .iter()
+            .take_while(|s| s.kind == Setup)
+            .any(|s| s.sql.contains("_wm_target")));
+        // txn markers are tagged, everything between them is Write
+        let begin = kidx(&plan, |s| s.kind == TxnBegin);
+        let commit = kidx(&plan, |s| s.kind == TxnCommit);
+        assert!(begin < commit);
+        assert!(plan.stmts[begin + 1..commit]
+            .iter()
+            .all(|s| s.kind == Write));
+        // bootstrap DDL is Write work (it targets the table, not the session)
+        let bootstrap = kidx(&plan, |s| s.sql.starts_with("CREATE TABLE IF NOT EXISTS"));
+        assert_eq!(plan.stmts[bootstrap].kind, Write);
+        // summary is last and carries the breakdown; checks ride along
+        let last = plan.stmts.last().unwrap();
+        assert_eq!(last.kind, Summary);
+        assert!(last.sql.contains("AS data_tests"));
+        assert_eq!(plan.checks.len(), 2);
+        assert!(plan.checks[0].name.contains("not_null(a)"));
+    }
+
+    #[test]
+    fn plan_whole_table_replace_has_no_txn_markers() {
+        use MaterializeStmtKind::*;
+        let plan = plan_for(MaterializeStrategy::Replace, false);
+        assert!(!plan.stmts.iter().any(|s| s.kind == TxnBegin));
+        assert!(!plan.stmts.iter().any(|s| s.kind == TxnCommit));
+        assert_eq!(
+            plan.stmts.iter().filter(|s| s.kind == Write).count(),
+            1,
+            "single atomic CREATE OR REPLACE"
+        );
+    }
+
+    #[test]
+    fn plan_scd2_captures_are_write_kind() {
+        use MaterializeStmtKind::*;
+        let plan = plan_for(
+            MaterializeStrategy::Scd2 { key: "a".into(), track: vec![], close_deleted: false },
+            false,
+        );
+        let capture = kidx(&plan, |s| s.sql.contains("TEMP TABLE _wm_scd2_changed"));
+        assert_eq!(plan.stmts[capture].kind, Write);
+        // no test declared ⇒ empty checks
+        let plain = ok("SELECT a FROM src");
+        let no_tests = build_wrap_blocks(
+            &plain,
+            "ATTACH 'ducklake:…' AS _wm_target;",
+            "orders",
+            "main/orders",
+            "_wm_partition",
+            "''",
+            false,
+            MaterializeStrategy::Append,
+            &[],
+        )
+        .unwrap();
+        assert!(no_tests.checks.is_empty());
+    }
+
    // -- data tests ---------------------------------------------------------

    fn ctx_partitioned() -> DataTestCtx<'static> {
@@ -1,12 +1,27 @@
-//! OSS fallback: pipeline partition backfills are an enterprise feature;
-//! their implementations live in windmill-ee-private (see
-//! `pipeline_advanced_ee`). In the public build the entry points report that
-//! the enterprise edition is required. (Freshness lives elsewhere: the
-//! fresh/stale badge is CE in the assets API, the active watchdog is
-//! windmill-queue's `freshness_watchdog`.)
+//! OSS fallback for enterprise pipeline features (implementations in
+//! windmill-ee-private, see `pipeline_advanced_ee`): partition backfill
+//! reports that the enterprise edition is required, and materialization-plan
+//! assembly runs the plan verbatim — dbt-like commit-then-test instead of the
+//! enterprise write-audit-publish. (Freshness lives elsewhere: the fresh/stale
+//! badge is CE in the assets API, the active watchdog is windmill-queue's
+//! `freshness_watchdog`.)

 use crate::error::Error;

 pub fn backfill_todo() -> Error {
    Error::internal_err("Pipeline partition backfill requires the enterprise edition".to_string())
 }
+
+/// Assemble a materialization plan into the statement list the DuckDB executor
+/// runs. The public build executes the plan verbatim — dbt-like
+/// commit-then-test: a failing `// data_test` still fails the run and stops
+/// the cascade, but the written slice stays live. The enterprise
+/// implementation (`pipeline_advanced_ee`) instead restructures the plan into
+/// write-audit-publish, where a failing test rolls the whole write back before
+/// anything is published.
+pub fn finalize_materialize_query(
+    plan: windmill_parser::sql_materialize::MaterializePlan,
+    _asset_path: &str,
+) -> Vec<String> {
+    plan.stmts.into_iter().map(|s| s.sql).collect()
+}
@@ -242,7 +242,7 @@ fn build_materialized_query(
        }
    }

-    let blocks = build_wrap_blocks(
+    let mat_plan = build_wrap_blocks(
        &plan,
        &synthetic_attach,
        table,
@@ -254,6 +254,11 @@ fn build_materialized_query(
        &resolved,
    )
    .map_err(Error::ExecutionErr)?;
+    // Enterprise seam: assembles the plan into the final statement list —
+    // verbatim on the public build (commit-then-test), restructured into
+    // write-audit-publish (guarded transaction, rollback on violation) on EE.
+    let blocks =
+        windmill_common::pipeline_advanced::finalize_materialize_query(mat_plan, &m.target_path);

    Ok(Some((Some(blocks.join("\n")), meta)))
 }
@@ -1240,6 +1245,11 @@ pub async fn do_duckdb(
            // committed (like dbt), so the slice is recorded `Failed` and the
            // cascade stops. The error lists *every* test so the user sees the
            // whole picture, not just the first failure.
+            // Under enterprise write-audit-publish an in-transaction guard
+            // (the enterprise `finalize_materialize_query` restructure) already aborted a failing run before
+            // COMMIT — that surfaces on the Err path above with the same
+            // breakdown in the error string, and nothing was published; this
+            // post-commit path then only ever sees passing counts.
            let tests = extract_data_tests(&result);
            // Captured output schema (gap #2a) — recorded only on the successful
            // path below, not on the failure paths (a failed run shouldn't
@@ -2498,6 +2508,37 @@ mod tests {
        );
    }

+    // The rewritten SQL is the plan assembled by
+    // `pipeline_advanced::finalize_materialize_query`, so what it contains is
+    // build-dependent: the public assembly runs the plan verbatim (tests only
+    // in the post-commit summary breakdown), the enterprise assembly adds the
+    // in-transaction write-audit-publish guard (whose shape/placement is
+    // tested next to its implementation in windmill-common's
+    // `pipeline_advanced_ee`).
+    #[test]
+    fn materialize_rewrite_carries_data_test_summary() {
+        let script = "-- materialize ducklake://main/orders\n\
+                      -- data_test not_null id\n\
+                      SELECT id FROM dl.src";
+        let (rewritten, meta) =
+            build_materialized_query(script, None, &std::collections::HashMap::new())
+                .expect("materialize builds")
+                .expect("materialize present");
+        let rewritten = rewritten.expect("managed mode rewrites the query");
+        assert!(rewritten.contains("AS data_tests"));
+        assert_eq!(meta.n_data_tests, 1);
+        #[cfg(not(feature = "private"))]
+        assert!(
+            !rewritten.contains("error("),
+            "public assembly is commit-then-test (no guard)"
+        );
+        #[cfg(all(feature = "private", feature = "enterprise"))]
+        assert!(
+            rewritten.contains("error("),
+            "enterprise assembly places the WAP guard"
+        );
+    }
+
    // Tests for parse_attach_db_resource function
    #[test]
    fn test_parse_attach_db_resource_postgres_res_prefix() {
@@ -354,10 +354,22 @@ load-bearing.
  repeats the natural key across closed versions, so an unscoped
  `unique(<key>)` would fail the run on the second change of any key. Custom
  tests see the raw history and scope themselves.
- **Commit-then-test.** Like dbt, the write commits before tests run; a failed
-  test fails the *run* (and records `Failed`, so downstream cascade stops) but
-  does not roll back the committed snapshot. Time-travel still lets you inspect
-  exactly what failed.
+- **Commit-then-test (public) / write-audit-publish (enterprise).** In the
+  public build, like dbt, the write commits before tests run; a failed test
+  fails the *run* (and records `Failed`, so downstream cascade stops) but does
+  not roll back the committed snapshot — time-travel still lets you inspect
+  exactly what failed. Enterprise upgrades this to write-audit-publish: the
+  same checks also run *inside* the write transaction as a guard statement
+  that raises on any violation, aborting the run before `COMMIT` — a failing
+  slice is never published, readers keep the previous version, and no snapshot
+  is created (even a *first* run of a new asset rolls back to "no table").
+  The whole mechanism lives in the enterprise repo:
+  `sql_materialize::build_wrap_blocks` produces a typed statement plan
+  (`MaterializePlan` — statements plus structural kinds and the compiled
+  checks), and `pipeline_advanced::finalize_materialize_query` assembles it —
+  verbatim on the public build, restructured into the guarded transaction on
+  EE. The public build thus carries only plan metadata, not the
+  write-audit-publish transform itself.
 - **Custom = DuckDB SQL, server worker.** The escape hatch fetches the deployed
  script's content (a single DuckDB `SELECT`/CTE returning the violating rows —
  it's embedded as a subquery, so a multi-statement body is rejected with a
@@ -47,3 +47,19 @@ cargo check --features enterprise,private
 # EE code that also requires license validation
 cargo check --features enterprise,private,license
 ```
+
+## Troubleshooting
+
+**`E0583: file not found for module <x>_ee` on `cargo check --features enterprise,private`**:
+the EE worktree is behind the OSS code it must satisfy. Fast-forward it to EE
+`origin/main` — EE worktrees are shallow clones, so run `git fetch --unshallow origin`
+first or the merge fails with "refusing to merge unrelated histories". After the
+fast-forward, any `*_ee.rs` file that is *new* since the worktree was created still
+needs its OSS symlink made by hand (worktree setup only links files that existed then).
+The EE repo has no `backend/` prefix — crates sit at its root. Use an absolute target
+(a relative one would resolve against the link's directory, not your cwd):
+
+```bash
+# from the windmill repo root
+ln -s ~/windmill-ee-private/<crate>/src/<x>_ee.rs backend/<crate>/src/<x>_ee.rs
+```
@@ -75,6 +75,12 @@ the annotation→verifier pattern that column-lineage will reuse. The keyword is
 `data_test`, not `test`, to stay clear of the unrelated `// test:` CI-test
 annotation.

+Enterprise goes one step *beyond* dbt here: write-audit-publish. dbt commits
+the model then tests it, so a failing test leaves the bad table live for BI
+readers; the enterprise build runs the same probes inside the write
+transaction and rolls back on violation, so a failing slice is never
+published (see `ducklake-materialization.md` §"Scoping decisions").
+
 ### 2. Incremental materializations

 See [Incremental deep-dive](#incremental-deep-dive) below.
@@ -10,6 +10,7 @@ After making changes, run the appropriate checks and fix all errors before consi
 | Enterprise code (`*_ee.rs`) | `cargo check --features enterprise,private` | Also do EE PR workflow (see `docs/enterprise.md`) |
 | Enterprise + license-gated code | `cargo check --features enterprise,private,license` | When the feature requires a valid license key |
 | Kafka trigger code | `cargo check --features kafka` | |
+| DuckDB executor code | `cargo check -p windmill-worker --features duckdb` | `duckdb_executor.rs` (and its `#[cfg(test)]` tests) only compile with this flag — a plain check/test silently skips them |
 | Native trigger code | `cargo check --features native_trigger` | |
 | Parquet code | `cargo check --features parquet` | |
 | Multiple gated modules | `cargo check --features enterprise,parquet` | Combine only the flags you need |