Integrate JiT compilation into CLI by rafal-hawrylak · Pull Request #2110 · dataform-co/dataform

rafal-hawrylak · 2026-03-09T20:08:25Z

No description provided.

kolina

Let's test compatibility with different Dataform core versions.

Also this change seems worth it to bump a minor version of Dataform core

ikholopov-omni

Blocked on review of https://github.com/dataform-co/dataform/pull/2109/changes

* feat: add worker process management for JiT compilation - Introduce base worker and JiT-specific child process logic - Implement RPC bridge for database access during JiT - Add VM scripts and loader for isolated execution - Add unit tests for the RPC mechanism

- Enhance error coercion in common/errors/errors.ts - Sync protos/execution.proto and protos/jit.proto with new fields

* feat: add worker process management for JiT compilation - Introduce base worker and JiT-specific child process logic - Implement RPC bridge for database access during JiT - Add VM scripts and loader for isolated execution - Add unit tests for the RPC mechanism

* feat: add action descriptor support to JiT compilation results - Implement action and column descriptor logic in core/jit_compiler.ts - Support overrides for descriptions, labels, and tags in JiT results - Enhance error coercion in common/errors/errors.ts - Add comprehensive JiT metadata tests in core/jit_compiler_test.ts - Sync protos/execution.proto and protos/jit.proto with new result fields

kolina · 2026-04-11T18:32:15Z

+    try {
+      if (action.jitCode) {
+        await this.compileJitAction(action, actionResult, this.dbadapter);
+        if ((actionResult.status as dataform.ActionResult.ExecutionStatus) === dataform.ActionResult.ExecutionStatus.FAILED) {


Do we need casts here? Here and below

Yes, to override TypeScript narrowing the type to exactly RUNNING (literal types). It believes the type is ExecutionStatus.FAILED and assumes it is not mutated.

It seems that compileJitAction only updates actionResult in the exception block, so you can consider not passing actionResult there and moving the exception block into this function

kolina · 2026-04-11T18:40:53Z

+    try {
+      const jitResponse = await this.executionOptions.jitCompiler(
+        jitRequest,
+        this.executionOptions.projectDir,


Reponding to the older comment

In the CLI, projectDir allows the JiT worker to resolve local require() calls and find @dataform/core. While this is convenient for local development, it does differ from the cloud environment where source code isn't available during execution. We may want to consider if JiT code should be strictly self-contained for better parity. This would make the project to be unportable from local to GCP and we don't want it, right?

I think it should be consistent. Allowing JiT compilation to access other files in a project seems ok to me, but if we don't support it in JiT, it'd make sense to not allow it in CLI as well

In both AoT compilation:

dataform/cli/vm/compile.ts

Line 47 in 38cd3b9

root: compileConfig.projectDir,

and JiT compilation: https://github.com/dataform-co/dataform/pull/2110/changes#diff-c88d00fcc6836e8a1bbecf72502bf3cd2f54cf0a585c58b1ba234bcf8b00bb4cR90 there importing files is allowed and in both there is restriction so that only files within projectDir can be imported.
What exactly you would align in AoT and JiT compilation?

My comment is about "GCP vs CLI": if my JiT task will be able to import other files locally, but then will break in GCP because JiT compilation in GCP doesn't support it, it'll be not a great experience.

kolina · 2026-04-11T21:58:58Z

+          return moduleName;
+        }
+      },
+      sourceExtensions: ["js", "json"]


add yaml here?

Hm, interesting - should yaml be really supported and is ever read in JiT compilation?

workflow_settings.yaml?

kolina · 2026-04-11T22:05:41Z

+          if (
+            argv[dryRunOptionName] &&
+            isJsonOutput &&
+            !executionGraph.actions.some(action => !!action.jitCode)


I think it makes sense to preserve printing graph here

If we preserve printing the executionGraph here, the JSON output for JiT actions will be incomplete (it will only contain the raw js jitCode, not the generated SQL). I added this condition so that a dry-run with JiT actions falls through to the Runner, forcing the JiT compilation to happen, and outputting a runResult that contains the actual compiled SQL. If we remove this, users lose the ability to inspect dynamically generated JiT queries via JSON dry-runs. Do you still want me to revert this and print the uncompiled graph?

kolina · 2026-04-11T22:12:24Z

@@ -1,2 +1,28 @@
-import { listenForCompileRequest } from "df/cli/vm/compile";


This function can be cleaned up?

I mean removing this function from listenForCompileRequest :_)

* feat: enable JiT compilation in CLI run, build and compile commands - Connect main CLI execution flow to the JiT worker process - Add comprehensive JiT test suite - Finalize AoT build logic to preserve JiT-specific fields - Update CLI console and entry points for integrated JiT support

kolina

No outstanding issues from my side except question about access of JiT compiled code to other files

kolina · 2026-04-22T01:18:29Z

+    ) {
+      return {
+        runResult: {
+          actions: [],


Don't need this if actions is present in executionOptions?

kolina · 2026-04-22T01:34:13Z

    private readonly dbadapter: dbadapters.IDbAdapter,
    private readonly graph: dataform.IExecutionGraph,
-    private readonly executionOptions: IExecutionOptions = {},
+    optionsOrProjectDir: RunOptionsOrProjectDir = ".",


verify what exactly are we passing here

kolina · 2026-04-23T16:47:49Z

+    try {
+      const jitResponse = await this.executionOptions.jitCompiler(
+        jitRequest,
+        this.executionOptions.projectDir,


My comment is about "GCP vs CLI": if my JiT task will be able to import other files locally, but then will break in GCP because JiT compilation in GCP doesn't support it, it'll be not a great experience.

kolina · 2026-04-23T16:56:32Z

+    try {
+      if (action.jitCode) {
+        await this.compileJitAction(action, actionResult, this.dbadapter);
+        if ((actionResult.status as dataform.ActionResult.ExecutionStatus) === dataform.ActionResult.ExecutionStatus.FAILED) {


It seems that compileJitAction only updates actionResult in the exception block, so you can consider not passing actionResult there and moving the exception block into this function

kolina · 2026-04-23T19:57:35Z

+          return moduleName;
+        }
+      },
+      sourceExtensions: ["js", "json"]


workflow_settings.yaml?

kolina · 2026-04-23T20:14:33Z

@@ -1,2 +1,28 @@
-import { listenForCompileRequest } from "df/cli/vm/compile";


I mean removing this function from listenForCompileRequest :_)

rafal-hawrylak self-assigned this Mar 9, 2026

rafal-hawrylak force-pushed the cli_jit_integration branch from fdccca4 to 371b080 Compare March 9, 2026 20:26

rafal-hawrylak requested review from andrzej-grudzien, apilaskowski, ikholopov-omni and kolina March 9, 2026 20:28

rafal-hawrylak marked this pull request as ready for review March 9, 2026 20:28

rafal-hawrylak requested a review from a team as a code owner March 9, 2026 20:28

rafal-hawrylak enabled auto-merge (squash) March 9, 2026 20:28

rafal-hawrylak force-pushed the cli_jit_integration branch from 371b080 to 18042dd Compare March 10, 2026 15:49

kolina requested changes Mar 11, 2026

View reviewed changes

ikholopov-omni requested changes Mar 11, 2026

View reviewed changes

rafal-hawrylak added 5 commits March 30, 2026 11:01

JiT: protos update to enable execution in CLI

d2f544f

- Enhance error coercion in common/errors/errors.ts - Sync protos/execution.proto and protos/jit.proto with new fields

feat: integrate JiT compilation and finalize test suite

e51e27c

rafal-hawrylak force-pushed the cli_jit_integration branch 13 times, most recently from 10bb4e6 to 824d7d9 Compare April 3, 2026 14:08

rafal-hawrylak force-pushed the cli_jit_integration branch 3 times, most recently from 2466f28 to 38cd3b9 Compare April 7, 2026 09:24

kolina reviewed Apr 11, 2026

View reviewed changes

rafal-hawrylak force-pushed the cli_jit_integration branch 4 times, most recently from 483f247 to 807dae4 Compare April 16, 2026 09:06

rafal-hawrylak force-pushed the cli_jit_integration branch from 807dae4 to 82adbdd Compare April 18, 2026 22:21

kolina reviewed Apr 23, 2026

View reviewed changes

		@@ -1,2 +1,28 @@
		import { listenForCompileRequest } from "df/cli/vm/compile";

Conversation

rafal-hawrylak commented Mar 9, 2026

Uh oh!

kolina left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ikholopov-omni left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kolina left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants