JiT context - Add Running execution metadata (#2132)

Extend metadata of JiT context with information
about the current execution run.

Author: ikholopov-omni

core/BUILD

@@ -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",
@@ -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",
@@ -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",
     ],
 )
 

core/contextables.ts

@@ -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,
 };
 
 /** 

core/jit_compiler_test.ts

@@ -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";
@@ -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");
+    });
+  });
 });

core/jit_context.ts

@@ -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>;
@@ -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 {

protos/jit.proto

@@ -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";
 
@@ -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.
@@ -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.