Skip to content

Commit 9f99149

Browse files
max-ostapenkomax-ostapenko
authored
Add BigQuery reservation configuration (#2096)
* feat: Add BigQuery reservation configuration for projects and individual actions. * lint Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * refactor: rename `bigqueryReservation` to `reservation` across configurations, protos, and related code Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * docs: clarify BigQuery reservation fallback behavior in protos and CLI Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * refactor: update default reservation test constant and expand project config assertions Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * fix test reservation name Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * added documentation for BigQuery reservation configuration Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * fix project name Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * fix reservation location Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * Add reservation to a profile schema Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * revert reservation attribute Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com> * Note support status --------- Signed-off-by: Max Ostapenko <1611259+max-ostapenko@users.noreply.github.com>
1 parent d276374 commit 9f99149

18 files changed

Lines changed: 376 additions & 84 deletions

‎cli/api/commands/run.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -348,7 +348,10 @@ export class Runner {
348348
},
349349
actionRetryLimit: this.executionOptions?.bigquery?.actionRetryLimit,
350350
jobPrefix: this.executionOptions?.bigquery?.jobPrefix,
351-
dryRun: this.executionOptions?.bigquery?.dryRun
351+
dryRun: this.executionOptions?.bigquery?.dryRun,
352+
reservation:
353+
action.actionDescriptor?.reservation ||
354+
this.graph.projectConfig?.defaultReservation
352355
}
353356
});
354357
if (taskStatus === dataform.TaskResult.ExecutionStatus.FAILED) {
@@ -375,7 +378,7 @@ export class Runner {
375378
actionResult.status === dataform.ActionResult.ExecutionStatus.RUNNING &&
376379
!(this.graph.runConfig && this.graph.runConfig.disableSetMetadata) &&
377380
// Only set metadata if not using BigQuery dry run
378-
!this.executionOptions.bigquery?.dryRun &&
381+
!this.executionOptions.bigquery?.dryRun &&
379382
action.type === "table"
380383
) {
381384
try {
@@ -420,9 +423,9 @@ export class Runner {
420423
};
421424
parentAction.tasks.push(taskResult);
422425
this.notifyListeners();
423-
if(options.bigquery?.dryRun && task.type === "assertion") {
426+
if (options.bigquery?.dryRun && task.type === "assertion") {
424427
taskResult.status = dataform.TaskResult.ExecutionStatus.SUCCESSFUL;
425-
}
428+
}
426429
else {
427430
try {
428431
// Retry this function a given number of times, configurable by user
@@ -469,7 +472,7 @@ class Timer {
469472
public static start(existingTiming?: dataform.ITiming) {
470473
return new Timer(existingTiming?.startTimeMillis.toNumber() || new Date().valueOf());
471474
}
472-
private constructor(readonly startTimeMillis: number) {}
475+
private constructor(readonly startTimeMillis: number) { }
473476

474477
public current(): dataform.ITiming {
475478
return {

‎cli/api/dbadapters/bigquery.ts‎

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ export interface IBigQueryExecutionOptions {
2727
location?: string;
2828
jobPrefix?: string;
2929
dryRun?: boolean;
30+
reservation?: string;
3031
}
3132

3233
export class BigQueryDbAdapter implements IDbAdapter {
@@ -69,23 +70,24 @@ export class BigQueryDbAdapter implements IDbAdapter {
6970
generator: () =>
7071
options?.interactive
7172
? this.runQuery(
72-
statement,
73-
options?.params,
74-
options?.rowLimit,
75-
options?.byteLimit,
76-
options.bigquery?.location
77-
)
73+
statement,
74+
options?.params,
75+
options?.rowLimit,
76+
options?.byteLimit,
77+
options.bigquery?.location
78+
)
7879
: this.createQueryJob(
79-
statement,
80-
options?.params,
81-
options?.rowLimit,
82-
options?.byteLimit,
83-
options?.onCancel,
84-
options?.bigquery?.labels,
85-
options?.bigquery?.location,
86-
options?.bigquery?.jobPrefix,
87-
options?.bigquery?.dryRun
88-
)
80+
statement,
81+
options?.params,
82+
options?.rowLimit,
83+
options?.byteLimit,
84+
options?.onCancel,
85+
options?.bigquery?.labels,
86+
options?.bigquery?.location,
87+
options?.bigquery?.jobPrefix,
88+
options?.bigquery?.dryRun,
89+
options?.bigquery?.reservation
90+
)
8991
})
9092
.promise();
9193
}
@@ -203,8 +205,8 @@ export class BigQueryDbAdapter implements IDbAdapter {
203205
metadata.type === "TABLE"
204206
? dataform.TableMetadata.Type.TABLE
205207
: metadata.type === "VIEW"
206-
? dataform.TableMetadata.Type.VIEW
207-
: dataform.TableMetadata.Type.UNKNOWN,
208+
? dataform.TableMetadata.Type.VIEW
209+
: dataform.TableMetadata.Type.UNKNOWN,
208210
target,
209211
fields: metadata.schema.fields?.map(field => convertField(field)),
210212
lastUpdatedMillis: Long.fromString(metadata.lastModifiedTime),
@@ -321,7 +323,8 @@ export class BigQueryDbAdapter implements IDbAdapter {
321323
labels?: { [label: string]: string },
322324
location?: string,
323325
jobPrefix?: string,
324-
dryRun?: boolean
326+
dryRun?: boolean,
327+
reservation?: string
325328
) {
326329
let isCancelled = false;
327330
onCancel?.(() => (isCancelled = true));
@@ -336,8 +339,9 @@ export class BigQueryDbAdapter implements IDbAdapter {
336339
params,
337340
labels,
338341
location,
339-
dryRun
340-
});
342+
dryRun,
343+
reservation
344+
} as any);
341345
const resultStream = job[0].getQueryResultsStream();
342346
return new Promise<IExecutionResult>((resolve, reject) => {
343347
if (isCancelled) {

‎cli/index.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -878,6 +878,16 @@ class ProjectConfigOptions {
878878
}
879879
};
880880

881+
public static defaultReservation: INamedOption<yargs.Options> = {
882+
name: "default-reservation",
883+
option: {
884+
describe:
885+
"The default BigQuery reservation to use for execution. If unset, the value from " +
886+
"workflow_settings.yaml is used. If neither is set, default BigQuery behavior applies.",
887+
type: "string"
888+
}
889+
};
890+
881891
public static allYargsOptions = [
882892
ProjectConfigOptions.defaultDatabase,
883893
ProjectConfigOptions.defaultSchema,
@@ -887,7 +897,8 @@ class ProjectConfigOptions {
887897
ProjectConfigOptions.databaseSuffix,
888898
ProjectConfigOptions.schemaSuffix,
889899
ProjectConfigOptions.tablePrefix,
890-
ProjectConfigOptions.disableAssertions
900+
ProjectConfigOptions.disableAssertions,
901+
ProjectConfigOptions.defaultReservation
891902
];
892903

893904
public static constructProjectConfigOverride(
@@ -922,6 +933,9 @@ class ProjectConfigOptions {
922933
if (argv[ProjectConfigOptions.disableAssertions.name]) {
923934
projectConfigOptions.disableAssertions = argv[ProjectConfigOptions.disableAssertions.name];
924935
}
936+
if (argv[ProjectConfigOptions.defaultReservation.name]) {
937+
projectConfigOptions.defaultReservation = argv[ProjectConfigOptions.defaultReservation.name];
938+
}
925939
return projectConfigOptions;
926940
}
927941
}

‎cli/index_test.ts‎

Lines changed: 102 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import { TmpDirFixture } from "df/testing/fixtures";
2525

2626
const DEFAULT_DATABASE = "dataform-open-source";
2727
const DEFAULT_LOCATION = "US";
28+
const DEFAULT_RESERVATION = "projects/dataform-open-source/locations/us/reservations/dataform-test";
2829
const CREDENTIALS_PATH = path.resolve(process.env.RUNFILES, "df/test_credentials/bigquery.json");
2930

3031
suite("@dataform/cli", ({ afterEach }) => {
@@ -800,7 +801,7 @@ SELECT 1 WHERE FALSE
800801
fs.writeFileSync(
801802
tableFilePath,
802803
`
803-
config {
804+
config {
804805
type: "table",
805806
assertions: {
806807
uniqueKey: ["id"]
@@ -1062,4 +1063,104 @@ SELECT 1 as id
10621063
});
10631064
});
10641065
});
1066+
1067+
suite("--default-reservation flag", ({ beforeEach }) => {
1068+
const projectDir = tmpDirFixture.createNewTmpDir();
1069+
1070+
beforeEach("setup test project", async () => {
1071+
const npmCacheDir = tmpDirFixture.createNewTmpDir();
1072+
const workflowSettingsPath = path.join(projectDir, "workflow_settings.yaml");
1073+
const packageJsonPath = path.join(projectDir, "package.json");
1074+
1075+
await getProcessResult(
1076+
execFile(nodePath, [cliEntryPointPath, "init", projectDir, DEFAULT_DATABASE, DEFAULT_LOCATION])
1077+
);
1078+
1079+
// Remove dataformCoreVersion so we can use the local package.
1080+
const workflowSettings = dataform.WorkflowSettings.create(
1081+
loadYaml(fs.readFileSync(workflowSettingsPath, "utf8"))
1082+
);
1083+
delete workflowSettings.dataformCoreVersion;
1084+
fs.writeFileSync(workflowSettingsPath, dumpYaml(workflowSettings));
1085+
1086+
fs.writeFileSync(
1087+
packageJsonPath,
1088+
`{
1089+
"dependencies":{
1090+
"@dataform/core": "${version}"
1091+
}
1092+
}`
1093+
);
1094+
await getProcessResult(
1095+
execFile(npmPath, [
1096+
"install",
1097+
"--prefix",
1098+
projectDir,
1099+
"--cache",
1100+
npmCacheDir,
1101+
corePackageTarPath
1102+
])
1103+
);
1104+
1105+
const tableFilePath = path.join(projectDir, "definitions", "example_table.sqlx");
1106+
fs.ensureFileSync(tableFilePath);
1107+
fs.writeFileSync(
1108+
tableFilePath,
1109+
`
1110+
config { type: "table" }
1111+
SELECT 1 as id
1112+
`
1113+
);
1114+
});
1115+
1116+
test("--default-reservation flag is applied to projectConfig in compile output", async () => {
1117+
const compileResult = await getProcessResult(
1118+
execFile(nodePath, [
1119+
cliEntryPointPath,
1120+
"compile",
1121+
projectDir,
1122+
"--json",
1123+
`--default-reservation=${DEFAULT_RESERVATION}`
1124+
])
1125+
);
1126+
1127+
expect(compileResult.exitCode).equals(0);
1128+
const compiledGraph = JSON.parse(compileResult.stdout);
1129+
expect(compiledGraph.projectConfig).deep.equals({
1130+
warehouse: "bigquery",
1131+
defaultSchema: "dataform",
1132+
assertionSchema: "dataform_assertions",
1133+
defaultDatabase: DEFAULT_DATABASE,
1134+
defaultLocation: DEFAULT_LOCATION,
1135+
defaultReservation: DEFAULT_RESERVATION
1136+
});
1137+
});
1138+
1139+
test("--default-reservation flag is applied to projectConfig in run (dry-run) output", async () => {
1140+
const runResult = await getProcessResult(
1141+
execFile(nodePath, [
1142+
cliEntryPointPath,
1143+
"run",
1144+
projectDir,
1145+
"--credentials",
1146+
CREDENTIALS_PATH,
1147+
"--dry-run",
1148+
"--json",
1149+
`--default-reservation=${DEFAULT_RESERVATION}`,
1150+
"--actions=example_table"
1151+
])
1152+
);
1153+
1154+
expect(runResult.exitCode).equals(0);
1155+
const executionGraph = JSON.parse(runResult.stdout);
1156+
expect(executionGraph.projectConfig).deep.equals({
1157+
warehouse: "bigquery",
1158+
defaultSchema: "dataform",
1159+
assertionSchema: "dataform_assertions",
1160+
defaultDatabase: DEFAULT_DATABASE,
1161+
defaultLocation: DEFAULT_LOCATION,
1162+
defaultReservation: DEFAULT_RESERVATION
1163+
});
1164+
});
1165+
});
10651166
});

‎core/actions/assertion.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,12 @@ export class Assertion extends ActionBuilder<dataform.Assertion> {
149149
if (config.filename) {
150150
this.proto.fileName = config.filename;
151151
}
152+
if (config.reservation) {
153+
if (!this.proto.actionDescriptor) {
154+
this.proto.actionDescriptor = {};
155+
}
156+
this.proto.actionDescriptor.reservation = config.reservation;
157+
}
152158
return this;
153159
}
154160

‎core/actions/assertion_test.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,9 @@ actions:
4646
- tagB
4747
disabled: true,
4848
description: description
49-
hermetic: true,
50-
dependOnDependencyAssertions: true`
49+
hermetic: true
50+
dependOnDependencyAssertions: true
51+
reservation: reservation`
5152
);
5253
fs.writeFileSync(path.join(projectDir, "definitions/action.sql"), "SELECT 1");
5354
fs.writeFileSync(path.join(projectDir, "definitions/operation.sqlx"), "SELECT 1");
@@ -69,7 +70,8 @@ actions:
6970
name: "name"
7071
},
7172
actionDescriptor: {
72-
description: "description"
73+
description: "description",
74+
reservation: "reservation"
7375
},
7476
disabled: true,
7577
fileName: "definitions/action.sql",

‎core/actions/incremental_table.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,13 @@ export class IncrementalTable extends ActionBuilder<dataform.Table> {
218218
});
219219
this.proto.onSchemaChange = this.mapOnSchemaChange(config.onSchemaChange);
220220

221+
if (config.reservation) {
222+
if (!this.proto.actionDescriptor) {
223+
this.proto.actionDescriptor = {};
224+
}
225+
this.proto.actionDescriptor.reservation = config.reservation;
226+
}
227+
221228
return this;
222229
}
223230

‎core/actions/incremental_table_test.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ actions:
101101
dependOnDependencyAssertions: true,
102102
${exampleBuiltInAssertions.inputAssertionBlock}
103103
hermetic: true,
104+
reservation: "reservation",
104105
onSchemaChange: "SYNCHRONIZE",
105106
metadata: {
106107
overview: "incremental table overview",
@@ -186,6 +187,7 @@ SELECT 1`
186187
incrementalQuery: "\n\n\nSELECT 1",
187188
actionDescriptor: {
188189
...exampleActionDescriptor.outputActionDescriptor,
190+
reservation: "reservation",
189191
// sqlxConfig.bigquery.labels are placed as bigqueryLabels.
190192
bigqueryLabels: {
191193
key: "val"
@@ -294,6 +296,7 @@ actions:
294296
dependOnDependencyAssertions: true
295297
${exampleBuiltInAssertionsAsYaml.inputActionConfigBlock}
296298
hermetic: true
299+
reservation: reservation
297300
onSchemaChange: FAIL
298301
`
299302
);
@@ -349,7 +352,8 @@ actions:
349352
bigqueryLabels: {
350353
key: "val"
351354
},
352-
description: "description"
355+
description: "description",
356+
reservation: "reservation"
353357
}
354358
}
355359
]);

‎core/actions/operation.ts‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,12 @@ export class Operation extends ActionBuilder<dataform.Operation> {
159159
if (config.filename) {
160160
this.proto.fileName = config.filename;
161161
}
162+
if (config.reservation) {
163+
if (!this.proto.actionDescriptor) {
164+
this.proto.actionDescriptor = {};
165+
}
166+
this.proto.actionDescriptor.reservation = config.reservation;
167+
}
162168
return this;
163169
}
164170

@@ -174,13 +180,13 @@ export class Operation extends ActionBuilder<dataform.Operation> {
174180
}
175181

176182
public jitCode(jitCode: JitContextable<IActionContext, JitOperationResult>) {
177-
if (!this.proto.actionDescriptor) {
178-
this.proto.actionDescriptor = {};
179-
}
180-
this.proto.actionDescriptor.compilationMode = dataform.ActionCompilationMode.ACTION_COMPILATION_MODE_JIT;
181-
this.contextableJitCode = jitCode;
182-
return this;
183+
if (!this.proto.actionDescriptor) {
184+
this.proto.actionDescriptor = {};
183185
}
186+
this.proto.actionDescriptor.compilationMode = dataform.ActionCompilationMode.ACTION_COMPILATION_MODE_JIT;
187+
this.contextableJitCode = jitCode;
188+
return this;
189+
}
184190

185191
/**
186192
* @deprecated Deprecated in favor of

0 commit comments

Comments
 (0)