diff --git a/core/jit_compiler.ts b/core/jit_compiler.ts index c79935283..780b076e3 100644 --- a/core/jit_compiler.ts +++ b/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 { + const mainBody = makeMainBody(request.jitCode); + + const jctx: JitContext = 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}`); } diff --git a/core/jit_compiler_test.ts b/core/jit_compiler_test.ts index 9d252e051..a45958f53 100644 --- a/core/jit_compiler_test.ts +++ b/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({ diff --git a/protos/jit.proto b/protos/jit.proto index e86e6ba2d..ca0d0707e 100644 --- a/protos/jit.proto +++ b/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; +}