diff --git a/cli/api/dbadapters/execution_sql.ts b/cli/api/dbadapters/execution_sql.ts index c54fd6143..fa0035af3 100644 --- a/cli/api/dbadapters/execution_sql.ts +++ b/cli/api/dbadapters/execution_sql.ts @@ -206,7 +206,12 @@ from (${query}) as insertions`; table.bigquery && table.bigquery.clusterBy && table.bigquery.clusterBy.length > 0 ? `cluster by ${table.bigquery.clusterBy.join(", ")} ` : "" - }${options.length > 0 ? `OPTIONS(${options.join(",")})` : ""}as ${table.query}`; + }${ + table.bigquery && table.bigquery.withConnection + ? `WITH CONNECTION ${table.bigquery.withConnection} ` + : "" + } + ${options.length > 0 ? `OPTIONS(${options.join(",")})` : ""}as ${table.query}`; } private createOrReplaceView(target: dataform.ITarget, query: string) { diff --git a/core/actions/incremental_table.ts b/core/actions/incremental_table.ts index 2a3e7a19e..91e1d2e87 100644 --- a/core/actions/incremental_table.ts +++ b/core/actions/incremental_table.ts @@ -45,6 +45,7 @@ interface ILegacyIncrementalTableBigqueryConfig { labels?: { [key: string]: string }; partitionExpirationDays?: number; requirePartitionFilter?: boolean; + withConnection?: string; additionalOptions?: { [key: string]: string }; } @@ -160,6 +161,7 @@ export class IncrementalTable extends ActionBuilder { labels: config.labels, partitionExpirationDays: config.partitionExpirationDays, requirePartitionFilter: config.requirePartitionFilter, + withConnection: config.withConnection, additionalOptions: config.additionalOptions }); if (config.filename) { @@ -477,6 +479,10 @@ export class IncrementalTable extends ActionBuilder { unverifiedConfig.requirePartitionFilter = unverifiedConfig.bigquery.requirePartitionFilter; } + if (!!unverifiedConfig.bigquery.withConnection) { + unverifiedConfig.withConnection = + unverifiedConfig.bigquery.withConnection; + } if (!!unverifiedConfig.bigquery.additionalOptions) { unverifiedConfig.additionalOptions = unverifiedConfig.bigquery.additionalOptions; } diff --git a/core/actions/table.ts b/core/actions/table.ts index 804822d03..272390cc1 100644 --- a/core/actions/table.ts +++ b/core/actions/table.ts @@ -96,6 +96,11 @@ export interface IBigQueryOptions { */ requirePartitionFilter?: boolean; + /** + * + */ + withConnection?: string; + /** * Key-value pairs for options [table](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#table_option_list), [view](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#view_option_list), [materialized view](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#materialized_view_option_list). * @@ -114,6 +119,7 @@ const IBigQueryOptionsProperties = () => "labels", "partitionExpirationDays", "requirePartitionFilter", + "withConnection", "additionalOptions" ]); @@ -319,6 +325,7 @@ export class Table extends ActionBuilder { config.requirePartitionFilter || config.clusterBy.length || Object.keys(config.labels).length || + config.withConnection || Object.keys(config.additionalOptions).length ? {} : undefined; @@ -338,6 +345,9 @@ export class Table extends ActionBuilder { if (Object.keys(config.labels).length) { bigqueryOptions.labels = config.labels; } + if (config.withConnection) { + bigqueryOptions.withConnection = config.withConnection; + } if (Object.keys(config.additionalOptions).length) { bigqueryOptions.additionalOptions = config.additionalOptions; } @@ -402,6 +412,7 @@ export class Table extends ActionBuilder { config.updatePartitionFilter || config.clusterBy.length || Object.keys(config.labels).length || + config.withConnection || Object.keys(config.additionalOptions).length ? {} : undefined; @@ -421,6 +432,9 @@ export class Table extends ActionBuilder { if (config.clusterBy.length) { bigqueryOptions.clusterBy = config.clusterBy; } + if (config.withConnection) { + bigqueryOptions.withConnection = config.withConnection + } if (Object.keys(config.labels).length) { bigqueryOptions.labels = config.labels; } diff --git a/core/main_test.ts b/core/main_test.ts index db46db07a..b387052df 100644 --- a/core/main_test.ts +++ b/core/main_test.ts @@ -1860,6 +1860,7 @@ ${exampleActionDescriptor.inputSqlxConfigBlock} partitionExpirationDays: 1, requirePartitionFilter: true, clusterBy: ["clusterBy"], + withConnection: "US.external_service_connection", labels: {"key": "val"}, additionalOptions: { option1Key: "option1", @@ -1925,7 +1926,8 @@ SELECT 1` }, partitionBy: "partitionBy", partitionExpirationDays: 1, - requirePartitionFilter: true + requirePartitionFilter: true, + withConnection: "US.external_service_connection" }, tags: ["tag1", "tag2"], dependencyTargets: [ @@ -2073,6 +2075,7 @@ SELECT 1` updatePartitionFilter: "updatePartitionFilter", clusterBy: ["clusterBy"], labels: {"key": "val"}, + withConnection: "US.external_service_connection", additionalOptions: { option1Key: "option1", option2Key: "option2", @@ -2139,7 +2142,8 @@ SELECT 1` partitionBy: "partitionBy", partitionExpirationDays: 1, requirePartitionFilter: true, - updatePartitionFilter: "updatePartitionFilter" + updatePartitionFilter: "updatePartitionFilter", + withConnection: "US.external_service_connection" }, tags: ["tag1", "tag2"], uniqueKey: ["key1", "key2"], diff --git a/core/session.ts b/core/session.ts index cea1746b6..94c7642d5 100644 --- a/core/session.ts +++ b/core/session.ts @@ -585,7 +585,9 @@ export class Session { (table.bigquery.partitionBy || table.bigquery.clusterBy?.length || table.bigquery.partitionExpirationDays || - table.bigquery.requirePartitionFilter) && + table.bigquery.requirePartitionFilter || + table.bigquery.withConnection + ) && table.enumType === dataform.TableType.VIEW && !table.materialized ) { diff --git a/protos/configs.proto b/protos/configs.proto index 802d93424..83814d4ed 100644 --- a/protos/configs.proto +++ b/protos/configs.proto @@ -188,6 +188,8 @@ message ActionConfig { // the action depends on data from a source which has not been declared as // a dependency. bool hermetic = 20; + + string withConnection = 21; } message ViewConfig { @@ -346,6 +348,8 @@ message ActionConfig { // the action depends on data from a source which has not been declared as // a dependency. bool hermetic = 23; + + string withConnection = 24; } message AssertionConfig { diff --git a/protos/core.proto b/protos/core.proto index c0f53c228..fa86e093a 100644 --- a/protos/core.proto +++ b/protos/core.proto @@ -59,6 +59,7 @@ message BigQueryOptions { int32 partition_expiration_days = 5; bool require_partition_filter = 6; map additional_options = 7; + string withConnection = 8; } message GraphErrors { diff --git a/tests/api/api.spec.ts b/tests/api/api.spec.ts index 8fd7bbae7..a9fe03419 100644 --- a/tests/api/api.spec.ts +++ b/tests/api/api.spec.ts @@ -703,6 +703,73 @@ suite("@dataform/api", () => { ); }); + test("bigquery_with_connection", () => { + const testGraph: dataform.ICompiledGraph = dataform.CompiledGraph.create({ + projectConfig: { warehouse: "bigquery", defaultDatabase: "deeb", defaultLocation: "US" }, + tables: [ + { + target: { + schema: "schema", + name: "with_connection" + }, + type: "table", + query: "select 1 as test", + bigquery: { + withConnection: "with_connection" + } + }, + { + target: { + schema: "schema", + name: "plain" + }, + type: "table", + query: "select 1 as test" + } + ] + }); + const expectedExecutionActions: dataform.IExecutionAction[] = [ + { + type: "table", + tableType: "table", + target: { + schema: "schema", + name: "with_connection" + }, + tasks: [ + { + type: "statement", + statement: + 'create or replace table `deeb.schema.additional_options` WITH CONNECTION with_connection as select 1 as test' + } + ], + dependencyTargets: [], + hermeticity: dataform.ActionHermeticity.HERMETIC + }, + { + type: "table", + tableType: "table", + target: { + schema: "schema", + name: "plain" + }, + tasks: [ + { + type: "statement", + statement: "create or replace table `deeb.schema.plain` as select 1 as test" + } + ], + dependencyTargets: [], + hermeticity: dataform.ActionHermeticity.HERMETIC + } + ]; + const executionGraph = new Builder(testGraph, {}, dataform.WarehouseState.create({})).build(); + expect(asPlainObject(executionGraph.actions)).deep.equals( + asPlainObject(expectedExecutionActions) + ); + }); + }); + test("bigquery_additional_options", () => { const testGraph: dataform.ICompiledGraph = dataform.CompiledGraph.create({ projectConfig: { warehouse: "bigquery", defaultDatabase: "deeb", defaultLocation: "US" },