Set labels on BQ jobs. (#1016)

Author: BenBirt

api/commands/run.ts

@@ -385,7 +385,9 @@ export class Runner {
           actionResult.status === dataform.ActionResult.ExecutionStatus.RUNNING &&
           !this.cancelled
         ) {
-          const taskStatus = await this.executeTask(client, task, actionResult);
+          const taskStatus = await this.executeTask(client, task, actionResult, {
+            bigquery: { labels: action.actionDescriptor?.bigqueryLabels }
+          });
           if (taskStatus === dataform.TaskResult.ExecutionStatus.FAILED) {
             actionResult.status = dataform.ActionResult.ExecutionStatus.FAILED;
           } else if (taskStatus === dataform.TaskResult.ExecutionStatus.CANCELLED) {
@@ -442,7 +444,8 @@ export class Runner {
   private async executeTask(
     client: dbadapters.IDbClient,
     task: dataform.IExecutionTask,
-    parentAction: dataform.IActionResult
+    parentAction: dataform.IActionResult,
+    options: { bigquery: { labels: { [label: string]: string } } }
   ): Promise<dataform.TaskResult.ExecutionStatus> {
     const timer = Timer.start();
     const taskResult: dataform.ITaskResult = {
@@ -458,7 +461,8 @@ export class Runner {
         () =>
           client.execute(task.statement, {
             onCancel: handleCancel => this.eEmitter.on(CANCEL_EVENT, handleCancel),
-            rowLimit: 1
+            rowLimit: 1,
+            bigquery: options.bigquery
           }),
         task.type === "operation" ? 0 : this.graph.projectConfig.idempotentActionRetries || 0
       );

api/dbadapters/bigquery.ts

@@ -60,8 +60,14 @@ export class BigQueryDbAdapter implements IDbAdapter {
       interactive?: boolean;
       rowLimit?: number;
       byteLimit?: number;
+      bigquery?: {
+        labels?: { [label: string]: string };
+      };
     } = { interactive: false, rowLimit: 1000, byteLimit: 1024 * 1024 }
   ): Promise<IExecutionResult> {
+    if (options?.interactive && options?.bigquery?.labels) {
+      throw new Error("BigQuery job labels may not be set for interactive queries.");
+    }
     return this.pool
       .addSingleTask({
         generator: () =>
@@ -72,7 +78,8 @@ export class BigQueryDbAdapter implements IDbAdapter {
                 options?.params,
                 options?.rowLimit,
                 options?.byteLimit,
-                options?.onCancel
+                options?.onCancel,
+                options?.bigquery?.labels
               )
       })
       .promise();
@@ -419,7 +426,8 @@ DELETE \`${CACHED_STATE_TABLE_NAME}\` WHERE target IN (${allActions
     params?: { [name: string]: any },
     rowLimit?: number,
     byteLimit?: number,
-    onCancel?: OnCancel
+    onCancel?: OnCancel,
+    labels?: { [label: string]: string }
   ) {
     let isCancelled = false;
     onCancel?.(() => (isCancelled = true));
@@ -429,7 +437,8 @@ DELETE \`${CACHED_STATE_TABLE_NAME}\` WHERE target IN (${allActions
         useLegacySql: false,
         jobPrefix: "dataform-",
         query,
-        params
+        params,
+        labels
       });
       const resultStream = job[0].getQueryResultsStream();
       return new Promise<IExecutionResult>((resolve, reject) => {

api/dbadapters/index.ts

@@ -28,6 +28,9 @@ export interface IDbClient {
       interactive?: boolean;
       rowLimit?: number;
       byteLimit?: number;
+      bigquery?: {
+        labels?: { [label: string]: string };
+      };
     }
   ): Promise<IExecutionResult>;
 }

package.json

@@ -9,7 +9,7 @@
     "@blueprintjs/core": "^3.22.3",
     "@blueprintjs/icons": "^3.13.0",
     "@blueprintjs/select": "^3.8.0",
-    "@google-cloud/bigquery": "^5.1.0",
+    "@google-cloud/bigquery": "^5.3.0",
     "@google-cloud/storage": "^5.1.1",
     "@mapbox/rehype-prism": "^0.4.0",
     "@nivo/bar": "^0.61.1",

yarn.lock

@@ -1090,10 +1090,10 @@
     classnames "^2.2"
     tslib "~1.10.0"
 
-"@google-cloud/bigquery@^5.1.0":
-  version "5.1.0"
-  resolved "https://registry.yarnpkg.com/@google-cloud/bigquery/-/bigquery-5.1.0.tgz#d960fe3317ba3d5e69cccc60b6230747f54727ac"
-  integrity sha512-n5RSnMaYATkfbAS+lMKQ5ELmCiA6j0nRQzXfcNS8X5pqC9ZANQ9jVpTcXyKXjuRh8vAi0Pp4+Sv6NJaYsMhVIg==
+"@google-cloud/bigquery@^5.3.0":
+  version "5.3.0"
+  resolved "https://registry.yarnpkg.com/@google-cloud/bigquery/-/bigquery-5.3.0.tgz#87db58ed6538d0be5dd146901ffd636d6b9fead0"
+  integrity sha512-IxHd7dWJygQgExdnwS7D8zMswe3VG5akKSLTkdl5ky1Xn6VHiXg30c8zK99jqRafA7kU023W/+L9F9zIQp1qiA==
   dependencies:
     "@google-cloud/common" "^3.1.0"
     "@google-cloud/paginator" "^3.0.0"