Assertion JiT integration

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",

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,

cli/api/dbadapters/bigquery.ts

@@ -154,10 +154,6 @@ export class BigQueryDbAdapter implements IDbAdapter {
       .promise();
   }
 
-  public async withClientLock<T>(callback: (client: IDbClient) => Promise<T>) {
-    return await callback(this);
-  }
-
   public async evaluate(queryOrAction: QueryOrAction) {
     const validationQueries = collectEvaluationQueries(queryOrAction, true);
 

cli/api/dbadapters/index.ts

@@ -51,8 +51,6 @@ export interface IDbClient {
 }
 
 export interface IDbAdapter extends IDbClient {
-  withClientLock<T>(callback: (client: IDbClient) => Promise<T>): Promise<T>;
-
   evaluate(queryOrAction: QueryOrAction): Promise<dataform.IQueryEvaluation[]>;
 
   schemas(database: string): Promise<string[]>;

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<void> {
+  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"; })`
+  );
+}

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");
+  });
+});