Skip to content

Commit dbcce3b

Browse files
ikholopov-omniikholopov-omni
committed
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.
1 parent 7356e28 commit dbcce3b

File tree

5 files changed

+148
-8
lines changed

5 files changed

+148
-8
lines changed

‎cli/index_test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,7 @@ select 1 as \${dataform.projectConfig.vars.testVar2}
683683
schemaSuffix: "test_schema_suffix"
684684
},
685685
graphErrors: {},
686+
jitData: {},
686687
dataformCoreVersion: version,
687688
targets: [
688689
{
@@ -872,6 +873,7 @@ SELECT 1 as id
872873
],
873874
dataformCoreVersion: version,
874875
graphErrors: {},
876+
jitData: {},
875877
projectConfig: {
876878
assertionSchema: "dataform_assertions",
877879
defaultDatabase: DEFAULT_DATABASE,

‎core/main.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ function dataformCompile(compileRequest: dataform.ICompileExecutionRequest, sess
229229
globalAny.declare = session.declare.bind(session);
230230
globalAny.notebook = session.notebook.bind(session);
231231
globalAny.test = session.test.bind(session);
232+
globalAny.jitData = session.jitData.bind(session);
232233

233234
loadActionConfigs(session, compileRequest.compileConfig.filePaths);
234235

‎core/main_test.ts‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { dump as dumpYaml } from "js-yaml";
55
import * as path from "path";
66

77
import { version } from "df/core/version";
8-
import { dataform } from "df/protos/ts";
8+
import { dataform, google } from "df/protos/ts";
99
import { asPlainObject, suite, test } from "df/testing";
1010
import { TmpDirFixture } from "df/testing/fixtures";
1111
import {
@@ -708,6 +708,7 @@ select 1 AS \${dataform.projectConfig.vars.selectVar}`
708708
asPlainObject({
709709
dataformCoreVersion: version,
710710
graphErrors: {},
711+
jitData: {},
711712
projectConfig: {
712713
defaultDatabase: "defaultProject",
713714
defaultLocation: "locationInOverride",
@@ -884,6 +885,7 @@ select 1 AS \${dataform.projectConfig.vars.columnVar}`
884885
],
885886
dataformCoreVersion: version,
886887
graphErrors: {},
888+
jitData: {},
887889
projectConfig: {
888890
defaultLocation: "us",
889891
vars: {
@@ -1593,6 +1595,93 @@ assert("name", {
15931595
});
15941596
});
15951597

1598+
suite("jitData", () => {
1599+
test("jitData is added to the compiled graph", () => {
1600+
const projectDir = tmpDirFixture.createNewTmpDir();
1601+
fs.writeFileSync(
1602+
path.join(projectDir, "workflow_settings.yaml"),
1603+
VALID_WORKFLOW_SETTINGS_YAML
1604+
);
1605+
fs.mkdirSync(path.join(projectDir, "definitions"));
1606+
fs.writeFileSync(
1607+
path.join(projectDir, "definitions/jit.js"),
1608+
`
1609+
dataform.jitData("key", {
1610+
"number": 123,
1611+
"string": "value",
1612+
"boolean": true,
1613+
"struct": {
1614+
"nestedKey": "nestedValue"
1615+
},
1616+
"list": [
1617+
"a",
1618+
"b",
1619+
"c"
1620+
],
1621+
"null": null,
1622+
});`
1623+
);
1624+
const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
1625+
1626+
expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]);
1627+
expect(result.compile.compiledGraph.jitData).to.deep.equal(
1628+
google.protobuf.Struct.create({
1629+
fields: {
1630+
key: google.protobuf.Value.create({
1631+
structValue: google.protobuf.Struct.create({
1632+
fields: {
1633+
number: google.protobuf.Value.create({ numberValue: 123 }),
1634+
string: google.protobuf.Value.create({ stringValue: "value" }),
1635+
boolean: google.protobuf.Value.create({ boolValue: true }),
1636+
struct: google.protobuf.Value.create({
1637+
structValue: google.protobuf.Struct.create({
1638+
fields: {
1639+
nestedKey: google.protobuf.Value.create({ stringValue: "nestedValue" })
1640+
}
1641+
})
1642+
}),
1643+
list: google.protobuf.Value.create({
1644+
listValue: google.protobuf.ListValue.create({
1645+
values: [
1646+
google.protobuf.Value.create({ stringValue: "a" }),
1647+
google.protobuf.Value.create({ stringValue: "b" }),
1648+
google.protobuf.Value.create({ stringValue: "c" })
1649+
]
1650+
})
1651+
}),
1652+
null: google.protobuf.Value.create({
1653+
nullValue: google.protobuf.NullValue.NULL_VALUE
1654+
})
1655+
}
1656+
})
1657+
})
1658+
}
1659+
})
1660+
);
1661+
});
1662+
1663+
test("jitData with duplicate key throws error", () => {
1664+
const projectDir = tmpDirFixture.createNewTmpDir();
1665+
fs.writeFileSync(
1666+
path.join(projectDir, "workflow_settings.yaml"),
1667+
VALID_WORKFLOW_SETTINGS_YAML
1668+
);
1669+
fs.mkdirSync(path.join(projectDir, "definitions"));
1670+
fs.writeFileSync(
1671+
path.join(projectDir, "definitions/jit.js"),
1672+
`
1673+
dataform.jitData("key", 1);
1674+
dataform.jitData("key", 2);
1675+
`
1676+
);
1677+
const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
1678+
1679+
expect(
1680+
result.compile.compiledGraph.graphErrors.compilationErrors.map(e => e.message)
1681+
).to.deep.equal(["JiT context data with key key already exists."]);
1682+
});
1683+
});
1684+
15961685
suite("invalid options", () => {
15971686
[
15981687
{

‎core/session.ts‎

Lines changed: 54 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import { targetAsReadableString, targetStringifier } from "df/core/targets";
2020
import * as utils from "df/core/utils";
2121
import { toResolvable } from "df/core/utils";
2222
import { version as dataformCoreVersion } from "df/core/version";
23-
import { dataform } from "df/protos/ts";
23+
import { dataform, google } from "df/protos/ts";
2424

2525
const DEFAULT_CONFIG = {
2626
defaultSchema: "dataform",
@@ -58,6 +58,9 @@ export class Session {
5858

5959
public graphErrors: dataform.IGraphErrors;
6060

61+
// jit_context.data, avilable at jit stage.
62+
public jitContextData: google.protobuf.Struct | undefined;
63+
6164
constructor(
6265
rootDir?: string,
6366
projectConfig?: dataform.ProjectConfig,
@@ -79,6 +82,7 @@ export class Session {
7982
this.actions = [];
8083
this.tests = {};
8184
this.graphErrors = { compilationErrors: [] };
85+
this.jitContextData = new google.protobuf.Struct();
8286
}
8387

8488
public compilationSql(): CompilationSql {
@@ -408,6 +412,49 @@ export class Session {
408412
return notebook;
409413
}
410414

415+
public jitData(key: string, data: unknown): void {
416+
function unknownToValue(raw: unknown): google.protobuf.Value {
417+
if (raw === null) {
418+
return google.protobuf.Value.create({ nullValue: google.protobuf.NullValue.NULL_VALUE });
419+
}
420+
if (typeof raw === "string") {
421+
return google.protobuf.Value.create({ stringValue: raw as string });
422+
}
423+
if (typeof raw === "number") {
424+
return google.protobuf.Value.create({ numberValue: raw as number });
425+
}
426+
if (typeof raw === "boolean") {
427+
return google.protobuf.Value.create({ boolValue: raw as boolean });
428+
}
429+
if (typeof raw === "object" && raw instanceof Array) {
430+
return google.protobuf.Value.create({
431+
listValue: google.protobuf.ListValue.create({
432+
values: (raw as unknown[]).map(unknownToValue)
433+
})
434+
});
435+
}
436+
if (typeof raw === "object") {
437+
return google.protobuf.Value.create({
438+
structValue: google.protobuf.Struct.create({
439+
fields: Object.fromEntries(Object.entries(raw).map(
440+
([field_key, field_value]) => ([
441+
field_key,
442+
unknownToValue(field_value)
443+
])
444+
))
445+
})
446+
})
447+
}
448+
throw new Error(`Unsupported context object: ${raw}`);
449+
}
450+
451+
if (this.jitContextData.fields[key] !== undefined) {
452+
throw new Error(`JiT context data with key ${key} already exists.`);
453+
}
454+
455+
this.jitContextData.fields[key] = unknownToValue(data);
456+
457+
}
411458
public compileError(err: Error | string, path?: string, actionTarget?: dataform.ITarget) {
412459
const fileName = path || utils.getCallerFile(this.rootDir) || __filename;
413460

@@ -461,7 +508,8 @@ export class Session {
461508
),
462509
graphErrors: this.graphErrors,
463510
dataformCoreVersion,
464-
targets: this.actions.map(action => action.getTarget())
511+
targets: this.actions.map(action => action.getTarget()),
512+
jitData: this.jitContextData,
465513
});
466514

467515
this.fullyQualifyDependencies(
@@ -579,9 +627,9 @@ export class Session {
579627
.find(dependency)
580628
.forEach(
581629
assertion =>
582-
(fullyQualifiedDependencies[
583-
targetAsReadableString(assertion.getTarget())
584-
] = assertion.getTarget())
630+
(fullyQualifiedDependencies[
631+
targetAsReadableString(assertion.getTarget())
632+
] = assertion.getTarget())
585633
);
586634
}
587635
} else {
@@ -668,7 +716,7 @@ export class Session {
668716
actions.forEach(action => {
669717
// Declarations cannot have dependencies.
670718
const cleanedDependencies = (action instanceof dataform.Declaration ||
671-
!action.dependencyTargets
719+
!action.dependencyTargets
672720
? []
673721
: action.dependencyTargets
674722
).filter(

‎protos/core.proto‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -404,7 +404,7 @@ message CompiledGraph {
404404

405405
repeated Target targets = 11;
406406

407-
google.protobuf.Struct jit_context = 15;
407+
google.protobuf.Struct jit_data = 15;
408408

409409
reserved 5, 6;
410410
}

0 commit comments

Comments
 (0)