From 27641f4bd7c95fa997161cbed6ce2b655b1672d6 Mon Sep 17 00:00:00 2001 From: Rafal Hawrylak Date: Thu, 28 May 2026 05:40:53 +0000 Subject: [PATCH 1/3] JiT test suite --- cli/BUILD | 5 + cli/api/dbadapters/bigquery.ts | 4 - cli/api/dbadapters/index.ts | 2 - cli/index_test_base.ts | 54 ++- cli/tests/jit/index_jit_advanced_test.ts | 249 +++++++++++ cli/tests/jit/index_jit_dependency_test.ts | 93 +++++ cli/tests/jit/index_jit_main_test.ts | 274 +++++++++++++ cli/tests/jit/index_jit_runtime_test.ts | 183 +++++++++ cli/tests/jit/jit_run_test.ts | 455 +++++++++++++++++++++ tests/api/api.spec.ts | 19 +- tests/integration/BUILD | 4 + tests/integration/bigquery.spec.ts | 125 ++++++ 12 files changed, 1442 insertions(+), 25 deletions(-) create mode 100644 cli/tests/jit/index_jit_advanced_test.ts create mode 100644 cli/tests/jit/index_jit_dependency_test.ts create mode 100644 cli/tests/jit/index_jit_main_test.ts create mode 100644 cli/tests/jit/index_jit_runtime_test.ts create mode 100644 cli/tests/jit/jit_run_test.ts diff --git a/cli/BUILD b/cli/BUILD index 780da5613..053df52d4 100644 --- a/cli/BUILD +++ b/cli/BUILD @@ -65,8 +65,13 @@ ts_test_suite( "index_project_test.ts", "index_compile_test.ts", "index_run_e2e_test.ts", + "tests/jit/index_jit_main_test.ts", + "tests/jit/index_jit_advanced_test.ts", + "tests/jit/index_jit_dependency_test.ts", + "tests/jit/index_jit_runtime_test.ts", "util_test.ts", "tests/jit/jit_build_test.ts", + "tests/jit/jit_run_test.ts", ], data = [ ":node_modules", diff --git a/cli/api/dbadapters/bigquery.ts b/cli/api/dbadapters/bigquery.ts index 80dfe39d0..303e08024 100644 --- a/cli/api/dbadapters/bigquery.ts +++ b/cli/api/dbadapters/bigquery.ts @@ -154,10 +154,6 @@ export class BigQueryDbAdapter implements IDbAdapter { .promise(); } - public async withClientLock(callback: (client: IDbClient) => Promise) { - return await callback(this); - } - public async evaluate(queryOrAction: QueryOrAction) { const validationQueries = collectEvaluationQueries(queryOrAction, true); diff --git a/cli/api/dbadapters/index.ts b/cli/api/dbadapters/index.ts index 380f1dbf6..ff55d1fff 100644 --- a/cli/api/dbadapters/index.ts +++ b/cli/api/dbadapters/index.ts @@ -51,8 +51,6 @@ export interface IDbClient { } export interface IDbAdapter extends IDbClient { - withClientLock(callback: (client: IDbClient) => Promise): Promise; - evaluate(queryOrAction: QueryOrAction): Promise; schemas(database: string): Promise; diff --git a/cli/index_test_base.ts b/cli/index_test_base.ts index 564d7f435..5c150404b 100644 --- a/cli/index_test_base.ts +++ b/cli/index_test_base.ts @@ -1,7 +1,14 @@ // tslint:disable tsr-detect-non-literal-fs-filename -import * as fs from "fs"; +import { execFile } from "child_process"; +import * as fs from "fs-extra"; +import { dump as dumpYaml, load as loadYaml } from "js-yaml"; import * as path from "path"; +import { version } from "df/core/version"; +import { dataform } from "df/protos/ts"; +import { corePackageTarPath, getProcessResult, nodePath, npmPath } from "df/testing"; +import { TmpDirFixture } from "df/testing/fixtures"; + export const DEFAULT_DATABASE = "dataform-open-source"; export const DEFAULT_LOCATION = "US"; export const DEFAULT_RESERVATION = "projects/dataform-open-source/locations/us/reservations/dataform-test"; @@ -15,3 +22,48 @@ if (!fs.existsSync(path.resolve(runfilesDir, "df"))) { export const CREDENTIALS_PATH = path.resolve(runfilesDir, workspaceName, "test_credentials/bigquery.json"); export const cliEntryPointPath = "cli/node_modules/@dataform/cli/bundle.js"; + +export async function setupJitProject( + tmpDirFixture: TmpDirFixture, + projectDir: string +): Promise { + const npmCacheDir = tmpDirFixture.createNewTmpDir(); + const packageJsonPath = path.join(projectDir, "package.json"); + + await getProcessResult( + execFile(nodePath, [cliEntryPointPath, "init", projectDir, DEFAULT_DATABASE, DEFAULT_LOCATION]) + ); + + const workflowSettingsPath = path.join(projectDir, "workflow_settings.yaml"); + const workflowSettings = dataform.WorkflowSettings.create( + loadYaml(fs.readFileSync(workflowSettingsPath, "utf8")) + ); + delete workflowSettings.dataformCoreVersion; + fs.writeFileSync(workflowSettingsPath, dumpYaml(workflowSettings)); + + fs.writeFileSync( + packageJsonPath, + `{ + "dependencies":{ + "@dataform/core": "${version}" + } +}` + ); + await getProcessResult( + execFile(npmPath, [ + "install", + "--prefix", + projectDir, + "--cache", + npmCacheDir, + corePackageTarPath + ]) + ); + + const jitTablePath = path.join(projectDir, "definitions", "jit_table.js"); + fs.ensureFileSync(jitTablePath); + fs.writeFileSync( + jitTablePath, + `publish("jit_table", {type: "table"}).jitCode(async (ctx) => { return "SELECT 1 as id"; })` + ); +} diff --git a/cli/tests/jit/index_jit_advanced_test.ts b/cli/tests/jit/index_jit_advanced_test.ts new file mode 100644 index 000000000..01c274fa4 --- /dev/null +++ b/cli/tests/jit/index_jit_advanced_test.ts @@ -0,0 +1,249 @@ +import { expect } from "chai"; +import { execFile } from "child_process"; +import * as fs from "fs-extra"; +import * as path from "path"; + +import { + cliEntryPointPath, + CREDENTIALS_PATH, + setupJitProject +} from "df/cli/index_test_base"; +import { getProcessResult, nodePath, suite, test } from "df/testing"; +import { TmpDirFixture } from "df/testing/fixtures"; + +suite("JiT support advanced", ({ afterEach }) => { + const tmpDirFixture = new TmpDirFixture(afterEach); + + test("JiT preOps and postOps support", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const prePostPath = path.join(projectDir, "definitions", "pre_post_jit.js"); + fs.writeFileSync( + prePostPath, + `publish("pre_post_jit", { type: "table" }).jitCode(async (jctx) => { + return { + query: "SELECT 1 as id", + preOps: ["SELECT 'pre' as p"], + postOps: ["SELECT 'post' as p"] + }; + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=pre_post_jit" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + const prePostAction = executedGraph.actions.find((a: any) => a.target.name === "pre_post_jit"); + const statement = prePostAction.tasks[0].compiledSql; + expect(statement).to.include("SELECT 'pre' as p"); + expect(statement).to.include("SELECT 1 as id"); + expect(statement).to.include("SELECT 'post' as p"); + }); + + test({ name: "JiT incremental pre/post ops support", timeout: 60000 }, async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const incPrePostPath = path.join(projectDir, "definitions", "inc_pre_post_jit.js"); + fs.writeFileSync( + incPrePostPath, + `publish("inc_pre_post_jit", { type: "incremental" }).jitCode(async (jctx) => { + if (jctx.incremental()) { + return { + query: "SELECT 'inc_path_query' as q", + preOps: ["SELECT 'inc_path_pre' as p"] + }; + } else { + return { + query: "SELECT 'reg_path_query' as q", + preOps: ["SELECT 'reg_path_pre' as p"] + }; + } + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=inc_pre_post_jit", + "--full-refresh" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + const incAction = executedGraph.actions.find((a: any) => a.target.name === "inc_pre_post_jit"); + const statement = incAction.tasks[0].compiledSql; + expect(statement).to.include("SELECT 'reg_path_pre' as p"); + expect(statement).to.include("SELECT 'reg_path_query' as q"); + + // Also validate when not using full-refresh. + // Since the table doesn't exist, jctx.incremental() should still be false. + const runResultIncremental = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=inc_pre_post_jit" + ]) + ); + + expect(runResultIncremental.exitCode).equals(0); + const executedGraphInc = JSON.parse(runResultIncremental.stdout); + const incActionInc = executedGraphInc.actions.find((a: any) => a.target.name === "inc_pre_post_jit"); + const statementInc = incActionInc.tasks[0].compiledSql; + expect(statementInc).to.include("SELECT 'reg_path_pre' as p"); + expect(statementInc).to.include("SELECT 'reg_path_query' as q"); + }); + + test({ name: "JiT incremental mode validation with consecutive runs", timeout: 60000 }, async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const incPath = path.join(projectDir, "definitions", "inc_jit.js"); + fs.writeFileSync( + incPath, + `publish("inc_jit", { type: "incremental" }).jitCode(async (jctx) => { + if (jctx.incremental()) { + return { + query: "SELECT 'inc_query' as q", + preOps: ["SELECT 'inc_pre' as p"] + }; + } else { + return { + query: "SELECT 'reg_query' as q", + preOps: ["SELECT 'reg_pre' as p"] + }; + } + })` + ); + + // 1. Initial run with full-refresh to create the table. + const firstRun = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--actions=inc_jit", + "--full-refresh" + ]) + ); + expect(firstRun.exitCode).equals(0); + + // 2. Second run without full-refresh. + // The table now exists, so it should use the incremental path. + const secondRun = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=inc_jit" + ]) + ); + + expect(secondRun.exitCode).equals(0); + const secondGraph = JSON.parse(secondRun.stdout); + const secondAction = secondGraph.actions.find((a: any) => a.target.name === "inc_jit"); + // Assert second run is INCREMENTAL + expect(secondAction.tasks[0].compiledSql).to.include("SELECT 'inc_pre' as p"); + expect(secondAction.tasks[0].compiledSql).to.include("SELECT 'inc_query' as q"); + }); + + test("JiT project-level data support", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + fs.writeFileSync( + path.join(projectDir, "definitions", "project_data.js"), + "const { session } = require('@dataform/core');\nsession.jitData('app_secret', 'e2e_secret_value');" + ); + fs.writeFileSync( + path.join(projectDir, "definitions", "jit_data_test.js"), + `publish("jit_data_test", { type: "table" }).jitCode(async (jctx) => { + const secret = jctx.data.app_secret; + return "SELECT '" + secret + "' as val"; + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_data_test" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + const dataAction = executedGraph.actions.find((a: any) => a.target.name === "jit_data_test"); + expect(dataAction.tasks[0].compiledSql).to.include("e2e_secret_value"); + }); + + test("JiT complex session data support", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + fs.writeFileSync( + path.join(projectDir, "definitions", "complex_project_data.js"), + "const { session } = require('@dataform/core');\n" + + "session.jitData('app_config', {\n" + + " env: 'test-env',\n" + + " version: 1.2,\n" + + " tags: ['t1', 't2']\n" + + "});" + ); + fs.writeFileSync( + path.join(projectDir, "definitions", "jit_complex_data_test.js"), + "publish('jit_complex_data_test', { type: 'table' }).jitCode(async (jctx) => {\n" + + " const config = jctx.data.app_config;\n" + + " return 'SELECT \\'' + config.env + '\\' as env, ' + config.version + ' as ver, \\'' + config.tags[0] + '\\' as tag';\n" + + "})" + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_complex_data_test" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + const dataAction = executedGraph.actions.find((a: any) => a.target.name === "jit_complex_data_test"); + expect(dataAction.tasks[0].compiledSql).to.include("SELECT 'test-env' as env, 1.2 as ver, 't1' as tag"); + }); +}); diff --git a/cli/tests/jit/index_jit_dependency_test.ts b/cli/tests/jit/index_jit_dependency_test.ts new file mode 100644 index 000000000..b3f1393d3 --- /dev/null +++ b/cli/tests/jit/index_jit_dependency_test.ts @@ -0,0 +1,93 @@ +import { expect } from "chai"; +import { execFile } from "child_process"; +import * as fs from "fs-extra"; +import * as path from "path"; + +import { + cliEntryPointPath, + CREDENTIALS_PATH, + setupJitProject +} from "df/cli/index_test_base"; +import { getProcessResult, nodePath, suite, test } from "df/testing"; +import { TmpDirFixture } from "df/testing/fixtures"; + +suite("JiT support dependencies", ({ afterEach }) => { + const tmpDirFixture = new TmpDirFixture(afterEach); + + test("JiT transitive dependency pruning", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + // A (AoT) -> B (JiT) + fs.writeFileSync( + path.join(projectDir, "definitions", "table_a.sqlx"), + "config { type: 'table' } SELECT 1 as val" + ); + fs.writeFileSync( + path.join(projectDir, "definitions", "table_b.js"), + `publish("table_b", { type: "table", dependencies: ["table_a"] }).jitCode(async (jctx) => { + const upstream = jctx.ref("table_a"); + return "SELECT '" + upstream + "' as ref_name"; + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=table_b", + "--include-deps" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + // Should have BOTH tables because of --include-deps + expect(executedGraph.actions.length).to.equal(2); + expect(executedGraph.actions.some((a: any) => a.target.name === "table_a")).to.equal(true); + const actionB = executedGraph.actions.find((a: any) => a.target.name === "table_b"); + expect(actionB).to.not.equal(undefined); + expect(actionB.tasks[0].compiledSql).to.include("SELECT '`dataform-open-source.dataform.table_a`' as ref_name"); + }); + + test("JiT to JiT dependency chain", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + // Action A (JiT) -> Action B (JiT) + fs.writeFileSync( + path.join(projectDir, "definitions", "jit_a.js"), + 'publish("jit_a", { type: "table" }).jitCode(async () => "SELECT 1 as val")' + ); + fs.writeFileSync( + path.join(projectDir, "definitions", "jit_b.js"), + "publish('jit_b', { type: 'table', dependencies: ['jit_a'] }).jitCode(async (jctx) => {\n" + + " const upstream = jctx.ref('jit_a');\n" + + " return 'SELECT \\'' + upstream + '\\' as ref_name';\n" + + "})" + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_b", + "--include-deps" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + expect(executedGraph.actions.length).to.equal(2); + const actionB = executedGraph.actions.find((a: any) => a.target.name === "jit_b"); + expect(actionB.tasks[0].compiledSql).to.include("SELECT '`dataform-open-source.dataform.jit_a`' as ref_name"); + }); +}); diff --git a/cli/tests/jit/index_jit_main_test.ts b/cli/tests/jit/index_jit_main_test.ts new file mode 100644 index 000000000..015581b38 --- /dev/null +++ b/cli/tests/jit/index_jit_main_test.ts @@ -0,0 +1,274 @@ +import { expect } from "chai"; +import { execFile } from "child_process"; +import * as fs from "fs-extra"; +import * as path from "path"; + +import { + cliEntryPointPath, + CREDENTIALS_PATH, + DEFAULT_DATABASE, + setupJitProject +} from "df/cli/index_test_base"; +import { getProcessResult, nodePath, suite, test } from "df/testing"; +import { TmpDirFixture } from "df/testing/fixtures"; + +suite("JiT support main", ({ afterEach }) => { + const tmpDirFixture = new TmpDirFixture(afterEach); + + test("compile command includes jitCode in output", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const compileResult = await getProcessResult( + execFile(nodePath, [cliEntryPointPath, "compile", projectDir, "--json"]) + ); + + expect(compileResult.exitCode).equals(0); + const compiledGraph = JSON.parse(compileResult.stdout); + const jitTable = compiledGraph.tables.find((t: any) => t.target.name === "jit_table"); + expect(!!jitTable).to.equal(true); + expect(jitTable.type).to.equal("table"); + expect(jitTable.jitCode).to.contain("async (ctx) => { return \"SELECT 1 as id\"; }"); + expect(compiledGraph).to.have.property("jitData"); + }); + + test("fails if both query and jitCode are provided", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const conflictPath = path.join(projectDir, "definitions", "conflict.js"); + fs.writeFileSync( + conflictPath, + `publish("conflict", {type: "table"}).query("SELECT 1").jitCode(async (ctx) => "SELECT 2")` + ); + + const compileResult = await getProcessResult( + execFile(nodePath, [cliEntryPointPath, "compile", projectDir, "--json"]) + ); + + expect(compileResult.exitCode).equals(1); + expect(compileResult.stderr).to.include("Cannot mix AoT and JiT compilation in action"); + }); + + test("run command performs JiT compilation during execution", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_table" + ]) + ); + + expect(runResult.exitCode).equals(0); + + const executedGraph = JSON.parse(runResult.stdout); + const jitAction = executedGraph.actions.find((a: any) => a.target.name === "jit_table"); + expect(!!jitAction).to.equal(true); + // Tasks array should be populated by the JiT runner + expect(jitAction.tasks.length).to.be.greaterThan(0); + expect(jitAction.tasks[0].compiledSql).to.include("SELECT 1 as id"); + }); + + test("mixed AoT and JiT support", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const aotTablePath = path.join(projectDir, "definitions", "aot_table.sqlx"); + fs.writeFileSync(aotTablePath, "config { type: 'table' } SELECT 2 as id"); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json" + ]) + ); + + expect(runResult.exitCode).equals(0); + + const executedGraph = JSON.parse(runResult.stdout); + const aotAction = executedGraph.actions.find((a: any) => a.target.name === "aot_table"); + const jitAction = executedGraph.actions.find((a: any) => a.target.name === "jit_table"); + + expect(!!aotAction).to.equal(true); + expect(!!jitAction).to.equal(true); + expect(executedGraph.actions.length).to.equal(2); + + expect(aotAction.tasks[0].compiledSql).to.include("SELECT 2 as id"); + // JiT action should have its tasks populated dynamically + expect(jitAction.tasks.length).to.be.greaterThan(0); + expect(jitAction.tasks[0].compiledSql).to.include("SELECT 1 as id"); + }); + + test("JiT respects disabled flag", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + + const disabledPath = path.join(projectDir, "definitions", "disabled_jit.js"); + fs.writeFileSync( + disabledPath, + `publish("disabled_jit", { type: "table", disabled: true }).jitCode(async (jctx) => { + throw new Error("Should not be executed"); + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--actions=disabled_jit"], + { + env: { ...process.env, NO_COLOR: "1" } + } + ) + ); + + expect(runResult.exitCode).equals(0); + // When an action is disabled, it should print a "disabled" message. + expect(runResult.stdout).to.include("Dataset creation disabled: dataform.disabled_jit [table] [disabled]"); + }); + + test("JiT compilation failure reporting", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const failingJitPath = path.join(projectDir, "definitions", "failing_jit.js"); + fs.writeFileSync( + failingJitPath, + `publish("failing_jit", {type: "table"}).jitCode(async (ctx) => { throw new Error("JiT compilation failed!"); })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=failing_jit" + ]) + ); + + expect(runResult.exitCode).equals(1); + + const executedGraph = JSON.parse(runResult.stdout); + const failingAction = executedGraph.actions.find((a: any) => a.target.name === "failing_jit"); + + expect(!!failingAction).to.equal(true); + expect(failingAction.status).to.equal(3); // FAILED + expect(failingAction.tasks[0].status).to.equal(3); // FAILED + expect(failingAction.tasks[0].errorMessage).to.include("JiT compilation failed!"); + }); + + test("surfaces 'Table not found' RPC error during JiT compilation", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const rpcJitPath = path.join(projectDir, "definitions", "rpc_jit.js"); + fs.writeFileSync( + rpcJitPath, + `publish("rpc_jit", {type: "table"}).jitCode(async (jctx) => { + // This will fail because the table does not exist in the warehouse, + // and jctx.adapter.getTable throws an error in this case. + const table = await jctx.adapter.getTable({target: {database: "${DEFAULT_DATABASE}", schema: "sch", name: "tab"}}); + return "SELECT 1 as id"; + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=rpc_jit" + ]) + ); + + expect(runResult.exitCode).equals(1); + + const executedGraph = JSON.parse(runResult.stdout); + const rpcAction = executedGraph.actions.find((a: any) => a.target.name === "rpc_jit"); + + expect(!!rpcAction).to.equal(true); + expect(rpcAction.status).to.equal(3); + expect(rpcAction.tasks[0].status).to.equal(3); + expect(rpcAction.tasks[0].errorMessage).to.include("JiT compilation error"); + expect(rpcAction.tasks[0].errorMessage).to.include("Table not found"); + expect(rpcAction.tasks[0].errorMessage).to.include(DEFAULT_DATABASE); + expect(rpcAction.tasks[0].errorMessage).to.include('"schema":"sch"'); + expect(rpcAction.tasks[0].errorMessage).to.include('"name":"tab"'); + }); + + test("mixed support with AoT filtered out", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const aotTablePath = path.join(projectDir, "definitions", "aot_table.sqlx"); + fs.writeFileSync(aotTablePath, "config { type: 'table' } SELECT 2 as id"); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_table" + ]) + ); + + expect(runResult.exitCode).equals(0); + + const executedGraph = JSON.parse(runResult.stdout); + expect(executedGraph.actions.length).to.equal(1); + const jitAction = executedGraph.actions.find((a: any) => a.target.name === "jit_table"); + expect(!!jitAction).to.equal(true); + expect(jitAction.tasks.length).to.be.greaterThan(0); + expect(jitAction.tasks[0].compiledSql).to.include("SELECT 1 as id"); + }); + + test("mixed support with JiT filtered out", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + const aotTablePath = path.join(projectDir, "definitions", "aot_table.sqlx"); + fs.writeFileSync(aotTablePath, "config { type: 'table' } SELECT 2 as id"); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=aot_table" + ]) + ); + + expect(runResult.exitCode).equals(0); + + const executedGraph = JSON.parse(runResult.stdout); + expect(executedGraph.actions.length).to.equal(1); + const aotAction = executedGraph.actions.find((a: any) => a.target.name === "aot_table"); + expect(!!aotAction).to.equal(true); + expect(aotAction.tasks[0].statement).to.include("SELECT 2 as id"); + }); +}); diff --git a/cli/tests/jit/index_jit_runtime_test.ts b/cli/tests/jit/index_jit_runtime_test.ts new file mode 100644 index 000000000..f3fefdee7 --- /dev/null +++ b/cli/tests/jit/index_jit_runtime_test.ts @@ -0,0 +1,183 @@ +import { expect } from "chai"; +import { execFile } from "child_process"; +import * as fs from "fs-extra"; +import * as path from "path"; + +import { + cliEntryPointPath, + CREDENTIALS_PATH, + setupJitProject +} from "df/cli/index_test_base"; +import { getProcessResult, nodePath, suite, test } from "df/testing"; +import { TmpDirFixture } from "df/testing/fixtures"; + +suite("JiT support runtime", ({ afterEach }) => { + const tmpDirFixture = new TmpDirFixture(afterEach); + + test("JiT require() of local files is rejected (GCP parity)", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + // Add a helper JS file + fs.ensureDirSync(path.join(projectDir, "helpers")); + fs.writeFileSync( + path.join(projectDir, "helpers", "utils.js"), + "module.exports = { getValue: () => 'required_value' };" + ); + // Add a JiT table that requires it + fs.writeFileSync( + path.join(projectDir, "definitions", "jit_require_test.js"), + `publish("jit_require_test", { type: "table" }).jitCode(async (jctx) => { + const utils = require("../helpers/utils.js"); + return "SELECT '" + utils.getValue() + "' as val"; + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_require_test" + ]) + ); + + expect(runResult.exitCode).equals(1); + expect(runResult.stdout).to.match(/Cannot find module/i); + }); + + test("JiT worker timeout handling", { timeout: 60000 }, async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + + // Add a JiT table that hangs in an infinite loop + const hangPath = path.join(projectDir, "definitions", "hang_jit.js"); + fs.writeFileSync( + hangPath, + `publish("hang_jit", { type: "table" }).jitCode(async (jctx) => { + while(true) { /* loop */ } + return "SELECT 1"; + })` + ); + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=hang_jit", + "--jit-timeout=4s" + ], { timeout: 50000 }) + ); + + expect(runResult.exitCode).equals(1); + expect(runResult.stdout).to.include("Compilation timed out"); + }); + + test("Global --timeout cancels in-flight JiT worker", { timeout: 90000 }, async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + + fs.writeFileSync( + path.join(projectDir, "definitions", "hang_jit_global.js"), + `publish("hang_jit_global", { type: "table" }).jitCode(async (jctx) => { + while(true) { /* loop */ } + return "SELECT 1"; + })` + ); + // --timeout must exceed BQ schema-prep time; smaller values fire the + // timer before the JiT compile starts, leaving the action SKIPPED and + // defeating the assertions below. + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=hang_jit_global", + "--timeout=15s" + ], { timeout: 80000 }) + ); + + expect(runResult.exitCode).equals(1); + const executedGraph = JSON.parse(runResult.stdout); + expect(executedGraph.status).equals(5); // TIMED_OUT + const hangAction = executedGraph.actions.find( + (a: any) => a.target.name === "hang_jit_global" + ); + expect(hangAction.status).equals(3); // FAILED + expect(hangAction.tasks[0].errorMessage).to.include( + "Run cancelled while worker was in flight." + ); + }); + + test("JiT parallel execution robustness", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + // Add multiple JiT tables + for (let i = 0; i < 5; i++) { + fs.writeFileSync( + path.join(projectDir, "definitions", `jit_${i}.js`), + `publish("jit_${i}", { type: "table" }).jitCode(async (jctx) => "SELECT ${i} as val")` + ); + } + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + expect(executedGraph.actions.filter((a: any) => a.target.name.startsWith("jit_")).length).to.equal(6); // jit_table + 5 others + }); + + test("JiT handles hard worker crash", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + // Add a JiT table that crashes the process + const crashPath = path.join(projectDir, "definitions", "crash_jit.js"); + fs.writeFileSync( + crashPath, + `publish("crash_jit", { type: "table" }).jitCode(async (jctx) => { + setTimeout(() => { throw new Error("Hard crash"); }, 10); + return new Promise(() => {}); // Hang until crash + })` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=crash_jit" + ]) + ); + + expect(runResult.exitCode).equals(1); + const executedGraph = JSON.parse(runResult.stdout); + const crashAction = executedGraph.actions.find((a: any) => a.target.name === "crash_jit"); + expect(crashAction.status).to.equal(3); // FAILED + expect(crashAction.tasks[0].errorMessage).to.include("Worker exited with code 1"); + }); +}); diff --git a/cli/tests/jit/jit_run_test.ts b/cli/tests/jit/jit_run_test.ts new file mode 100644 index 000000000..64009121f --- /dev/null +++ b/cli/tests/jit/jit_run_test.ts @@ -0,0 +1,455 @@ +import { expect } from "chai"; +import { anything, capture, instance, mock, verify, when } from "ts-mockito"; + +import { handleDbRequest as handleRpc } from "df/cli/api/commands/jit/rpc"; +import { Runner } from "df/cli/api/commands/run"; +import { IDbAdapter, IDbClient } from "df/cli/api/dbadapters"; +import { jitCompile } from "df/core/jit_compiler"; +import { dataform } from "df/protos/ts"; +import { suite, test } from "df/testing"; + +suite("run", () => { + test("JiT compilation is performed for Table actions", async () => { + const { mockAdapter, adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "jit_table" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return 'SELECT 1'; }", + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + // RPC callback bridge for tests + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + // Verify overall run status + if (result.status !== dataform.RunResult.ExecutionStatus.SUCCESSFUL) { + process.stderr.write("Run failed with actions: " + JSON.stringify(result.actions, null, 2) + "\n"); + } + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + // Verify action result + const actionResult = result.actions[0]; + expect(actionResult.target.name).equals("jit_table"); + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.SUCCESSFUL); + + // Verify task results + expect(actionResult.tasks.length).equals(1); + expect(actionResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + + // Verify that the Runner executed the query statement returned by JiT compilation + verify(mockAdapter.execute(anything(), anything())).atLeast(1); + }); + + test("JiT compilation is performed for Operation actions", async () => { + const { mockAdapter, adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "jit_op" }, + type: "operation", + jitCode: "async (jctx) => { return ['SELECT 1', 'SELECT 2']; }", + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + // RPC callback bridge for tests + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + const actionResult = result.actions[0]; + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.SUCCESSFUL); + expect(actionResult.tasks.length).equals(2); + expect(actionResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + expect(actionResult.tasks[1].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + + verify(mockAdapter.execute("SELECT 1", anything())).once(); + verify(mockAdapter.execute("SELECT 2", anything())).once(); + }); + + test("Mixed run with JiT and AoT actions", async () => { + const { mockAdapter, adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "aot_table" }, + type: "table", + tableType: "table", + tasks: [dataform.ExecutionTask.create({ statement: "SELECT 'aot'", type: "statement" })] + }, + { + target: { database: "db", schema: "sch", name: "jit_table" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return 'SELECT \"jit\"'; }", + tasks: [], + dependencyTargets: [{ database: "db", schema: "sch", name: "aot_table" }] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + // RPC callback bridge for tests + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + expect(result.actions.length).equals(2); + + const aotResult = result.actions.find((a: dataform.IActionResult) => a.target.name === "aot_table"); + const jitResult = result.actions.find((a: dataform.IActionResult) => a.target.name === "jit_table"); + + expect(aotResult.status).equals(dataform.ActionResult.ExecutionStatus.SUCCESSFUL); + expect(aotResult.tasks.length).equals(1); + expect(aotResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + + expect(jitResult.status).equals(dataform.ActionResult.ExecutionStatus.SUCCESSFUL); + expect(jitResult.tasks.length).equals(1); + expect(jitResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + + // Verify that both actions resulted in database execution calls + verify(mockAdapter.execute(anything(), anything())).atLeast(2); + const [firstStatement] = capture(mockAdapter.execute).first(); + const [secondStatement] = capture(mockAdapter.execute).second(); + const allStatements = [firstStatement, secondStatement]; + expect(allStatements.some((s: string) => s.includes("SELECT 'aot'"))).to.equal(true); + expect(allStatements.some((s: string) => s.includes("SELECT \"jit\""))).to.equal(true); + }); + + test("Handles JiT compilation syntax error", async () => { + const { adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "bad_jit" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return syntax error; }", + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + // RPC callback bridge for tests + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.FAILED); + + const actionResult = result.actions[0]; + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.FAILED); + expect(actionResult.tasks.length).equals(1); + expect(actionResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.FAILED); + expect(actionResult.tasks[0].errorMessage).to.contain("JiT compilation error"); + }); + + test("Handles database error during JiT compilation (RPC failure)", async () => { + const { adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "jit_db_error" }, + type: "table", + tableType: "table", + // This code calls jctx.adapter.execute() which triggers our mockClient.execute + jitCode: "async (jctx) => { await jctx.adapter.execute({statement: 'SELECT fail'}); return 'SELECT 2'; }", + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + // RPC callback bridge for tests + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.FAILED); + + const actionResult = result.actions[0]; + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.FAILED); + expect(actionResult.tasks.length).equals(1); + expect(actionResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.FAILED); + expect(actionResult.tasks[0].errorMessage).to.contain("RPC DB Fail"); + }); + + test("Handles JiT incremental table compilation", async () => { + const target = { database: "db", schema: "sch", name: "incremental_jit" }; + const executionGraph = createGraph([ + { + target, + type: "table", + tableType: "incremental", + jitCode: `async (jctx) => { + return jctx.incremental() ? "SELECT 'inc' as t" : "SELECT 'full' as t"; + }`, + tasks: [] + } + ]); + + let runner: Runner; + + // 1. First run - empty warehouse, should use 'full' path + const { mockAdapter: mockAdapterFull, adapterInstance: adapterInstanceFull } = createMocks(); + runner = new Runner(adapterInstanceFull, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const fullResult = await runner.execute().result(); + expect(fullResult.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + verify(mockAdapterFull.execute(anything(), anything())).atLeast(1); + const [executedSqlFull] = capture(mockAdapterFull.execute).last(); + expect(executedSqlFull).to.contain("create or replace table `db.sch.incremental_jit` as"); + expect(executedSqlFull).to.contain("SELECT 'full' as t"); + + // 2. Mock that the table now exists in the warehouse + executionGraph.warehouseState.tables.push({ + target, + type: dataform.TableMetadata.Type.TABLE, + fields: [{ name: "t" }] + }); + + // 3. Second run - table exists, should use 'incremental' path + const { + mockAdapter: mockAdapterIncremental, + adapterInstance: adapterInstanceIncremental + } = createMocks(); + runner = new Runner(adapterInstanceIncremental, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const incrementalResult = await runner.execute().result(); + expect(incrementalResult.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + verify(mockAdapterIncremental.execute(anything(), anything())).atLeast(1); + const [executedSqlIncremental] = capture(mockAdapterIncremental.execute).last(); + expect(executedSqlIncremental).to.contain("SELECT 'inc' as t"); + }); + + test("Handles JiT incremental table compilation - incremental mode", async () => { + const { mockAdapter, adapterInstance } = createMocks(); + + const target = { database: "db", schema: "sch", name: "incremental_jit" }; + const executionGraph = createGraph([ + { + target, + type: "table", + tableType: "incremental", + jitCode: `async (jctx) => { + return jctx.incremental() ? "SELECT 'inc' as t" : "SELECT 'full' as t"; + }`, + tasks: [] + } + ]); + // Mock that the table already exists in the warehouse as a TABLE with a 't' field + executionGraph.warehouseState.tables.push({ + target, + type: dataform.TableMetadata.Type.TABLE, + fields: [{ name: "t" }] + }); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + // Verify it used the 'incremental' query path + verify(mockAdapter.execute(anything(), anything())).atLeast(1); + const [executedSql] = capture(mockAdapter.execute).last(); + // For BigQuery, it should be an 'insert into' because no uniqueKey was specified. + // We check for substrings without trailing spaces to avoid exact whitespace mismatches. + // tslint:disable: tsr-detect-sql-literal-injection + expect(executedSql).to.equal( + "insert into `db.sch.incremental_jit` \n" + + "(`t`) \n" + + "select `t` \n" + + "from (SELECT 'inc' as t) as insertions" + ); + // tslint:enable: tsr-detect-sql-literal-injection + }); + + test("JiT compilation with RPC calls (ListTables, GetTable, DeleteTable)", async () => { + const { mockAdapter, adapterInstance } = createMocks(); + + const target = { database: "db", schema: "sch", name: "existing_table" }; + when(mockAdapter.tables(anything(), anything())).thenResolve([{ target }]); + when(mockAdapter.table(anything())).thenResolve({ + target, + type: dataform.TableMetadata.Type.TABLE + } as any); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "jit_rpc_test" }, + type: "table", + tableType: "table", + jitCode: `async (jctx) => { + const list = await jctx.adapter.listTables({ database: "db", schema: "sch" }); + const table = await jctx.adapter.getTable({ target: list.tables[0].target }); + await jctx.adapter.deleteTable({ target: table.target }); + return "SELECT '" + table.target.name + "' as deleted_table"; + }`, + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + const actionResult = result.actions[0]; + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.SUCCESSFUL); + + verify(mockAdapter.deleteTable(anything())).once(); + const [deletedTarget] = capture(mockAdapter.deleteTable).last(); + expect(deletedTarget.name).equals("existing_table"); + + verify(mockAdapter.execute(anything(), anything())).once(); + const [executedSql] = capture(mockAdapter.execute).last(); + expect(executedSql).to.contain("SELECT 'existing_table' as deleted_table"); + }); + + test("Global timeout cancels in-flight JiT worker", async () => { + const { adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "hang_jit" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { /* hangs forever in mock */ }", + tasks: [] + } + ]); + executionGraph.runConfig.timeoutMillis = 200; + + let cancelCallbackInvoked = false; + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: (req, pdir, adapter, client, timeoutMs, opts, onCancel) => + new Promise((resolve, reject) => { + onCancel(() => { + cancelCallbackInvoked = true; + reject(new Error("Run cancelled while worker was in flight.")); + }); + }) + }); + + const result = await runner.execute().result(); + + expect(cancelCallbackInvoked).equals(true); + expect(result.status).equals(dataform.RunResult.ExecutionStatus.TIMED_OUT); + + const actionResult = result.actions[0]; + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.FAILED); + expect(actionResult.tasks.length).equals(1); + expect(actionResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.FAILED); + expect(actionResult.tasks[0].errorMessage).to.match( + /JiT compilation error.*Run cancelled while worker was in flight/ + ); + }); +}); + +function createMocks() { + const mockAdapter = mock(); + const mockClient = mock(); + + when(mockAdapter.schemas(anything())).thenResolve([]); + when(mockAdapter.execute(anything(), anything())).thenCall((statement: string) => { + if (statement.includes("fail") || statement.includes("nonexistent")) { + throw new Error("RPC DB Fail"); + } + return Promise.resolve({ + rows: [], + metadata: {} + }); + }); + when(mockClient.executeRaw(anything(), anything())).thenCall((statement: string) => { + if (statement.includes("fail") || statement.includes("nonexistent")) { + throw new Error("RPC DB Fail"); + } + return Promise.resolve({ + rows: [], + metadata: {} + }); + }); + when(mockClient.execute(anything(), anything())).thenCall((statement: string) => { + if (statement.includes("fail") || statement.includes("nonexistent")) { + throw new Error("RPC DB Fail"); + } + return Promise.resolve({ + rows: [], + metadata: {} + }); + }); + + const adapterInstance = instance(mockAdapter); + (adapterInstance as any).rpcImpl = (method: string, req: Uint8Array, callback: any) => { + handleRpc(instance(mockAdapter), instance(mockClient), method, req) + .then((res: Uint8Array) => callback(null, res)) + .catch((err: Error) => callback(err, null)); + }; + + return { mockAdapter, mockClient, adapterInstance }; +} + +function createGraph(actions: any[]): dataform.ExecutionGraph { + return dataform.ExecutionGraph.create({ + projectConfig: { warehouse: "bigquery" }, + runConfig: { fullRefresh: false, timeoutMillis: 30000 }, + warehouseState: { tables: [] }, + actions: actions.map(a => ({ + dependencyTargets: [], + ...a + })) + }); +} diff --git a/tests/api/api.spec.ts b/tests/api/api.spec.ts index 3dbdf98fe..696ab80b0 100644 --- a/tests/api/api.spec.ts +++ b/tests/api/api.spec.ts @@ -973,8 +973,6 @@ suite("@dataform/api", () => { ).thenReject(new Error("bad statement")); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const runner = new Runner(mockDbAdapterInstance, RUN_TEST_GRAPH); @@ -1018,8 +1016,6 @@ suite("@dataform/api", () => { ).thenReject(new Error("bad statement")); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); let runner = new Runner(mockDbAdapterInstance, RUN_TEST_GRAPH); runner.execute(); @@ -1041,7 +1037,7 @@ suite("@dataform/api", () => { }).toJSON() ); - runner = new Runner(mockDbAdapterInstance, RUN_TEST_GRAPH, undefined, result); + runner = Runner.resume(mockDbAdapterInstance, RUN_TEST_GRAPH, result); expect( dataform.RunResult.create(cleanTiming(await runner.execute().result())).toJSON() @@ -1079,8 +1075,6 @@ suite("@dataform/api", () => { .thenResolve({ rows: [], metadata: {} }); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const runner = new Runner(mockDbAdapterInstance, NEW_TEST_GRAPH, { bigquery: { actionRetryLimit: 1 } @@ -1119,8 +1113,6 @@ suite("@dataform/api", () => { .thenResolve({ rows: [], metadata: {} }); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const runner = new Runner(mockDbAdapterInstance, NEW_TEST_GRAPH, { bigquery: { actionRetryLimit: 2 } @@ -1183,8 +1175,6 @@ suite("@dataform/api", () => { .thenResolve({ rows: [], metadata: {} }); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const runner = new Runner(mockDbAdapterInstance, NEW_TEST_GRAPH_WITH_OPERATION, { bigquery: { actionRetryLimit: 3 } @@ -1235,7 +1225,6 @@ suite("@dataform/api", () => { reject(new Error("Run cancelled")); }); }), - withClientLock: callback => callback(mockDbAdapter), schemas: _ => Promise.resolve([]), createSchema: (_, __) => Promise.resolve(), table: _ => undefined @@ -1292,8 +1281,6 @@ suite("@dataform/api", () => { }); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const labels = { env: "testing", team: "dataform" }; const runner = new Runner(mockDbAdapterInstance, NEW_TEST_GRAPH, { @@ -1355,8 +1342,6 @@ suite("@dataform/api", () => { }); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const globalLabels = { env: "testing", team: "dataform" }; const runner = new Runner(mockDbAdapterInstance, NEW_TEST_GRAPH, { @@ -1434,8 +1419,6 @@ suite("@dataform/api", () => { ); const mockDbAdapterInstance = instance(mockedDbAdapter); - mockDbAdapterInstance.withClientLock = async callback => - await callback(mockDbAdapterInstance); const runner = new Runner(mockDbAdapterInstance, METADATA_TEST_GRAPH); diff --git a/tests/integration/BUILD b/tests/integration/BUILD index 5f1ee674c..fdccfff82 100644 --- a/tests/integration/BUILD +++ b/tests/integration/BUILD @@ -12,6 +12,8 @@ ts_test_suite( "//test_credentials:bigquery.json", "//tests/integration/bigquery_project:files", "//tests/integration/bigquery_project:node_modules", + "//packages/@dataform/core:bundle.js", + "//packages/@dataform/core:package.json", ], tags = ["integration"], deps = [ @@ -22,8 +24,10 @@ ts_test_suite( "//protos:ts", "//testing", "@npm//@types/chai", + "@npm//@types/fs-extra", "@npm//@types/long", "@npm//@types/node", "@npm//chai", + "@npm//fs-extra", ], ) diff --git a/tests/integration/bigquery.spec.ts b/tests/integration/bigquery.spec.ts index d6b2e95f3..0e4f278cb 100644 --- a/tests/integration/bigquery.spec.ts +++ b/tests/integration/bigquery.spec.ts @@ -1,5 +1,7 @@ import { expect } from "chai"; +import * as fs from "fs-extra"; import Long from "long"; +import * as path from "path"; import * as dfapi from "df/cli/api"; import * as dbadapters from "df/cli/api/dbadapters"; @@ -490,6 +492,129 @@ suite("@dataform/integration/bigquery", { parallel: true }, ({ before, after }) expect(partialSearch.length).equals(2); expect(columnSearch.length).greaterThan(0); }); + + test("JiT execution e2e", { timeout: 120000 }, async () => { + // Create a simple project with a JiT table + const projectDir = "tests/integration/jit_project"; + if (fs.existsSync(projectDir)) { + fs.removeSync(projectDir); + } + fs.mkdirpSync(path.join(projectDir, "definitions")); + fs.writeFileSync( + path.join(projectDir, "workflow_settings.yaml"), + ` +defaultProject: dataform-open-source +defaultLocation: US +defaultDataset: df_integration_test_jit +` + ); + fs.writeFileSync( + path.join(projectDir, "definitions/jit_table.js"), + `publish("jit_table", { type: "table" }).jitCode(async (jctx) => "SELECT 1 as id")` + ); + + // Mock @dataform/core to avoid npm install and 403 error + const nodeModulesDir = path.join(projectDir, "node_modules", "@dataform", "core"); + fs.mkdirpSync(nodeModulesDir); + const coreBundlePath = path.resolve("packages/@dataform/core/bundle.js"); + fs.copyFileSync(coreBundlePath, path.join(nodeModulesDir, "bundle.js")); + const corePackageJsonPath = path.resolve("packages/@dataform/core/package.json"); + fs.copyFileSync(corePackageJsonPath, path.join(nodeModulesDir, "package.json")); + // We also need a package.json in the project root to bypass the dataformCoreVersion check + fs.writeFileSync( + path.join(projectDir, "package.json"), + JSON.stringify({ dependencies: { "@dataform/core": "3.0.0-alpha.0" } }) + ); + + try { + const compiledGraph = await dfapi.compile({ projectDir }); + + // Drop dataset to start fresh + await dbadapter.execute( + "drop schema if exists `dataform-open-source.df_integration_test_jit` cascade" + ); + + const executionGraph = await dfapi.build(compiledGraph, {}, dbadapter); + const runResult = await dfapi.run(dbadapter, executionGraph, { projectDir }).result(); + + expect(dataform.RunResult.ExecutionStatus[runResult.status]).eql( + dataform.RunResult.ExecutionStatus[dataform.RunResult.ExecutionStatus.SUCCESSFUL] + ); + + const rows = await dbadapter.execute("SELECT * FROM `dataform-open-source.df_integration_test_jit.jit_table`").then(res => res.rows); + expect(rows).to.eql([{ id: 1 }]); + } finally { + if (fs.existsSync(projectDir)) { + fs.removeSync(projectDir); + } + } + }); + + test("JiT dry run integration", { timeout: 120000 }, async () => { + // Create a simple project with a JiT table + const projectDir = "tests/integration/jit_dry_run_project"; + if (fs.existsSync(projectDir)) { + fs.removeSync(projectDir); + } + fs.mkdirpSync(path.join(projectDir, "definitions")); + fs.writeFileSync( + path.join(projectDir, "workflow_settings.yaml"), + ` +defaultProject: dataform-open-source +defaultLocation: US +defaultDataset: df_integration_test_jit_dry_run +` + ); + fs.writeFileSync( + path.join(projectDir, "definitions/jit_table.js"), + `publish("jit_table_dry_run", { type: "table" }).jitCode(async (jctx) => "SELECT 1 as id")` + ); + + // Mock @dataform/core + const nodeModulesDir = path.join(projectDir, "node_modules", "@dataform", "core"); + fs.mkdirpSync(nodeModulesDir); + const coreBundlePath = path.resolve("packages/@dataform/core/bundle.js"); + fs.copyFileSync(coreBundlePath, path.join(nodeModulesDir, "bundle.js")); + const corePackageJsonPath = path.resolve("packages/@dataform/core/package.json"); + fs.copyFileSync(corePackageJsonPath, path.join(nodeModulesDir, "package.json")); + fs.writeFileSync( + path.join(projectDir, "package.json"), + JSON.stringify({ dependencies: { "@dataform/core": "3.0.0-alpha.0" } }) + ); + + try { + const compiledGraph = await dfapi.compile({ projectDir }); + + // Drop dataset to start fresh + await dbadapter.execute( + "drop schema if exists `dataform-open-source.df_integration_test_jit_dry_run` cascade" + ); + + const executionGraph = await dfapi.build(compiledGraph, {}, dbadapter); + + const runResult = await dfapi.run(dbadapter, executionGraph, { + projectDir, + bigquery: { dryRun: true } + }).result(); + + expect(dataform.RunResult.ExecutionStatus[runResult.status]).eql( + dataform.RunResult.ExecutionStatus[dataform.RunResult.ExecutionStatus.SUCCESSFUL] + ); + + // Verify that the table was NOT created + const tables = await dbadapter.schemas("dataform-open-source").then(schemas => { + if (!schemas.includes("df_integration_test_jit_dry_run")) { + return []; + } + return dbadapter.tables("dataform-open-source", "df_integration_test_jit_dry_run"); + }); + expect(tables.length).to.equal(0); + } finally { + if (fs.existsSync(projectDir)) { + fs.removeSync(projectDir); + } + } + }); }); async function cleanWarehouse( From 61ac2a35fd6e4ff723e80f80a9a31893a693b6ca Mon Sep 17 00:00:00 2001 From: Rafal Hawrylak Date: Thu, 16 Jul 2026 06:04:24 +0000 Subject: [PATCH 2/3] Add regression test for JiT CANCEL_EVENT listener leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runs 4 JiT actions and asserts the CANCEL_EVENT listener count returns to zero after the run completes. Fails today (leaks 4 listeners) — will pass once run.ts unregisters each per-compile cancel handler. --- cli/tests/jit/jit_run_test.ts | 57 +++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/cli/tests/jit/jit_run_test.ts b/cli/tests/jit/jit_run_test.ts index 64009121f..a5b453acb 100644 --- a/cli/tests/jit/jit_run_test.ts +++ b/cli/tests/jit/jit_run_test.ts @@ -1,4 +1,5 @@ import { expect } from "chai"; +import { EventEmitter } from "events"; import { anything, capture, instance, mock, verify, when } from "ts-mockito"; import { handleDbRequest as handleRpc } from "df/cli/api/commands/jit/rpc"; @@ -397,6 +398,62 @@ suite("run", () => { /JiT compilation error.*Run cancelled while worker was in flight/ ); }); + + test("JiT compilation cleans up CANCEL_EVENT listeners after each action", async () => { + const { adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "jit_a" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return 'SELECT 1'; }", + tasks: [] + }, + { + target: { database: "db", schema: "sch", name: "jit_b" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return 'SELECT 2'; }", + tasks: [] + }, + { + target: { database: "db", schema: "sch", name: "jit_c" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return 'SELECT 3'; }", + tasks: [] + }, + { + target: { database: "db", schema: "sch", name: "jit_d" }, + type: "table", + tableType: "table", + jitCode: "async (jctx) => { return 'SELECT 4'; }", + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter, client, timeoutMs, opts, onCancel) => { + // Mirror the real compile worker: register a cancel handler on every + // invocation so we can observe whether Runner unregisters it once the + // compile completes. + onCancel(() => undefined); + return await jitCompile(req, (method, internalReq, callback) => { + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + // Each JiT compile registers a cancel listener; those listeners must be + // removed once the compile finishes. Otherwise a run with many actions + // leaks listeners and eventually triggers MaxListenersExceededWarning. + const emitter = (runner as any).eEmitter as EventEmitter; + expect(emitter.listenerCount("jobCancel")).equals(0); + }); }); function createMocks() { From 160c6a76228b8eeec46b7c5b0c6d70be98056d0d Mon Sep 17 00:00:00 2001 From: Rafal Hawrylak Date: Wed, 15 Jul 2026 14:04:10 +0000 Subject: [PATCH 3/3] Wire assertion actions into the JiT compilation runtime Runner.compileJitAction was missing an assertion branch in the target-type selector, so assertions fell through to JIT_COMPILATION_TARGET_TYPE_UNSPECIFIED and the compiler rejected them. createTasksFromJitResponse also had no branch for jitResponse.assertion, so the compiled query never reached executionSql.createAssertionTasks. Wire both: select ASSERTION for assertion actions, and hydrate a dataform.Assertion from the JiT response before building tasks. Tests: - Runner-level unit test (jit_run_test.ts) drives a single JiT assertion action end-to-end and checks the resulting task set. - CLI e2e test (index_jit_runtime_test.ts) uses assert().jitCode() against dry-run BigQuery. --- cli/api/commands/run.ts | 8 ++++++ cli/tests/jit/index_jit_runtime_test.ts | 32 +++++++++++++++++++++ cli/tests/jit/jit_run_test.ts | 38 +++++++++++++++++++++++++ 3 files changed, 78 insertions(+) diff --git a/cli/api/commands/run.ts b/cli/api/commands/run.ts index 77eb49b6a..302ed013d 100644 --- a/cli/api/commands/run.ts +++ b/cli/api/commands/run.ts @@ -578,6 +578,8 @@ export class Runner { : dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_TABLE; } else if (action.type === "operation") { compilationTargetType = dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_OPERATION; + } else if (action.type === "assertion") { + compilationTargetType = dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_ASSERTION; } const jitRequest = dataform.JitCompilationRequest.create({ @@ -631,6 +633,12 @@ export class Runner { ...jitResponse.operation }); return this.executionSql.createOperationTasks(operation); + } else if (jitResponse.assertion) { + const assertion = dataform.Assertion.create({ + ...action, + ...jitResponse.assertion + }); + return this.executionSql.createAssertionTasks(assertion); } else if (jitResponse.incrementalTable) { const table = dataform.Table.create({ ...action, diff --git a/cli/tests/jit/index_jit_runtime_test.ts b/cli/tests/jit/index_jit_runtime_test.ts index f3fefdee7..ca587c41e 100644 --- a/cli/tests/jit/index_jit_runtime_test.ts +++ b/cli/tests/jit/index_jit_runtime_test.ts @@ -120,6 +120,38 @@ suite("JiT support runtime", ({ afterEach }) => { ); }); + test("JiT compilation of an Assertion action", async () => { + const projectDir = tmpDirFixture.createNewTmpDir(); + await setupJitProject(tmpDirFixture, projectDir); + + fs.writeFileSync( + path.join(projectDir, "definitions", "jit_assertion.js"), + `assert("jit_assertion").jitCode(async (jctx) => "SELECT 1 as row_count")` + ); + + const runResult = await getProcessResult( + execFile(nodePath, [ + cliEntryPointPath, + "run", + projectDir, + "--credentials", + CREDENTIALS_PATH, + "--dry-run", + "--json", + "--actions=jit_assertion" + ]) + ); + + expect(runResult.exitCode).equals(0); + const executedGraph = JSON.parse(runResult.stdout); + const assertionAction = executedGraph.actions.find( + (a: any) => a.target.name === "jit_assertion" + ); + expect(assertionAction.status).to.equal(2); // SUCCESSFUL + // createAssertionTasks emits 2 tasks: create-or-replace view + row-count check. + expect(assertionAction.tasks.length).equals(2); + }); + test("JiT parallel execution robustness", async () => { const projectDir = tmpDirFixture.createNewTmpDir(); await setupJitProject(tmpDirFixture, projectDir); diff --git a/cli/tests/jit/jit_run_test.ts b/cli/tests/jit/jit_run_test.ts index a5b453acb..a1f4d4ede 100644 --- a/cli/tests/jit/jit_run_test.ts +++ b/cli/tests/jit/jit_run_test.ts @@ -86,6 +86,44 @@ suite("run", () => { verify(mockAdapter.execute("SELECT 2", anything())).once(); }); + test("JiT compilation is performed for Assertion actions", async () => { + const { mockAdapter, adapterInstance } = createMocks(); + + const executionGraph = createGraph([ + { + target: { database: "db", schema: "sch", name: "jit_assertion" }, + type: "assertion", + jitCode: "async (jctx) => { return 'SELECT * FROM t WHERE bad'; }", + tasks: [] + } + ]); + + const runner = new Runner(adapterInstance, executionGraph, { + jitCompiler: async (req, pdir, adapter) => { + return await jitCompile(req, (method, internalReq, callback) => { + (adapter as any).rpcImpl(method, internalReq, callback); + }); + } + }); + const result = await runner.execute().result(); + + if (result.status !== dataform.RunResult.ExecutionStatus.SUCCESSFUL) { + process.stderr.write("Run failed with actions: " + JSON.stringify(result.actions, null, 2) + "\n"); + } + expect(result.status).equals(dataform.RunResult.ExecutionStatus.SUCCESSFUL); + + const actionResult = result.actions[0]; + expect(actionResult.target.name).equals("jit_assertion"); + expect(actionResult.status).equals(dataform.ActionResult.ExecutionStatus.SUCCESSFUL); + + // createAssertionTasks emits 2 tasks: create-or-replace view + row-count check. + expect(actionResult.tasks.length).equals(2); + expect(actionResult.tasks[0].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + expect(actionResult.tasks[1].status).equals(dataform.TaskResult.ExecutionStatus.SUCCESSFUL); + + verify(mockAdapter.execute(anything(), anything())).atLeast(1); + }); + test("Mixed run with JiT and AoT actions", async () => { const { mockAdapter, adapterInstance } = createMocks();