From 18ad8d57fb00f622e6a4d8e48f80023bf3d29fde Mon Sep 17 00:00:00 2001 From: ikholopov Date: Thu, 27 Nov 2025 15:07:31 +0100 Subject: [PATCH] Jit data in Session Adding a JiT data to the session. Renaming the proto field as well, it was submitted before the update of the design, but since it is not used yet, it should be fine to rename it directly. --- cli/index_test.ts | 2 + core/main.ts | 1 + core/main_test.ts | 115 +++++++++++++++++++++++++++++++++++++++++++++- core/session.ts | 60 +++++++++++++++++++++--- protos/core.proto | 2 +- 5 files changed, 172 insertions(+), 8 deletions(-) diff --git a/cli/index_test.ts b/cli/index_test.ts index 769998d28..f7c91fb93 100644 --- a/cli/index_test.ts +++ b/cli/index_test.ts @@ -683,6 +683,7 @@ select 1 as \${dataform.projectConfig.vars.testVar2} schemaSuffix: "test_schema_suffix" }, graphErrors: {}, + jitData: {}, dataformCoreVersion: version, targets: [ { @@ -872,6 +873,7 @@ SELECT 1 as id ], dataformCoreVersion: version, graphErrors: {}, + jitData: {}, projectConfig: { assertionSchema: "dataform_assertions", defaultDatabase: DEFAULT_DATABASE, diff --git a/core/main.ts b/core/main.ts index 9dbe1f413..16505dbc1 100644 --- a/core/main.ts +++ b/core/main.ts @@ -229,6 +229,7 @@ function dataformCompile(compileRequest: dataform.ICompileExecutionRequest, sess globalAny.declare = session.declare.bind(session); globalAny.notebook = session.notebook.bind(session); globalAny.test = session.test.bind(session); + globalAny.jitData = session.jitData.bind(session); loadActionConfigs(session, compileRequest.compileConfig.filePaths); diff --git a/core/main_test.ts b/core/main_test.ts index 4f6bf5d6e..d5f49190e 100644 --- a/core/main_test.ts +++ b/core/main_test.ts @@ -5,7 +5,7 @@ import { dump as dumpYaml } from "js-yaml"; import * as path from "path"; import { version } from "df/core/version"; -import { dataform } from "df/protos/ts"; +import { dataform, google } from "df/protos/ts"; import { asPlainObject, suite, test } from "df/testing"; import { TmpDirFixture } from "df/testing/fixtures"; import { @@ -708,6 +708,7 @@ select 1 AS \${dataform.projectConfig.vars.selectVar}` asPlainObject({ dataformCoreVersion: version, graphErrors: {}, + jitData: {}, projectConfig: { defaultDatabase: "defaultProject", defaultLocation: "locationInOverride", @@ -884,6 +885,7 @@ select 1 AS \${dataform.projectConfig.vars.columnVar}` ], dataformCoreVersion: version, graphErrors: {}, + jitData: {}, projectConfig: { defaultLocation: "us", vars: { @@ -1593,6 +1595,117 @@ assert("name", { }); }); + suite("jitData", () => { + test("jitData is added to the compiled graph", () => { + 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/jit.js"), + ` +dataform.jitData("key", { + "number": 123, + "string": "value", + "boolean": true, + "struct": { + "nestedKey": "nestedValue" + }, + "list": [ + "a", + "b", + "c" + ], + "null": null, + "undef": undefined, +});` + ); + const result = runMainInVm(coreExecutionRequestFromPath(projectDir)); + + expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]); + expect(result.compile.compiledGraph.jitData).to.deep.equal( + google.protobuf.Struct.create({ + fields: { + key: google.protobuf.Value.create({ + structValue: google.protobuf.Struct.create({ + fields: { + number: google.protobuf.Value.create({ numberValue: 123 }), + string: google.protobuf.Value.create({ stringValue: "value" }), + boolean: google.protobuf.Value.create({ boolValue: true }), + struct: google.protobuf.Value.create({ + structValue: google.protobuf.Struct.create({ + fields: { + nestedKey: google.protobuf.Value.create({ stringValue: "nestedValue" }) + } + }) + }), + list: google.protobuf.Value.create({ + listValue: google.protobuf.ListValue.create({ + values: [ + google.protobuf.Value.create({ stringValue: "a" }), + google.protobuf.Value.create({ stringValue: "b" }), + google.protobuf.Value.create({ stringValue: "c" }) + ] + }) + }), + null: google.protobuf.Value.create({ + nullValue: google.protobuf.NullValue.NULL_VALUE + }), + undef: google.protobuf.Value.create({ + nullValue: google.protobuf.NullValue.NULL_VALUE + }), + } + }) + }) + } + }) + ); + }); + + test("jitData with duplicate key throws error", () => { + 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/jit.js"), + ` +dataform.jitData("key", 1); +dataform.jitData("key", 2); +` + ); + const result = runMainInVm(coreExecutionRequestFromPath(projectDir)); + + expect( + result.compile.compiledGraph.graphErrors.compilationErrors.map(e => e.message) + ).to.deep.equal(["JiT context data with key key already exists."]); + }); + + test("jitData with unsupported type throws error", () => { + 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/jit.js"), + ` +dataform.jitData("key", {test: () => {}}); +` + ); + const result = runMainInVm(coreExecutionRequestFromPath(projectDir)); + + expect( + result.compile.compiledGraph.graphErrors.compilationErrors.map(e => e.message) + ).to.deep.equal(["Unsupported context object: () => {}"]); + }); + }); + suite("invalid options", () => { [ { diff --git a/core/session.ts b/core/session.ts index 8052de79f..6f3fe0574 100644 --- a/core/session.ts +++ b/core/session.ts @@ -20,7 +20,7 @@ import { targetAsReadableString, targetStringifier } from "df/core/targets"; import * as utils from "df/core/utils"; import { toResolvable } from "df/core/utils"; import { version as dataformCoreVersion } from "df/core/version"; -import { dataform } from "df/protos/ts"; +import { dataform, google } from "df/protos/ts"; const DEFAULT_CONFIG = { defaultSchema: "dataform", @@ -58,6 +58,9 @@ export class Session { public graphErrors: dataform.IGraphErrors; + // jit_context.data, avilable at jit stage. + public jitContextData: google.protobuf.Struct | undefined; + constructor( rootDir?: string, projectConfig?: dataform.ProjectConfig, @@ -79,6 +82,7 @@ export class Session { this.actions = []; this.tests = {}; this.graphErrors = { compilationErrors: [] }; + this.jitContextData = new google.protobuf.Struct(); } public compilationSql(): CompilationSql { @@ -408,6 +412,49 @@ export class Session { return notebook; } + public jitData(key: string, data: unknown): void { + function unknownToValue(raw: unknown): google.protobuf.Value { + if (raw === null || typeof raw === "undefined") { + return google.protobuf.Value.create({ nullValue: google.protobuf.NullValue.NULL_VALUE }); + } + if (typeof raw === "string") { + return google.protobuf.Value.create({ stringValue: raw as string }); + } + if (typeof raw === "number") { + return google.protobuf.Value.create({ numberValue: raw as number }); + } + if (typeof raw === "boolean") { + return google.protobuf.Value.create({ boolValue: raw as boolean }); + } + if (typeof raw === "object" && raw instanceof Array) { + return google.protobuf.Value.create({ + listValue: google.protobuf.ListValue.create({ + values: (raw as unknown[]).map(unknownToValue) + }) + }); + } + if (typeof raw === "object") { + return google.protobuf.Value.create({ + structValue: google.protobuf.Struct.create({ + fields: Object.fromEntries(Object.entries(raw).map( + ([fieldKey, fieldValue]) => ([ + fieldKey, + unknownToValue(fieldValue) + ]) + )) + }) + }) + } + throw new Error(`Unsupported context object: ${raw}`); + } + + if (this.jitContextData.fields[key] !== undefined) { + throw new Error(`JiT context data with key ${key} already exists.`); + } + + this.jitContextData.fields[key] = unknownToValue(data); + + } public compileError(err: Error | string, path?: string, actionTarget?: dataform.ITarget) { const fileName = path || utils.getCallerFile(this.rootDir) || __filename; @@ -461,7 +508,8 @@ export class Session { ), graphErrors: this.graphErrors, dataformCoreVersion, - targets: this.actions.map(action => action.getTarget()) + targets: this.actions.map(action => action.getTarget()), + jitData: this.jitContextData, }); this.fullyQualifyDependencies( @@ -579,9 +627,9 @@ export class Session { .find(dependency) .forEach( assertion => - (fullyQualifiedDependencies[ - targetAsReadableString(assertion.getTarget()) - ] = assertion.getTarget()) + (fullyQualifiedDependencies[ + targetAsReadableString(assertion.getTarget()) + ] = assertion.getTarget()) ); } } else { @@ -668,7 +716,7 @@ export class Session { actions.forEach(action => { // Declarations cannot have dependencies. const cleanedDependencies = (action instanceof dataform.Declaration || - !action.dependencyTargets + !action.dependencyTargets ? [] : action.dependencyTargets ).filter( diff --git a/protos/core.proto b/protos/core.proto index 31235cebc..82de48a4e 100644 --- a/protos/core.proto +++ b/protos/core.proto @@ -404,7 +404,7 @@ message CompiledGraph { repeated Target targets = 11; - google.protobuf.Struct jit_context = 15; + google.protobuf.Struct jit_data = 15; reserved 5, 6; }