JiT Worker in CLI and RPC infrastructure by rafal-hawrylak · Pull Request #2109 · dataform-co/dataform

rafal-hawrylak · 2026-03-09T19:43:35Z

This change establishes the foundation for executing JiT compilation in an isolated environment. It introduces the worker process management, the RPC bridge for database access during JiT, and the necessary Bazel targets.

cli/api/commands/base_worker.ts

cli/api/commands/jit_compile_child_process.ts

cli/api/commands/jit/compiler.ts

cli/api/commands/jit_compile_child_process.ts

cli/api/commands/base_worker.ts

cli/api/utils/constants.ts

cli/vm/jit_worker.ts

cli/api/commands/jit/rpc_test.ts

kolina · 2026-03-10T10:46:21Z

cli/api/commands/jit/rpc.ts

@@ -0,0 +1,101 @@
+import { IDbAdapter, IDbClient } from "df/cli/api/dbadapters";


To clarify: is this protobuf bridge protocol the same for GCP and CLI implementations?

Both the CLI and GCP implementations follow the dataform.DbAdapter service defined in jit.proto.

They both use the same binary "payloads" for the actual data (e.g., ExecuteRequest, ExecuteResponse).

In CLI all 4 methods are implemented (Execute, ListTables, GetTable, and DeleteTable) and in GCP only Execute as of now (https://source.corp.google.com/piper///depot/google3/cloud/dataform/compilation/jitrunner/worker/engine.cc;rcl=881961593;l=163)

See my comment below on the issue with "Execute" compatibility

PTAL on the approach with executeRaw

This one is waiting for:

this to be released: feat(bigquery): allow the user to ask for skipping parsing rows when quering googleapis/google-cloud-node#7848

this to be reviewed, merged and released: fix(bigquery): allow timestamp format overrides in createQueryStream googleapis/google-cloud-node#7905

cli/api/commands/jit_compile_child_process.ts

protos/jit.proto

ikholopov-omni · 2026-03-11T19:17:36Z

cli/api/commands/jit/rpc.ts

+  return new Uint8Array();
+}
+
+function mapRowToProto(row: { [key: string]: any }): google.protobuf.IStruct {


jit.proto declares that struct contains:

// Rows. For BigQuery, see // https://docs.cloud.google.com/bigquery/docs/reference/rest/v2/jobs/getQueryResults.

In other words, right now we expect a raw API result of "f,v" JSON struct, not bespoke conversion.

The BigQuery client for Node is strange in respect for this - it forcefully decodes those and removes rows from the original request. We could either:
a) Use googleapis package client instead for JiT
b) Come up with another protocol for encoding rows and implement it both here and in GCP. Let me know on chat if you need code pointers for GCP part.

PTAL on the approach with executeRaw

It has exactly the same problem. query() inside rawExecute already removes rows from response and only returns decoded ones as first component in the tuple [rows, _, response]

cli/api/commands/jit/rpc.ts

cli/api/dbadapters/bigquery.ts

cli/vm/jit_worker.ts

kolina · 2026-03-27T23:11:58Z

cli/api/commands/jit/rpc.ts

+
+async function handleListTables(dbadapter: IDbAdapter, request: Uint8Array): Promise<Uint8Array> {
+  const listTablesRequest = dataform.ListTablesRequest.decode(request);
+  const targets = await dbadapter.tables(listTablesRequest.database, listTablesRequest.schema);


I think the ListTables API could potentially return all needed metadata without the need to request it separately for each table?

Yes, we can hide it in bigquery adapter. Eventually, we can't avoid api requests per table to bigquery and that's because to fetch metadata there is no suitable api in bigquery client. We may get most of the metadata directly from INFORMATION_SCHEMA.TABLE and INFORMATION_SCHEMA.COLUMNS but we'd miss lastModifiedTime and labels. Labels though can be obtained from INFORMATION_SCHEMA.TABLE_OPTIONS but this is not stightforward, requires parsing and what the most important this is low level information for which we'd couple bigquery implementation knowledge into our code.

cli/api/commands/jit/rpc_test.ts

cli/api/dbadapters/bigquery.ts

cli/vm/jit_worker.ts

rafal-hawrylak · 2026-03-29T08:45:20Z

cli/api/dbadapters/bigquery.ts

+    options?: { concurrencyLimit?: number; bigqueryClient?: BigQuery }
+  ) {
    this.bigQueryCredentials = credentials;
+    this.bigqueryClient = options?.bigqueryClient;


@kolina like this? Honestly I don't like having the field in class just for tests. Maybe extract logic of creating BigQuery clients outside BigQueryDbAdapter to BigQueryClientProvider - it would be easily mockable and we'd use the benefits of dependency injection?

* 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

rafal-hawrylak requested review from andrzej-grudzien and apilaskowski March 9, 2026 19:43

rafal-hawrylak force-pushed the cli_jit_workers branch from a08d4da to 01fa6ba Compare March 9, 2026 19:55

rafal-hawrylak requested review from ikholopov-omni and kolina March 9, 2026 19:59

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

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

rafal-hawrylak enabled auto-merge (squash) March 9, 2026 19:59

rafal-hawrylak self-assigned this Mar 9, 2026

kolina requested changes Mar 10, 2026

View reviewed changes

rafal-hawrylak force-pushed the cli_jit_workers branch from 01fa6ba to 33b9f19 Compare March 10, 2026 14:39

apilaskowski reviewed Mar 11, 2026

View reviewed changes

cli/api/commands/jit_compile_child_process.ts Outdated Show resolved Hide resolved

ikholopov-omni requested changes Mar 11, 2026

View reviewed changes

rafal-hawrylak force-pushed the cli_jit_workers branch 10 times, most recently from a05cf58 to be56edb Compare March 12, 2026 15:50

kolina requested changes Mar 14, 2026

View reviewed changes

rafal-hawrylak force-pushed the cli_jit_workers branch 5 times, most recently from 9ca1f45 to 0fe90aa Compare March 27, 2026 18:00

kolina reviewed Mar 28, 2026

View reviewed changes

rafal-hawrylak force-pushed the cli_jit_workers branch 5 times, most recently from c6650a7 to 22d3661 Compare March 29, 2026 08:40

rafal-hawrylak commented Mar 29, 2026

View reviewed changes

rafal-hawrylak force-pushed the cli_jit_workers branch 2 times, most recently from d512e48 to 38433f2 Compare March 30, 2026 06:56

rafal-hawrylak force-pushed the cli_jit_workers branch from 38433f2 to e687190 Compare March 30, 2026 11:01

rafal-hawrylak requested review from apilaskowski, ikholopov-omni and kolina March 30, 2026 11:01

		@@ -0,0 +1,101 @@
		import { IDbAdapter, IDbClient } from "df/cli/api/dbadapters";

Conversation

rafal-hawrylak commented Mar 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

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!

ikholopov-omni Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

rafal-hawrylak Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ikholopov-omni Mar 12, 2026 •

edited

Loading

rafal-hawrylak Mar 30, 2026 •

edited

Loading