diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle index df1eee310ef6..23859ded50e0 100644 --- a/sdks/java/io/google-cloud-platform/build.gradle +++ b/sdks/java/io/google-cloud-platform/build.gradle @@ -283,7 +283,8 @@ task bigQueryEarlyRolloutIntegrationTest(type: Test, dependsOn: processTestResou include '**/StorageApiDirectWriteProtosIT.class' include '**/StorageApiSinkFailedRowsIT.class' include '**/StorageApiSinkRowUpdateIT.class' - include '**/StorageApiSinkSchemaUpdateIT.class' + include '**/StorageApiSinkSchemaUpdateWithInputSchemaIT.class' + include '**/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.class' include '**/TableRowToStorageApiProtoIT.class' // file loads include '**/BigQuerySchemaUpdateOptionsIT.class' diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java similarity index 94% rename from sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java rename to sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java index 3cb9897ada52..c4aeffd25aae 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java @@ -70,44 +70,32 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.joda.time.Duration; import org.joda.time.Instant; -import org.junit.AfterClass; -import org.junit.BeforeClass; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TestName; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -@RunWith(Parameterized.class) -public class StorageApiSinkSchemaUpdateIT { - @Parameterized.Parameters - public static Iterable data() { - return ImmutableList.of( - new Object[] {false, false}, - new Object[] {false, true}, - new Object[] {true, false}, - new Object[] {true, true}); - } - - @Parameterized.Parameter(0) - public boolean useInputSchema; +abstract class StorageApiSinkSchemaUpdateITBase { + private final boolean useInputSchema; + private final boolean changeTableSchema; + private final String bigQueryDatasetId; - @Parameterized.Parameter(1) - public boolean changeTableSchema; + StorageApiSinkSchemaUpdateITBase( + boolean useInputSchema, boolean changeTableSchema, String bigQueryDatasetId) { + this.useInputSchema = useInputSchema; + this.changeTableSchema = changeTableSchema; + this.bigQueryDatasetId = bigQueryDatasetId; + } @Rule public TestName testName = new TestName(); - private static final Logger LOG = LoggerFactory.getLogger(StorageApiSinkSchemaUpdateIT.class); + private static final Logger LOG = LoggerFactory.getLogger(StorageApiSinkSchemaUpdateITBase.class); private static final BigqueryClient BQ_CLIENT = new BigqueryClient("StorageApiSinkSchemaChangeIT"); private static final String PROJECT = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); - private static final String BIG_QUERY_DATASET_ID = - "storage_api_sink_schema_change_" + System.nanoTime(); - private static final String[] FIELDS = { "BOOL", "BOOLEAN", @@ -149,18 +137,17 @@ public static Iterable data() { // used when test suite specifies a particular GCP location for BigQuery operations private static String bigQueryLocation; - @BeforeClass - public static void setUpTestEnvironment() throws IOException, InterruptedException { + static void setUpTestEnvironment(String bigQueryDatasetId) + throws IOException, InterruptedException { // Create one BQ dataset for all test cases. bigQueryLocation = TestPipeline.testingPipelineOptions().as(TestBigQueryOptions.class).getBigQueryLocation(); - BQ_CLIENT.createNewDataset(PROJECT, BIG_QUERY_DATASET_ID, null, bigQueryLocation); + BQ_CLIENT.createNewDataset(PROJECT, bigQueryDatasetId, null, bigQueryLocation); } - @AfterClass - public static void cleanUp() { - LOG.info("Cleaning up dataset {} and tables.", BIG_QUERY_DATASET_ID); - BQ_CLIENT.deleteDataset(PROJECT, BIG_QUERY_DATASET_ID); + static void cleanUp(String bigQueryDatasetId) { + LOG.info("Cleaning up dataset {} and tables.", bigQueryDatasetId); + BQ_CLIENT.deleteDataset(PROJECT, bigQueryDatasetId); } private String createTable(TableSchema tableSchema) throws IOException, InterruptedException { @@ -178,16 +165,16 @@ private String createTable(TableSchema tableSchema, String suffix) } tableId += suffix; - BQ_CLIENT.deleteTable(PROJECT, BIG_QUERY_DATASET_ID, tableId); + BQ_CLIENT.deleteTable(PROJECT, bigQueryDatasetId, tableId); BQ_CLIENT.createNewTable( PROJECT, - BIG_QUERY_DATASET_ID, + bigQueryDatasetId, new Table() .setSchema(tableSchema) .setTableReference( new TableReference() .setTableId(tableId) - .setDatasetId(BIG_QUERY_DATASET_ID) + .setDatasetId(bigQueryDatasetId) .setProjectId(PROJECT))); return tableId; } @@ -395,7 +382,7 @@ private void runStreamingPipelineWithSchemaChange( makeTableSchemaFromTypes(fieldNamesWithExtra, ImmutableSet.of(extraField)); String tableId = createTable(bqTableSchema); - String tableSpec = PROJECT + ":" + BIG_QUERY_DATASET_ID + "." + tableId; + String tableSpec = PROJECT + ":" + bigQueryDatasetId + "." + tableId; // build write transform Write write = @@ -462,7 +449,7 @@ private void runStreamingPipelineWithSchemaChange( "Update Schema", ParDo.of( new UpdateSchemaDoFn( - PROJECT, BIG_QUERY_DATASET_ID, ImmutableMap.of(tableId, updatedSchema)))); + PROJECT, bigQueryDatasetId, ImmutableMap.of(tableId, updatedSchema)))); } WriteResult result = rows.apply("Stream to BigQuery", write); if (useIgnoreUnknownValues) { @@ -648,7 +635,7 @@ public void runDynamicDestinationsWithAutoSchemaUpdate(boolean useAtLeastOnce) t GenerateRowFunc generateRowFunc = new GenerateRowFunc(fieldNamesOrigin, fieldNamesWithExtra); String tableId = createTable(bqTableSchema, "_dynamic_" + i); - String tableSpec = PROJECT + ":" + BIG_QUERY_DATASET_ID + "." + tableId; + String tableSpec = PROJECT + ":" + bigQueryDatasetId + "." + tableId; rowFuncs.put((long) i, generateRowFunc); destinations.put((long) i, tableSpec); @@ -727,7 +714,7 @@ public void runDynamicDestinationsWithAutoSchemaUpdate(boolean useAtLeastOnce) t .apply("Add a dummy key", WithKeys.of(1)) .apply( "Update Schema", - ParDo.of(new UpdateSchemaDoFn(PROJECT, BIG_QUERY_DATASET_ID, updatedSchemas))); + ParDo.of(new UpdateSchemaDoFn(PROJECT, bigQueryDatasetId, updatedSchemas))); } WriteResult result = rows.apply("Stream to BigQuery", write); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java new file mode 100644 index 000000000000..cd29d9fb65a1 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import java.io.IOException; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +@RunWith(Parameterized.class) +public class StorageApiSinkSchemaUpdateWithInputSchemaIT extends StorageApiSinkSchemaUpdateITBase { + private static final String BIG_QUERY_DATASET_ID = + "storage_api_sink_schema_change_with_input_" + System.nanoTime(); + + @Parameterized.Parameters(name = "changeTableSchema={0}") + public static Iterable data() { + return ImmutableList.of(new Object[] {false}, new Object[] {true}); + } + + public StorageApiSinkSchemaUpdateWithInputSchemaIT(boolean changeTableSchema) { + super(true, changeTableSchema, BIG_QUERY_DATASET_ID); + } + + @BeforeClass + public static void setUpTestEnvironment() throws IOException, InterruptedException { + StorageApiSinkSchemaUpdateITBase.setUpTestEnvironment(BIG_QUERY_DATASET_ID); + } + + @AfterClass + public static void cleanUp() { + StorageApiSinkSchemaUpdateITBase.cleanUp(BIG_QUERY_DATASET_ID); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java new file mode 100644 index 000000000000..9b02e0be396b --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import java.io.IOException; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +@RunWith(Parameterized.class) +public class StorageApiSinkSchemaUpdateWithoutInputSchemaIT + extends StorageApiSinkSchemaUpdateITBase { + private static final String BIG_QUERY_DATASET_ID = + "storage_api_sink_schema_change_without_input_" + System.nanoTime(); + + @Parameterized.Parameters(name = "changeTableSchema={0}") + public static Iterable data() { + return ImmutableList.of(new Object[] {false}, new Object[] {true}); + } + + public StorageApiSinkSchemaUpdateWithoutInputSchemaIT(boolean changeTableSchema) { + super(false, changeTableSchema, BIG_QUERY_DATASET_ID); + } + + @BeforeClass + public static void setUpTestEnvironment() throws IOException, InterruptedException { + StorageApiSinkSchemaUpdateITBase.setUpTestEnvironment(BIG_QUERY_DATASET_ID); + } + + @AfterClass + public static void cleanUp() { + StorageApiSinkSchemaUpdateITBase.cleanUp(BIG_QUERY_DATASET_ID); + } +}