From b636dfebe9f3d3789e67bd710a8642236f9f3e4c Mon Sep 17 00:00:00 2001 From: liferoad Date: Thu, 24 Oct 2024 10:41:46 -0400 Subject: [PATCH 1/6] Added the queryTempDataset option for the BQ read options --- .../v2/transforms/BigQueryConverters.java | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java b/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java index 53de64259c..742f8d657f 100644 --- a/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java +++ b/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java @@ -229,6 +229,18 @@ public interface BigQueryReadOptions extends PipelineOptions { String getQueryLocation(); void setQueryLocation(String query); + + @TemplateParameter.Text( + order = 6, + optional = true, + description = "BigQuery temporary dataset reference when reading from the query.", + helpText = + "With this option, you can set an existing dataset to create the temporary table " + + "to store the results of the query.", + example = "temp_dataset") + String getQueryTempDataset(); + + void setQueryTempDataset(String queryTempDataset); } /** @@ -349,6 +361,7 @@ public PCollection expand(PBegin pipeline) { .fromQuery(options().getQuery()) .withTemplateCompatibility() .withQueryLocation(options().getQueryLocation()) + .withQueryTempDataset(options().getQueryTempDataset()) .usingStandardSql()); } else { @@ -358,7 +371,8 @@ public PCollection expand(PBegin pipeline) { readFunction() .fromQuery(options().getQuery()) .withTemplateCompatibility() - .withQueryLocation(options().getQueryLocation())); + .withQueryLocation(options().getQueryLocation()) + .withQueryTempDataset(options().getQueryTempDataset())); } } } From 94aab4816208d9dc99b05eb25601e701df0d82b4 Mon Sep 17 00:00:00 2001 From: liferoad Date: Thu, 24 Oct 2024 11:11:12 -0400 Subject: [PATCH 2/6] Fixed the spotless --- .../teleport/v2/transforms/BigQueryConverters.java | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java b/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java index 742f8d657f..5502aead88 100644 --- a/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java +++ b/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/BigQueryConverters.java @@ -231,13 +231,13 @@ public interface BigQueryReadOptions extends PipelineOptions { void setQueryLocation(String query); @TemplateParameter.Text( - order = 6, - optional = true, - description = "BigQuery temporary dataset reference when reading from the query.", - helpText = - "With this option, you can set an existing dataset to create the temporary table " - + "to store the results of the query.", - example = "temp_dataset") + order = 6, + optional = true, + description = "BigQuery temporary dataset reference when reading from the query.", + helpText = + "With this option, you can set an existing dataset to create the temporary table " + + "to store the results of the query.", + example = "temp_dataset") String getQueryTempDataset(); void setQueryTempDataset(String queryTempDataset); From d1c0277deb3c951d39e219079d403c55800b350b Mon Sep 17 00:00:00 2001 From: liferoad Date: Thu, 24 Oct 2024 13:05:39 -0400 Subject: [PATCH 3/6] Added one dummy test for BigQueryReadOptions --- .../v2/transforms/BigQueryConvertersTest.java | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java index cf7cb99a44..7591317bb6 100644 --- a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java +++ b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java @@ -346,6 +346,22 @@ public void testReadBigQueryInvalidInput() { pipeline.run(); } + /** + * Tests {@link BigQueryConverters.BigQueryReadOptions} works with some options. + */ + @Test + public void testBigQueryReadOptions() { + + BigQueryConverters.BigQueryReadOptions options = + PipelineOptionsFactory.create().as(BigQueryConverters.BigQueryReadOptions.class); + + options.setInputTableSpec(null); + options.setQuery("select * from sampledb.sample_table"); + options.setQueryTempDataset("temp_dataset"); + + } + + /** * Tests that {@link BigQueryConverters.TableRowToFailsafeJsonDocument} transform returns the * correct element. From d37cb4a17da73b6d641f6f214adb10ced8c26f3a Mon Sep 17 00:00:00 2001 From: liferoad Date: Thu, 24 Oct 2024 13:10:41 -0400 Subject: [PATCH 4/6] fixed the spotless --- .../teleport/v2/transforms/BigQueryConvertersTest.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java index 7591317bb6..ab0706a08a 100644 --- a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java +++ b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java @@ -346,22 +346,18 @@ public void testReadBigQueryInvalidInput() { pipeline.run(); } - /** - * Tests {@link BigQueryConverters.BigQueryReadOptions} works with some options. - */ + /** Tests {@link BigQueryConverters.BigQueryReadOptions} works with some options. */ @Test public void testBigQueryReadOptions() { BigQueryConverters.BigQueryReadOptions options = - PipelineOptionsFactory.create().as(BigQueryConverters.BigQueryReadOptions.class); + PipelineOptionsFactory.create().as(BigQueryConverters.BigQueryReadOptions.class); options.setInputTableSpec(null); options.setQuery("select * from sampledb.sample_table"); options.setQueryTempDataset("temp_dataset"); - } - /** * Tests that {@link BigQueryConverters.TableRowToFailsafeJsonDocument} transform returns the * correct element. From fab6cb69ace5f92669d2760c7885bfec85d0fed0 Mon Sep 17 00:00:00 2001 From: liferoad Date: Thu, 24 Oct 2024 13:26:27 -0400 Subject: [PATCH 5/6] fixed the test --- .../cloud/teleport/v2/transforms/BigQueryConvertersTest.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java index ab0706a08a..03d02d1311 100644 --- a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java +++ b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java @@ -356,6 +356,9 @@ public void testBigQueryReadOptions() { options.setInputTableSpec(null); options.setQuery("select * from sampledb.sample_table"); options.setQueryTempDataset("temp_dataset"); + + pipeline.apply( + BigQueryConverters.ReadBigQueryTableRows.newBuilder().setOptions(options).build()); } /** From a15c28ac54a2cadb017f8b03397a9d9cdec8f333 Mon Sep 17 00:00:00 2001 From: liferoad Date: Thu, 24 Oct 2024 14:31:12 -0400 Subject: [PATCH 6/6] fixed the test --- .../cloud/teleport/v2/transforms/BigQueryConvertersTest.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java index 03d02d1311..95aea4dd49 100644 --- a/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java +++ b/v2/common/src/test/java/com/google/cloud/teleport/v2/transforms/BigQueryConvertersTest.java @@ -357,8 +357,7 @@ public void testBigQueryReadOptions() { options.setQuery("select * from sampledb.sample_table"); options.setQueryTempDataset("temp_dataset"); - pipeline.apply( - BigQueryConverters.ReadBigQueryTableRows.newBuilder().setOptions(options).build()); + assertThat(options.getQueryTempDataset()).isEqualTo("temp_dataset"); } /**