Proto field additions for JiT flow

cli/api/dbadapters/bigquery.ts

@@ -37,19 +37,55 @@ export interface IBigQueryExecutionOptions {
   reservation?: string;
 }
 
+export interface IBigQueryClientProvider {
+  get(projectId?: string): BigQuery;
+}
+
+export class BigQueryClientProvider implements IBigQueryClientProvider {
+  private readonly clients = new Map<string, BigQuery>();
+
+  constructor(private readonly credentials: dataform.IBigQuery) {}
+
+  public get(projectId?: string): BigQuery {
+    projectId = projectId || this.credentials.projectId;
+    if (!this.clients.has(projectId)) {
+      this.clients.set(
+        projectId,
+        new BigQuery({
+          projectId,
+          scopes: EXTRA_GOOGLE_SCOPES,
+          location: this.credentials.location,
+          credentials: this.credentials.credentials && JSON.parse(this.credentials.credentials)
+        })
+      );
+    }
+    return this.clients.get(projectId);
+  }
+}
+
+export class StaticBigQueryClientProvider implements IBigQueryClientProvider {
+  constructor(private readonly client: BigQuery) {}
+
+  public get(projectId?: string): BigQuery {
+    return this.client;
+  }
+}
+
 export class BigQueryDbAdapter implements IDbAdapter {
   private bigQueryCredentials: dataform.IBigQuery;
   private pool: PromisePoolExecutor;
-
-  private readonly clients = new Map<string, BigQuery>();
-  private readonly bigqueryClient?: BigQuery;
+  private clientProvider: IBigQueryClientProvider;
 
   constructor(
     credentials: dataform.IBigQuery,
-    options?: { concurrencyLimit?: number; bigqueryClient?: BigQuery }
+    options?: {
+      concurrencyLimit?: number;
+      clientProvider?: IBigQueryClientProvider;
+    }
   ) {
     this.bigQueryCredentials = credentials;
-    this.bigqueryClient = options?.bigqueryClient;
+    this.clientProvider = options?.clientProvider || new BigQueryClientProvider(credentials);
+
     // Bigquery allows 50 concurrent queries, and a rate limit of 100/user/second by default.
     // These limits should be safely low enough for most projects.
     this.pool = new PromisePoolExecutor({
@@ -174,7 +210,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
       datasetIds.map(async datasetId => {
         const [tables] = await this.getClient(database)
           .dataset(datasetId)
-          .getTables();
+          .getTables({ autoPaginate: true, maxResults: 1000 });
         await Promise.all(
           tables.map(async table => {
             const metadata = await this.table({
@@ -272,14 +308,14 @@ export class BigQueryDbAdapter implements IDbAdapter {
   }
 
   public async schemas(database: string): Promise<string[]> {
-    const data = await this.getClient(database).getDatasets();
+    const data = await this.getClient(database).getDatasets({ autoPaginate: true, maxResults: 1000 });
     return data[0].map(dataset => dataset.id);
   }
 
   public async createSchema(database: string, schema: string): Promise<void> {
     await this.execute(
       `create schema if not exists \`${database || this.bigQueryCredentials.projectId}.${schema}\``,
-      { bigquery: { location: this.bigQueryCredentials.location } }
+      { bigquery: { location: this.bigQueryCredentials.location || "US" } }
     );
   }
 
@@ -320,23 +356,7 @@ export class BigQueryDbAdapter implements IDbAdapter {
   }
 
   private getClient(projectId?: string) {
-    if (this.bigqueryClient) {
-      return this.bigqueryClient;
-    }
-    projectId = projectId || this.bigQueryCredentials.projectId;
-    if (!this.clients.has(projectId)) {
-      this.clients.set(
-        projectId,
-        new BigQuery({
-          projectId,
-          scopes: EXTRA_GOOGLE_SCOPES,
-          location: this.bigQueryCredentials.location,
-          credentials:
-            this.bigQueryCredentials.credentials && JSON.parse(this.bigQueryCredentials.credentials)
-        })
-      );
-    }
-    return this.clients.get(projectId);
+    return this.clientProvider.get(projectId);
   }
 
   private async runQuery(
@@ -544,11 +564,14 @@ function convertFieldType(type: string) {
     case "INT64":
       return dataform.Field.Primitive.INTEGER;
     case "NUMERIC":
+    case "BIGNUMERIC":
       return dataform.Field.Primitive.NUMERIC;
     case "BOOL":
     case "BOOLEAN":
       return dataform.Field.Primitive.BOOLEAN;
     case "STRING":
+    case "JSON":
+    case "INTERVAL":
       return dataform.Field.Primitive.STRING;
     case "DATE":
       return dataform.Field.Primitive.DATE;

cli/api/dbadapters/bigquery_test.ts

@@ -2,7 +2,7 @@ import { Dataset, Table } from "@google-cloud/bigquery";
 import { expect } from "chai";
 import { anything, instance, mock, verify, when } from "ts-mockito";
 
-import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
+import { BigQueryDbAdapter, StaticBigQueryClientProvider } from "df/cli/api/dbadapters/bigquery";
 import { dataform } from "df/protos/ts";
 import { suite, test } from "df/testing";
 
@@ -17,12 +17,12 @@ suite("BigQueryDbAdapter", () => {
     const projectId = "project1";
 
     const credentials = dataform.BigQuery.create({ projectId, location: "US" });
-    const adapter = new BigQueryDbAdapter(credentials, { bigqueryClient: instance(mockBigQuery) });
+    const adapter = new BigQueryDbAdapter(credentials, { clientProvider: new StaticBigQueryClientProvider(instance(mockBigQuery)) });
 
     when(mockBigQuery.dataset(schemaName)).thenReturn(instance(mockDataset));
     // getTables returns an array where the first element is an array of tables.
     // Each table object needs an 'id' property.
-    when(mockDataset.getTables()).thenReturn(Promise.resolve([[{ id: tableName }]] as any));
+    when(mockDataset.getTables(anything())).thenReturn(Promise.resolve([[{ id: tableName }]] as any));
     when(mockDataset.table(tableName)).thenReturn(instance(mockTable));
     when(mockTable.getMetadata()).thenReturn(
       Promise.resolve([
@@ -54,10 +54,10 @@ suite("BigQueryDbAdapter", () => {
     const projectId = "project";
 
     const credentials = dataform.BigQuery.create({ projectId, location: "US" });
-    const adapter = new BigQueryDbAdapter(credentials, { bigqueryClient: instance(mockBigQuery) });
+    const adapter = new BigQueryDbAdapter(credentials, { clientProvider: new StaticBigQueryClientProvider(instance(mockBigQuery)) });
 
     when(mockBigQuery.dataset(schemaName)).thenReturn(instance(mockDataset));
-    when(mockDataset.getTables()).thenReturn(Promise.resolve([[{ id: tableName }]] as any));
+    when(mockDataset.getTables(anything())).thenReturn(Promise.resolve([[{ id: tableName }]] as any));
     when(mockDataset.table(tableName)).thenReturn(instance(mockTable));
     when(mockTable.getMetadata()).thenReturn(
       Promise.resolve([
@@ -70,7 +70,7 @@ suite("BigQueryDbAdapter", () => {
       ] as any)
     );
 
-    when(mockBigQuery.getDatasets()).thenReturn(Promise.resolve([[{ id: schemaName }]] as any));
+    when(mockBigQuery.getDatasets(anything())).thenReturn(Promise.resolve([[{ id: schemaName }]] as any));
 
     const result = await adapter.tables(projectId);
 
@@ -79,4 +79,67 @@ suite("BigQueryDbAdapter", () => {
     expect(result[0].target.schema).to.equal(schemaName);
     expect(result[0].target.name).to.equal(tableName);
   });
+
+  test("setMetadata handles action without columns", async () => {
+    // Partial mock for BigQuery client to avoid real network calls
+    const mockBigQuery: any = {
+      dataset: () => ({
+        table: () => ({
+          getMetadata: () => Promise.resolve([{ schema: { fields: [] } }]),
+          setMetadata: (metadata: any) => {
+            expect(metadata.description).to.equal("test");
+            return Promise.resolve([]);
+          }
+        })
+      })
+    };
+
+    const credentials = dataform.BigQuery.create({ projectId: "p", location: "US" });
+    const adapter = new BigQueryDbAdapter(credentials, {
+      concurrencyLimit: 1,
+      clientProvider: { get: () => mockBigQuery }
+    });
+
+    const action = dataform.ExecutionAction.create({
+      target: { database: "db", schema: "sch", name: "tab" },
+      actionDescriptor: { description: "test" }
+      // columns is missing/null in this action
+    });
+
+    // This should not throw "cannot read property 'find' of undefined"
+    await adapter.setMetadata(action);
+  });
+
+  test("setMetadata correctly maps column descriptions", async () => {
+    const mockBigQuery: any = {
+      dataset: () => ({
+        table: () => ({
+          getMetadata: () => Promise.resolve([{
+            schema: {
+              fields: [{ name: "id", type: "INTEGER" }]
+            }
+          }]),
+          setMetadata: (metadata: any) => {
+            expect(metadata.schema[0].description).to.equal("id desc");
+            return Promise.resolve([]);
+          }
+        })
+      })
+    };
+
+    const credentials = dataform.BigQuery.create({ projectId: "p", location: "US" });
+    const adapter = new BigQueryDbAdapter(credentials, {
+      concurrencyLimit: 1,
+      clientProvider: { get: () => mockBigQuery }
+    });
+
+    const action = dataform.ExecutionAction.create({
+      target: { database: "db", schema: "sch", name: "tab" },
+      actionDescriptor: {
+        columns: [{ path: ["id"], description: "id desc" }]
+      }
+    });
+
+    await adapter.setMetadata(action);
+  });
 });

protos/core.proto

@@ -51,9 +51,16 @@ message CompileConfig {
   // Project configuration overrides.
   ProjectConfig project_config_override = 3;
 
-  // Override compilation timeout settings.
+  // Override top-level compilation timeout (the single compile worker that
+  // runs once per invocation). Does not bound per-model JiT compilation;
+  // see jit_timeout_millis.
   int32 timeout_millis = 6;
 
+  // Per-model JiT compilation worker timeout. Each model with JiT code gets
+  // its own fresh budget; this is per-model, not a shared budget across the
+  // compile.
+  int32 jit_timeout_millis = 13;
+
   Extension extension = 10;
 
   // Whether to emit compilation debug logs.

protos/execution.proto

@@ -14,8 +14,16 @@ message RunConfig {
   bool include_dependencies = 3;
   bool include_dependents = 11;
   bool full_refresh = 2;
+
+  // Wall-clock deadline for the entire run. When hit, all in-flight actions are
+  // cancelled. Does not bound per-model JiT compilation; see jit_timeout_millis.
   int32 timeout_millis = 7;
 
+  // Per-model JiT compilation worker timeout. Each action with JiT code gets
+  // its own fresh budget; this is per-model, not a shared budget across the
+  // run. Falls back to DEFAULT_JIT_COMPILATION_TIMEOUT_MILLIS when unset.
+  int32 jit_timeout_millis = 10;
+
   // For internal use only, will be removed at a later date.
   bool disable_set_metadata = 9;
 
@@ -46,6 +54,8 @@ message ExecutionAction {
 
   string jit_code = 12;
 
+  bool disabled = 13;
+
   reserved 1, 3, 7;
 }
 
@@ -118,6 +128,7 @@ message TaskResult {
   string error_message = 2;
   Timing timing = 3;
   ExecutionMetadata metadata = 4;
+  string compiled_sql = 5;
 }
 
 message TestResult {

protos/jit.proto

@@ -126,6 +126,8 @@ message JitCompilationRequest {
   repeated string file_paths = 6;
   // Current execution information for introspection.
   RunningExecutionData execution_data = 7;
+  // File name where the target is defined.
+  string file_name = 8;
 }
 
 // JiT compilation response.