Add JiT compilation support for assertion actions

core/jit_compiler.ts

@@ -1,5 +1,6 @@
 import * as $protobuf from "protobufjs";
 
+import { JitAssertionResult } from "df/core/actions/assertion";
 import { JitOperationResult } from "df/core/actions/operation";
 import { JitTableResult } from "df/core/actions/table";
 import { IActionContext, ITableContext, JitContext } from "df/core/contextables";
@@ -63,6 +64,18 @@ function jitCompileTable(
   return mainBody(jctx).then(makeJitTableResult);
 }
 
+function jitCompileAssertion(
+  request: dataform.IJitCompilationRequest,
+  adapter: dataform.DbAdapter,
+): Promise<dataform.IJitAssertionResult> {
+  const mainBody = makeMainBody<IActionContext, JitAssertionResult>(request.jitCode);
+
+  const jctx: JitContext<IActionContext> = new SqlActionJitContext(
+    adapter, request,
+  );
+  return mainBody(jctx).then(query => dataform.JitAssertionResult.create({ query }));
+}
+
 function jitCompileIncrementalTable(
   request: dataform.IJitCompilationRequest,
   adapter: dataform.DbAdapter,
@@ -110,6 +123,9 @@ export function jitCompile(request: dataform.IJitCompilationRequest, rpcCallback
     case dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_INCREMENTAL_TABLE:
       return jitCompileIncrementalTable(request, dbAdapter).then(
         incrementalTable => dataform.JitCompilationResponse.create({ incrementalTable }));
+    case dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_ASSERTION:
+      return jitCompileAssertion(request, dbAdapter).then(
+        assertion => dataform.JitCompilationResponse.create({ assertion }));
     default:
       throw new Error(`Unrecognized compilation target type: ${request.compilationTargetType}`);
   }

core/jit_compiler_test.ts

@@ -120,6 +120,35 @@ suite("jit_compiler", () => {
     });
   });
 
+  suite("jitCompileAssertion", () => {
+    test("compiles assertion returning string", async () => {
+      const request = dataform.JitCompilationRequest.create({
+        jitCode: `async (ctx) => "SELECT * FROM t WHERE invalid"`,
+        target,
+        jitData: {},
+        compilationTargetType: dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_ASSERTION,
+      });
+      const result = await jitCompile(request, rpcCallback);
+      expect(result.assertion.query).to.equal("SELECT * FROM t WHERE invalid");
+    });
+
+    test("compiles assertion using context", async () => {
+      const request = dataform.JitCompilationRequest.create({
+        jitCode: `async (ctx) => \`SELECT * FROM \${ctx.ref('other')} WHERE invalid\``,
+        target,
+        jitData: {},
+        dependencies: [dataform.Target.create({
+          database: "db",
+          schema: "schema",
+          name: "other",
+        })],
+        compilationTargetType: dataform.JitCompilationTargetType.JIT_COMPILATION_TARGET_TYPE_ASSERTION,
+      });
+      const result = await jitCompile(request, rpcCallback);
+      expect(result.assertion.query).to.equal("SELECT * FROM `db.schema.other` WHERE invalid");
+    });
+  });
+
   suite("jitCompileIncrementalTable", () => {
     test("compiles incremental table", async () => {
       const request = dataform.JitCompilationRequest.create({

protos/jit.proto

@@ -108,6 +108,8 @@ enum JitCompilationTargetType {
   JIT_COMPILATION_TARGET_TYPE_OPERATION = 2;
   // Incremental table target.
   JIT_COMPILATION_TARGET_TYPE_INCREMENTAL_TABLE = 3;
+  // Assertion target.
+  JIT_COMPILATION_TARGET_TYPE_ASSERTION = 4;
 }
 
 // JiT compilation request.
@@ -134,6 +136,7 @@ message JitCompilationResponse {
     JitTableResult table = 1;
     JitOperationResult operation = 2;
     JitIncrementalTableResult incremental_table = 3;
+    JitAssertionResult assertion = 4;
   }
 }
 
@@ -161,3 +164,9 @@ message JitOperationResult {
   // Sequence of SQL operations.
   repeated string queries = 1;
 }
+
+// JiT compilation result for assertion actions.
+message JitAssertionResult {
+  // SQL Select query that returns rows iff the assertion fails.
+  string query = 1;
+}