From 1bb7e3228edf4a12af960cba9395ec7bf3e65c8c Mon Sep 17 00:00:00 2001 From: Rafal Hawrylak Date: Thu, 21 May 2026 12:45:39 +0000 Subject: [PATCH] feat: implement .jitCode() support for Assertion actions - 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 --- core/actions/assertion.ts | 28 +++++++++++++++++++++++++--- core/main_test.ts | 37 +++++++++++++++++++++++++++++++++++++ core/session.ts | 25 +++++++++++++++++-------- protos/core.proto | 2 ++ 4 files changed, 81 insertions(+), 11 deletions(-) diff --git a/core/actions/assertion.ts b/core/actions/assertion.ts index e638e8b37..e1ab9be90 100644 --- a/core/actions/assertion.ts +++ b/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 | ((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 { /** @hidden We delay contextification until the final compile step, so hold these here for now. */ private contextableQuery: AContextable; + /** @hidden */ + private contextableJitCode: JitContextable | undefined; + /** @hidden */ constructor(session?: Session, unverifiedConfig?: any, configPath?: string) { super(session); @@ -166,6 +172,15 @@ export class Assertion extends ActionBuilder { return this; } + public jitCode(jitCode: JitContextable) { + 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 { 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, diff --git a/core/main_test.ts b/core/main_test.ts index 1b53e0e91..6f74f89f9 100644 --- a/core/main_test.ts +++ b/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( diff --git a/core/session.ts b/core/session.ts index b29bc971f..5f9416d17 100644 --- a/core/session.ts +++ b/core/session.ts @@ -2,7 +2,7 @@ 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, @@ -10,12 +10,12 @@ import { 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 + | JitContextable + | JitContextable + | JitContextable + // Fallback for context types without explicit JIT results (e.g. DataPreparation) + | JitContextable; }) { 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); + } + this.actions.push(assertion); break; case "dataPreparation": const dataPreparation = new DataPreparation(this, sqlxConfig).query( diff --git a/protos/core.proto b/protos/core.proto index 663332019..91657f3cd 100644 --- a/protos/core.proto +++ b/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;