This is an automated email from the ASF dual-hosted git repository.
reuvenlax pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new de8d83c107d Merge pull request #40259 from
reuvenlax/improve_gcp_presubmit_speed
de8d83c107d is described below
commit de8d83c107d6af6df3cb24769650b537c98cf812
Author: Reuven Lax <[email protected]>
AuthorDate: Fri Sep 25 15:48:48 2026 -0700
Merge pull request #40259 from reuvenlax/improve_gcp_presubmit_speed
Speed up Beam presubmits
---
.../beam_PreCommit_Java_GCP_IO_Direct.yml | 48 ++++----
sdks/java/io/google-cloud-platform/build.gradle | 128 +++++++++++++--------
.../BigQueryIOWriteStorageApiBatchTest.java | 35 ++++++
.../BigQueryIOWriteStorageApiStreamTest.java | 35 ++++++
.../sdk/io/gcp/bigquery/BigQueryIOWriteTest.java | 9 +-
.../dofn/ReadChangeStreamPartitionDoFnTest.java | 14 ++-
.../beam/sdk/io/gcp/spanner/SpannerReadIT.java | 52 +++++----
.../beam/sdk/io/gcp/spanner/SpannerWriteIT.java | 57 +++++----
.../changestreams/it/IntegrationTestEnv.java | 80 +++++++------
9 files changed, 289 insertions(+), 169 deletions(-)
diff --git a/.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml
b/.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml
index 9ebb56b76d5..3e2c9c16248 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/[email protected]
- 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 ff648b4f903..b1ad17c006d 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 00000000000..1fe6bd350fc
--- /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<Object[]> 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 00000000000..3fff3bd31ae
--- /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<Object[]> 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 fdfd2699236..62e57ec848c 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
@@ -229,14 +229,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<Object[]> 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 49540e20594..f157eea476a 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 class ReadChangeStreamPartitionDoFnTest {
@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 class ReadChangeStreamPartitionDoFnTest {
"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 class ReadChangeStreamPartitionDoFnTest {
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 247f403e252..988a56aa113 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.PCollectionView;
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 class SpannerReadIT {
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 class SpannerReadIT {
+ " 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<Database, CreateDatabaseMetadata> 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 class SpannerReadIT {
p.run();
}
- private void makeTestData() {
+ private static void makeTestData() {
DatabaseClient databaseClient = getDatabaseClient();
DatabaseClient pgDatabaseClient = getPgDatabaseClient();
@@ -496,23 +496,27 @@ public class SpannerReadIT {
.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 89c2a194b93..318237e1118 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.Database;
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.Predicate
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 class SpannerWriteIT {
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 class SpannerWriteIT {
+ " Key INT64,"
+ " Value STRING(MAX) NOT NULL,"
+ ") PRIMARY KEY (Key)"));
- op.get();
- databaseAdminClient
- .createDatabase(
+ OperationFuture<Database, CreateDatabaseMetadata> 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 class SpannerWriteIT {
.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 class SpannerWriteIT {
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<Long, Mutation> {
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 b00568f8d40..3571a417cf0 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 @@ public class IntegrationTestEnv extends ExternalResource {
@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);
}
}
}
