From 6648e93436b3d2f59eb927b37f3a8a2add80d7c9 Mon Sep 17 00:00:00 2001 From: Rahul Bhonsale Date: Wed, 3 Jun 2026 00:07:32 +0200 Subject: [PATCH] Include notebooks and data preparations in tag/action selection prune() only considered tables, operations and assertions when computing the set of actions selected by --tags / --actions. Notebooks and data preparations were excluded from the selection union, so a tag selector could never match them, and they leaked through the returned graph unfiltered (passed via the spread but never filtered). This meant a notebook registered with a tag via the JS API, e.g. notebook({ name: "run_notebook", filename: "...", tags: ["reconciliation"] }) would not appear when executing by that tag in the CLI or UI. Add notebooks and data preparations to both the selection union (so tag and action selectors match them and dependency/dependent traversal includes them) and the filtered output graph. Co-Authored-By: Claude Opus 4.8 (1M context) --- cli/api/commands/prune.ts | 19 ++++++-- tests/api/api.spec.ts | 92 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 3 deletions(-) diff --git a/cli/api/commands/prune.ts b/cli/api/commands/prune.ts index 71358ea0b..1336583a3 100644 --- a/cli/api/commands/prune.ts +++ b/cli/api/commands/prune.ts @@ -2,7 +2,12 @@ import { targetAsReadableString } from "df/core/targets"; import * as utils from "df/core/utils"; import { dataform } from "df/protos/ts"; -type CompileAction = dataform.ITable | dataform.IOperation | dataform.IAssertion; +type CompileAction = + | dataform.ITable + | dataform.IOperation + | dataform.IAssertion + | dataform.INotebook + | dataform.IDataPreparation; export function prune( compiledGraph: dataform.ICompiledGraph, @@ -20,6 +25,12 @@ export function prune( ), operations: compiledGraph.operations.filter(action => includedActionNames.has(targetAsReadableString(action.target)) + ), + notebooks: compiledGraph.notebooks.filter(action => + includedActionNames.has(targetAsReadableString(action.target)) + ), + dataPreparations: compiledGraph.dataPreparations.filter(action => + includedActionNames.has(targetAsReadableString(action.target)) ) }; } @@ -28,11 +39,13 @@ function computeIncludedActionNames( compiledGraph: dataform.ICompiledGraph, runConfig: dataform.IRunConfig ): Set { - // Union all tables, operations, assertions. + // Union all tables, operations, assertions, notebooks and data preparations. const allActions: CompileAction[] = [].concat( compiledGraph.tables, compiledGraph.operations, - compiledGraph.assertions + compiledGraph.assertions, + compiledGraph.notebooks, + compiledGraph.dataPreparations ); const allActionNames = new Set( diff --git a/tests/api/api.spec.ts b/tests/api/api.spec.ts index b29936140..544d1f445 100644 --- a/tests/api/api.spec.ts +++ b/tests/api/api.spec.ts @@ -377,6 +377,98 @@ suite("@dataform/api", () => { expect(actionNames).not.includes("schema.b"); expect(actionNames).not.includes("schema.d"); }); + + // Mirrors a project that registers a notebook and a data preparation with a tag via the + // JS API, e.g. notebook({ name: "nb_a", filename: "...", tags: ["tag1"] }), alongside + // operations/tables that share the same tag. "tag1" actions depend on an untagged + // operation (op_dep) to exercise dependency traversal; "tag2" actions are unrelated. + const TEST_GRAPH_WITH_NOTEBOOK_TAGS: dataform.ICompiledGraph = dataform.CompiledGraph.create({ + projectConfig: { warehouse: "bigquery", defaultLocation: "US" }, + operations: [ + { + target: { schema: "schema", name: "op_a" }, + tags: ["tag1"], + queries: ["create or replace view schema.someview as select 1 as test"] + }, + { + target: { schema: "schema", name: "op_b" }, + tags: ["tag2"], + queries: ["create or replace view schema.someview as select 1 as test"] + }, + { + target: { schema: "schema", name: "op_dep" }, + tags: [], + queries: ["create or replace view schema.someview as select 1 as test"] + } + ], + notebooks: [ + { + target: { schema: "schema", name: "nb_a" }, + tags: ["tag1"], + dependencyTargets: [{ schema: "schema", name: "op_dep" }] + }, + { + target: { schema: "schema", name: "nb_b" }, + tags: ["tag2"] + } + ], + dataPreparations: [ + { + target: { schema: "schema", name: "dp_a" }, + tags: ["tag1"] + }, + { + target: { schema: "schema", name: "dp_b" }, + tags: ["tag2"] + } + ] + }); + + const prunedActionNames = (prunedGraph: dataform.ICompiledGraph) => [ + ...prunedGraph.tables.map(action => targetAsReadableString(action.target)), + ...prunedGraph.operations.map(action => targetAsReadableString(action.target)), + ...prunedGraph.assertions.map(action => targetAsReadableString(action.target)), + ...prunedGraph.notebooks.map(action => targetAsReadableString(action.target)), + ...prunedGraph.dataPreparations.map(action => targetAsReadableString(action.target)) + ]; + + test("prune notebooks and data preparations with --tags", () => { + const prunedGraph = prune(TEST_GRAPH_WITH_NOTEBOOK_TAGS, { + tags: ["tag1"], + includeDependencies: false, + includeDependents: false + }); + const actionNames = prunedActionNames(prunedGraph); + + // Operations are filtered by tag correctly: the "tag1" operation is kept and the + // "tag2" operation is dropped. + expect(actionNames).includes("schema.op_a"); + expect(actionNames).not.includes("schema.op_b"); + + // Notebooks behave identically: a notebook tagged "tag1" is selected, and a notebook + // tagged "tag2" is dropped. + expect(actionNames).includes("schema.nb_a"); + expect(actionNames).not.includes("schema.nb_b"); + + // Data preparations behave identically. + expect(actionNames).includes("schema.dp_a"); + expect(actionNames).not.includes("schema.dp_b"); + }); + + test("prune notebooks with --tags pulls in dependencies", () => { + const prunedGraph = prune(TEST_GRAPH_WITH_NOTEBOOK_TAGS, { + tags: ["tag1"], + includeDependencies: true, + includeDependents: false + }); + const actionNames = prunedActionNames(prunedGraph); + + // The tag-selected notebook is included, along with its (untagged) dependency. + expect(actionNames).includes("schema.nb_a"); + expect(actionNames).includes("schema.op_dep"); + // Unrelated "tag2" actions remain excluded. + expect(actionNames).not.includes("schema.nb_b"); + }); }); suite("sql_generating", () => {