Fix JiT assertion compilation to accept object result shape

core/actions/assertion.ts

@@ -33,7 +33,7 @@ interface ILegacyAssertionConfig extends dataform.ActionConfig.AssertionConfig {
 export type AContextable<T> = 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

core/jit_compiler.ts

@@ -73,7 +73,9 @@ function jitCompileAssertion(
   const jctx: JitContext<IActionContext> = 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(

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", () => {