feat: implement .jitCode() support for Assertion actions

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,15 @@ 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) {
+      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,43 @@ 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("assert API returns disabled assertions when disableAssertions is true", () => {
         const projectDir = tmpDirFixture.createNewTmpDir();
         fs.writeFileSync(

core/session.ts

@@ -2,20 +2,20 @@ import { default as TarjanGraphConstructor, Graph as TarjanGraph } from "tarjan-
 
 import { encode64, unknownToValue, verifyObjectMatchesProto, VerifyProtoErrorBehaviour } from "df/common/protos";
 import { Action, ActionProto, ILegacyTableConfig, TableType } from "df/core/actions";
-import { AContextable, Assertion, AssertionContext } from "df/core/actions/assertion";
+import { AContextable, Assertion, AssertionContext, JitAssertionResult } from "df/core/actions/assertion";
 import {
   DataPreparation,
   DataPreparationContext,
 } from "df/core/actions/data_preparation";
 import { Declaration } from "df/core/actions/declaration";
 import { IncrementalTable } from "df/core/actions/incremental_table";
 import { Notebook } from "df/core/actions/notebook";
-import { Operation, OperationContext } from "df/core/actions/operation";
-import { Table, TableContext } from "df/core/actions/table";
+import { JitOperationResult, Operation, OperationContext } from "df/core/actions/operation";
+import { JitTableResult, Table, TableContext } from "df/core/actions/table";
 import { Test } from "df/core/actions/test";
-import { View } from "df/core/actions/view";
+import { JitViewResult, View } from "df/core/actions/view";
 import { CompilationSql } from "df/core/compilation_sql";
-import { Contextable, IActionContext, ITableContext, Resolvable } from "df/core/contextables";
+import { Contextable, IActionContext, ITableContext, JitContextable, Resolvable } from "df/core/contextables";
 import { targetAsReadableString, targetStringifier } from "df/core/targets";
 import * as utils from "df/core/utils";
 import { ResolvableMap, toResolvable } from "df/core/utils";
@@ -113,6 +113,13 @@ export class Session {
         contextable: (ctx: IActionContext) => string;
       }
     ];
+    jitContextable?:
+      | JitContextable<TableContext, JitTableResult>
+      | JitContextable<TableContext, JitViewResult>
+      | JitContextable<AssertionContext, JitAssertionResult>
+      | JitContextable<OperationContext, JitOperationResult>
+      // Fallback for context types without explicit JIT results (e.g. DataPreparation)
+      | JitContextable<any, any>;
   }) {
     const { sqlxConfig } = actionOptions;
     const actionType = sqlxConfig.hasOwnProperty("type") ? sqlxConfig.type : "operations";
@@ -190,9 +197,11 @@ export class Session {
         this.actions.push(table);
         break;
       case "assertion":
-        this.actions.push(
-          new Assertion(this, sqlxConfig).query(ctx => actionOptions.sqlContextable(ctx)[0])
-        );
+        const assertion = new Assertion(this, sqlxConfig).query(ctx => actionOptions.sqlContextable(ctx)[0]);
+        if (actionOptions.jitContextable) {
+          assertion.jitCode(actionOptions.jitContextable as JitContextable<AssertionContext, JitAssertionResult>);
+        }
+        this.actions.push(assertion);
         break;
       case "dataPreparation":
         const dataPreparation = new DataPreparation(this, sqlxConfig).query(

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;