feat: support per-job reservation assignment (#1477)

alvarowolfx · web-flow · commit 8151e72bb1e1 · 2025-05-27T16:59:19.000-04:00
diff --git a/src/bigquery.ts b/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;
@@ -1562,6 +1564,11 @@ export class BigQuery extends Service {
       delete query.jobId;
     }
 
+    if (query.reservation) {
+      reqOpts.configuration.reservation = query.reservation;
+      delete query.reservation;
+    }
+
     this.createJob(reqOpts, callback!);
   }
 
@@ -1731,11 +1738,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 +2339,13 @@ export class BigQuery extends Service {
       useLegacySql: false,
       requestId: randomUUID(),
       jobCreationMode: 'JOB_CREATION_OPTIONAL',
+      reservation: queryObj.reservation,
+      continuous: queryObj.continuous,
+      destinationEncryptionConfiguration:
+        queryObj.destinationEncryptionConfiguration,
+      writeIncrementalResults: queryObj.writeIncrementalResults,
+      connectionProperties: queryObj.connectionProperties,
+      preserveNulls: queryObj.preserveNulls,
     };
     if (!this._enableQueryPreview) {
       delete req.jobCreationMode;
diff --git a/src/model.ts b/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!);
   }
 
diff --git a/src/table.ts b/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;
diff --git a/test/bigquery.ts b/test/bigquery.ts
@@ -2644,6 +2644,20 @@ describe('BigQuery', () => {
       bq.createQueryJob(options, assert.ifError);
     });
 
+    it('should accept a reservation id', done => {
+      const options = {
+        query: QUERY_STRING,
+        reservation: 'reservation/1',
+      };
+
+      bq.createJob = (reqOpts: JobOptions) => {
+        assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
+        done();
+      };
+
+      bq.createQueryJob(options, assert.ifError);
+    });
+
     it('should accept a location', done => {
       const options = {
         query: QUERY_STRING,
@@ -3300,6 +3314,29 @@ describe('BigQuery', () => {
 
       bq.query(QUERY_STRING, fakeOptions, assert.ifError);
     });
+
+    it('should accept a reservation id', done => {
+      const query: Query = {
+        query: QUERY_STRING,
+        reservation: 'reservation/1',
+      };
+      const fakeJob = {
+        getQueryResults: (options: {}) => {
+          done();
+        },
+      };
+
+      bq.createJob = (reqOpts: JobOptions, callback: Function) => {
+        assert(reqOpts.configuration?.reservation, 'reservation/1');
+        callback(null, fakeJob, FAKE_RESPONSE);
+      };
+
+      bq.buildQueryRequest_ = (query: {}, opts: {}) => {
+        return undefined;
+      };
+
+      bq.query(query, assert.ifError);
+    });
   });
 
   describe('buildQueryRequest_', () => {
diff --git a/test/model.ts b/test/model.ts
@@ -257,6 +257,25 @@ describe('BigQuery/Model', () => {
         model.createExtractJob(URI, options, done);
       });
 
+      it('should accept a reservation id', done => {
+        const options = {
+          reservation: 'reservation/1',
+        };
+
+        model.bigQuery.createJob = (
+          reqOpts: JobOptions,
+          callback: Function,
+        ) => {
+          assert.strictEqual(
+            reqOpts.configuration?.reservation,
+            'reservation/1',
+          );
+          callback(); // the done fn
+        };
+
+        model.createExtractJob(URI, options, done);
+      });
+
       it('should accept a job id', done => {
         const jobId = 'job-id';
         const options = {jobId};
diff --git a/test/table.ts b/test/table.ts
@@ -767,6 +767,19 @@ describe('BigQuery/Table', () => {
       table.createCopyJob(DEST_TABLE, options, done);
     });
 
+    it('should accept a reservation id', done => {
+      const options = {
+        reservation: 'reservation/1',
+      };
+
+      table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
+        assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
+        callback(); // the done fn
+      };
+
+      table.createCopyJob(DEST_TABLE, options, done);
+    });
+
     it('should use the default location', done => {
       table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
         assert.strictEqual(reqOpts.location, LOCATION);
@@ -911,6 +924,19 @@ describe('BigQuery/Table', () => {
       table.createCopyFromJob(SOURCE_TABLE, options, done);
     });
 
+    it('should accept a reservation id', done => {
+      const options = {
+        reservation: 'reservation/1',
+      };
+
+      table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
+        assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
+        callback(); // the done fn
+      };
+
+      table.createCopyFromJob(SOURCE_TABLE, options, done);
+    });
+
     it('should use the default location', done => {
       table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
         assert.strictEqual(reqOpts.location, LOCATION);
@@ -1195,6 +1221,19 @@ describe('BigQuery/Table', () => {
       table.createExtractJob(FILE, options, done);
     });
 
+    it('should accept a reservation id', done => {
+      const options = {
+        reservation: 'reservation/1',
+      };
+
+      table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
+        assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
+        callback(); // the done fn
+      };
+
+      table.createExtractJob(FILE, options, done);
+    });
+
     it('should use the default location', done => {
       const table = new Table(DATASET, TABLE_ID, {location: LOCATION});
 
@@ -1425,6 +1464,19 @@ describe('BigQuery/Table', () => {
       );
     });
 
+    it('should set the job reservation', async () => {
+      const reservation = 'reservation/1';
+      await table.createLoadJob(FILE, {reservation});
+      assert(bqCreateJobStub.calledOnce);
+      assert(
+        bqCreateJobStub.calledWithMatch({
+          configuration: {
+            reservation,
+          },
+        }),
+      );
+    });
+
     it('should use the default location', async () => {
       const table = new Table(DATASET, TABLE_ID, {location: LOCATION});
       await table.createLoadJob(FILE);
