diff --git a/core/actions/assertion.ts b/core/actions/assertion.ts index 954293a1e..c18b5143d 100644 --- a/core/actions/assertion.ts +++ b/core/actions/assertion.ts @@ -33,7 +33,7 @@ interface ILegacyAssertionConfig extends dataform.ActionConfig.AssertionConfig { export type AContextable = T | ((ctx: AssertionContext) => T); /** JiT compilation stage result for assertions. */ -export type JitAssertionResult = string; +export type JitAssertionResult = string | { query: string }; /** * An assertion is a data quality test query that finds rows that violate one or more conditions diff --git a/core/jit_compiler.ts b/core/jit_compiler.ts index 780b076e3..086906670 100644 --- a/core/jit_compiler.ts +++ b/core/jit_compiler.ts @@ -73,7 +73,9 @@ function jitCompileAssertion( const jctx: JitContext = new SqlActionJitContext( adapter, request, ); - return mainBody(jctx).then(query => dataform.JitAssertionResult.create({ query })); + return mainBody(jctx).then(result => + typeof result === "string" ? { query: result } : result + ); } function jitCompileIncrementalTable( diff --git a/core/jit_compiler_test.ts b/core/jit_compiler_test.ts index a45958f53..52d4c1407 100644 --- a/core/jit_compiler_test.ts +++ b/core/jit_compiler_test.ts @@ -147,6 +147,29 @@ suite("jit_compiler", () => { const result = await jitCompile(request, rpcCallback); expect(result.assertion.query).to.equal("SELECT * FROM `db.schema.other` WHERE invalid"); }); + + test("compiles assertion returning bare object", async () => { + const request = dataform.JitCompilationRequest.create({ + jitCode: `async (ctx) => ({ query: "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 returning object with extra fields", async () => { + const request = dataform.JitCompilationRequest.create({ + jitCode: + `async (ctx) => ({ query: "SELECT * FROM t WHERE invalid", preOps: [], postOps: [] })`, + 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"); + }); }); suite("jitCompileIncrementalTable", () => {