feat: support per-job reservation assignment

Author: alvarowolfx

src/bigquery.ts

@@ -76,6 +76,7 @@ export type JobRequest<J> = J & {
   jobPrefix?: string;
   location?: string;
   projectId?: string;
+  reservation?: string;
 };
 
 export type PagedRequest<P> = P & {
@@ -114,6 +115,7 @@ export type Query = JobRequest<bigquery.IJobConfigurationQuery> & {
   job?: Job;
   maxResults?: number;
   jobTimeoutMs?: number;
+  reservation?: string;
   pageToken?: string;
   wrapIntegers?: boolean | IntegerTypeCastOptions;
   parseJSON?: boolean;
@@ -1731,11 +1733,16 @@ export class BigQuery extends Service {
       location: this.location,
     };
 
-    if (options.location) {
-      reqOpts.jobReference.location = options.location;
+    if (reqOpts.location) {
+      reqOpts.jobReference.location = reqOpts.location;
       delete reqOpts.location;
     }
 
+    if (reqOpts.configuration && reqOpts.reservation) {
+      reqOpts.configuration.reservation = reqOpts.reservation;
+      delete reqOpts.reservation;
+    }
+
     const job = this.job(jobId!, {
       location: reqOpts.jobReference.location,
     });
@@ -2327,6 +2334,7 @@ export class BigQuery extends Service {
       useLegacySql: false,
       requestId: randomUUID(),
       jobCreationMode: 'JOB_CREATION_OPTIONAL',
+      reservation: queryObj.reservation,
     };
     if (!this._enableQueryPreview) {
       delete req.jobCreationMode;

src/model.ts

@@ -26,7 +26,7 @@ import {
   RequestCallback,
   JobRequest,
 } from '.';
-import {JobMetadata} from './job';
+import {JobMetadata, JobOptions} from './job';
 import bigquery from './types';
 
 // This is supposed to be a @google-cloud/storage `File` type. The storage npm
@@ -424,8 +424,7 @@ class Model extends ServiceObject {
       }
     }
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const body: any = {
+    const body: JobOptions = {
       configuration: {
         extract: extend(true, options, {
           sourceModel: {
@@ -447,6 +446,11 @@ class Model extends ServiceObject {
       delete options.jobId;
     }
 
+    if (body.configuration && options.reservation) {
+      body.configuration.reservation = options.reservation;
+      delete options.reservation;
+    }
+
     this.bigQuery.createJob(body, callback!);
   }
 

src/table.ts

@@ -52,7 +52,7 @@ import {
 } from '.';
 import {GoogleErrorBody} from '@google-cloud/common/build/src/util';
 import {Duplex, Writable} from 'stream';
-import {JobMetadata} from './job';
+import {JobMetadata, JobOptions} from './job';
 import bigquery from './types';
 import {IntegerTypeCastOptions} from './bigquery';
 import {RowQueue} from './rowQueue';
@@ -923,8 +923,7 @@ class Table extends ServiceObject {
     const callback =
       typeof metadataOrCallback === 'function' ? metadataOrCallback : cb;
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const body: any = {
+    const body: JobOptions = {
       configuration: {
         copy: extend(true, metadata, {
           destinationTable: {
@@ -955,6 +954,11 @@ class Table extends ServiceObject {
       delete metadata.jobId;
     }
 
+    if (body.configuration && metadata.reservation) {
+      body.configuration.reservation = metadata.reservation;
+      delete metadata.reservation;
+    }
+
     this.bigQuery.createJob(body, callback!);
   }
 
@@ -1045,8 +1049,7 @@ class Table extends ServiceObject {
     const callback =
       typeof metadataOrCallback === 'function' ? metadataOrCallback : cb;
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const body: any = {
+    const body: JobOptions = {
       configuration: {
         copy: extend(true, metadata, {
           destinationTable: {
@@ -1080,6 +1083,11 @@ class Table extends ServiceObject {
       delete metadata.jobId;
     }
 
+    if (body.configuration && metadata.reservation) {
+      body.configuration.reservation = metadata.reservation;
+      delete metadata.reservation;
+    }
+
     this.bigQuery.createJob(body, callback!);
   }
 
@@ -1218,8 +1226,7 @@ class Table extends ServiceObject {
       delete options.gzip;
     }
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const body: any = {
+    const body: JobOptions = {
       configuration: {
         extract: extend(true, options, {
           sourceTable: {
@@ -1245,6 +1252,11 @@ class Table extends ServiceObject {
       delete options.jobId;
     }
 
+    if (body.configuration && options.reservation) {
+      body.configuration.reservation = options.reservation;
+      delete options.reservation;
+    }
+
     this.bigQuery.createJob(body, callback!);
   }
 
@@ -1399,8 +1411,7 @@ class Table extends ServiceObject {
       return [jobResponse, jobResponse.metadata];
     }
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const body: any = {
+    const body: JobOptions = {
       configuration: {
         load: {
           destinationTable: {
@@ -1427,7 +1438,12 @@ class Table extends ServiceObject {
       delete metadata.jobId;
     }
 
-    extend(true, body.configuration.load, metadata, {
+    if (body.configuration && metadata.reservation) {
+      body.configuration.reservation = metadata.reservation;
+      delete metadata.reservation;
+    }
+
+    extend(true, body.configuration?.load, metadata, {
       sourceUris: toArray(source).map(src => {
         if (!util.isCustomType(src, 'storage/file')) {
           throw new Error('Source must be a File object.');
@@ -1437,7 +1453,12 @@ class Table extends ServiceObject {
         // the file's extension. If no match, don't set, and default upstream
         // to CSV.
         const format = FORMATS[path.extname(src.name).substr(1).toLowerCase()];
-        if (!metadata.sourceFormat && format) {
+        if (
+          !metadata.sourceFormat &&
+          format &&
+          body.configuration &&
+          body.configuration.load
+        ) {
           body.configuration.load.sourceFormat = format;
         }
         return 'gs://' + src.bucket.name + '/' + src.name;