From f069821055851479789922e2407d7b1815b52684 Mon Sep 17 00:00:00 2001 From: hiracky16 Date: Sun, 15 Dec 2024 15:23:57 +0900 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20add=20bigquery=20opt?= =?UTF-8?q?ion=20of=20with=20connection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cli/api/dbadapters/execution_sql.ts | 7 ++- core/actions/incremental_table.ts | 6 +++ core/actions/table.ts | 14 ++++++ core/main_test.ts | 4 +- core/session.ts | 4 +- protos/configs.proto | 4 ++ protos/core.proto | 1 + tests/api/api.spec.ts | 67 +++++++++++++++++++++++++++++ 8 files changed, 104 insertions(+), 3 deletions(-) diff --git a/cli/api/dbadapters/execution_sql.ts b/cli/api/dbadapters/execution_sql.ts index c54fd6143..c54dd87b8 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 e9291d4d4..ca4eb9e17 100644 --- a/core/actions/incremental_table.ts +++ b/core/actions/incremental_table.ts @@ -44,6 +44,7 @@ interface ILegacyIncrementalTableBigqueryConfig { labels?: { [key: string]: string }; partitionExpirationDays?: number; requirePartitionFilter?: boolean; + withConnection?: string; additionalOptions?: { [key: string]: string }; } @@ -159,6 +160,7 @@ export class IncrementalTable extends ActionBuilder { labels: config.labels, partitionExpirationDays: config.partitionExpirationDays, requirePartitionFilter: config.requirePartitionFilter, + withConnection: config.withConnection, additionalOptions: config.additionalOptions }); if (config.filename) { @@ -488,6 +490,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 3f639b177..f7400b5b8 100644 --- a/core/main_test.ts +++ b/core/main_test.ts @@ -1892,6 +1892,7 @@ ${exampleActionDescriptor.inputSqlxConfigBlock} partitionExpirationDays: 1, requirePartitionFilter: true, clusterBy: ["clusterBy"], + withConnection: "US.external_service_connection", labels: {"key": "val"}, additionalOptions: { option1Key: "option1", @@ -1928,7 +1929,8 @@ SELECT 1`; }, partitionBy: "partitionBy", partitionExpirationDays: 1, - requirePartitionFilter: true + requirePartitionFilter: true, + withConnection: "US.external_service_connection" }, tags: ["tag1", "tag2"], dependencyTargets: [ 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" }, From cc443a76faafa432a6d387a45c737875bbb1a211 Mon Sep 17 00:00:00 2001 From: hiracky16 Date: Sun, 15 Dec 2024 15:47:05 +0900 Subject: [PATCH 2/3] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20add=20test=20case?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/main_test.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/core/main_test.ts b/core/main_test.ts index f7400b5b8..92535044e 100644 --- a/core/main_test.ts +++ b/core/main_test.ts @@ -2096,6 +2096,7 @@ config { updatePartitionFilter: "updatePartitionFilter", clusterBy: ["clusterBy"], labels: {"key": "val"}, + withConnection: "US.external_service_connection", additionalOptions: { option1Key: "option1", option2Key: "option2", @@ -2133,7 +2134,8 @@ SELECT 1` partitionBy: "partitionBy", partitionExpirationDays: 1, requirePartitionFilter: true, - updatePartitionFilter: "updatePartitionFilter" + updatePartitionFilter: "updatePartitionFilter", + withConnection: "US.external_service_connection" }, tags: ["tag1", "tag2"], uniqueKey: ["key1", "key2"], From f22de260b0177fb54968a6486826ed28362527b2 Mon Sep 17 00:00:00 2001 From: hiracky16 Date: Wed, 18 Dec 2024 06:27:27 +0900 Subject: [PATCH 3/3] =?UTF-8?q?fix:=20=F0=9F=90=9B=20upper=20case?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cli/api/dbadapters/execution_sql.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/api/dbadapters/execution_sql.ts b/cli/api/dbadapters/execution_sql.ts index c54dd87b8..fa0035af3 100644 --- a/cli/api/dbadapters/execution_sql.ts +++ b/cli/api/dbadapters/execution_sql.ts @@ -208,7 +208,7 @@ from (${query}) as insertions`; : "" }${ table.bigquery && table.bigquery.withConnection - ? `with connection ${table.bigquery.withConnection} ` + ? `WITH CONNECTION ${table.bigquery.withConnection} ` : "" } ${options.length > 0 ? `OPTIONS(${options.join(",")})` : ""}as ${table.query}`;