From 882a72d64456633394cd140291bbe250fd774ad5 Mon Sep 17 00:00:00 2001 From: reuvenlax Date: Fri, 25 Sep 2026 15:48:48 -0700 Subject: [PATCH] Merge pull request #40259 from reuvenlax/improve_gcp_presubmit_speed Speed up Beam presubmits --- .../beam_PreCommit_Java_GCP_IO_Direct.yml | 48 +++---- .../io/google-cloud-platform/build.gradle | 128 +++++++++++------- .../BigQueryIOWriteStorageApiBatchTest.java | 35 +++++ .../BigQueryIOWriteStorageApiStreamTest.java | 35 +++++ .../io/gcp/bigquery/BigQueryIOWriteTest.java | 9 +- .../ReadChangeStreamPartitionDoFnTest.java | 14 +- .../sdk/io/gcp/spanner/SpannerReadIT.java | 52 +++---- .../sdk/io/gcp/spanner/SpannerWriteIT.java | 57 +++++--- .../changestreams/it/IntegrationTestEnv.java | 80 ++++++----- 9 files changed, 289 insertions(+), 169 deletions(-) create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiBatchTest.java create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiStreamTest.java diff --git a/.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml b/.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml index 9ebb56b76d5e..3e2c9c162489 100644 --- a/.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml +++ b/.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml @@ -82,9 +82,26 @@ jobs: beam_PreCommit_Java_GCP_IO_Direct: name: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) strategy: + fail-fast: false matrix: - job_name: ["beam_PreCommit_Java_GCP_IO_Direct"] - job_phrase: ["Run Java_GCP_IO_Direct PreCommit"] + include: + - job_name: "beam_PreCommit_Java_GCP_IO_Direct" + job_phrase: "Run Java_GCP_IO_Direct PreCommit" + gradle_command: | + :sdks:java:io:google-cloud-platform:build \ + :sdks:java:io:google-cloud-platform:expansion-service:build \ + extra_arguments: "" + is_build_job: true + - job_name: "beam_PreCommit_Java_GCP_IO_Direct_BigQuery" + job_phrase: "Run Java_GCP_IO_Direct PreCommit" + gradle_command: ":sdks:java:io:google-cloud-platform:postCommitBigQuery" + extra_arguments: "-x test" + is_build_job: false + - job_name: "beam_PreCommit_Java_GCP_IO_Direct_OtherGCP" + job_phrase: "Run Java_GCP_IO_Direct PreCommit" + gradle_command: ":sdks:java:io:google-cloud-platform:postCommitOtherGcp" + extra_arguments: "-x test" + is_build_job: false timeout-minutes: 180 if: | github.event_name == 'push' || @@ -108,19 +125,16 @@ jobs: - name: run PreCommit Java GCP IO Direct script uses: ./.github/actions/gradle-command-self-hosted-action with: - gradle-command: | - :sdks:java:io:google-cloud-platform:build \ - :sdks:java:io:google-cloud-platform:expansion-service:build \ - :sdks:java:io:google-cloud-platform:postCommit \ + gradle-command: ${{ matrix.gradle_command }} arguments: | -PdisableSpotlessCheck=true \ -PdisableCheckStyle=true \ - -PenableJacocoReport \ + ${{ matrix.extra_arguments }} - name: Archive JUnit Test Results uses: actions/upload-artifact@v7 if: ${{ !success() }} with: - name: JUnit Test Results + name: JUnit Test Results (${{ matrix.job_name }}) path: "**/build/reports/tests/" - name: Publish JUnit Test Results uses: EnricoMi/publish-unit-test-result-action@v2 @@ -132,27 +146,13 @@ jobs: large_files: true - name: Archive SpotBugs Results uses: actions/upload-artifact@v7 - if: always() + if: ${{ always() && matrix.is_build_job }} with: name: SpotBugs Results path: '**/build/reports/spotbugs/*.html' - name: Publish SpotBugs Results uses: jwgmeligmeyling/spotbugs-github-action@v1.2 - if: always() + if: ${{ always() && matrix.is_build_job }} with: name: Publish SpotBugs path: '**/build/reports/spotbugs/*.html' - - name: Set Jacoco report path - id: jacoco_report_path - run: | - echo "path=sdks/java/io/google-cloud-platform/build/reports/jacoco/test/jacocoTestReport.xml" >> $GITHUB_OUTPUT - - name: Archive Jacoco Results - uses: actions/upload-artifact@v7 - with: - name: Jacoco Results - path: '**/build/jacoco/report/**' - - name: Upload coverage reports to Codecov - uses: codecov/codecov-action@v3 - with: - file: ${{ steps.jacoco_report_path.outputs.path }} - flags: java diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle index ff648b4f9033..b1ad17c006df 100644 --- a/sdks/java/io/google-cloud-platform/build.gradle +++ b/sdks/java/io/google-cloud-platform/build.gradle @@ -33,6 +33,14 @@ ext.summary = "IO library to read and write Google Cloud Platform systems from B tasks.withType(Test).configureEach { // Open java.nio for Apache Arrow direct memory access under Java 17+ jvmArgs '--add-opens=java.base/java.nio=ALL-UNNAMED' + maxParallelForks 4 + jacoco.enabled = false +} + +gradle.taskGraph.whenReady { + tasks.withType(Test).configureEach { + jacoco.enabled = false + } } dependencies { @@ -328,62 +336,86 @@ task bigQueryEarlyRolloutIntegrationTest(type: Test, dependsOn: processTestResou testClassesDirs = sourceSets.test.output.classesDirs } -// path(s) for Cloud Spanner related classes -def spannerIncludes = [ - '**/org/apache/beam/sdk/io/gcp/spanner/**', -] - -// exclude auto-generated classes and integration tests -def jacocoExcludes = [ - '**/AutoValue_*', - '**/*IT*', -] - -task spannerCodeCoverageReport(type: JacocoReport, dependsOn: test) { - group = "Reporting" - description = "Generates code coverage report for Cloud Spanner related classes" - classDirectories.setFrom(files(files(project.sourceSets.main.output).collect { - project.fileTree( - dir: it, - includes: spannerIncludes, - excludes: jacocoExcludes) - })) - sourceDirectories.setFrom(files(project.sourceSets.main.allSource.srcDirs)) - executionData.setFrom(file("${buildDir}/jacoco/test.exec")) - reports { - html.getRequired().set(true) - html.getOutputLocation().set(file("${buildDir}/reports/jacoco/spanner/")) +task integrationTestBigQuery(type: Test, dependsOn: processTestResources) { + group = "Verification" + systemProperty "beamTestPipelineOptions", integrationTest.systemProperties["beamTestPipelineOptions"] + outputs.upToDateWhen { false } + include '**/bigquery/**/*IT.class' + exclude '**/BigQueryIOIcebergManagedTableIT.class' + exclude '**/BigQueryIOReadIT.class' + exclude '**/BigQueryIOStorageQueryIT.class' + exclude '**/BigQueryIOStorageReadIT.class' + exclude '**/BigQueryIOStorageWriteIT.class' + exclude '**/BigQueryIODynamicQueryIT.class' + exclude '**/BigQueryIODynamicReadIT.class' + exclude '**/BigQueryToTableIT.class' + maxParallelForks 4 + classpath = sourceSets.test.runtimeClasspath + testClassesDirs = sourceSets.test.output.classesDirs + useJUnit { + excludeCategories "org.apache.beam.sdk.testing.UsesKms" + } +} + +task integrationTestBigQueryKms(type: Test, dependsOn: processTestResources) { + group = "Verification" + systemProperty "beamTestPipelineOptions", integrationTestKms.systemProperties["beamTestPipelineOptions"] + outputs.upToDateWhen { false } + include '**/bigquery/**/*IT.class' + maxParallelForks 4 + classpath = sourceSets.test.runtimeClasspath + testClassesDirs = sourceSets.test.output.classesDirs + useJUnit { + includeCategories "org.apache.beam.sdk.testing.UsesKms" } } -task spannerCodeCoverageVerification(type: JacocoCoverageVerification, dependsOn: spannerCodeCoverageReport) { +task integrationTestOtherGcp(type: Test, dependsOn: processTestResources) { group = "Verification" - description = "Enforces code coverage verification for Cloud Spanner related classes" - classDirectories.setFrom(files(files(project.sourceSets.main.output).collect { - project.fileTree( - dir: it, - includes: spannerIncludes, - excludes: jacocoExcludes) - })) - sourceDirectories.setFrom(files(project.sourceSets.main.allSource.srcDirs)) - executionData.setFrom(file("${buildDir}/jacoco/test.exec")) - violationRules { - failOnViolation = true - rule { - element = 'BUNDLE' - - limit { - value = 'COVEREDRATIO' - counter = 'INSTRUCTION' - minimum = 0.60 - } + systemProperty "beamTestPipelineOptions", integrationTest.systemProperties["beamTestPipelineOptions"] + outputs.upToDateWhen { false } + include '**/*IT.class' + exclude '**/bigquery/**/*IT.class' + maxParallelForks 4 + classpath = sourceSets.test.runtimeClasspath + testClassesDirs = sourceSets.test.output.classesDirs + useJUnit { + excludeCategories "org.apache.beam.sdk.testing.UsesKms" + filter { + // https://github.com/apache/beam/issues/32071 + excludeTestsMatching 'org.apache.beam.sdk.io.gcp.bigtable.BigtableReadIT.testE2EBigtableSegmentRead' } } } -// make the check task depend on spannerCodeCoverageVerification so that the build -// fails when the code coverage threshold is violated. -project.check.dependsOn "spannerCodeCoverageVerification" +task integrationTestOtherGcpKms(type: Test, dependsOn: processTestResources) { + group = "Verification" + systemProperty "beamTestPipelineOptions", integrationTestKms.systemProperties["beamTestPipelineOptions"] + outputs.upToDateWhen { false } + include '**/*IT.class' + exclude '**/bigquery/**/*IT.class' + maxParallelForks 4 + classpath = sourceSets.test.runtimeClasspath + testClassesDirs = sourceSets.test.output.classesDirs + useJUnit { + includeCategories "org.apache.beam.sdk.testing.UsesKms" + } +} + +task postCommitBigQuery { + group = "Verification" + description = "Integration tests of BigQuery connectors using the DirectRunner." + dependsOn integrationTestBigQuery + dependsOn integrationTestBigQueryKms + dependsOn bigQueryManagedIcebergIntegrationTest +} + +task postCommitOtherGcp { + group = "Verification" + description = "Integration tests of non-BigQuery GCP connectors using the DirectRunner." + dependsOn integrationTestOtherGcp + dependsOn integrationTestOtherGcpKms +} task postCommit { group = "Verification" diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiBatchTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiBatchTest.java new file mode 100644 index 000000000000..1fe6bd350fcc --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiBatchTest.java @@ -0,0 +1,35 @@ +/* + * 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 org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; + +/** + * Runs {@link BigQueryIOWriteTest} for Storage Write API batch mode in a separate test class so + * Gradle {@code maxParallelForks} can execute it in parallel with other modes. + */ +@RunWith(Parameterized.class) +public class BigQueryIOWriteStorageApiBatchTest extends BigQueryIOWriteTest { + @Parameters(name = "useStorageApi={0}, useStorageApiApproximate={1}, useStreaming={2}") + public static Iterable data() { + return ImmutableList.of(new Object[] {true, false, false}); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiStreamTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiStreamTest.java new file mode 100644 index 000000000000..3fff3bd31ae3 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiStreamTest.java @@ -0,0 +1,35 @@ +/* + * 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 org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; + +/** + * Runs {@link BigQueryIOWriteTest} for Storage Write API streaming modes in a separate test class + * so Gradle {@code maxParallelForks} can execute them in parallel with other modes. + */ +@RunWith(Parameterized.class) +public class BigQueryIOWriteStorageApiStreamTest extends BigQueryIOWriteTest { + @Parameters(name = "useStorageApi={0}, useStorageApiApproximate={1}, useStreaming={2}") + public static Iterable data() { + return ImmutableList.of(new Object[] {true, false, true}, new Object[] {true, true, true}); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 7310e2f850e5..3f61e61fbd66 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -228,14 +228,9 @@ public class BigQueryIOWriteTest implements Serializable { private transient TemporaryFolder testFolder = new TemporaryFolder(); private transient TestPipeline p; - @Parameters + @Parameters(name = "useStorageApi={0}, useStorageApiApproximate={1}, useStreaming={2}") public static Iterable data() { - return ImmutableList.of( - new Object[] {false, false, false}, - new Object[] {false, false, true}, - new Object[] {true, false, false}, - new Object[] {true, false, true}, - new Object[] {true, true, true}); + return ImmutableList.of(new Object[] {false, false, false}, new Object[] {false, false, true}); } @Parameter(0) diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java index 49540e20594a..f157eea476aa 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java @@ -103,8 +103,13 @@ public void setup() throws IOException { @Test public void testProcessElementAndGetSize() throws IOException, InterruptedException { + Instant now = Instant.now(); + doFn = + new ReadChangeStreamPartitionDoFn( + daoFactory, actionFactory, metrics, Duration.ZERO, () -> now); + doFn.setSizeEstimator(sizeEstimator); long watermarkLag = 10; - Instant tenSecondsAgo = Instant.now().minus(Duration.standardSeconds(watermarkLag)); + Instant tenSecondsAgo = now.minus(Duration.standardSeconds(watermarkLag)); Range.ByteStringRange partitionRange = Range.ByteStringRange.create("", ""); ChangeStreamContinuationToken testToken = ChangeStreamContinuationToken.create(partitionRange, "test"); @@ -115,7 +120,7 @@ public void testProcessElementAndGetSize() throws IOException, InterruptedExcept "uid-a", tenSecondsAgo, Collections.emptyList(), - Instant.now().plus(Duration.standardSeconds(60))); + now.plus(Duration.standardSeconds(60))); long mutationSize = 100L; when(sizeEstimator.sizeOf(any())).thenReturn(mutationSize); ReadChangeStreamPartitionProgressTracker restrictionTracker = @@ -149,11 +154,10 @@ public void testProcessElementAndGetSize() throws IOException, InterruptedExcept doFn.processElement(partition, restrictionTracker, receiver, watermarkEstimator); double sizeEstimate = doFn.getSize( - new StreamProgress( - testToken, tenSecondsAgo, BigDecimal.valueOf(20), Instant.now(), false)); + new StreamProgress(testToken, tenSecondsAgo, BigDecimal.valueOf(20), now, false)); // we should have output 2 100B mutations in the past 10s long bytesPerSecond = (mutationSize * 2) / 10; - assertEquals(sizeEstimate, (double) (bytesPerSecond * watermarkLag), 10); + assertEquals((double) (bytesPerSecond * watermarkLag), sizeEstimate, 10); verify(receiver, times(2)).outputWithTimestamp(KV.of(rowKey, mockMutation), Instant.EPOCH); } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java index 247f403e2523..988a56aa1131 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java @@ -51,8 +51,8 @@ import org.apache.beam.sdk.values.TypeDescriptor; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; -import org.junit.After; -import org.junit.Before; +import org.junit.AfterClass; +import org.junit.BeforeClass; import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; @@ -96,15 +96,15 @@ public interface SpannerTestPipelineOptions extends TestPipelineOptions { void setTable(String value); } - private Spanner spanner; - private DatabaseAdminClient databaseAdminClient; - private SpannerTestPipelineOptions options; - private String databaseName; - private String pgDatabaseName; - private String project; + private static Spanner spanner; + private static DatabaseAdminClient databaseAdminClient; + private static SpannerTestPipelineOptions options; + private static String databaseName; + private static String pgDatabaseName; + private static String project; - @Before - public void setUp() throws Exception { + @BeforeClass + public static void setUpTestEnvironment() throws Exception { PipelineOptionsFactory.register(SpannerTestPipelineOptions.class); options = TestPipeline.testingPipelineOptions().as(SpannerTestPipelineOptions.class); @@ -142,16 +142,16 @@ public void setUp() throws Exception { + " Key INT64," + " Value STRING(MAX)," + ") PRIMARY KEY (Key)")); - op.get(); // PG-dialect databases need 2-steps to create: Create DB then update DDL. - databaseAdminClient - .createDatabase( + OperationFuture pgOp = + databaseAdminClient.createDatabase( databaseAdminClient .newDatabaseBuilder(DatabaseId.of(project, options.getInstanceId(), pgDatabaseName)) .setDialect(Dialect.POSTGRESQL) .build(), - Collections.emptyList()) - .get(); + Collections.emptyList()); + op.get(); + pgOp.get(); databaseAdminClient .updateDatabaseDdl( options.getInstanceId(), @@ -461,7 +461,7 @@ public void testReadAllRecordsInDb() throws Exception { p.run(); } - private void makeTestData() { + private static void makeTestData() { DatabaseClient databaseClient = getDatabaseClient(); DatabaseClient pgDatabaseClient = getPgDatabaseClient(); @@ -496,23 +496,27 @@ private SpannerConfig createPgSpannerConfig() { .withDatabaseId(pgDatabaseName)); } - private DatabaseClient getDatabaseClient() { + private static DatabaseClient getDatabaseClient() { return spanner.getDatabaseClient(DatabaseId.of(project, options.getInstanceId(), databaseName)); } - private DatabaseClient getPgDatabaseClient() { + private static DatabaseClient getPgDatabaseClient() { return spanner.getDatabaseClient( DatabaseId.of(project, options.getInstanceId(), pgDatabaseName)); } - @After - public void tearDown() throws Exception { - databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName); - databaseAdminClient.dropDatabase(options.getInstanceId(), pgDatabaseName); - spanner.close(); + @AfterClass + public static void tearDown() throws Exception { + if (databaseAdminClient != null) { + databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName); + databaseAdminClient.dropDatabase(options.getInstanceId(), pgDatabaseName); + } + if (spanner != null) { + spanner.close(); + } } - private String generateDatabaseName() { + private static String generateDatabaseName() { String random = RandomUtils.randomAlphaNumeric( MAX_DB_NAME_LENGTH - 4 - options.getDatabaseIdPrefix().length()); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java index 89c2a194b937..318237e1118b 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java @@ -28,6 +28,7 @@ import com.google.cloud.spanner.DatabaseAdminClient; import com.google.cloud.spanner.DatabaseId; import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.KeySet; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.SessionPoolOptions; @@ -61,8 +62,9 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; import org.checkerframework.checker.nullness.qual.Nullable; import org.hamcrest.TypeSafeMatcher; -import org.junit.After; +import org.junit.AfterClass; import org.junit.Before; +import org.junit.BeforeClass; import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; @@ -105,15 +107,15 @@ public interface SpannerTestPipelineOptions extends TestPipelineOptions { void setTable(String value); } - private Spanner spanner; - private DatabaseAdminClient databaseAdminClient; - private SpannerTestPipelineOptions options; - private String databaseName; - private String pgDatabaseName; - private String project; + private static Spanner spanner; + private static DatabaseAdminClient databaseAdminClient; + private static SpannerTestPipelineOptions options; + private static String databaseName; + private static String pgDatabaseName; + private static String project; - @Before - public void setUp() throws Exception { + @BeforeClass + public static void setUpTestEnvironment() throws Exception { PipelineOptionsFactory.register(SpannerTestPipelineOptions.class); options = TestPipeline.testingPipelineOptions().as(SpannerTestPipelineOptions.class); @@ -152,15 +154,15 @@ public void setUp() throws Exception { + " Key INT64," + " Value STRING(MAX) NOT NULL," + ") PRIMARY KEY (Key)")); - op.get(); - databaseAdminClient - .createDatabase( + OperationFuture pgOp = + databaseAdminClient.createDatabase( databaseAdminClient .newDatabaseBuilder(DatabaseId.of(project, options.getInstanceId(), pgDatabaseName)) .setDialect(Dialect.POSTGRESQL) .build(), - Collections.emptyList()) - .get(); + Collections.emptyList()); + op.get(); + pgOp.get(); databaseAdminClient .updateDatabaseDdl( options.getInstanceId(), @@ -176,7 +178,18 @@ public void setUp() throws Exception { .get(); } - private String generateDatabaseName() { + @Before + public void setUp() { + Mutation deleteAll = Mutation.delete(options.getTable(), KeySet.all()); + spanner + .getDatabaseClient(DatabaseId.of(project, options.getInstanceId(), databaseName)) + .writeAtLeastOnce(Collections.singletonList(deleteAll)); + spanner + .getDatabaseClient(DatabaseId.of(project, options.getInstanceId(), pgDatabaseName)) + .writeAtLeastOnce(Collections.singletonList(deleteAll)); + } + + private static String generateDatabaseName() { String random = RandomUtils.randomAlphaNumeric( MAX_DB_NAME_LENGTH - 4 - options.getDatabaseIdPrefix().length()); @@ -390,11 +403,15 @@ public void testPgFailFast() throws Exception { result.waitUntilFinish(); } - @After - public void tearDown() throws Exception { - databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName); - databaseAdminClient.dropDatabase(options.getInstanceId(), pgDatabaseName); - spanner.close(); + @AfterClass + public static void tearDown() throws Exception { + if (databaseAdminClient != null) { + databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName); + databaseAdminClient.dropDatabase(options.getInstanceId(), pgDatabaseName); + } + if (spanner != null) { + spanner.close(); + } } private static class GenerateMutations extends DoFn { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/IntegrationTestEnv.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/IntegrationTestEnv.java index b00568f8d40c..3571a417cf0a 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/IntegrationTestEnv.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/IntegrationTestEnv.java @@ -119,52 +119,50 @@ protected void before() throws Throwable { @Override protected void after() { - for (String changeStream : changeStreams) { - try { - if (this.isPostgres) { - databaseAdminClient - .updateDatabaseDdl( - instanceId, - databaseId, - Collections.singletonList("DROP CHANGE STREAM \"" + changeStream + "\""), - null) - .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); - } else { - databaseAdminClient - .updateDatabaseDdl( - instanceId, - databaseId, - Collections.singletonList("DROP CHANGE STREAM " + changeStream), - null) - .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); + if (isPlacementTable) { + for (String changeStream : changeStreams) { + try { + if (this.isPostgres) { + databaseAdminClient + .updateDatabaseDdl( + instanceId, + databaseId, + Collections.singletonList("DROP CHANGE STREAM \"" + changeStream + "\""), + null) + .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); + } else { + databaseAdminClient + .updateDatabaseDdl( + instanceId, + databaseId, + Collections.singletonList("DROP CHANGE STREAM " + changeStream), + null) + .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); + } + } catch (Exception e) { + LOG.error("Failed to drop change stream {}. Skipping...", changeStream, e); } - } catch (Exception e) { - LOG.error("Failed to drop change stream {}. Skipping...", changeStream, e); } - } - for (String table : tables) { - try { - if (this.isPostgres) { - databaseAdminClient - .updateDatabaseDdl( - instanceId, - databaseId, - Collections.singletonList("DROP TABLE \"" + table + "\""), - null) - .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); - } else { - databaseAdminClient - .updateDatabaseDdl( - instanceId, databaseId, Collections.singletonList("DROP TABLE " + table), null) - .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); - } - } catch (Exception e) { - if (isPlacementTable) { + for (String table : tables) { + try { + if (this.isPostgres) { + databaseAdminClient + .updateDatabaseDdl( + instanceId, + databaseId, + Collections.singletonList("DROP TABLE \"" + table + "\""), + null) + .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); + } else { + databaseAdminClient + .updateDatabaseDdl( + instanceId, databaseId, Collections.singletonList("DROP TABLE " + table), null) + .get(TIMEOUT_MINUTES, TimeUnit.MINUTES); + } + } catch (Exception e) { // Drop placement table requires all rows deleted and garbage collected. LOG.info("Failed to drop table {}. Skipping...", table, e); - } else { - LOG.error("Failed to drop table {}. Skipping...", table, e); } } }