Skip to content

Commit 7fc8c87

Browse files
ikholopov-omniikholopov-omni
authored
Jit data in Session (#2080)
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 bb35744 commit 7fc8c87

5 files changed

Lines changed: 172 additions & 8 deletions

File tree

‎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: 114 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,117 @@ 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+
"undef": undefined,
1623+
});`
1624+
);
1625+
const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
1626+
1627+
expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]);
1628+
expect(result.compile.compiledGraph.jitData).to.deep.equal(
1629+
google.protobuf.Struct.create({
1630+
fields: {
1631+
key: google.protobuf.Value.create({
1632+
structValue: google.protobuf.Struct.create({
1633+
fields: {
1634+
number: google.protobuf.Value.create({ numberValue: 123 }),
1635+
string: google.protobuf.Value.create({ stringValue: "value" }),
1636+
boolean: google.protobuf.Value.create({ boolValue: true }),
1637+
struct: google.protobuf.Value.create({
1638+
structValue: google.protobuf.Struct.create({
1639+
fields: {
1640+
nestedKey: google.protobuf.Value.create({ stringValue: "nestedValue" })
1641+
}
1642+
})
1643+
}),
1644+
list: google.protobuf.Value.create({
1645+
listValue: google.protobuf.ListValue.create({
1646+
values: [
1647+
google.protobuf.Value.create({ stringValue: "a" }),
1648+
google.protobuf.Value.create({ stringValue: "b" }),
1649+
google.protobuf.Value.create({ stringValue: "c" })
1650+
]
1651+
})
1652+
}),
1653+
null: google.protobuf.Value.create({
1654+
nullValue: google.protobuf.NullValue.NULL_VALUE
1655+
}),
1656+
undef: google.protobuf.Value.create({
1657+
nullValue: google.protobuf.NullValue.NULL_VALUE
1658+
}),
1659+
}
1660+
})
1661+
})
1662+
}
1663+
})
1664+
);
1665+
});
1666+
1667+
test("jitData with duplicate key throws error", () => {
1668+
const projectDir = tmpDirFixture.createNewTmpDir();
1669+
fs.writeFileSync(
1670+
path.join(projectDir, "workflow_settings.yaml"),
1671+
VALID_WORKFLOW_SETTINGS_YAML
1672+
);
1673+
fs.mkdirSync(path.join(projectDir, "definitions"));
1674+
fs.writeFileSync(
1675+
path.join(projectDir, "definitions/jit.js"),
1676+
`
1677+
dataform.jitData("key", 1);
1678+
dataform.jitData("key", 2);
1679+
`
1680+
);
1681+
const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
1682+
1683+
expect(
1684+
result.compile.compiledGraph.graphErrors.compilationErrors.map(e => e.message)
1685+
).to.deep.equal(["JiT context data with key key already exists."]);
1686+
});
1687+
1688+
test("jitData with unsupported type throws error", () => {
1689+
const projectDir = tmpDirFixture.createNewTmpDir();
1690+
fs.writeFileSync(
1691+
path.join(projectDir, "workflow_settings.yaml"),
1692+
VALID_WORKFLOW_SETTINGS_YAML
1693+
);
1694+
fs.mkdirSync(path.join(projectDir, "definitions"));
1695+
fs.writeFileSync(
1696+
path.join(projectDir, "definitions/jit.js"),
1697+
`
1698+
dataform.jitData("key", {test: () => {}});
1699+
`
1700+
);
1701+
const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
1702+
1703+
expect(
1704+
result.compile.compiledGraph.graphErrors.compilationErrors.map(e => e.message)
1705+
).to.deep.equal(["Unsupported context object: () => {}"]);
1706+
});
1707+
});
1708+
15961709
suite("invalid options", () => {
15971710
[
15981711
{

‎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 || typeof raw === "undefined") {
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+
([fieldKey, fieldValue]) => ([
441+
fieldKey,
442+
unknownToValue(fieldValue)
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)