Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions core/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,6 @@ ts_library(
ts_test_suite(
name = "tests",
srcs = [
"main_test.ts",
"utils_test.ts",
"compilers_test.ts",
"actions/assertion_test.ts",
"actions/data_preparation_test.ts",
"actions/declaration_test.ts",
Expand All @@ -79,8 +76,11 @@ ts_test_suite(
"actions/table_test.ts",
"actions/test_test.ts",
"actions/view_test.ts",
"compilers_test.ts",
"jit_compiler_test.ts",
"jit_context_test.ts",
"main_test.ts",
"utils_test.ts",
],
data = [
":node_modules",
Expand All @@ -94,10 +94,12 @@ ts_test_suite(
"@npm//@types/chai",
"@npm//@types/fs-extra",
"@npm//@types/js-yaml",
"@npm//@types/long",
"@npm//@types/node",
"@npm//chai",
"@npm//fs-extra",
"@npm//js-yaml",
"@npm//long",
],
)

Expand Down
2 changes: 2 additions & 0 deletions core/contextables.ts
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ export type JitContext<T> = T & {
data?: { [k: string]: any },
/** Original JiT compilation request. */
request: dataform.IJitCompilationRequest,
/** Current execution information for introspection. */
executionData: dataform.IRunningExecutionData,
};

/**
Expand Down
41 changes: 38 additions & 3 deletions core/jit_compiler_test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { expect } from "chai";
import Long from "long";

import { jitCompile } from "df/core/jit_compiler";
import { dataform } from "df/protos/ts";
Expand Down Expand Up @@ -121,20 +122,54 @@ suite("jit_compiler", () => {

suite("jitCompileIncrementalTable", () => {
test("compiles incremental table", async () => {
const request = dataform.JitCompilationRequest.create({
const request = dataform.JitCompilationRequest.create({
jitCode: `async (ctx) => {
if (ctx.incremental()) {
return { query: "SELECT INC" };
}
return { query: "SELECT REG" };
}`,
target,
jitData: {},
compilationTargetType: dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_INCREMENTAL_TABLE,
jitData: {},
compilationTargetType: dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_INCREMENTAL_TABLE,
});
const result = await jitCompile(request, rpcCallback);
expect(result.incrementalTable.incremental?.query).to.equal("SELECT INC");
expect(result.incrementalTable.regular?.query).to.equal("SELECT REG");
});
});

suite("jitCompileContext", () => {
test("can reference self and other tables", async () => {
const request = dataform.JitCompilationRequest.create({
jitCode: `async (jctx) => \`$\{jctx.self()\}\n$\{jctx.ref('other')\}\``,
target,
jitData: {},
dependencies: [dataform.Target.create({
database: "db",
schema: "schema",
name: "other",
})],
compilationTargetType: dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_TABLE,

});
const result = await jitCompile(request, rpcCallback);
expect(result.table.query).to.equal("`db.schema.name`\n`db.schema.other`");
});

test("can reference execution info data", async () => {
const request = dataform.JitCompilationRequest.create({
jitCode: `async (jctx) => \`$\{jctx.executionData.executionStartTime.seconds\}\n$\{jctx.executionData.executionId\}\``,
target,
jitData: {},
compilationTargetType: dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_TABLE,
executionData: {
executionId: "test-id",
executionStartTime: {seconds: Long.fromNumber(1774974514), nanos: 481},
}
});
const result = await jitCompile(request, rpcCallback);
expect(result.table.query).to.equal("1774974514\ntest-id");
});
});
});
2 changes: 2 additions & 0 deletions core/jit_context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ function canonicalTargetValue(target: dataform.ITarget): string {
/** Generate SQL action JiT context. */
export class SqlActionJitContext implements JitContext<IActionContext> {
public readonly data: { [k: string]: any } | undefined;
public readonly executionData: dataform.IRunningExecutionData;

private readonly target: dataform.ITarget;
private readonly resolvableMap: ResolvableMap<string>;
Expand All @@ -25,6 +26,7 @@ export class SqlActionJitContext implements JitContext<IActionContext> {
value: canonicalTargetValue(dep)
})));
this.data = jitDataToJsValue(request.jitData);
this.executionData = request.executionData;
}

public self(): string {
Expand Down
13 changes: 13 additions & 0 deletions protos/jit.proto
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ package dataform;

import "google/protobuf/empty.proto";
import "google/protobuf/struct.proto";
import "google/protobuf/timestamp.proto";
import "protos/core.proto";
import "protos/db_adapter.proto";

Expand Down Expand Up @@ -83,6 +84,16 @@ message DeleteTableRequest {
Target target = 1;
}

// Execution information for introspection.
message RunningExecutionData {
// Targets that are being run in the current execution.
repeated Target run_targets = 1;
// ID of the current execution.
string execution_id = 2;
// Start time of the current execution.
google.protobuf.Timestamp execution_start_time = 3;
}

// JiT compilation target type.
enum JitCompilationTargetType {
// Unspecified - will result in error.
Expand All @@ -109,6 +120,8 @@ message JitCompilationRequest {
JitCompilationTargetType compilation_target_type = 5;
// List of additional file paths accessible at compilation.
repeated string file_paths = 6;
// Current execution information for introspection.
RunningExecutionData execution_data = 7;
}

// JiT compilation response.
Expand Down
Loading