Skip to content

Commit 8151e72

Browse files
alvarowolfxalvarowolfx
authored
feat: support per-job reservation assignment (#1477)
1 parent 5fc2a41 commit 8151e72

File tree

6 files changed

+168
-16
lines changed

6 files changed

+168
-16
lines changed

‎src/bigquery.ts‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ export type JobRequest<J> = J & {
7676
jobPrefix?: string;
7777
location?: string;
7878
projectId?: string;
79+
reservation?: string;
7980
};
8081

8182
export type PagedRequest<P> = P & {
@@ -114,6 +115,7 @@ export type Query = JobRequest<bigquery.IJobConfigurationQuery> & {
114115
job?: Job;
115116
maxResults?: number;
116117
jobTimeoutMs?: number;
118+
reservation?: string;
117119
pageToken?: string;
118120
wrapIntegers?: boolean | IntegerTypeCastOptions;
119121
parseJSON?: boolean;
@@ -1562,6 +1564,11 @@ export class BigQuery extends Service {
15621564
delete query.jobId;
15631565
}
15641566

1567+
if (query.reservation) {
1568+
reqOpts.configuration.reservation = query.reservation;
1569+
delete query.reservation;
1570+
}
1571+
15651572
this.createJob(reqOpts, callback!);
15661573
}
15671574

@@ -1731,11 +1738,16 @@ export class BigQuery extends Service {
17311738
location: this.location,
17321739
};
17331740

1734-
if (options.location) {
1735-
reqOpts.jobReference.location = options.location;
1741+
if (reqOpts.location) {
1742+
reqOpts.jobReference.location = reqOpts.location;
17361743
delete reqOpts.location;
17371744
}
17381745

1746+
if (reqOpts.configuration && reqOpts.reservation) {
1747+
reqOpts.configuration.reservation = reqOpts.reservation;
1748+
delete reqOpts.reservation;
1749+
}
1750+
17391751
const job = this.job(jobId!, {
17401752
location: reqOpts.jobReference.location,
17411753
});
@@ -2327,6 +2339,13 @@ export class BigQuery extends Service {
23272339
useLegacySql: false,
23282340
requestId: randomUUID(),
23292341
jobCreationMode: 'JOB_CREATION_OPTIONAL',
2342+
reservation: queryObj.reservation,
2343+
continuous: queryObj.continuous,
2344+
destinationEncryptionConfiguration:
2345+
queryObj.destinationEncryptionConfiguration,
2346+
writeIncrementalResults: queryObj.writeIncrementalResults,
2347+
connectionProperties: queryObj.connectionProperties,
2348+
preserveNulls: queryObj.preserveNulls,
23302349
};
23312350
if (!this._enableQueryPreview) {
23322351
delete req.jobCreationMode;

‎src/model.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ import {
2626
RequestCallback,
2727
JobRequest,
2828
} from '.';
29-
import {JobMetadata} from './job';
29+
import {JobMetadata, JobOptions} from './job';
3030
import bigquery from './types';
3131

3232
// This is supposed to be a @google-cloud/storage `File` type. The storage npm
@@ -424,8 +424,7 @@ class Model extends ServiceObject {
424424
}
425425
}
426426

427-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
428-
const body: any = {
427+
const body: JobOptions = {
429428
configuration: {
430429
extract: extend(true, options, {
431430
sourceModel: {
@@ -447,6 +446,11 @@ class Model extends ServiceObject {
447446
delete options.jobId;
448447
}
449448

449+
if (body.configuration && options.reservation) {
450+
body.configuration.reservation = options.reservation;
451+
delete options.reservation;
452+
}
453+
450454
this.bigQuery.createJob(body, callback!);
451455
}
452456

‎src/table.ts‎

Lines changed: 32 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ import {
5252
} from '.';
5353
import {GoogleErrorBody} from '@google-cloud/common/build/src/util';
5454
import {Duplex, Writable} from 'stream';
55-
import {JobMetadata} from './job';
55+
import {JobMetadata, JobOptions} from './job';
5656
import bigquery from './types';
5757
import {IntegerTypeCastOptions} from './bigquery';
5858
import {RowQueue} from './rowQueue';
@@ -923,8 +923,7 @@ class Table extends ServiceObject {
923923
const callback =
924924
typeof metadataOrCallback === 'function' ? metadataOrCallback : cb;
925925

926-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
927-
const body: any = {
926+
const body: JobOptions = {
928927
configuration: {
929928
copy: extend(true, metadata, {
930929
destinationTable: {
@@ -955,6 +954,11 @@ class Table extends ServiceObject {
955954
delete metadata.jobId;
956955
}
957956

957+
if (body.configuration && metadata.reservation) {
958+
body.configuration.reservation = metadata.reservation;
959+
delete metadata.reservation;
960+
}
961+
958962
this.bigQuery.createJob(body, callback!);
959963
}
960964

@@ -1045,8 +1049,7 @@ class Table extends ServiceObject {
10451049
const callback =
10461050
typeof metadataOrCallback === 'function' ? metadataOrCallback : cb;
10471051

1048-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1049-
const body: any = {
1052+
const body: JobOptions = {
10501053
configuration: {
10511054
copy: extend(true, metadata, {
10521055
destinationTable: {
@@ -1080,6 +1083,11 @@ class Table extends ServiceObject {
10801083
delete metadata.jobId;
10811084
}
10821085

1086+
if (body.configuration && metadata.reservation) {
1087+
body.configuration.reservation = metadata.reservation;
1088+
delete metadata.reservation;
1089+
}
1090+
10831091
this.bigQuery.createJob(body, callback!);
10841092
}
10851093

@@ -1218,8 +1226,7 @@ class Table extends ServiceObject {
12181226
delete options.gzip;
12191227
}
12201228

1221-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1222-
const body: any = {
1229+
const body: JobOptions = {
12231230
configuration: {
12241231
extract: extend(true, options, {
12251232
sourceTable: {
@@ -1245,6 +1252,11 @@ class Table extends ServiceObject {
12451252
delete options.jobId;
12461253
}
12471254

1255+
if (body.configuration && options.reservation) {
1256+
body.configuration.reservation = options.reservation;
1257+
delete options.reservation;
1258+
}
1259+
12481260
this.bigQuery.createJob(body, callback!);
12491261
}
12501262

@@ -1399,8 +1411,7 @@ class Table extends ServiceObject {
13991411
return [jobResponse, jobResponse.metadata];
14001412
}
14011413

1402-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1403-
const body: any = {
1414+
const body: JobOptions = {
14041415
configuration: {
14051416
load: {
14061417
destinationTable: {
@@ -1427,7 +1438,12 @@ class Table extends ServiceObject {
14271438
delete metadata.jobId;
14281439
}
14291440

1430-
extend(true, body.configuration.load, metadata, {
1441+
if (body.configuration && metadata.reservation) {
1442+
body.configuration.reservation = metadata.reservation;
1443+
delete metadata.reservation;
1444+
}
1445+
1446+
extend(true, body.configuration?.load, metadata, {
14311447
sourceUris: toArray(source).map(src => {
14321448
if (!util.isCustomType(src, 'storage/file')) {
14331449
throw new Error('Source must be a File object.');
@@ -1437,7 +1453,12 @@ class Table extends ServiceObject {
14371453
// the file's extension. If no match, don't set, and default upstream
14381454
// to CSV.
14391455
const format = FORMATS[path.extname(src.name).substr(1).toLowerCase()];
1440-
if (!metadata.sourceFormat && format) {
1456+
if (
1457+
!metadata.sourceFormat &&
1458+
format &&
1459+
body.configuration &&
1460+
body.configuration.load
1461+
) {
14411462
body.configuration.load.sourceFormat = format;
14421463
}
14431464
return 'gs://' + src.bucket.name + '/' + src.name;

‎test/bigquery.ts‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2644,6 +2644,20 @@ describe('BigQuery', () => {
26442644
bq.createQueryJob(options, assert.ifError);
26452645
});
26462646

2647+
it('should accept a reservation id', done => {
2648+
const options = {
2649+
query: QUERY_STRING,
2650+
reservation: 'reservation/1',
2651+
};
2652+
2653+
bq.createJob = (reqOpts: JobOptions) => {
2654+
assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
2655+
done();
2656+
};
2657+
2658+
bq.createQueryJob(options, assert.ifError);
2659+
});
2660+
26472661
it('should accept a location', done => {
26482662
const options = {
26492663
query: QUERY_STRING,
@@ -3300,6 +3314,29 @@ describe('BigQuery', () => {
33003314

33013315
bq.query(QUERY_STRING, fakeOptions, assert.ifError);
33023316
});
3317+
3318+
it('should accept a reservation id', done => {
3319+
const query: Query = {
3320+
query: QUERY_STRING,
3321+
reservation: 'reservation/1',
3322+
};
3323+
const fakeJob = {
3324+
getQueryResults: (options: {}) => {
3325+
done();
3326+
},
3327+
};
3328+
3329+
bq.createJob = (reqOpts: JobOptions, callback: Function) => {
3330+
assert(reqOpts.configuration?.reservation, 'reservation/1');
3331+
callback(null, fakeJob, FAKE_RESPONSE);
3332+
};
3333+
3334+
bq.buildQueryRequest_ = (query: {}, opts: {}) => {
3335+
return undefined;
3336+
};
3337+
3338+
bq.query(query, assert.ifError);
3339+
});
33033340
});
33043341

33053342
describe('buildQueryRequest_', () => {

‎test/model.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,25 @@ describe('BigQuery/Model', () => {
257257
model.createExtractJob(URI, options, done);
258258
});
259259

260+
it('should accept a reservation id', done => {
261+
const options = {
262+
reservation: 'reservation/1',
263+
};
264+
265+
model.bigQuery.createJob = (
266+
reqOpts: JobOptions,
267+
callback: Function,
268+
) => {
269+
assert.strictEqual(
270+
reqOpts.configuration?.reservation,
271+
'reservation/1',
272+
);
273+
callback(); // the done fn
274+
};
275+
276+
model.createExtractJob(URI, options, done);
277+
});
278+
260279
it('should accept a job id', done => {
261280
const jobId = 'job-id';
262281
const options = {jobId};

‎test/table.ts‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -767,6 +767,19 @@ describe('BigQuery/Table', () => {
767767
table.createCopyJob(DEST_TABLE, options, done);
768768
});
769769

770+
it('should accept a reservation id', done => {
771+
const options = {
772+
reservation: 'reservation/1',
773+
};
774+
775+
table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
776+
assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
777+
callback(); // the done fn
778+
};
779+
780+
table.createCopyJob(DEST_TABLE, options, done);
781+
});
782+
770783
it('should use the default location', done => {
771784
table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
772785
assert.strictEqual(reqOpts.location, LOCATION);
@@ -911,6 +924,19 @@ describe('BigQuery/Table', () => {
911924
table.createCopyFromJob(SOURCE_TABLE, options, done);
912925
});
913926

927+
it('should accept a reservation id', done => {
928+
const options = {
929+
reservation: 'reservation/1',
930+
};
931+
932+
table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
933+
assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
934+
callback(); // the done fn
935+
};
936+
937+
table.createCopyFromJob(SOURCE_TABLE, options, done);
938+
});
939+
914940
it('should use the default location', done => {
915941
table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
916942
assert.strictEqual(reqOpts.location, LOCATION);
@@ -1195,6 +1221,19 @@ describe('BigQuery/Table', () => {
11951221
table.createExtractJob(FILE, options, done);
11961222
});
11971223

1224+
it('should accept a reservation id', done => {
1225+
const options = {
1226+
reservation: 'reservation/1',
1227+
};
1228+
1229+
table.bigQuery.createJob = (reqOpts: JobOptions, callback: Function) => {
1230+
assert.strictEqual(reqOpts.configuration?.reservation, 'reservation/1');
1231+
callback(); // the done fn
1232+
};
1233+
1234+
table.createExtractJob(FILE, options, done);
1235+
});
1236+
11981237
it('should use the default location', done => {
11991238
const table = new Table(DATASET, TABLE_ID, {location: LOCATION});
12001239

@@ -1425,6 +1464,19 @@ describe('BigQuery/Table', () => {
14251464
);
14261465
});
14271466

1467+
it('should set the job reservation', async () => {
1468+
const reservation = 'reservation/1';
1469+
await table.createLoadJob(FILE, {reservation});
1470+
assert(bqCreateJobStub.calledOnce);
1471+
assert(
1472+
bqCreateJobStub.calledWithMatch({
1473+
configuration: {
1474+
reservation,
1475+
},
1476+
}),
1477+
);
1478+
});
1479+
14281480
it('should use the default location', async () => {
14291481
const table = new Table(DATASET, TABLE_ID, {location: LOCATION});
14301482
await table.createLoadJob(FILE);

0 commit comments

Comments
 (0)