feat: implement .jitCode() support for Assertion actions (#2170)

- Add `jit_code` field to the `Assertion` message in `protos/core.proto`
- Fix missing protobuf imports and configuration in `protos/BUILD`
- Introduce `JitAssertionResult` type alias and `jitCode` method to `Assertion` builder (`core/actions/assertion.ts`)
- Add strictly typed `jitContextable` parameter to `Session.sqlxAction` in `core/session.ts`
- Update assertion `compile()` to handle `jitCode` execution bypassing standard string evaluation
- Parse and apply `jitContextable` safely for assertions in `core/session.ts`
- Add unit test coverage in `core/main_test.ts` to verify compilation output

Author: rafal-hawrylak

core/actions/assertion.ts

@@ -1,6 +1,6 @@
 import { verifyObjectMatchesProto, VerifyProtoErrorBehaviour } from "df/common/protos";
 import { ActionBuilder } from "df/core/actions";
-import { IActionContext, Resolvable } from "df/core/contextables";
+import { IActionContext, JitContextable, Resolvable } from "df/core/contextables";
 import * as Path from "df/core/path";
 import { Session } from "df/core/session";
 import {
@@ -32,6 +32,9 @@ interface ILegacyAssertionConfig extends dataform.ActionConfig.AssertionConfig {
 /** @hidden */
 export type AContextable<T> = T | ((ctx: AssertionContext) => T);
 
+/** JiT compilation stage result for assertions. */
+export type JitAssertionResult = string;
+
 /**
  * An assertion is a data quality test query that finds rows that violate one or more conditions
  * specified in the query. If the query returns any rows, the assertion fails.
@@ -90,6 +93,9 @@ export class Assertion extends ActionBuilder<dataform.Assertion> {
   /** @hidden We delay contextification until the final compile step, so hold these here for now. */
   private contextableQuery: AContextable<string>;
 
+  /** @hidden */
+  private contextableJitCode: JitContextable<AssertionContext, JitAssertionResult> | undefined;
+
   /** @hidden */
   constructor(session?: Session, unverifiedConfig?: any, configPath?: string) {
     super(session);
@@ -166,6 +172,15 @@ export class Assertion extends ActionBuilder<dataform.Assertion> {
     return this;
   }
 
+  public jitCode(jitCode: JitContextable<AssertionContext, JitAssertionResult>) {
+    if (!this.proto.actionDescriptor) {
+      this.proto.actionDescriptor = {};
+    }
+    this.proto.actionDescriptor.compilationMode = dataform.ActionCompilationMode.ACTION_COMPILATION_MODE_JIT;
+    this.contextableJitCode = jitCode;
+    return this;
+  }
+
   /**
    * @deprecated Deprecated in favor of
    * [AssertionConfig.dependencies](configs#dataform-ActionConfig-AssertionConfig).
@@ -293,8 +308,24 @@ export class Assertion extends ActionBuilder<dataform.Assertion> {
   public compile() {
     const context = new AssertionContext(this);
 
-    this.proto.query = context.apply(this.contextableQuery);
-    validateQueryString(this.session, this.proto.query, this.proto.fileName);
+    if (this.contextableJitCode && this.contextableQuery) {
+      this.session.compileError(
+        new Error("Assertion may set either .jitCode() or .query(), but not both."),
+        this.proto.fileName,
+        this.proto.target
+      );
+      return this.proto;
+    }
+
+    if (this.contextableJitCode) {
+      if (!this.proto.actionDescriptor) {
+        this.proto.actionDescriptor = {};
+      }
+      this.proto.jitCode = this.contextableJitCode.toString();
+    } else {
+      this.proto.query = context.apply(this.contextableQuery);
+      validateQueryString(this.session, this.proto.query, this.proto.fileName);
+    }
 
     return verifyObjectMatchesProto(
       dataform.Assertion,

core/main_test.ts

@@ -1537,6 +1537,65 @@ assert("name", {
         ]);
       });
 
+      test("jitCode correctly populates the jitCode field", () => {
+        const projectDir = tmpDirFixture.createNewTmpDir();
+        fs.writeFileSync(
+          path.join(projectDir, "workflow_settings.yaml"),
+          VALID_WORKFLOW_SETTINGS_YAML
+        );
+        fs.mkdirSync(path.join(projectDir, "definitions"));
+        fs.writeFileSync(
+          path.join(projectDir, "definitions/assert.js"),
+          `
+assert("name").jitCode(ctx => "jit");`
+        );
+
+        const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
+
+        expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]);
+        expect(asPlainObject(result.compile.compiledGraph.assertions)).deep.equals([
+          {
+            actionDescriptor: {
+              compilationMode: "ACTION_COMPILATION_MODE_JIT"
+            },
+            canonicalTarget: {
+              database: "defaultProject",
+              name: "name",
+              schema: "defaultDataset"
+            },
+            fileName: "definitions/assert.js",
+            jitCode: 'ctx => "jit"',
+            target: {
+              database: "defaultProject",
+              name: "name",
+              schema: "defaultDataset"
+            }
+          }
+        ]);
+      });
+
+      test("fails when both jitCode and query are set on an assertion", () => {
+        const projectDir = tmpDirFixture.createNewTmpDir();
+        fs.writeFileSync(
+          path.join(projectDir, "workflow_settings.yaml"),
+          VALID_WORKFLOW_SETTINGS_YAML
+        );
+        fs.mkdirSync(path.join(projectDir, "definitions"));
+        fs.writeFileSync(
+          path.join(projectDir, "definitions/assert.js"),
+          `
+assert("name").query("SELECT 1").jitCode(ctx => "jit");`
+        );
+
+        const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
+
+        expect(
+          result.compile.compiledGraph.graphErrors.compilationErrors?.map(error => error.message)
+        ).deep.equals([
+          "Assertion may set either .jitCode() or .query(), but not both."
+        ]);
+      });
+
       test("assert API returns disabled assertions when disableAssertions is true", () => {
         const projectDir = tmpDirFixture.createNewTmpDir();
         fs.writeFileSync(

protos/core.proto

@@ -260,6 +260,8 @@ message Assertion {
   // Only present for auto assertions.
   Target parent_action = 15;
 
+  string jit_code = 16;
+
   // Generated.
   string file_name = 7;