From 202d339ad9f13b13c03bafab5509c2ed64487f81 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Thu, 25 Jun 2026 19:30:03 +0200 Subject: [PATCH 01/29] Fix OpenSearch env fallback for username/password in examples --- .../notebooks/beam-ml/rag_usecase/opensearch_connector.py | 8 ++++---- .../beam-ml/rag_usecase/opensearch_enrichment.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py b/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py index fc83c8d443cc..e9f07daef101 100644 --- a/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py +++ b/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py @@ -88,8 +88,8 @@ def __init__(self, """ self.host = host self.port = port - self.username = username | os.getenv("OPENSEARCH_USERNAME") - self.password = password | os.getenv("OPENSEARCH_PASSWORD") + self.username = username or os.getenv("OPENSEARCH_USERNAME") + self.password = password or os.getenv("OPENSEARCH_PASSWORD") self._batch_size = batch_size if not self.username or not self.password: @@ -247,8 +247,8 @@ def __init__(self, """ self.host = host self.port = port - self.username = username | os.getenv("OPENSEARCH_USERNAME") - self.password = password | os.getenv("OPENSEARCH_PASSWORD") + self.username = username or os.getenv("OPENSEARCH_USERNAME") + self.password = password or os.getenv("OPENSEARCH_PASSWORD") self.batch_size = batch_size self.embedded_columns = embedded_columns diff --git a/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py b/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py index 70397550241f..c43d9af181ba 100644 --- a/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py +++ b/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py @@ -77,8 +77,8 @@ def __init__( """ self.opensearch_host = opensearch_host self.opensearch_port = opensearch_port - self.username = username | os.getenv("OPENSEARCH_USERNAME") - self.password = password | os.getenv("OPENSEARCH_PASSWORD") + self.username = username or os.getenv("OPENSEARCH_USERNAME") + self.password = password or os.getenv("OPENSEARCH_PASSWORD") self.index_name = index_name self.vector_field = vector_field self.k = k From 22b2f385a8a9a82ad5c2aa92217068ce01c8500a Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Tue, 7 Jul 2026 15:38:43 +0400 Subject: [PATCH 02/29] Refactor LoadTestBase --- ...m_PostCommit_Java_IO_Performance_Tests.yml | 2 +- it/build.gradle | 6 +- it/common/build.gradle | 10 + .../bigquery/BigQueryResourceManager.java | 2 +- .../BigQueryResourceManagerException.java | 2 +- .../BigQueryResourceManagerUtils.java | 2 +- .../beam/it/common/bigquery/package-info.java | 20 ++ .../dataflow/AbstractPipelineLauncher.java | 2 +- .../dataflow/DefaultPipelineLauncher.java | 4 +- .../beam/it/common/dataflow/package-info.java | 20 ++ .../common}/monitoring/MonitoringClient.java | 2 +- .../it/common}/monitoring/package-info.java | 2 +- .../beam/it/common}/IOLoadTestBase.java | 7 +- .../beam/it/common}/IOStressTestBase.java | 2 +- .../apache/beam/it/common}/LoadTestBase.java | 13 +- it/google-cloud-platform/build.gradle | 3 + .../conditions/BigQueryRowsCheck.java | 2 +- .../gcp/dataflow/ClassicTemplateClient.java | 2 + .../it/gcp/dataflow/FlexTemplateClient.java | 2 + .../beam/it/gcp/bigquery/BigQueryIOLT.java | 3 +- .../beam/it/gcp/bigquery/BigQueryIOST.java | 3 +- .../bigquery/BigQueryResourceManagerTest.java | 2 + .../BigQueryResourceManagerUtilsTest.java | 4 +- .../it/gcp/bigquery/BigQueryStreamingLT.java | 2 +- .../beam/it/gcp/bigtable/BigTableIOLT.java | 2 +- .../beam/it/gcp/bigtable/BigTableIOST.java | 2 +- .../AbstractPipelineLauncherTest.java | 1 + .../dataflow/ClassicTemplateClientTest.java | 9 +- .../dataflow/DefaultPipelineLauncherTest.java | 1 + .../gcp/dataflow/FlexTemplateClientTest.java | 8 +- .../apache/beam/it/gcp/pubsub/PubSubIOST.java | 2 +- .../apache/beam/it/gcp/pubsub/PubsubIOLT.java | 2 +- .../beam/it/gcp/spanner/SpannerIOLT.java | 2 +- .../beam/it/gcp/spanner/SpannerIOST.java | 2 +- .../beam/it/gcp/storage/FileBasedIOLT.java | 2 +- it/iceberg/build.gradle | 172 ++++++++++++++++++ .../src/main/resources/test-artifact.json | 1 + .../apache/beam/it/iceberg/IcebergIOLT.java | 7 + settings.gradle.kts | 1 + 39 files changed, 291 insertions(+), 42 deletions(-) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/bigquery/BigQueryResourceManager.java (99%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/bigquery/BigQueryResourceManagerException.java (96%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/bigquery/BigQueryResourceManagerUtils.java (98%) create mode 100644 it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/dataflow/AbstractPipelineLauncher.java (99%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/dataflow/DefaultPipelineLauncher.java (99%) create mode 100644 it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/monitoring/MonitoringClient.java (99%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/monitoring/package-info.java (94%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/IOLoadTestBase.java (96%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/IOStressTestBase.java (99%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/LoadTestBase.java (98%) create mode 100644 it/iceberg/build.gradle create mode 100644 it/iceberg/src/main/resources/test-artifact.json create mode 100644 it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java diff --git a/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml b/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml index db43441f2cda..0bc0a37ea4ea 100644 --- a/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml +++ b/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml @@ -64,7 +64,7 @@ jobs: matrix: job_name: ["beam_PostCommit_Java_IO_Performance_Tests"] job_phrase: ["Run Java PostCommit IO Performance Tests"] - test_case: ["GCSPerformanceTest", "BigTablePerformanceTest", "BigQueryStorageApiStreamingPerformanceTest"] + test_case: ["IcebergPerformanceTest"] steps: - uses: actions/checkout@v7 - name: Setup repository diff --git a/it/build.gradle b/it/build.gradle index 42a9ad9f4ee8..c550e26d305f 100644 --- a/it/build.gradle +++ b/it/build.gradle @@ -34,4 +34,8 @@ tasks.register('BigTablePerformanceTest') { tasks.register('BigQueryStorageApiStreamingPerformanceTest') { dependsOn(":it:google-cloud-platform:BigQueryStorageApiStreamingPerformanceTest") -} \ No newline at end of file +} + +tasks.register('IcebergPerformanceTest') { + dependsOn(":it:iceberg:IcebergPerformanceTest") +} diff --git a/it/common/build.gradle b/it/common/build.gradle index c92ec551cb25..7b8559aa6e1f 100644 --- a/it/common/build.gradle +++ b/it/common/build.gradle @@ -28,9 +28,16 @@ ext.summary = "Code used by all integration test utilities." dependencies { implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation project(path: ":runners:google-cloud-dataflow-java") + implementation project(path: ":sdks:java:testing:test-utils") + implementation 'com.google.cloud:google-cloud-bigquery' + implementation 'com.google.cloud:google-cloud-monitoring' + provided 'com.google.api.grpc:proto-google-cloud-monitoring-v3' + implementation library.java.gax implementation library.java.google_api_services_dataflow implementation library.java.google_auth_library_credentials implementation library.java.google_auth_library_oauth2_http + implementation library.java.google_api_services_bigquery implementation library.java.vendored_guava_32_1_2_jre implementation library.java.slf4j_api implementation library.java.commons_lang3 @@ -38,6 +45,8 @@ dependencies { implementation library.java.google_code_gson implementation library.java.google_http_client implementation library.java.guava + implementation library.java.protobuf_java_util + implementation library.java.protobuf_java // TODO: excluding Guava until Truth updates it to >32.1.x testImplementation(library.java.truth) { @@ -46,4 +55,5 @@ dependencies { testImplementation library.java.junit testImplementation library.java.mockito_inline testRuntimeOnly library.java.slf4j_simple + testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") } diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java rename to it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java index 5ca6a8c4655c..728a47b224b4 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import com.google.api.gax.paging.Page; import com.google.auth.Credentials; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java similarity index 96% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java rename to it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java index 61ef27e45260..ce95dc891551 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; /** Custom exception for {@link BigQueryResourceManager} implementations. */ public class BigQueryResourceManagerException extends RuntimeException { diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java rename to it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java index f2b6849caa61..56cb82013885 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import static org.apache.beam.it.common.utils.ResourceManagerUtils.generateResourceId; diff --git a/it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java new file mode 100644 index 000000000000..f2d5978fb4eb --- /dev/null +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java @@ -0,0 +1,20 @@ +/* + * 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 for managing BigQuery common resources within integration tests. */ +package org.apache.beam.it.common.bigquery; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java rename to it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java index 9d00fa67f145..a6efaeca5afc 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; import static org.apache.beam.it.common.PipelineLauncher.JobState.ACTIVE_STATES; import static org.apache.beam.it.common.PipelineLauncher.JobState.FAILED; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java rename to it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java index 11a09c4ba749..a872868c12bb 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; import static org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS; @@ -41,7 +41,7 @@ import java.util.stream.StreamSupport; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.utils.PipelineUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowPipelineJob; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.DistributionResult; diff --git a/it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java new file mode 100644 index 000000000000..eb10897523bc --- /dev/null +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java @@ -0,0 +1,20 @@ +/* + * 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 for managing Dataflow jobs from integration tests. */ +package org.apache.beam.it.common.dataflow; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java b/it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java rename to it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java index 66550b5a5073..6b447d84512d 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java +++ b/it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.monitoring; +package org.apache.beam.it.common.monitoring; import com.google.api.gax.core.CredentialsProvider; import com.google.cloud.monitoring.v3.MetricServiceClient; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java similarity index 94% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java index 9cd79ea9cf12..731bbc3a6b25 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java @@ -17,4 +17,4 @@ */ /** Package for querying metrics from cloud monitoring. */ -package org.apache.beam.it.gcp.monitoring; +package org.apache.beam.it.common.monitoring; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOLoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java similarity index 96% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOLoadTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java index 14770a429731..fb9feb6663b0 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOLoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp; +package org.apache.beam.it.common; import com.google.cloud.Timestamp; import java.io.IOException; @@ -23,9 +23,8 @@ import java.util.Collection; import java.util.Map; import java.util.UUID; -import org.apache.beam.it.common.PipelineLauncher; -import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.dataflow.DefaultPipelineLauncher; + +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.testutils.NamedTestResult; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOStressTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOStressTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java index 5c2fb74cd2fb..a15e98d65d81 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOStressTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp; +package org.apache.beam.it.common; import java.io.Serializable; import java.time.Duration; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/LoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/LoadTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java index cd9ef52ed835..7d7adbbae0ae 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/LoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java @@ -15,10 +15,10 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp; +package org.apache.beam.it.common; import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.RUNNER_V2; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.RUNNER_V2; import com.google.api.gax.core.CredentialsProvider; import com.google.api.gax.core.FixedCredentialsProvider; @@ -40,12 +40,11 @@ import java.util.Map; import java.util.Map.Entry; import java.util.regex.Pattern; -import org.apache.beam.it.common.PipelineLauncher; + import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; -import org.apache.beam.it.common.PipelineOperator; -import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.bigquery.BigQueryResourceManager; -import org.apache.beam.it.gcp.monitoring.MonitoringClient; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; +import org.apache.beam.it.common.monitoring.MonitoringClient; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.After; diff --git a/it/google-cloud-platform/build.gradle b/it/google-cloud-platform/build.gradle index 3a46f2b94d83..dd8376e35693 100644 --- a/it/google-cloud-platform/build.gradle +++ b/it/google-cloud-platform/build.gradle @@ -27,6 +27,8 @@ applyJavaNature( description = "Apache Beam :: IT :: Google Cloud Platform" ext.summary = "Integration test utilities for Google Cloud Platform." +evaluationDependsOn(":it:common") + dependencies { implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) implementation project(path: ":sdks:java:core", configuration: "shadow") @@ -84,6 +86,7 @@ dependencies { implementation 'com.google.cloud:google-cloud-secretmanager' provided 'com.google.api.grpc:proto-google-cloud-secretmanager-v1' + testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:google-cloud-platform") testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:synthetic") diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java index 28626f7edd18..57b7d5ac99c9 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java @@ -20,7 +20,7 @@ import com.google.auto.value.AutoValue; import com.google.cloud.bigquery.TableId; import org.apache.beam.it.conditions.ConditionCheck; -import org.apache.beam.it.gcp.bigquery.BigQueryResourceManager; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; import org.checkerframework.checker.nullness.qual.Nullable; /** ConditionCheck to validate if BigQuery has received a certain number of rows. */ diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java index 00aa4003817f..296f7092cfaa 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java @@ -30,6 +30,8 @@ import com.google.auth.http.HttpCredentialsAdapter; import dev.failsafe.Failsafe; import java.io.IOException; + +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java index 0a9731b13c32..3e7270dbd34f 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java @@ -32,6 +32,8 @@ import com.google.auth.http.HttpCredentialsAdapter; import dev.failsafe.Failsafe; import java.io.IOException; + +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java index f123ffa5954c..9a4af4082e23 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java @@ -43,8 +43,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java index 14c83e48308a..bbe4df611916 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java @@ -41,8 +41,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; +import org.apache.beam.it.common.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java index 53987dbe2670..3a633f00c6fd 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java @@ -41,6 +41,8 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableInfo; import com.google.cloud.bigquery.TimePartitioning; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; +import org.apache.beam.it.common.bigquery.BigQueryResourceManagerException; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.junit.Before; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java index 432b7ef7362f..7248fed0e878 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java @@ -17,10 +17,12 @@ */ package org.apache.beam.it.gcp.bigquery; -import static org.apache.beam.it.gcp.bigquery.BigQueryResourceManagerUtils.checkValidTableId; +import static org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils.checkValidTableId; import static org.junit.Assert.assertThrows; import java.util.Arrays; + +import org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils; import org.junit.Test; /** Unit tests for {@link BigQueryResourceManagerUtils}. */ diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java index 6e511bd8e5c6..7bf759b42f90 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java @@ -41,7 +41,7 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowRunner; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java index 410c992fe2d6..aea26f2f0a49 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java @@ -35,7 +35,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; import org.apache.beam.sdk.testing.TestPipeline; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java index 052a0839adf5..85950772bce0 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; +import org.apache.beam.it.common.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java index f2c33d70116b..df1c5cea6406 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java @@ -34,6 +34,7 @@ import java.io.IOException; import java.net.SocketTimeoutException; import org.apache.beam.it.common.PipelineLauncher.JobState; +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java index 88c35589f2be..a50c0951005f 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java @@ -18,10 +18,10 @@ package org.apache.beam.it.gcp.dataflow; import static com.google.common.truth.Truth.assertThat; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; @@ -49,6 +49,7 @@ import org.apache.beam.it.common.PipelineLauncher.JobState; import org.apache.beam.it.common.PipelineLauncher.LaunchConfig; import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java index b6c0f8cdc58f..52ff5c0a230d 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java @@ -23,6 +23,7 @@ import java.io.IOException; import java.time.Instant; import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher; import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.it.gcp.IOLoadTestBase.PipelineMetricsType; import org.apache.beam.sdk.io.GenerateSequence; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java index 06f44437414a..b94d52e45afb 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java @@ -18,10 +18,10 @@ package org.apache.beam.it.gcp.dataflow; import static com.google.common.truth.Truth.assertThat; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java index cb96db40b749..74bd994055a7 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; +import org.apache.beam.it.common.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java index 477c79dd3bc7..a9c7b63e64cd 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java index 8f3de1b58f40..a29cf32eae1c 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; import org.apache.beam.sdk.io.synthetic.SyntheticSourceOptions; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java index 76312b3e10fb..8ce34d06de88 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; +import org.apache.beam.it.common.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java index 58ba88d900e8..febcfc01f3f3 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java @@ -34,7 +34,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.sdk.io.Compression; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.TextIO; diff --git a/it/iceberg/build.gradle b/it/iceberg/build.gradle new file mode 100644 index 000000000000..1a61f86367b2 --- /dev/null +++ b/it/iceberg/build.gradle @@ -0,0 +1,172 @@ +/* + * 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. + */ + +import org.apache.beam.gradle.IoPerformanceTestUtilities + +plugins { id 'org.apache.beam.module' } +applyJavaNature( + automaticModuleName: 'org.apache.beam.it.iceberg', + exportJavadoc: false, + // iceberg ended support for Java 8 in 1.7.0 + requireJavaVersion: JavaVersion.VERSION_11, +) + +description = "Apache Beam :: IT :: Iceberg" +ext.summary = "Integration test utilities for Iceberg." + +def iceberg_version = "1.10.0" +def parquet_version = "1.16.0" +def orc_version = "1.9.6" +def hive_version = "3.1.3" + +def hadoopVersions = [ + "336": "3.3.6", +] + +hadoopVersions.each {kv -> configurations.create("hadoopVersion$kv.key")} + +evaluationDependsOn(":it:common") + +dependencies { + implementation library.java.vendored_guava_32_1_2_jre + implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation project(path: ":model:pipeline", configuration: "shadow") + implementation library.java.avro + implementation library.java.slf4j_api + implementation library.java.joda_time + implementation "org.apache.parquet:parquet-column:$parquet_version" + implementation "org.apache.parquet:parquet-hadoop:$parquet_version" + implementation "org.apache.parquet:parquet-common:$parquet_version" + implementation project(":sdks:java:io:parquet") + implementation "org.apache.orc:orc-core:$orc_version" + implementation "org.apache.iceberg:iceberg-core:$iceberg_version" + implementation "org.apache.iceberg:iceberg-api:$iceberg_version" + implementation "org.apache.iceberg:iceberg-parquet:$iceberg_version" + implementation "org.apache.iceberg:iceberg-orc:$iceberg_version" + implementation "org.apache.iceberg:iceberg-data:$iceberg_version" + implementation "org.apache.hadoop:hadoop-common:3.3.6" + + provided "org.immutables:value:2.8.8" + permitUnusedDeclared "org.immutables:value:2.8.8" + implementation library.java.vendored_calcite_1_40_0 + runtimeOnly "org.apache.iceberg:iceberg-gcp:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-aws:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-aws-bundle:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-azure:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-azure-bundle:$iceberg_version" + runtimeOnly library.java.bigdataoss_gcs_connector + runtimeOnly library.java.bigdataoss_util_hadoop + runtimeOnly library.java.hadoop_client + + testImplementation project(":sdks:java:managed") + testImplementation library.java.bigdataoss_gcsio + testImplementation library.java.bigdataoss_util_hadoop + testImplementation "org.apache.parquet:parquet-avro:$parquet_version" + testImplementation "org.apache.iceberg:iceberg-data:$iceberg_version" + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation library.java.junit + testImplementation library.java.hamcrest + testImplementation project(path: ":sdks:java:extensions:avro") + testImplementation 'org.awaitility:awaitility:4.2.0' + + // Hive catalog test dependencies + testImplementation project(path: ":sdks:java:io:iceberg:hive") + testImplementation "org.apache.iceberg:iceberg-common:$iceberg_version" + testImplementation ("org.apache.iceberg:iceberg-hive-metastore:$iceberg_version") + testImplementation ("org.apache.hive:hive-metastore:$hive_version") + testImplementation "org.assertj:assertj-core:3.11.1" + testRuntimeOnly ("org.apache.hive.hcatalog:hive-hcatalog-core:$hive_version") { + exclude group: "org.apache.hive", module: "hive-exec" + exclude group: "org.apache.parquet", module: "parquet-hadoop-bundle" + } + + // BigQueryMetastore catalog dep + testImplementation project(path: ":sdks:java:io:iceberg:bqms", configuration: "shadow") + testImplementation project(":sdks:java:io:google-cloud-platform") + testImplementation library.java.google_api_services_bigquery + testImplementation 'com.google.cloud:google-cloud-storage' + + testImplementation library.java.google_auth_library_oauth2_http + testRuntimeOnly library.java.slf4j_jdk14 + testImplementation project(path: ":runners:direct-java", configuration: "shadow") + testRuntimeOnly project(path: ":runners:google-cloud-dataflow-java") + testRuntimeOnly project(path: ":sdks:java:harness") + hadoopVersions.each {kv -> + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-client:$kv.value" + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-minicluster:$kv.value" + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-hdfs-client:$kv.value" + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-mapreduce-client-core:$kv.value" + } + + implementation project(path: ":runners:google-cloud-dataflow-java") + implementation project(path: ":it:common") + implementation project(path: ":it:conditions") + implementation project(path: ":it:truthmatchers") + implementation project(path: ":sdks:java:testing:test-utils") + implementation library.java.failsafe + implementation library.java.google_api_services_dataflow + implementation(library.java.truth) { + exclude group: 'com.google.guava', module: 'guava' + } + implementation library.java.jackson_core + implementation library.java.jackson_databind + implementation 'org.apache.commons:commons-lang3:3.9' + implementation library.java.google_api_client + implementation library.java.google_http_client + implementation library.java.guava + implementation library.java.protobuf_java_util + implementation library.java.protobuf_java + implementation library.java.threetenbp + + testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") + testImplementation project(path: ":sdks:java:io:synthetic") + testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") + testImplementation library.java.mockito_inline + testImplementation project(path: ":sdks:java:extensions:google-cloud-platform-core", configuration: "testRuntimeMigration") + testRuntimeOnly library.java.slf4j_simple +} + +configurations.all { + // iceberg-core needs avro:1.12.0 + resolutionStrategy.force 'org.apache.avro:avro:1.12.0' + // TODO(https://github.com/apache/beam/issues/38515): + // Remove below pins when parquet-hadoop upgrades to hadoop-common:3.4.2 + resolutionStrategy.force 'org.apache.hadoop:hadoop-common:3.3.6' + resolutionStrategy.force 'org.apache.hadoop:hadoop-client:3.3.6' + resolutionStrategy.force 'org.apache.hadoop:hadoop-hdfs:3.3.6' + resolutionStrategy.force 'org.apache.hadoop:hadoop-hdfs-client:3.3.6' +} + +hadoopVersions.each {kv -> + configurations."hadoopVersion$kv.key" { + resolutionStrategy { + force "org.apache.hadoop:hadoop-client:$kv.value" + force "org.apache.hadoop:hadoop-common:$kv.value" + force "org.apache.hadoop:hadoop-mapreduce-client-core:$kv.value" + force "org.apache.hadoop:hadoop-minicluster:$kv.value" + force "org.apache.hadoop:hadoop-hdfs:$kv.value" + force "org.apache.hadoop:hadoop-hdfs-client:$kv.value" + } + } +} + +tasks.register( + "IcebergPerformanceTest", IoPerformanceTestUtilities.IoPerformanceTest, project, 'iceberg', 'IcebergIOLT', + ['configuration':'large','project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'] + + System.properties +) diff --git a/it/iceberg/src/main/resources/test-artifact.json b/it/iceberg/src/main/resources/test-artifact.json new file mode 100644 index 000000000000..551c80d14a66 --- /dev/null +++ b/it/iceberg/src/main/resources/test-artifact.json @@ -0,0 +1 @@ +["This is a test artifact."] \ No newline at end of file diff --git a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java new file mode 100644 index 000000000000..d2e342ccecb4 --- /dev/null +++ b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java @@ -0,0 +1,7 @@ +package org.apache.beam.it.iceberg; + +import org.apache.beam.it.common.IOLoadTestBase; + +public class IcebergIOLT extends IOLoadTestBase { + +} diff --git a/settings.gradle.kts b/settings.gradle.kts index d9dbbf9021ef..5ee8f8a4d394 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -122,6 +122,7 @@ include(":it:conditions") include(":it:datadog") include(":it:elasticsearch") include(":it:google-cloud-platform") +include(":it:iceberg") include(":it:jdbc") include(":it:kafka") include(":it:testcontainers") From 875cfd11cedd5c7b192fbdb98eeb2484e6c01123 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Tue, 7 Jul 2026 15:58:46 +0400 Subject: [PATCH 03/29] Refactor DefaultPipelineLauncher --- .../dataflow/DefaultPipelineLauncher.java | 45 ++++++++++++------- .../bigquery/BigQueryResourceManagerTest.java | 4 +- .../BigQueryResourceManagerUtilsTest.java | 3 +- .../AbstractPipelineLauncherTest.java | 3 +- .../dataflow/DefaultPipelineLauncherTest.java | 9 ++-- .../common/{ => dataflow}/IOLoadTestBase.java | 23 +++------- .../{ => dataflow}/IOStressTestBase.java | 2 +- .../common/{ => dataflow}/LoadTestBase.java | 6 ++- .../it/gcp/datagenerator/package-info.java | 20 --------- .../org/apache/beam/it/gcp/WordCountIT.java | 2 + .../beam/it/gcp/bigquery/BigQueryIOLT.java | 2 +- .../beam/it/gcp/bigquery/BigQueryIOST.java | 2 +- .../it/gcp/bigquery/BigQueryStreamingLT.java | 2 +- .../beam/it/gcp/bigtable/BigTableIOLT.java | 2 +- .../beam/it/gcp/bigtable/BigTableIOST.java | 2 +- .../it/gcp/datagenerator/DataGenerator.java | 2 +- .../apache/beam/it/gcp/pubsub/PubSubIOST.java | 2 +- .../apache/beam/it/gcp/pubsub/PubsubIOLT.java | 2 +- .../beam/it/gcp/spanner/SpannerIOLT.java | 2 +- .../beam/it/gcp/spanner/SpannerIOST.java | 2 +- .../beam/it/gcp/storage/FileBasedIOLT.java | 2 +- .../apache/beam/it/iceberg/IcebergIOLT.java | 2 +- it/kafka/build.gradle | 3 +- .../org/apache/beam/it/kafka/KafkaIOLT.java | 3 +- .../org/apache/beam/it/kafka/KafkaIOST.java | 3 +- 25 files changed, 65 insertions(+), 85 deletions(-) rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/bigquery/BigQueryResourceManagerTest.java (98%) rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/bigquery/BigQueryResourceManagerUtilsTest.java (94%) rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/dataflow/AbstractPipelineLauncherTest.java (98%) rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/dataflow/DefaultPipelineLauncherTest.java (93%) rename it/common/src/test/java/org/apache/beam/it/common/{ => dataflow}/IOLoadTestBase.java (88%) rename it/common/src/test/java/org/apache/beam/it/common/{ => dataflow}/IOStressTestBase.java (99%) rename it/common/src/test/java/org/apache/beam/it/common/{ => dataflow}/LoadTestBase.java (99%) delete mode 100644 it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java rename it/google-cloud-platform/src/{main => test}/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java (99%) diff --git a/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java index a872868c12bb..32a4eb6eb1b6 100644 --- a/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java @@ -18,7 +18,6 @@ package org.apache.beam.it.common.dataflow; import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; -import static org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.fasterxml.jackson.databind.ObjectMapper; @@ -41,7 +40,6 @@ import java.util.stream.StreamSupport; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.utils.PipelineUtils; -import org.apache.beam.it.common.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowPipelineJob; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.DistributionResult; @@ -73,6 +71,8 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher { private static final String READ_PIPELINE_NAME_OVERWRITE = "readPipelineNameOverride"; private static final String WRITE_PIPELINE_NAME_OVERWRITE = "writePipelineNameOverride"; private static final Pattern JOB_ID_PATTERN = Pattern.compile("Submitted job: (\\S+)"); + /** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */ + public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS"; // For unsupported runners (other than dataflow), implement launcher methods by operating with // PipelineResult. @@ -101,6 +101,18 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher { .put(PipelineResult.State.UNRECOGNIZED, JobState.UNKNOWN) .build(); + // To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and + // pipeline defined + // metrics, here we wrap Beam provided metrics as a pre-defined metrics name + // [name_space:metric_type:metric_name + // which will be recognized by getMetric method + public enum PipelineMetricsType { + COUNTER, + STARTTIME, + ENDTIME, + RUNTIME, + } + private DefaultPipelineLauncher(Builder builder) { super( new Dataflow( @@ -169,7 +181,7 @@ private static void checkIfMetricResultIsUnique( resultCount <= 1, "More than one metric result matches name: %s in namespace %s. Metric results count: %s", name, - IOLoadTestBase.BEAM_METRICS_NAMESPACE, + BEAM_METRICS_NAMESPACE, resultCount); } @@ -181,14 +193,14 @@ private static Iterable> getDistributions( .queryMetrics( MetricsFilter.builder() .addNameFilter( - MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName)) + MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) .build()); return metrics.getDistributions(); } /** Pull Beam pipeline defined metrics given the jobId. */ public Long getBeamMetric( - String jobId, IOLoadTestBase.PipelineMetricsType metricType, String metricName) { + String jobId, PipelineMetricsType metricType, String metricName) { PipelineResult pipelineResult = MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null)); if (pipelineResult != null) { @@ -198,7 +210,7 @@ public Long getBeamMetric( .queryMetrics( MetricsFilter.builder() .addNameFilter( - MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName)) + MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) .build()); switch (metricType) { @@ -212,7 +224,7 @@ public Long getBeamMetric( LOG.error( "Failed to get metric {}, from namespace {}", metricName, - IOLoadTestBase.BEAM_METRICS_NAMESPACE); + BEAM_METRICS_NAMESPACE); } return UNKNOWN_METRIC_VALUE; case STARTTIME: @@ -230,9 +242,9 @@ public Long getBeamMetric( .map(element -> Objects.requireNonNull(element.getAttempted()).getMax()) .max(Long::compareTo) .orElse(UNKNOWN_METRIC_VALUE); - if (metricType == IOLoadTestBase.PipelineMetricsType.STARTTIME) { + if (metricType == PipelineMetricsType.STARTTIME) { return lowestMin; - } else if (metricType == IOLoadTestBase.PipelineMetricsType.ENDTIME) { + } else if (metricType == PipelineMetricsType.ENDTIME) { return greatestMax; } else { if (lowestMin != UNKNOWN_METRIC_VALUE && greatestMax != UNKNOWN_METRIC_VALUE) { @@ -254,15 +266,15 @@ public Long getBeamMetric( @Override public Double getMetric(String project, String region, String jobId, String metricName) throws IOException { - if (metricName.startsWith(IOLoadTestBase.BEAM_METRICS_NAMESPACE)) { + if (metricName.startsWith(BEAM_METRICS_NAMESPACE)) { String[] nameSpacedMetrics = metricName.split(":", 3); Preconditions.checkState( nameSpacedMetrics.length == 3, String.format( "Invalid Beam metrics name: %s, expected: '%s:metric_type:metric_name'", - metricName, IOLoadTestBase.BEAM_METRICS_NAMESPACE)); - IOLoadTestBase.PipelineMetricsType metricType = - IOLoadTestBase.PipelineMetricsType.valueOf(nameSpacedMetrics[1]); + metricName, BEAM_METRICS_NAMESPACE)); + PipelineMetricsType metricType = + PipelineMetricsType.valueOf(nameSpacedMetrics[1]); // Pipeline defined metrics are long values. Have to cast to double that is what the base // class defined. @@ -437,16 +449,15 @@ private List extractOptions(String project, String region, LaunchConfig // add pipeline options from beamTestPipelineOptions system property to preserve the // pipeline options already set in TestPipeline. @Nullable - String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS); + String beamTestPipelineOptions = System.getProperty(org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS); if (!Strings.isNullOrEmpty(beamTestPipelineOptions)) { try { additionalOptions.addAll(MAPPER.readValue(beamTestPipelineOptions, List.class)); } catch (IOException e) { throw new RuntimeException( "Unable to instantiate test options from system property " - + PROPERTY_BEAM_TEST_PIPELINE_OPTIONS - + ":" - + System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS), + + org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS + + ":" + beamTestPipelineOptions, e); } } diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java similarity index 98% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java rename to it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java index 3a633f00c6fd..62aac93b4eed 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -41,8 +41,6 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableInfo; import com.google.cloud.bigquery.TimePartitioning; -import org.apache.beam.it.common.bigquery.BigQueryResourceManager; -import org.apache.beam.it.common.bigquery.BigQueryResourceManagerException; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.junit.Before; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java similarity index 94% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java rename to it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java index 7248fed0e878..eaff80740727 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java @@ -15,14 +15,13 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import static org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils.checkValidTableId; import static org.junit.Assert.assertThrows; import java.util.Arrays; -import org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils; import org.junit.Test; /** Unit tests for {@link BigQueryResourceManagerUtils}. */ diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java similarity index 98% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java index df1c5cea6406..72d5514a6e3c 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -34,7 +34,6 @@ import java.io.IOException; import java.net.SocketTimeoutException; import org.apache.beam.it.common.PipelineLauncher.JobState; -import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java similarity index 93% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java index 52ff5c0a230d..333549c621be 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java @@ -15,17 +15,16 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; +import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import java.io.IOException; import java.time.Instant; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.PipelineLauncher; -import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher; -import org.apache.beam.it.gcp.IOLoadTestBase; -import org.apache.beam.it.gcp.IOLoadTestBase.PipelineMetricsType; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testutils.metrics.TimeMonitor; @@ -49,7 +48,7 @@ public void testPipelineMetrics() throws IOException { pipeline .apply(GenerateSequence.from(0).to(numElements)) - .apply(ParDo.of(new TimeMonitor<>(IOLoadTestBase.BEAM_METRICS_NAMESPACE, timeMetrics))) + .apply(ParDo.of(new TimeMonitor<>(BEAM_METRICS_NAMESPACE, timeMetrics))) .apply(ParDo.of(new IOLoadTestBase.CountingFn<>(counterMetrics))); PipelineLauncher.LaunchConfig options = diff --git a/it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java similarity index 88% rename from it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java index fb9feb6663b0..8a1a94df1aec 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/IOLoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.common; +package org.apache.beam.it.common.dataflow; import com.google.cloud.Timestamp; import java.io.IOException; @@ -24,7 +24,9 @@ import java.util.Map; import java.util.UUID; -import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.testutils.NamedTestResult; @@ -38,6 +40,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE; + /** Base class for IO Load tests. */ @RunWith(JUnit4.class) @SuppressWarnings({ @@ -89,21 +93,6 @@ public void processElement(ProcessContext ctx) { } } - // To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and - // pipeline defined - // metrics, here we wrap Beam provided metrics as a pre-defined metrics name - // [name_space:metric_type:metric_name - // which will be recognized by getMetric method - public enum PipelineMetricsType { - COUNTER, - STARTTIME, - ENDTIME, - RUNTIME, - } - - /** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */ - public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS"; - /** Given a metrics name, return Beam metrics name. */ public static String getBeamMetricsName(PipelineMetricsType metricstype, String metricsName) { return BEAM_METRICS_NAMESPACE + ":" + metricstype + ":" + metricsName; diff --git a/it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java similarity index 99% rename from it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java index a15e98d65d81..e08eec031eb5 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/IOStressTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.common; +package org.apache.beam.it.common.dataflow; import java.io.Serializable; import java.time.Duration; diff --git a/it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java similarity index 99% rename from it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java index 7d7adbbae0ae..360712835062 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/LoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.common; +package org.apache.beam.it.common.dataflow; import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.RUNNER_V2; @@ -41,9 +41,11 @@ import java.util.Map.Entry; import java.util.regex.Pattern; +import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; +import org.apache.beam.it.common.PipelineOperator; +import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.bigquery.BigQueryResourceManager; -import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.apache.beam.it.common.monitoring.MonitoringClient; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.checkerframework.checker.nullness.qual.Nullable; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java deleted file mode 100644 index d563025ebeb4..000000000000 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * 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. - */ - -/** Data generator for load tests. */ -package org.apache.beam.it.gcp.datagenerator; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java index b561c0c71fb6..9719caa59f1b 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java @@ -25,6 +25,8 @@ import java.time.Duration; import java.util.Arrays; import java.util.List; + +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.PipelineLauncher.LaunchConfig; import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; import org.apache.beam.it.common.PipelineLauncher.Sdk; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java index 9a4af4082e23..b734f93478ef 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java @@ -45,7 +45,7 @@ import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.bigquery.BigQueryResourceManager; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java index bbe4df611916..4c86eb139bf9 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java @@ -43,7 +43,7 @@ import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.bigquery.BigQueryResourceManager; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOStressTestBase; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java index 7bf759b42f90..75a7be50a508 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java @@ -41,7 +41,7 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowRunner; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java index aea26f2f0a49..03efb42dafb6 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java @@ -35,7 +35,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; import org.apache.beam.sdk.testing.TestPipeline; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java index 85950772bce0..5e70a0e2d582 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOStressTestBase; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java rename to it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java index 99016b5dd3a4..0b838bddeb6d 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java @@ -17,7 +17,7 @@ */ package org.apache.beam.it.gcp.datagenerator; -import static org.apache.beam.it.gcp.LoadTestBase.createConfig; +import static org.apache.beam.it.common.dataflow.LoadTestBase.createConfig; import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatPipeline; import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java index 74bd994055a7..98092152ad22 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOStressTestBase; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java index a9c7b63e64cd..360073ac0e9e 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java index a29cf32eae1c..99682746d2b3 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; import org.apache.beam.sdk.io.synthetic.SyntheticSourceOptions; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java index 8ce34d06de88..01d8deef1773 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java @@ -38,7 +38,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOStressTestBase; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java index febcfc01f3f3..abec875981fc 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java @@ -34,7 +34,7 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.sdk.io.Compression; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.TextIO; diff --git a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java index d2e342ccecb4..3f486ec3f5d0 100644 --- a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java +++ b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java @@ -1,6 +1,6 @@ package org.apache.beam.it.iceberg; -import org.apache.beam.it.common.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; public class IcebergIOLT extends IOLoadTestBase { diff --git a/it/kafka/build.gradle b/it/kafka/build.gradle index 0496c423bebd..32ef0623fea9 100644 --- a/it/kafka/build.gradle +++ b/it/kafka/build.gradle @@ -42,8 +42,7 @@ dependencies { permitUsedUndeclared library.java.guava testImplementation library.java.mockito_core testRuntimeOnly library.java.slf4j_simple - // TODO: Remove the below dependencies after KafkaIOLT has been moved to the right directory. - testImplementation project(path: ":it:google-cloud-platform") + testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:kafka") testImplementation project(path: ":sdks:java:io:synthetic") testImplementation project(path: ":runners:direct-java") diff --git a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java index c20699d61a55..4b0ecbfd3433 100644 --- a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java +++ b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java @@ -27,11 +27,12 @@ import java.time.format.DateTimeFormatter; import java.util.Map; import java.util.UUID; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.io.kafka.KafkaIO; import org.apache.beam.sdk.io.synthetic.SyntheticOptions; diff --git a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java index bb1cb37a11b6..cfc11127c34f 100644 --- a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java +++ b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java @@ -31,10 +31,11 @@ import java.util.Map; import java.util.Objects; import java.util.UUID; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.IOStressTestBase; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.kafka.KafkaIO; From 5312b03502c4bd02e77fe8293f2fb31766d8ff35 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Tue, 7 Jul 2026 17:04:31 +0400 Subject: [PATCH 04/29] Add IcebergIOLT --- it/clickhouse/build.gradle | 2 +- it/common/build.gradle | 23 +- .../beam/it/common}/artifacts/Artifact.java | 2 +- .../it/common}/artifacts/ArtifactClient.java | 2 +- .../it/common}/artifacts/GcsArtifact.java | 2 +- .../it/common}/artifacts/package-info.java | 2 +- .../artifacts/utils/ArtifactUtils.java | 2 +- .../common}/artifacts/utils/AvroTestUtil.java | 2 +- .../common}/artifacts/utils/JsonTestUtil.java | 2 +- .../artifacts/utils/ParquetTestUtil.java | 2 +- .../common}/artifacts/utils/package-info.java | 2 +- .../dataflow/DefaultPipelineLauncher.java | 23 +- .../common}/storage/GcsResourceManager.java | 10 +- .../beam/it/common}/storage/package-info.java | 2 +- .../src/main/resources/test-artifact.json | 1 + .../it/common}/artifacts/GcsArtifactTest.java | 4 +- .../artifacts/utils/ArtifactUtilsTest.java | 4 +- .../BigQueryResourceManagerUtilsTest.java | 1 - .../dataflow/DefaultPipelineLauncherTest.java | 2 +- .../it/common/dataflow/IOLoadTestBase.java | 5 +- .../beam/it/common/dataflow/LoadTestBase.java | 3 +- .../storage/GcsResourceManagerTest.java | 19 +- it/google-cloud-platform/build.gradle | 18 - .../gcp/artifacts/matchers/package-info.java | 20 - .../conditions/BigQueryRowsCheck.java | 2 +- .../gcp/dataflow/ClassicTemplateClient.java | 1 - .../it/gcp/dataflow/FlexTemplateClient.java | 1 - .../gcp/spanner/matchers/SpannerAsserts.java | 2 +- .../org/apache/beam/it/gcp/WordCountIT.java | 3 +- .../beam/it/gcp/bigquery/BigQueryIOLT.java | 3 +- .../beam/it/gcp/bigquery/BigQueryIOST.java | 3 +- .../beam/it/gcp/bigtable/BigTableIOLT.java | 3 +- .../beam/it/gcp/bigtable/BigTableIOST.java | 3 +- .../apache/beam/it/gcp/pubsub/PubSubIOST.java | 3 +- .../apache/beam/it/gcp/pubsub/PubsubIOLT.java | 3 +- .../beam/it/gcp/spanner/SpannerIOLT.java | 3 +- .../beam/it/gcp/spanner/SpannerIOST.java | 3 +- .../beam/it/gcp/storage/FileBasedIOLT.java | 4 +- it/iceberg/build.gradle | 42 +-- .../apache/beam/it/iceberg/IcebergIOLT.java | 351 +++++++++++++++++- .../org/apache/beam/it/kafka/KafkaIOLT.java | 4 +- .../org/apache/beam/it/kafka/KafkaIOST.java | 2 +- it/truthmatchers/build.gradle | 6 +- .../it/truthmatchers}/ArtifactAsserts.java | 15 +- .../it/truthmatchers}/ArtifactsSubject.java | 21 +- 45 files changed, 459 insertions(+), 174 deletions(-) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/Artifact.java (97%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/ArtifactClient.java (99%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/GcsArtifact.java (96%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/package-info.java (95%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/utils/ArtifactUtils.java (98%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/utils/AvroTestUtil.java (98%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/utils/JsonTestUtil.java (99%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/utils/ParquetTestUtil.java (98%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/artifacts/utils/package-info.java (94%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/storage/GcsResourceManager.java (97%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp => common/src/main/java/org/apache/beam/it/common}/storage/package-info.java (95%) create mode 100644 it/common/src/main/resources/test-artifact.json rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/artifacts/GcsArtifactTest.java (94%) rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/artifacts/utils/ArtifactUtilsTest.java (94%) rename it/{google-cloud-platform/src/test/java/org/apache/beam/it/gcp => common/src/test/java/org/apache/beam/it/common}/storage/GcsResourceManagerTest.java (95%) delete mode 100644 it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers => truthmatchers/src/main/java/org/apache/beam/it/truthmatchers}/ArtifactAsserts.java (79%) rename it/{google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers => truthmatchers/src/main/java/org/apache/beam/it/truthmatchers}/ArtifactsSubject.java (87%) diff --git a/it/clickhouse/build.gradle b/it/clickhouse/build.gradle index 858d0dea9207..9fa1c8634a9f 100644 --- a/it/clickhouse/build.gradle +++ b/it/clickhouse/build.gradle @@ -52,4 +52,4 @@ dependencies { permitUnusedDeclared "com.clickhouse:client-v2:$clickhouse_java_client_version:all" -} \ No newline at end of file +} diff --git a/it/common/build.gradle b/it/common/build.gradle index 7b8559aa6e1f..62dd45ddacf5 100644 --- a/it/common/build.gradle +++ b/it/common/build.gradle @@ -29,31 +29,40 @@ dependencies { implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) implementation project(path: ":sdks:java:core", configuration: "shadow") implementation project(path: ":runners:google-cloud-dataflow-java") - implementation project(path: ":sdks:java:testing:test-utils") implementation 'com.google.cloud:google-cloud-bigquery' implementation 'com.google.cloud:google-cloud-monitoring' + implementation 'com.google.cloud:google-cloud-storage' + implementation 'org.apache.hadoop:hadoop-common:3.3.6' + implementation 'org.apache.avro:avro:1.11.1' + implementation 'org.apache.parquet:parquet-avro:1.15.2' + implementation 'org.apache.parquet:parquet-common:1.15.2' + implementation 'org.apache.parquet:parquet-hadoop:1.15.2' provided 'com.google.api.grpc:proto-google-cloud-monitoring-v3' + implementation library.java.jackson_core + implementation library.java.jackson_databind implementation library.java.gax implementation library.java.google_api_services_dataflow implementation library.java.google_auth_library_credentials implementation library.java.google_auth_library_oauth2_http - implementation library.java.google_api_services_bigquery implementation library.java.vendored_guava_32_1_2_jre implementation library.java.slf4j_api implementation library.java.commons_lang3 implementation library.java.failsafe implementation library.java.google_code_gson + implementation library.java.google_cloud_core + implementation library.java.google_api_client implementation library.java.google_http_client implementation library.java.guava implementation library.java.protobuf_java_util implementation library.java.protobuf_java - + implementation library.java.junit + testImplementation library.java.mockito_inline + testRuntimeOnly library.java.slf4j_simple // TODO: excluding Guava until Truth updates it to >32.1.x testImplementation(library.java.truth) { - exclude group: 'com.google.guava', module: 'guava' + exclude group: 'com.google.guava', module: 'guava' } - testImplementation library.java.junit - testImplementation library.java.mockito_inline - testRuntimeOnly library.java.slf4j_simple + testImplementation project(path: ":runners:direct-java", configuration: "shadow") + testImplementation project(path: ":sdks:java:testing:test-utils") testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") } diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java similarity index 97% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java index 830f44598be7..6371d6eedb1e 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; /** * Represents a single artifact. diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java index e463baffeaf1..b275d35703e1 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; import java.io.IOException; import java.nio.file.Path; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java similarity index 96% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java index c60e3a0fe53a..5f4360ddbb2b 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; import com.google.cloud.storage.Blob; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java similarity index 95% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java index de2a2541d558..44b54f82aeba 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java @@ -17,4 +17,4 @@ */ /** Package for working with test artifacts. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java index 854651ae1f14..3a3bcb2d7ec4 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import static java.util.Arrays.stream; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java index 5cfe4c137e35..58241f9a1db3 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import java.io.ByteArrayOutputStream; import java.io.IOException; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java index 9a83558f7bfc..08c3cdfb0c5a 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.JsonNode; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java index d6ebf825e2c1..fb4326e1b4a7 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import java.io.ByteArrayInputStream; import java.io.File; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java similarity index 94% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java index 53cb4ed1d3af..aaf820b25103 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java @@ -17,4 +17,4 @@ */ /** Package for artifact utilities. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; diff --git a/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java index 32a4eb6eb1b6..af6438a561aa 100644 --- a/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java @@ -192,15 +192,13 @@ private static Iterable> getDistributions( .metrics() .queryMetrics( MetricsFilter.builder() - .addNameFilter( - MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) + .addNameFilter(MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) .build()); return metrics.getDistributions(); } /** Pull Beam pipeline defined metrics given the jobId. */ - public Long getBeamMetric( - String jobId, PipelineMetricsType metricType, String metricName) { + public Long getBeamMetric(String jobId, PipelineMetricsType metricType, String metricName) { PipelineResult pipelineResult = MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null)); if (pipelineResult != null) { @@ -209,8 +207,7 @@ public Long getBeamMetric( .metrics() .queryMetrics( MetricsFilter.builder() - .addNameFilter( - MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) + .addNameFilter(MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) .build()); switch (metricType) { @@ -222,9 +219,7 @@ public Long getBeamMetric( return metricResult.getAttempted(); } catch (NoSuchElementException e) { LOG.error( - "Failed to get metric {}, from namespace {}", - metricName, - BEAM_METRICS_NAMESPACE); + "Failed to get metric {}, from namespace {}", metricName, BEAM_METRICS_NAMESPACE); } return UNKNOWN_METRIC_VALUE; case STARTTIME: @@ -273,8 +268,7 @@ public Double getMetric(String project, String region, String jobId, String metr String.format( "Invalid Beam metrics name: %s, expected: '%s:metric_type:metric_name'", metricName, BEAM_METRICS_NAMESPACE)); - PipelineMetricsType metricType = - PipelineMetricsType.valueOf(nameSpacedMetrics[1]); + PipelineMetricsType metricType = PipelineMetricsType.valueOf(nameSpacedMetrics[1]); // Pipeline defined metrics are long values. Have to cast to double that is what the base // class defined. @@ -449,7 +443,9 @@ private List extractOptions(String project, String region, LaunchConfig // add pipeline options from beamTestPipelineOptions system property to preserve the // pipeline options already set in TestPipeline. @Nullable - String beamTestPipelineOptions = System.getProperty(org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS); + String beamTestPipelineOptions = + System.getProperty( + org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS); if (!Strings.isNullOrEmpty(beamTestPipelineOptions)) { try { additionalOptions.addAll(MAPPER.readValue(beamTestPipelineOptions, List.class)); @@ -457,7 +453,8 @@ private List extractOptions(String project, String region, LaunchConfig throw new RuntimeException( "Unable to instantiate test options from system property " + org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS - + ":" + beamTestPipelineOptions, + + ":" + + beamTestPipelineOptions, e); } } diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/GcsResourceManager.java b/it/common/src/main/java/org/apache/beam/it/common/storage/GcsResourceManager.java similarity index 97% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/GcsResourceManager.java rename to it/common/src/main/java/org/apache/beam/it/common/storage/GcsResourceManager.java index 730ca23ee54b..85f2c8b0b018 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/GcsResourceManager.java +++ b/it/common/src/main/java/org/apache/beam/it/common/storage/GcsResourceManager.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.storage; +package org.apache.beam.it.common.storage; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; @@ -39,10 +39,10 @@ import java.util.function.Consumer; import java.util.regex.Pattern; import org.apache.beam.it.common.ResourceManager; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.gcp.artifacts.ArtifactClient; -import org.apache.beam.it.gcp.artifacts.GcsArtifact; -import org.apache.beam.it.gcp.artifacts.utils.ArtifactUtils; +import org.apache.beam.it.common.artifacts.Artifact; +import org.apache.beam.it.common.artifacts.ArtifactClient; +import org.apache.beam.it.common.artifacts.GcsArtifact; +import org.apache.beam.it.common.artifacts.utils.ArtifactUtils; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.junit.rules.TestName; import org.slf4j.Logger; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/storage/package-info.java similarity index 95% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/storage/package-info.java index efcea3aba267..1c5b3b3ce21e 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/storage/package-info.java @@ -17,4 +17,4 @@ */ /** Package for managing Google Cloud Storage resources within integration tests. */ -package org.apache.beam.it.gcp.storage; +package org.apache.beam.it.common.storage; diff --git a/it/common/src/main/resources/test-artifact.json b/it/common/src/main/resources/test-artifact.json new file mode 100644 index 000000000000..551c80d14a66 --- /dev/null +++ b/it/common/src/main/resources/test-artifact.json @@ -0,0 +1 @@ +["This is a test artifact."] \ No newline at end of file diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/GcsArtifactTest.java b/it/common/src/test/java/org/apache/beam/it/common/artifacts/GcsArtifactTest.java similarity index 94% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/GcsArtifactTest.java rename to it/common/src/test/java/org/apache/beam/it/common/artifacts/GcsArtifactTest.java index 59b7931b603a..470741b399ad 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/GcsArtifactTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/artifacts/GcsArtifactTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; import static com.google.common.truth.Truth.assertThat; import static org.mockito.Mockito.when; @@ -30,7 +30,7 @@ import org.mockito.junit.MockitoJUnit; import org.mockito.junit.MockitoRule; -/** Unit tests for {@link GcsArtifact}. */ +/** Unit tests for {@link org.apache.beam.it.common.artifacts.GcsArtifact}. */ @RunWith(JUnit4.class) public class GcsArtifactTest { @Rule public final MockitoRule mockito = MockitoJUnit.rule(); diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtilsTest.java b/it/common/src/test/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtilsTest.java similarity index 94% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtilsTest.java rename to it/common/src/test/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtilsTest.java index 491d32d0ce6a..137c35aaf4bb 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtilsTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtilsTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -24,7 +24,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Artifacts for {@link ArtifactUtils}. */ +/** Artifacts for {@link org.apache.beam.it.common.artifacts.utils.ArtifactUtils}. */ @RunWith(JUnit4.class) public final class ArtifactUtilsTest { diff --git a/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java index eaff80740727..7fc848b45292 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java @@ -21,7 +21,6 @@ import static org.junit.Assert.assertThrows; import java.util.Arrays; - import org.junit.Test; /** Unit tests for {@link BigQueryResourceManagerUtils}. */ diff --git a/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java index 333549c621be..35510a0f1d1c 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java @@ -23,8 +23,8 @@ import java.io.IOException; import java.time.Instant; -import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testutils.metrics.TimeMonitor; diff --git a/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java index 8a1a94df1aec..c35f402bd8ac 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java @@ -17,13 +17,14 @@ */ package org.apache.beam.it.common.dataflow; +import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE; + import com.google.cloud.Timestamp; import java.io.IOException; import java.util.ArrayList; import java.util.Collection; import java.util.Map; import java.util.UUID; - import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; @@ -40,8 +41,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE; - /** Base class for IO Load tests. */ @RunWith(JUnit4.class) @SuppressWarnings({ diff --git a/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java index 360712835062..cd1b71dc7552 100644 --- a/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java @@ -17,8 +17,8 @@ */ package org.apache.beam.it.common.dataflow; -import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.RUNNER_V2; +import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; import com.google.api.gax.core.CredentialsProvider; import com.google.api.gax.core.FixedCredentialsProvider; @@ -40,7 +40,6 @@ import java.util.Map; import java.util.Map.Entry; import java.util.regex.Pattern; - import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; import org.apache.beam.it.common.PipelineOperator; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/GcsResourceManagerTest.java b/it/common/src/test/java/org/apache/beam/it/common/storage/GcsResourceManagerTest.java similarity index 95% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/GcsResourceManagerTest.java rename to it/common/src/test/java/org/apache/beam/it/common/storage/GcsResourceManagerTest.java index 0153573feaed..70dc6dcbc9c4 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/GcsResourceManagerTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/storage/GcsResourceManagerTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.storage; +package org.apache.beam.it.common.storage; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -35,6 +35,7 @@ import com.google.cloud.storage.Storage; import com.google.cloud.storage.Storage.BlobListOption; import com.google.cloud.storage.Storage.BucketListOption; +import com.google.common.truth.Truth; import java.io.IOException; import java.net.URISyntaxException; import java.nio.file.Files; @@ -43,8 +44,8 @@ import java.util.List; import java.util.UUID; import java.util.regex.Pattern; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.gcp.artifacts.GcsArtifact; +import org.apache.beam.it.common.artifacts.Artifact; +import org.apache.beam.it.common.artifacts.GcsArtifact; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Resources; import org.junit.After; @@ -128,7 +129,7 @@ public void testCreateArtifactInRunDir() { verify(client).create(blobInfoCaptor.capture(), contentsCaptor.capture()); BlobInfo actualInfo = blobInfoCaptor.getValue(); - assertThat(actual.blob).isSameInstanceAs(blob); + Truth.assertThat(actual.blob).isSameInstanceAs(blob); assertThat(actualInfo.getBucket()).isEqualTo(BUCKET); assertThat(actualInfo.getName()) .isEqualTo(String.format("%s/%s/%s", TEST_CLASS, gcsClient.runId(), artifactName)); @@ -144,7 +145,7 @@ public void testUploadArtifact() throws IOException { verify(client).create(blobInfoCaptor.capture(), contentsCaptor.capture()); BlobInfo actualInfo = blobInfoCaptor.getValue(); - assertThat(actual.blob).isSameInstanceAs(blob); + Truth.assertThat(actual.blob).isSameInstanceAs(blob); assertThat(actualInfo.getBucket()).isEqualTo(BUCKET); assertThat(actualInfo.getName()) .isEqualTo(String.format("%s/%s/%s", TEST_CLASS, gcsClient.runId(), ARTIFACT_NAME)); @@ -185,8 +186,8 @@ public void testListArtifactsInMethodDirSinglePage() { BlobListOption actualOptions = listOptionsCaptor.getValue(); assertThat(actual).hasSize(2); - assertThat(actual.get(0).name()).isEqualTo(name1); - assertThat(actual.get(1).name()).isEqualTo(name3); + Truth.assertThat(actual.get(0).name()).isEqualTo(name1); + Truth.assertThat(actual.get(1).name()).isEqualTo(name3); assertThat(actualBucket).isEqualTo(BUCKET); assertThat(actualOptions) .isEqualTo( @@ -221,8 +222,8 @@ public void testListArtifactsInMethodDirMultiplePages() { BlobListOption actualOptions = listOptionsCaptor.getValue(); assertThat(actual).hasSize(2); - assertThat(actual.get(0).name()).isEqualTo(name1); - assertThat(actual.get(1).name()).isEqualTo(name3); + Truth.assertThat(actual.get(0).name()).isEqualTo(name1); + Truth.assertThat(actual.get(1).name()).isEqualTo(name3); assertThat(actualBucket).isEqualTo(BUCKET); assertThat(actualOptions) .isEqualTo( diff --git a/it/google-cloud-platform/build.gradle b/it/google-cloud-platform/build.gradle index dd8376e35693..164a75c06ba0 100644 --- a/it/google-cloud-platform/build.gradle +++ b/it/google-cloud-platform/build.gradle @@ -31,47 +31,29 @@ evaluationDependsOn(":it:common") dependencies { implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) - implementation project(path: ":sdks:java:core", configuration: "shadow") - implementation project(path: ":runners:google-cloud-dataflow-java") implementation project(path: ":it:common") implementation project(path: ":it:conditions") implementation project(path: ":it:truthmatchers") - implementation project(path: ":sdks:java:testing:test-utils") implementation library.java.failsafe implementation library.java.google_api_services_dataflow implementation library.java.junit implementation library.java.slf4j_api - implementation(library.java.truth) { - exclude group: 'com.google.guava', module: 'guava' - } implementation library.java.vendored_guava_32_1_2_jre implementation library.java.jackson_core implementation library.java.jackson_databind - implementation 'org.apache.hadoop:hadoop-common:3.3.5' - implementation 'org.apache.avro:avro:1.11.1' - implementation 'org.apache.parquet:parquet-avro:1.15.2' - implementation 'org.apache.parquet:parquet-common:1.15.2' - implementation 'org.apache.parquet:parquet-hadoop:1.15.2' implementation 'org.apache.commons:commons-lang3:3.9' implementation library.java.gax implementation library.java.google_api_common implementation library.java.google_api_client - implementation library.java.google_api_services_bigquery implementation library.java.google_auth_library_credentials implementation library.java.google_auth_library_oauth2_http - implementation library.java.google_http_client implementation library.java.guava - implementation library.java.protobuf_java_util implementation library.java.protobuf_java implementation library.java.threetenbp implementation 'org.awaitility:awaitility:4.2.0' - implementation 'joda-time:joda-time:2.10.10' // Google Cloud Dependencies implementation library.java.google_cloud_core - implementation 'com.google.cloud:google-cloud-storage' implementation 'com.google.cloud:google-cloud-bigquery' - implementation 'com.google.cloud:google-cloud-monitoring' - provided 'com.google.api.grpc:proto-google-cloud-monitoring-v3' implementation 'com.google.cloud:google-cloud-bigtable' implementation 'com.google.cloud:google-cloud-spanner' implementation 'com.google.cloud:google-cloud-pubsub' diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java deleted file mode 100644 index 9759071b397e..000000000000 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * 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 for Artifact Truth matchers / subjects to have reusable assertions. */ -package org.apache.beam.it.gcp.artifacts.matchers; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java index 57b7d5ac99c9..da0cf5e55dbe 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java @@ -19,8 +19,8 @@ import com.google.auto.value.AutoValue; import com.google.cloud.bigquery.TableId; -import org.apache.beam.it.conditions.ConditionCheck; import org.apache.beam.it.common.bigquery.BigQueryResourceManager; +import org.apache.beam.it.conditions.ConditionCheck; import org.checkerframework.checker.nullness.qual.Nullable; /** ConditionCheck to validate if BigQuery has received a certain number of rows. */ diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java index 296f7092cfaa..7908246dce13 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java @@ -30,7 +30,6 @@ import com.google.auth.http.HttpCredentialsAdapter; import dev.failsafe.Failsafe; import java.io.IOException; - import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java index 3e7270dbd34f..496c4b0e980e 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java @@ -32,7 +32,6 @@ import com.google.auth.http.HttpCredentialsAdapter; import dev.failsafe.Failsafe; import java.io.IOException; - import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java index 5a101e08d375..1778d70cf60a 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java @@ -17,7 +17,7 @@ */ package org.apache.beam.it.gcp.spanner.matchers; -import static org.apache.beam.it.gcp.artifacts.utils.JsonTestUtil.parseJsonString; +import static org.apache.beam.it.common.artifacts.utils.JsonTestUtil.parseJsonString; import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatRecords; import com.google.cloud.spanner.Mutation; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java index 9719caa59f1b..16cf2b8f4d31 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java @@ -25,12 +25,11 @@ import java.time.Duration; import java.util.Arrays; import java.util.List; - -import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.PipelineLauncher.LaunchConfig; import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; import org.apache.beam.it.common.PipelineLauncher.Sdk; import org.apache.beam.it.common.PipelineOperator.Result; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowRunner; import org.apache.beam.sdk.io.TextIO; import org.apache.beam.sdk.options.PipelineOptions; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java index b734f93478ef..9ac03db1c741 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java @@ -44,8 +44,9 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.bigquery.BigQueryResourceManager; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java index 4c86eb139bf9..495a4d4485d8 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java @@ -42,8 +42,9 @@ import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.bigquery.BigQueryResourceManager; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOStressTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java index 03efb42dafb6..27773121e79d 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java @@ -34,8 +34,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; import org.apache.beam.sdk.testing.TestPipeline; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java index 5e70a0e2d582..b86f13a0de10 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOStressTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java index 98092152ad22..47058c75e399 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOStressTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java index 360073ac0e9e..da43fca2eaa3 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java index 99682746d2b3..4441b8dcd6b6 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; import org.apache.beam.sdk.io.synthetic.SyntheticSourceOptions; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java index 01d8deef1773..4eb7030c23c3 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOStressTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java index abec875981fc..ac1a7fc103c6 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java @@ -33,8 +33,10 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.storage.GcsResourceManager; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.sdk.io.Compression; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.TextIO; diff --git a/it/iceberg/build.gradle b/it/iceberg/build.gradle index 1a61f86367b2..f11374a1ca4f 100644 --- a/it/iceberg/build.gradle +++ b/it/iceberg/build.gradle @@ -31,7 +31,6 @@ ext.summary = "Integration test utilities for Iceberg." def iceberg_version = "1.10.0" def parquet_version = "1.16.0" -def orc_version = "1.9.6" def hive_version = "3.1.3" def hadoopVersions = [ @@ -43,27 +42,8 @@ hadoopVersions.each {kv -> configurations.create("hadoopVersion$kv.key")} evaluationDependsOn(":it:common") dependencies { - implementation library.java.vendored_guava_32_1_2_jre - implementation project(path: ":sdks:java:core", configuration: "shadow") - implementation project(path: ":model:pipeline", configuration: "shadow") - implementation library.java.avro - implementation library.java.slf4j_api - implementation library.java.joda_time - implementation "org.apache.parquet:parquet-column:$parquet_version" - implementation "org.apache.parquet:parquet-hadoop:$parquet_version" - implementation "org.apache.parquet:parquet-common:$parquet_version" - implementation project(":sdks:java:io:parquet") - implementation "org.apache.orc:orc-core:$orc_version" - implementation "org.apache.iceberg:iceberg-core:$iceberg_version" - implementation "org.apache.iceberg:iceberg-api:$iceberg_version" - implementation "org.apache.iceberg:iceberg-parquet:$iceberg_version" - implementation "org.apache.iceberg:iceberg-orc:$iceberg_version" - implementation "org.apache.iceberg:iceberg-data:$iceberg_version" - implementation "org.apache.hadoop:hadoop-common:3.3.6" - provided "org.immutables:value:2.8.8" permitUnusedDeclared "org.immutables:value:2.8.8" - implementation library.java.vendored_calcite_1_40_0 runtimeOnly "org.apache.iceberg:iceberg-gcp:$iceberg_version" runtimeOnly "org.apache.iceberg:iceberg-aws:$iceberg_version" runtimeOnly "org.apache.iceberg:iceberg-aws-bundle:$iceberg_version" @@ -113,28 +93,10 @@ dependencies { "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-mapreduce-client-core:$kv.value" } - implementation project(path: ":runners:google-cloud-dataflow-java") - implementation project(path: ":it:common") - implementation project(path: ":it:conditions") - implementation project(path: ":it:truthmatchers") - implementation project(path: ":sdks:java:testing:test-utils") - implementation library.java.failsafe - implementation library.java.google_api_services_dataflow - implementation(library.java.truth) { - exclude group: 'com.google.guava', module: 'guava' - } - implementation library.java.jackson_core - implementation library.java.jackson_databind - implementation 'org.apache.commons:commons-lang3:3.9' - implementation library.java.google_api_client - implementation library.java.google_http_client - implementation library.java.guava - implementation library.java.protobuf_java_util - implementation library.java.protobuf_java - implementation library.java.threetenbp - testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:synthetic") + testImplementation project(path: ":sdks:java:io:iceberg") + testImplementation project(path: ":it:truthmatchers") testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") testImplementation library.java.mockito_inline testImplementation project(path: ":sdks:java:extensions:google-cloud-platform-core", configuration: "testRuntimeMigration") diff --git a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java index 3f486ec3f5d0..683ff4529b68 100644 --- a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java +++ b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java @@ -1,7 +1,356 @@ +/* + * 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.it.iceberg; +import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; +import static org.junit.Assert.assertEquals; + +import com.google.auto.value.AutoValue; +import java.io.IOException; +import java.text.ParseException; +import java.time.Duration; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.Map; +import java.util.Random; +import java.util.UUID; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.PipelineOperator; +import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.storage.GcsResourceManager; +import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.sdk.io.GenerateSequence; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergIO; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestPipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.catalog.TableIdentifier; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; + +/** + * IcebergIO performance tests using a Hadoop catalog backed by GCS. + * + *

The test writes generated rows to an Iceberg table, reads the table back, validates the number + * of rows, and exports write and read performance metrics. + * + *

Example trigger command: + * + * 

+ * ./gradlew :it:iceberg:integrationTest \
+ *   --tests "org.apache.beam.it.iceberg.IcebergIOLT" \
+ *   -Dconfiguration=local \
+ *   -Dproject=[gcpProject] \
+ *   -DartifactBucket=[artifactBucket]
+ *
+ */ +public final class IcebergIOLT extends IOLoadTestBase { + + private static final String READ_ELEMENT_METRIC_NAME = "read_count"; + private static final String NAMESPACE = "default"; + + private static final Schema ROW_SCHEMA = + Schema.builder().addInt64Field("id").addByteArrayField("data").build(); + + private static final Map TEST_CONFIGS = + ImmutableMap.of( + "local", + Configuration.of(1_000L, 5, "DirectRunner", 1_000), // approximately 1 MB + "small", + Configuration.of(1_000_000L, 20, "DataflowRunner", 1_000), // approximately 1 GB + "medium", + Configuration.of(10_000_000L, 40, "DataflowRunner", 1_000), // approximately 10 GB + "large", + Configuration.of(100_000_000L, 100, "DataflowRunner", 1_000) // approximately 100 GB + ); + + private static GcsResourceManager resourceManager; + + private Configuration configuration; + private IcebergCatalogConfig catalogConfig; + private TableIdentifier tableIdentifier; + + @Rule public TestPipeline writePipeline = TestPipeline.create(); + + @Rule public TestPipeline readPipeline = TestPipeline.create(); + + @BeforeClass + public static void beforeClass() { + resourceManager = + GcsResourceManager.builder(TestProperties.artifactBucket(), "icebergiolt", CREDENTIALS) + .build(); + } + + @AfterClass + public static void tearDownClass() { + ResourceManagerUtils.cleanResources(resourceManager); + } + + @Before + public void setup() { + configuration = getTestConfiguration(); + + String uniqueSuffix = + DateTimeFormatter.ofPattern("MMddHHmmssSSS") + .withZone(ZoneOffset.UTC) + .format(java.time.Instant.now()) + + "-" + + UUID.randomUUID().toString().substring(0, 10); + + String warehouseDirectory = "icebergiolt-" + uniqueSuffix; + resourceManager.registerTempDir(warehouseDirectory); + + String warehouseLocation = + String.format("gs://%s/%s", TestProperties.artifactBucket(), warehouseDirectory); + + tableIdentifier = TableIdentifier.of(NAMESPACE, "table_" + uniqueSuffix.replace("-", "_")); + + Map catalogProperties = + ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation); + + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("iceberg-load-test") + .setCatalogProperties(catalogProperties) + .build(); + + setTempLocation(writePipeline); + setTempLocation(readPipeline); + } + + /** Writes generated rows to Iceberg and verifies that all rows can be read back. */ + @Test + public void testIcebergWriteAndRead() throws IOException { + PipelineLauncher.LaunchInfo writeInfo = testWrite(); + + PipelineOperator.Result writeResult = + pipelineOperator.waitUntilDone( + createConfig(writeInfo, Duration.ofMinutes(configuration.getPipelineTimeout()))); + + assertThatResult(writeResult).isLaunchFinished(); + assertEquals( + PipelineLauncher.JobState.DONE, + pipelineLauncher.getJobStatus(project, region, writeInfo.jobId())); + + PipelineLauncher.LaunchInfo readInfo = testRead(); + + PipelineOperator.Result readResult = + pipelineOperator.waitUntilDone( + createConfig(readInfo, Duration.ofMinutes(configuration.getPipelineTimeout()))); + + assertThatResult(readResult).isLaunchFinished(); + assertEquals( + PipelineLauncher.JobState.DONE, + pipelineLauncher.getJobStatus(project, region, readInfo.jobId())); + + double numRecords = + pipelineLauncher.getMetric( + project, + region, + readInfo.jobId(), + getBeamMetricsName(PipelineMetricsType.COUNTER, READ_ELEMENT_METRIC_NAME)); + + assertEquals((double) configuration.getNumRows(), numRecords, 0.5); + + exportMetrics(writeInfo, readInfo); + } + + private PipelineLauncher.LaunchInfo testWrite() throws IOException { + writePipeline + .apply("Generate records", GenerateSequence.from(0).to(configuration.getNumRows())) + .apply("Map records", ParDo.of(new MapToIcebergFormat(configuration.getValueSizeBytes()))) + .setRowSchema(ROW_SCHEMA) + .apply("Write to Iceberg", IcebergIO.writeRows(catalogConfig).to(tableIdentifier)); + + PipelineLauncher.LaunchConfig options = + PipelineLauncher.LaunchConfig.builder("write-iceberg") + .setSdk(PipelineLauncher.Sdk.JAVA) + .setPipeline(writePipeline) + .addParameter("runner", configuration.getRunner()) + .addParameter( + "maxNumWorkers", + TestProperties.getProperty("maxNumWorkers", "10", TestProperties.Type.PROPERTY)) + .build(); + + return pipelineLauncher.launch(project, region, options); + } + + private PipelineLauncher.LaunchInfo testRead() throws IOException { + readPipeline + .apply("Read from Iceberg", IcebergIO.readRows(catalogConfig).from(tableIdentifier)) + .apply("Counting element", ParDo.of(new CountingFn<>(READ_ELEMENT_METRIC_NAME))); + + PipelineLauncher.LaunchConfig options = + PipelineLauncher.LaunchConfig.builder("read-iceberg") + .setSdk(PipelineLauncher.Sdk.JAVA) + .setPipeline(readPipeline) + .addParameter("runner", configuration.getRunner()) + .addParameter( + "maxNumWorkers", + TestProperties.getProperty("maxNumWorkers", "10", TestProperties.Type.PROPERTY)) + .build(); + + return pipelineLauncher.launch(project, region, options); + } + + private void exportMetrics( + PipelineLauncher.LaunchInfo writeInfo, PipelineLauncher.LaunchInfo readInfo) { + + /* + * Only use PCollection names belonging to transforms defined directly in this test. + * Internal IcebergIO transform names may change when the implementation changes. + */ + MetricsConfiguration writeMetricsConfig = + MetricsConfiguration.builder() + .setInputPCollection("Map records.out0") + .setInputPCollectionV2("Map records/ParMultiDo(MapToIcebergFormat).out0") + .build(); + + MetricsConfiguration readMetricsConfig = + MetricsConfiguration.builder() + .setOutputPCollection("Counting element.out0") + .setOutputPCollectionV2("Counting element/ParMultiDo(Counting).out0") + .build(); + + try { + exportMetricsToBigQuery(writeInfo, getMetrics(writeInfo, writeMetricsConfig)); + + exportMetricsToBigQuery(readInfo, getMetrics(readInfo, readMetricsConfig)); + } catch (IOException | ParseException | InterruptedException e) { + throw new RuntimeException("Unable to collect or export IcebergIO load-test metrics.", e); + } + } + + private Configuration getTestConfiguration() { + String testConfig = + TestProperties.getProperty("configuration", "local", TestProperties.Type.PROPERTY); + + Configuration selectedConfiguration = TEST_CONFIGS.get(testConfig); + + if (selectedConfiguration == null) { + throw new IllegalArgumentException( + String.format( + "Unknown test configuration: [%s]. Known configurations: %s", + testConfig, TEST_CONFIGS.keySet())); + } + + return selectedConfiguration; + } + + private static void setTempLocation(TestPipeline pipeline) { + String tempBucketName = TestProperties.artifactBucket(); + + if (Strings.isNullOrEmpty(tempBucketName)) { + return; + } + + String tempLocation = String.format("gs://%s/temp/", tempBucketName); + + pipeline.getOptions().as(TestPipelineOptions.class).setTempRoot(tempLocation); + + pipeline.getOptions().setTempLocation(tempLocation); + } + + /** Maps a generated sequence value to a deterministic Iceberg-compatible Beam row. */ + private static class MapToIcebergFormat extends DoFn { + + private final int valueSizeBytes; + + private MapToIcebergFormat(int valueSizeBytes) { + if (valueSizeBytes <= 0) { + throw new IllegalArgumentException("valueSizeBytes must be positive."); + } + + this.valueSizeBytes = valueSizeBytes; + } + + @ProcessElement + public void processElement(@Element Long index, OutputReceiver output) { + + byte[] value = new byte[valueSizeBytes]; + + /* + * Using the row index as the seed makes retries deterministic. Random data also prevents + * compression from making the generated payload significantly smaller than configured. + */ + new Random(index).nextBytes(value); + + output.output(Row.withSchema(ROW_SCHEMA).addValues(index, value).build()); + } + } + + /** Options for the IcebergIO load test. */ + @AutoValue + abstract static class Configuration { + + abstract long getNumRows(); + + abstract int getPipelineTimeout(); + + abstract String getRunner(); + + abstract int getValueSizeBytes(); + + static Configuration of(long numRows, int pipelineTimeout, String runner, int valueSizeBytes) { + + if (numRows <= 0) { + throw new IllegalArgumentException("numRows must be positive."); + } + + if (pipelineTimeout <= 0) { + throw new IllegalArgumentException("pipelineTimeout must be positive."); + } + + if (valueSizeBytes <= 0) { + throw new IllegalArgumentException("valueSizeBytes must be positive."); + } + + return new AutoValue_IcebergIOLT_Configuration.Builder() + .setNumRows(numRows) + .setPipelineTimeout(pipelineTimeout) + .setRunner(runner) + .setValueSizeBytes(valueSizeBytes) + .build(); + } + + @AutoValue.Builder + abstract static class Builder { + + abstract Builder setNumRows(long numRows); + + abstract Builder setPipelineTimeout(int pipelineTimeout); + + abstract Builder setRunner(String runner); -public class IcebergIOLT extends IOLoadTestBase { + abstract Builder setValueSizeBytes(int valueSizeBytes); + abstract Configuration build(); + } + } } diff --git a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java index 4b0ecbfd3433..1607ca1cd674 100644 --- a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java +++ b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java @@ -27,12 +27,12 @@ import java.time.format.DateTimeFormatter; import java.util.Map; import java.util.UUID; -import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.io.kafka.KafkaIO; import org.apache.beam.sdk.io.synthetic.SyntheticOptions; diff --git a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java index cfc11127c34f..d6231a69f1e7 100644 --- a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java +++ b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java @@ -31,10 +31,10 @@ import java.util.Map; import java.util.Objects; import java.util.UUID; -import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; diff --git a/it/truthmatchers/build.gradle b/it/truthmatchers/build.gradle index fd54ad591a0a..2e21c1b5df04 100644 --- a/it/truthmatchers/build.gradle +++ b/it/truthmatchers/build.gradle @@ -27,11 +27,15 @@ ext.summary = "Truth matchers for integration tests." dependencies { implementation project(path: ":it:common") + implementation 'org.apache.avro:avro:1.11.1' // TODO: excluding Guava until Truth updates it to >32.1.x implementation(library.java.truth) { exclude group: 'com.google.guava', module: 'guava' } + implementation library.java.jackson_core + implementation library.java.jackson_databind implementation library.java.guava + implementation library.java.vendored_guava_32_1_2_jre permitUsedUndeclared library.java.guava implementation library.java.google_code_gson -} \ No newline at end of file +} diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactAsserts.java b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactAsserts.java similarity index 79% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactAsserts.java rename to it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactAsserts.java index 3b1251307d0d..92407e5a2fc9 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactAsserts.java +++ b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactAsserts.java @@ -15,15 +15,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.matchers; - -import static com.google.common.truth.Truth.assertAbout; -import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatRecords; +package org.apache.beam.it.truthmatchers; +import com.google.common.truth.Truth; import java.util.List; import org.apache.avro.generic.GenericRecord; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.truthmatchers.RecordsSubject; +import org.apache.beam.it.common.artifacts.Artifact; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; public class ArtifactAsserts { @@ -36,7 +33,7 @@ public class ArtifactAsserts { * @return Truth Subject to chain assertions. */ public static ArtifactsSubject assertThatArtifacts(List artifacts) { - return assertAbout(ArtifactsSubject.records()).that(artifacts); + return Truth.assertAbout(ArtifactsSubject.records()).that(artifacts); } /** @@ -47,7 +44,7 @@ public static ArtifactsSubject assertThatArtifacts(List artifacts) { * @return Truth Subject to chain assertions. */ public static ArtifactsSubject assertThatArtifact(Artifact artifact) { - return assertAbout(ArtifactsSubject.records()).that(ImmutableList.of(artifact)); + return Truth.assertAbout(ArtifactsSubject.records()).that(ImmutableList.of(artifact)); } /** @@ -57,6 +54,6 @@ public static ArtifactsSubject assertThatArtifact(Artifact artifact) { * @return Truth Subject to chain assertions. */ public static RecordsSubject assertThatGenericRecords(List records) { - return assertThatRecords(ArtifactsSubject.genericRecordToRecords(records)); + return PipelineAsserts.assertThatRecords(ArtifactsSubject.genericRecordToRecords(records)); } } diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactsSubject.java b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactsSubject.java similarity index 87% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactsSubject.java rename to it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactsSubject.java index ab5b49699484..e2518280345d 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactsSubject.java +++ b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactsSubject.java @@ -15,11 +15,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.matchers; +package org.apache.beam.it.truthmatchers; -import static org.apache.beam.it.gcp.artifacts.matchers.ArtifactAsserts.assertThatGenericRecords; -import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatRecords; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing.sha256; +import static org.apache.beam.it.truthmatchers.ArtifactAsserts.assertThatGenericRecords; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; @@ -31,11 +29,11 @@ import java.util.Map; import org.apache.avro.Schema; import org.apache.avro.generic.GenericRecord; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.gcp.artifacts.utils.AvroTestUtil; -import org.apache.beam.it.gcp.artifacts.utils.JsonTestUtil; -import org.apache.beam.it.gcp.artifacts.utils.ParquetTestUtil; -import org.apache.beam.it.truthmatchers.RecordsSubject; +import org.apache.beam.it.common.artifacts.Artifact; +import org.apache.beam.it.common.artifacts.utils.AvroTestUtil; +import org.apache.beam.it.common.artifacts.utils.JsonTestUtil; +import org.apache.beam.it.common.artifacts.utils.ParquetTestUtil; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; /** * Subject that has assertion operations for artifact lists (GCS files), usually coming from the @@ -98,7 +96,8 @@ public void hasContent(String content) { */ public void hasHash(String hash) { if (actual.stream() - .noneMatch(artifact -> sha256().hashBytes(artifact.contents()).toString().equals(hash))) { + .noneMatch( + artifact -> Hashing.sha256().hashBytes(artifact.contents()).toString().equals(hash))) { failWithActual("expected to contain hash", hash); } } @@ -167,6 +166,6 @@ public RecordsSubject asJsonRecords() { throw new RuntimeException("Error reading " + artifact.name() + " as JSON.", e); } } - return assertThatRecords(allRecords); + return PipelineAsserts.assertThatRecords(allRecords); } } From 7fb07784bf2c7b01480d65cc1d8319599982f58b Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Thu, 9 Jul 2026 14:42:45 +0400 Subject: [PATCH 05/29] Add looker metrics for IcebergIO --- ...m_PostCommit_Java_IO_Performance_Tests.yml | 2 +- .test-infra/tools/refresh_looker_metrics.py | 2 + .../www/site/content/en/performance/_index.md | 1 + .../en/performance/icebergio/_index.md | 50 +++++++++++++++++++ website/www/site/data/performance.yaml | 31 ++++++++++++ 5 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 website/www/site/content/en/performance/icebergio/_index.md diff --git a/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml b/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml index 0bc0a37ea4ea..9bda3d1bf31f 100644 --- a/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml +++ b/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml @@ -64,7 +64,7 @@ jobs: matrix: job_name: ["beam_PostCommit_Java_IO_Performance_Tests"] job_phrase: ["Run Java PostCommit IO Performance Tests"] - test_case: ["IcebergPerformanceTest"] + test_case: ["IcebergPerformanceTest", "GCSPerformanceTest", "BigTablePerformanceTest", "BigQueryStorageApiStreamingPerformanceTest"] steps: - uses: actions/checkout@v7 - name: Setup repository diff --git a/.test-infra/tools/refresh_looker_metrics.py b/.test-infra/tools/refresh_looker_metrics.py index c8d66f4a4bd3..573cee0d4e2b 100644 --- a/.test-infra/tools/refresh_looker_metrics.py +++ b/.test-infra/tools/refresh_looker_metrics.py @@ -34,6 +34,8 @@ ("33", ["21", "70", "116", "69", "115"]), # BigTableIO_Write ("34", ["22", "56", "96", "55", "95"]), # TextIO_Read ("35", ["23", "64", "110", "63", "109"]), # TextIO_Write + ("113", ["386", "388", "390", "392", "394"]), # IcebergIO_Read + ("114", ["387", "389", "391", "393", "395"]), # IcebergIO_Write ("75", ["258", "259", "260", "261", "262"]), # TensorFlow MNIST ("76", ["233", "234", "235", "236", "237"]), # PyTorch BERT base uncased ("77", ["238", "239", "240", "241", "242"]), # PyTorch BERT large uncased diff --git a/website/www/site/content/en/performance/_index.md b/website/www/site/content/en/performance/_index.md index a0eaba2aa0ec..0181965fc105 100644 --- a/website/www/site/content/en/performance/_index.md +++ b/website/www/site/content/en/performance/_index.md @@ -38,6 +38,7 @@ writing to various Beam IOs. - [BigQuery](/performance/bigquery) - [BigTable](/performance/bigtable) - [TextIO](/performance/textio) +- [IcebergIO](/performance/icebergio) # Measured Beam Python ML Pipelines diff --git a/website/www/site/content/en/performance/icebergio/_index.md b/website/www/site/content/en/performance/icebergio/_index.md new file mode 100644 index 000000000000..e618865097b4 --- /dev/null +++ b/website/www/site/content/en/performance/icebergio/_index.md @@ -0,0 +1,50 @@ +--- +title: "IcebergIO Performance" +--- + + + +# IcebergIO Performance + +The following graphs show various metrics when reading from or writing to Apache Iceberg tables using +IcebergIO. See the [glossary](/performance/glossary) for definitions. + +## Read + +### What is the estimated cost of reading from Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="read" section="test_name" >}} + +### How has various metrics changed when reading from Apache Iceberg tables using IcebergIO for different Beam SDK versions? + +{{< performance_looks io="icebergio" read_or_write="read" section="version" >}} + +### How has various metrics changed over time when reading from Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="read" section="date" >}} + +## Write + +### What is the estimated cost of writing to Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="write" section="test_name" >}} + +### How has various metrics changed when writing to Apache Iceberg tables using IcebergIO for different Beam SDK versions? + +{{< performance_looks io="icebergio" read_or_write="write" section="version" >}} + +### How has various metrics changed over time when writing to Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="write" section="date" >}} diff --git a/website/www/site/data/performance.yaml b/website/www/site/data/performance.yaml index 42fb5bee92f5..e996ea262911 100644 --- a/website/www/site/data/performance.yaml +++ b/website/www/site/data/performance.yaml @@ -107,6 +107,37 @@ looks: title: AvgInputThroughputBytesPerSec by Version - id: fVVHhXCrHNgBG52TJsTjR8VbmWCCQnVN title: AvgInputThroughputElementsPerSec by Version + icebergio: + read: + folder: 113 + test_name: + - id: x5R7w8VnzkccgzrBTWQ82J7JBNtkGxH2 + title: Read IcebergIO RunTime and EstimatedCost + date: + - id: HQsKkmn7cS7m5bCYzPBYqnnFM6MXvJgV + title: AvgOutputThroughputBytesPerSec by Date + - id: MNYqD2ZCwV2XfYjwDg4NDVT7vr4Ytpz2 + title: AvgOutputThroughputElementsPerSec by Date + version: + - id: Cf7wGQfrM4hn8Wdf64ZjwshrHHdZKmKs + title: AvgOutputThroughputBytesPerSec by Version + - id: MGSX85MhCnCFFWYzSCGS36JRWxcBccqd + title: AvgOutputThroughputElementsPerSec by Version + write: + folder: 114 + test_name: + - id: yvvWB2JbbbFQ4R5MTwF2hbVW2X4bDxnY + title: Write IcebergIO RunTime and EstimatedCost + date: + - id: Gcw6jfhvvkkHKDwqqTkhZtYYNNH2Z8DZ + title: AvgInputThroughputBytesPerSec by Date + - id: RYtRGKsxBRkyRgWk6RGjshFnZPk76vhM + title: AvgInputThroughputElementsPerSec by Date + version: + - id: B5q25ktnXNhHxpYrGmqzgcGMgX26zvM3 + title: AvgInputThroughputBytesPerSec by Version + - id: VfMWk6rk26JQNr8wRpsK5Gk6XQPfd7vW + title: AvgInputThroughputElementsPerSec by Version pytorchbertbase: write: folder: 76 From 9432acebb2b9b34a852da4e1bdf1de9f43e94b53 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Tue, 14 Jul 2026 16:01:52 +0400 Subject: [PATCH 06/29] Refactoring --- .../src/main/resources/test-artifact.json | 1 - it/iceberg/src/main/resources/test-artifact.json | 1 - .../java/org/apache/beam/it/iceberg/IcebergIOLT.java | 9 ++++++++- 3 files changed, 8 insertions(+), 3 deletions(-) delete mode 100644 it/google-cloud-platform/src/main/resources/test-artifact.json delete mode 100644 it/iceberg/src/main/resources/test-artifact.json diff --git a/it/google-cloud-platform/src/main/resources/test-artifact.json b/it/google-cloud-platform/src/main/resources/test-artifact.json deleted file mode 100644 index 551c80d14a66..000000000000 --- a/it/google-cloud-platform/src/main/resources/test-artifact.json +++ /dev/null @@ -1 +0,0 @@ -["This is a test artifact."] \ No newline at end of file diff --git a/it/iceberg/src/main/resources/test-artifact.json b/it/iceberg/src/main/resources/test-artifact.json deleted file mode 100644 index 551c80d14a66..000000000000 --- a/it/iceberg/src/main/resources/test-artifact.json +++ /dev/null @@ -1 +0,0 @@ -["This is a test artifact."] \ No newline at end of file diff --git a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java index 683ff4529b68..2a6098a1518e 100644 --- a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java +++ b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java @@ -281,6 +281,7 @@ private static void setTempLocation(TestPipeline pipeline) { private static class MapToIcebergFormat extends DoFn { private final int valueSizeBytes; + private transient Random random; private MapToIcebergFormat(int valueSizeBytes) { if (valueSizeBytes <= 0) { @@ -290,6 +291,11 @@ private MapToIcebergFormat(int valueSizeBytes) { this.valueSizeBytes = valueSizeBytes; } + @Setup + public void setup() { + random = new Random(); + } + @ProcessElement public void processElement(@Element Long index, OutputReceiver output) { @@ -299,7 +305,8 @@ public void processElement(@Element Long index, OutputReceiver output) { * Using the row index as the seed makes retries deterministic. Random data also prevents * compression from making the generated payload significantly smaller than configured. */ - new Random(index).nextBytes(value); + random.setSeed(index); + random.nextBytes(value); output.output(Row.withSchema(ROW_SCHEMA).addValues(index, value).build()); } From ab142e93fb6421eca06a6136bd770b95445782b7 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Tue, 14 Jul 2026 16:31:18 +0400 Subject: [PATCH 07/29] Fix conflict --- it/clickhouse/build.gradle | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/it/clickhouse/build.gradle b/it/clickhouse/build.gradle index 9fa1c8634a9f..0727f0328e45 100644 --- a/it/clickhouse/build.gradle +++ b/it/clickhouse/build.gradle @@ -29,6 +29,12 @@ ext.summary = "Integration test utilities for ClickHouse." def clickhouse_java_client_version = "0.9.6" +configurations.configureEach { + resolutionStrategy.capabilitiesResolution.withCapability("org.lz4:lz4-java") { + select("at.yawk.lz4:lz4-java:1.10.2") + } +} + dependencies { implementation project(path: ":it:common") implementation project(path: ":it:testcontainers") From f80294498cf6ebc2da24b825445dea83b5c51767 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 07:07:33 -0400 Subject: [PATCH 08/29] Bump actions/checkout from 4 to 7 (#39399) Bumps [actions/checkout](https://github.com/actions/checkout) from 4 to 7. - [Release notes](https://github.com/actions/checkout/releases) - [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md) - [Commits](https://github.com/actions/checkout/compare/v4...v7) --- updated-dependencies: - dependency-name: actions/checkout dependency-version: '7' dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/codeql.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index be534203eaed..7072b9a8da4a 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -36,7 +36,7 @@ jobs: matrix: ${{ steps.set-matrix.outputs.matrix }} steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v7 - uses: dorny/paths-filter@fbd0ab8f3e69293af611ebaee6363fc25e6d187d id: filter with: From 7735b396afb4f56625b4fc22dfba732c8ef64dc2 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 07:09:02 -0400 Subject: [PATCH 09/29] Bump pillow (#39395) Bumps [pillow](https://github.com/python-pillow/Pillow) from 10.2.0 to 12.3.0. - [Release notes](https://github.com/python-pillow/Pillow/releases) - [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst) - [Commits](https://github.com/python-pillow/Pillow/compare/10.2.0...12.3.0) --- updated-dependencies: - dependency-name: pillow dependency-version: 12.3.0 dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../ml-orchestration/kfp/components/train/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt b/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt index 5b996fdfc70a..2fc217218387 100644 --- a/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt +++ b/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt @@ -15,4 +15,4 @@ torch==2.12.0 numpy==1.22.4 -Pillow==10.2.0 \ No newline at end of file +Pillow==12.3.0 \ No newline at end of file From ab36569659ac2423c39c5fe9da6ef200aa153837 Mon Sep 17 00:00:00 2001 From: Abdelrahman Ibrahim Date: Tue, 21 Jul 2026 16:30:11 +0300 Subject: [PATCH 10/29] fix GCP_PATH expansion (#39400) Co-authored-by: Abdelrahman Ibrahim --- .github/workflows/build_wheels.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml index 6236888521a5..02b223943fea 100644 --- a/.github/workflows/build_wheels.yml +++ b/.github/workflows/build_wheels.yml @@ -37,7 +37,8 @@ concurrency: cancel-in-progress: true env: - GCP_PATH: "gs://${{ secrets.GCP_PYTHON_WHEELS_BUCKET }}/${GITHUB_REF##*/}/${GITHUB_SHA}-${GITHUB_RUN_ID}/" + # Use github.* context; workflow env is not bash-evaluated. + GCP_PATH: "gs://${{ secrets.GCP_PYTHON_WHEELS_BUCKET }}/${{ github.ref_name }}/${{ github.sha }}-${{ github.run_id }}/" jobs: From dca1253c7e3886d125c526aa830b7945daef16ce Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 21 Jul 2026 09:47:05 -0400 Subject: [PATCH 11/29] Fix CdapIO dependency (#39393) --- .../groovy/org/apache/beam/gradle/BeamModulePlugin.groovy | 2 +- sdks/java/io/cdap/build.gradle | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy index d73f2e7a2bae..70a22e397724 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy @@ -613,7 +613,7 @@ class BeamModulePlugin implements Plugin { def autoservice_version = "1.0.1" def aws_java_sdk2_version = "2.20.162" def cassandra_driver_version = "3.10.2" - def cdap_version = "6.5.1" + def cdap_version = "6.11.4" def checkerframework_version = "3.42.0" def classgraph_version = "4.8.162" def delta_lake_version = "4.2.0" diff --git a/sdks/java/io/cdap/build.gradle b/sdks/java/io/cdap/build.gradle index e0af1717e162..cc7e3ba12cde 100644 --- a/sdks/java/io/cdap/build.gradle +++ b/sdks/java/io/cdap/build.gradle @@ -90,3 +90,8 @@ test { // Open java.lang for Gson reflection on StackTraceElement under Java 17+ jvmArgs '--add-opens=java.base/java.lang=ALL-UNNAMED' } + +// spark3_streaming depends on old 'org.lz4:lz4-java'; conflict with kafka-clients:3.9.2 +configurations.all { + resolveCapabilitiesConflict(it, 'org.lz4:lz4-java', 'at.yawk.lz4') +} \ No newline at end of file From 7bc2fb49f8dcadb035b4a6caad3da4a6e6f00ab8 Mon Sep 17 00:00:00 2001 From: tvalentyn Date: Tue, 21 Jul 2026 07:54:39 -0700 Subject: [PATCH 12/29] Fix a memory leak in the failed instruction id cache. (#39405) * Sanitize the exception message stored for failed instructions to avoid inadvertently capturing the content of stack frames in the heap and limit RAM growth due to cache size. * Use Py3.10 syntax --- .../apache_beam/runners/worker/sdk_worker.py | 12 +++++- .../runners/worker/sdk_worker_test.py | 38 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index a79cb0e8de6e..db25a40a405c 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -80,7 +80,7 @@ MAX_KNOWN_NOT_RUNNING_INSTRUCTIONS = 1000 # The number of ProcessBundleRequest instruction ids that BundleProcessorCache # will remember for failed instructions. -MAX_FAILED_INSTRUCTIONS = 10000 +MAX_FAILED_INSTRUCTIONS = 1000 # retry on transient UNAVAILABLE grpc error from state channels. _GRPC_SERVICE_CONFIG = json.dumps({ @@ -559,7 +559,15 @@ def discard(self, instruction_id, exception): """ processor = None with self._lock: - self.failed_instruction_ids[instruction_id] = exception + tb_str = "".join(traceback.format_exception(exception)) + if len(tb_str) > 10240: + tb_str = ( + tb_str[:5000] + "\n... [traceback truncated] ...\n" + + tb_str[-5000:]) + clean_exception = RuntimeError( + f"Original Exception: {type(exception).__name__}: {str(exception)[:2000]}\n{tb_str}" + ) + self.failed_instruction_ids[instruction_id] = clean_exception while len(self.failed_instruction_ids) > MAX_FAILED_INSTRUCTIONS: self.failed_instruction_ids.popitem(last=False) if instruction_id in self.active_bundle_processors: diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py index 76e428f06464..bea313a4d2fd 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -296,6 +296,44 @@ def test_failed_bundle_processor_returns_failed_split_response(self): worker.do_instruction(split_request).error, hc.contains_string('test message')) + def test_failed_instruction_id_cache_size_is_capped(self): + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + with mock.patch( + 'apache_beam.runners.worker.sdk_worker.MAX_FAILED_INSTRUCTIONS', 10): + for i in range(15): + bundle_processor_cache.discard(f'inst_{i}', RuntimeError(f'error {i}')) + + for i in range(5): + self.assertNotIn( + f'inst_{i}', bundle_processor_cache.failed_instruction_ids) + for i in range(5, 15): + self.assertIn( + f'inst_{i}', bundle_processor_cache.failed_instruction_ids) + + def test_failed_instruction_tracebacks_are_truncated_when_too_long(self): + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + long_message = "x" * 15000 + bundle_processor_cache.discard('instruction_id', RuntimeError(long_message)) + + stored_exception = bundle_processor_cache.failed_instruction_ids[ + 'instruction_id'] + tb_str = str(stored_exception) + + self.assertLessEqual(len(tb_str), 13000) + self.assertIn('[traceback truncated]', tb_str) + def test_data_sampling_response(self): # Create a data sampler with some fake sampled data. This data will be seen # in the sample response. From 5b441f270828cdc2796058e623e352da5a13a153 Mon Sep 17 00:00:00 2001 From: tvalentyn Date: Tue, 21 Jul 2026 07:58:00 -0700 Subject: [PATCH 13/29] Update CHANGES.md to document a memory growth issue. (#39407) * Update CHANGES.md * Fix AI review comments --- CHANGES.md | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index d39a62f1a780..991c341460a0 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -84,7 +84,7 @@ * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)). * Improved Java pipeline performance by avoiding repeated `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)). -* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Python) Fixed a memory leak in Python SDK caused by storing exceptions with potentially large stack frames in a cache ([#39406](https://github.com/apache/beam/issues/39406)). ## Security Fixes @@ -138,6 +138,7 @@ ## Known Issues * (Java) Projects using the Flink runner with Flink 2.1 or later alongside libraries requiring `org.lz4:lz4-java` (e.g., Kafka clients) may encounter a Gradle capability conflict, because Flink 2.1+ ships `at.yawk.lz4:lz4-java` which declares the same capability. To resolve, add a `capabilitiesResolution` rule to your `build.gradle` that selects `at.yawk.lz4:lz4-java` ([#38947](https://github.com/apache/beam/issues/38947)). +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.74.0] - 2026-06-02 @@ -179,6 +180,10 @@ * Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). * Added `max_batch_duration_secs` passthrough support in Python Enrichment BigQuery and CloudSQL handlers so batching duration can be forwarded to `BatchElements` ([#38243](https://github.com/apache/beam/issues/38243)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.73.0] - 2026-04-29 ## Highlights @@ -215,6 +220,10 @@ * Fixed [CVE-2023-46604](https://www.cve.org/CVERecord?id=CVE-2023-46604) (CVSS 10.0) and [CVE-2022-41678](https://www.cve.org/CVERecord?id=CVE-2022-41678) by upgrading ActiveMQ from 5.14.5 to 5.19.2 (Java) ([#37943](https://github.com/apache/beam/issues/37943)). * Fixed [CVE-2024-1597](https://www.cve.org/CVERecord?id=CVE-2024-1597), [CVE-2022-31197](https://www.cve.org/CVERecord?id=CVE-2022-31197), and [CVE-2022-21724](https://www.cve.org/CVERecord?id=CVE-2022-21724) by upgrading PostgreSQL JDBC Driver from 42.2.16 to 42.6.2 (Java) ([#37942](https://github.com/apache/beam/issues/37942)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.72.0] - 2026-03-30 ## Highlights @@ -248,6 +257,10 @@ * Fixed [CVE-2024-28397](https://www.cve.org/CVERecord?id=CVE-2024-28397) by switching from js2py to pythonmonkey (Yaml) ([#37560](https://github.com/apache/beam/issues/37560)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.71.0] - 2026-01-22 ## I/Os @@ -269,6 +282,7 @@ ## Known Issues +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.70.0] - 2025-12-16 @@ -291,6 +305,10 @@ Now Beam has full support for Milvus integration including Milvus enrichment and * (Python) Python 3.9 reached EOL in October 2025 and support for the language version has been removed. ([#36665](https://github.com/apache/beam/issues/36665)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.69.0] - 2025-10-28 ## Highlights @@ -350,6 +368,10 @@ Now Beam has full support for Milvus integration including Milvus enrichment and ([#36141](https://github.com/apache/beam/issues/36141)). * Fixed Spanner Change Stream reading stuck issue due to watermark of partition moving backwards ([#36470](https://github.com/apache/beam/issues/36470)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.68.0] - 2025-09-22 ## Highlights @@ -404,6 +426,7 @@ Now Beam has full support for Milvus integration including Milvus enrichment and ## Known Issues * ([#36470](https://github.com/apache/beam/issues/36470)). Spanner Change Stream reading stuck issue due to watermark of partition moving backwards. This issue exists in 2.67.0 and 2.68.0. To mitigate the issue, either use old version 2.66.0 or go to 2.69.0. +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.67.0] - 2025-08-12 @@ -452,6 +475,7 @@ Now Beam has full support for Milvus integration including Milvus enrichment and * ([#35666](https://github.com/apache/beam/issues/35666)). YAML Flatten incorrectly drops fields when input PCollections' schema are different. This issue exists for all versions since 2.52.0. * ([#36470](https://github.com/apache/beam/issues/36470)). Spanner Change Stream reading stuck issue due to watermark of partition moving backwards. This issue exists in 2.67.0 and 2.68.0. To mitigate the issue, either use old version 2.66.0 or go to 2.69.0. +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.66.0] - 2025-07-01 From 9baa19ad5caa7f334d2d6ecbce62a1f28ca6a2ff Mon Sep 17 00:00:00 2001 From: Abdelrahman Ibrahim Date: Tue, 21 Jul 2026 19:40:41 +0300 Subject: [PATCH 14/29] Increase CI timeouts for Python 3.14 multiarch container jobs (#39401) Co-authored-by: Abdelrahman Ibrahim --- .github/workflows/beam_PostCommit_Python_Arm.yml | 3 ++- .../workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml | 2 ++ .github/workflows/update_python_dependencies.yml | 1 + 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/beam_PostCommit_Python_Arm.yml b/.github/workflows/beam_PostCommit_Python_Arm.yml index 0cb03894b255..961dfbb2c53a 100644 --- a/.github/workflows/beam_PostCommit_Python_Arm.yml +++ b/.github/workflows/beam_PostCommit_Python_Arm.yml @@ -54,7 +54,8 @@ jobs: beam_PostCommit_Python_Arm: name: ${{ matrix.job_name }} ${{ matrix.python_version }} runs-on: [self-hosted, ubuntu-24.04, main] - timeout-minutes: 480 + # Python 3.14 multiarch container builds take much longer than other versions. + timeout-minutes: ${{ matrix.python_version == '3.14' && 600 || 480 }} strategy: fail-fast: false matrix: diff --git a/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml b/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml index d7e012e9f036..6636f70197e6 100644 --- a/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml +++ b/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml @@ -62,6 +62,8 @@ jobs: startsWith(github.event.comment.body, 'Run Python ValidatesContainer Dataflow ARM') runs-on: ubuntu-22.04 + # Python 3.14 multiarch container builds compile many deps from source on arm64. + timeout-minutes: ${{ matrix.python_version == '3.14' && 300 || 120 }} steps: - uses: actions/checkout@v7 with: diff --git a/.github/workflows/update_python_dependencies.yml b/.github/workflows/update_python_dependencies.yml index 9da85090c4c7..f8630ebd0aec 100644 --- a/.github/workflows/update_python_dependencies.yml +++ b/.github/workflows/update_python_dependencies.yml @@ -51,6 +51,7 @@ jobs: runs-on: [self-hosted, ubuntu-24.04, main] needs: set-properties name: Update Python Dependencies + timeout-minutes: 180 steps: - name: Checkout code uses: actions/checkout@v7 From 0c151d8ddb23df45b96d660bb8b7e644bece7068 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 21 Jul 2026 12:41:46 -0400 Subject: [PATCH 15/29] Fix gRPC stream observer leak on ProcessBundleHandler shutdown (#39390) * Fix gRPC stream observer leak on ProcessBundleHandler shutdown * tears down active gRPC multiplexers on ProcessBundleHandler shutdown * Update sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../fn/harness/control/ProcessBundleHandler.java | 1 + .../beam/fn/harness/data/BeamFnDataClient.java | 9 ++++++++- .../beam/fn/harness/data/BeamFnDataGrpcClient.java | 12 ++++++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 449afd6a0243..b744da10cf6f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -767,6 +767,7 @@ public BeamFnApi.InstructionResponse.Builder trySplit(InstructionRequest request /** Shutdown the bundles, running the tearDown() functions. */ public void shutdown() throws Exception { bundleProcessorCache.shutdown(); + beamFnDataClient.close(); } @VisibleForTesting diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java index 1a50f5b448c5..458ed1209f43 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java @@ -17,6 +17,8 @@ */ package org.apache.beam.fn.harness.data; +import java.io.Closeable; +import java.io.IOException; import java.util.List; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.model.pipeline.v1.Endpoints; @@ -30,7 +32,7 @@ * provide a receiver of outbound elements. Callers can register themselves as receivers for inbound * elements or can get a handle for a receiver of outbound elements. */ -public interface BeamFnDataClient { +public interface BeamFnDataClient extends Closeable { /** * Registers a receiver for the provided instruction id. * @@ -72,4 +74,9 @@ void unregisterReceiver( /** Get the outbound observer for the specified apiServiceDescriptor and dataStreamId. */ StreamObserver getOutboundObserver( Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId); + + @Override + default void close() throws IOException { + // Default to no-op + } } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java index 2f2a6b0fc660..79e34fe5765f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java @@ -118,6 +118,18 @@ public void poisonInstructionId(String instructionId) { } } + @Override + public void close() { + for (BeamFnDataGrpcMultiplexer client : multiplexerCache.values()) { + try { + client.close(); + } catch (Exception e) { + LOG.warn("Failed to close multiplexer", e); + } + } + multiplexerCache.clear(); + } + @Override public StreamObserver getOutboundObserver( ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) { From 414b48ad644d7fda3db4d2b81f9931f6a40d0394 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Tue, 21 Jul 2026 13:23:12 -0400 Subject: [PATCH 16/29] add Apache third party allow list check and doc (#39392) * add third party allow list check and doc * add python action and shift validate further down * fix two workflow/action issues caught during test run * reuse setup environment step per comment --- .github/actions/setup-k8s-access/action.yml | 2 +- .github/workflows/README.md | 9 +++++++++ .../beam_Infrastructure_UsersPermissions.yml | 2 +- .github/workflows/beam_PreCommit_GHA.yml | 12 ++++++++---- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/.github/actions/setup-k8s-access/action.yml b/.github/actions/setup-k8s-access/action.yml index cb00c853738b..264aa77cfc88 100644 --- a/.github/actions/setup-k8s-access/action.yml +++ b/.github/actions/setup-k8s-access/action.yml @@ -68,7 +68,7 @@ runs: run: | kubectl config set-context --current --namespace=${{ steps.replace_namespace.outputs.TEST_NAMESPACE }} - name: Post cleanup - uses: pyTooling/Actions/with-post-step@v4.2.2 + uses: pytooling/actions/with-post-step@42e17fae05f224e5ac3d79d021a4e3577878efe5 # v4.2.2 with: main: echo "Post Cleanup" post: | diff --git a/.github/workflows/README.md b/.github/workflows/README.md index dcc885feec9d..41287872ee25 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -58,6 +58,7 @@ On top of normal Actions workflow steps, all new CI workflows (excluding release 4) A set of GitHub token permissions 5) Concurrency Groups 6) Comment Triggering Support +7) Allowed Actions (See below) Each of these is described in more detail below. ## Name and Phrase @@ -165,6 +166,14 @@ In order to make it easier for non-committers to interact with workflows, workfl **Note:** Comment triggering is found not scalable ([#28909](https://github.com/apache/beam/issues/28909)) and is currently limited to a subset of suites. For more information see the [Running Workflows Manually](#running-workflows-manually) section. +## Allowed Actions + +GitHub Actions under the trusted namespaces `actions/*`, `github/*`, and `apache/*` are automatically approved. Any third-party actions outside of these namespaces must be explicitly approved and added to the ASF-wide allowlist before use. + +Additionally, third-party actions must be pinned to a specific Git commit SHA (e.g., `zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02`) instead of a version tag or branch name to comply with ASF security policy. + +For details on the allowlist and how to request new actions, see the [ASF Infrastructure Actions Allowlist Check Documentation](https://github.com/apache/infrastructure-actions/blob/main/allowlist-check/README.md). + # Testing new workflows or workflow updates ## Testing New Workflows diff --git a/.github/workflows/beam_Infrastructure_UsersPermissions.yml b/.github/workflows/beam_Infrastructure_UsersPermissions.yml index 0e806034b5ca..b188e5ff875b 100644 --- a/.github/workflows/beam_Infrastructure_UsersPermissions.yml +++ b/.github/workflows/beam_Infrastructure_UsersPermissions.yml @@ -52,7 +52,7 @@ jobs: - name: Setup gcloud uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db - name: Install Terraform - uses: hashicorp/setup-terraform@v4 + uses: hashicorp/setup-terraform@dfe3c3f87815947d99a8997f908cb6525fc44e9e # v4.0.1 with: terraform_version: 1.12.2 - name: Initialize Terraform diff --git a/.github/workflows/beam_PreCommit_GHA.yml b/.github/workflows/beam_PreCommit_GHA.yml index 7ad6b5e67583..fa3d2c9605b4 100644 --- a/.github/workflows/beam_PreCommit_GHA.yml +++ b/.github/workflows/beam_PreCommit_GHA.yml @@ -80,15 +80,19 @@ jobs: comment_phrase: ${{ matrix.job_phrase }} github_token: ${{ secrets.GITHUB_TOKEN }} github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) - - name: Run zizmor - uses: zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02 # v0.6.0 - with: - advanced-security: true - name: Setup environment uses: ./.github/actions/setup-environment-action with: java-version: default go-version: default + python-version: default + # For details on the ASF GHA allowlist, see: https://github.com/apache/infrastructure-actions/blob/main/allowlist-check/README.md + - name: Validate GHA Allowlist + uses: apache/infrastructure-actions/allowlist-check@main # zizmor: ignore[unpinned-uses] + - name: Run zizmor + uses: zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02 # v0.6.0 + with: + advanced-security: true - name: run GHA PreCommit script uses: ./.github/actions/gradle-command-self-hosted-action with: From 83fcb2f9cdfa630f63743bde1c48c88c49812e59 Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:34:07 -0400 Subject: [PATCH 17/29] [Gemini] Add Java GeminiModelHandler class (#39245) * [Gemini] Add Java GeminiModelHandler class * Remove model scratch files * address review comments * checkstyle issue * unit test fix * Fix weird coder UX * fix altered base class signatures --- sdks/java/ml/inference/gemini/build.gradle | 42 +++++ .../inference/gemini/GeminiImageResponse.java | 33 ++++ .../gemini/GeminiInferenceFunctions.java | 67 +++++++ .../inference/gemini/GeminiModelHandler.java | 86 +++++++++ .../gemini/GeminiModelParameters.java | 60 +++++++ .../gemini/GeminiRequestFunction.java | 31 ++++ .../inference/gemini/GeminiStringInput.java | 33 ++++ .../gemini/GeminiStringResponse.java | 33 ++++ .../sdk/ml/inference/gemini/package-info.java | 20 +++ .../gemini/GeminiModelHandlerIT.java | 163 ++++++++++++++++++ .../gemini/GeminiModelHandlerTest.java | 133 ++++++++++++++ .../remote/PredictionResultCoder.java | 74 ++++++++ .../ml/inference/remote/RemoteInference.java | 34 +++- .../inference/remote/RemoteInferenceTest.java | 4 +- settings.gradle.kts | 1 + 15 files changed, 809 insertions(+), 5 deletions(-) create mode 100644 sdks/java/ml/inference/gemini/build.gradle create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java create mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java create mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java create mode 100644 sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java diff --git a/sdks/java/ml/inference/gemini/build.gradle b/sdks/java/ml/inference/gemini/build.gradle new file mode 100644 index 000000000000..f62984eba886 --- /dev/null +++ b/sdks/java/ml/inference/gemini/build.gradle @@ -0,0 +1,42 @@ +/* + * 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. + */ +plugins { + id 'org.apache.beam.module' +} + +applyJavaNature( + automaticModuleName: 'org.apache.beam.sdk.ml.inference.gemini', +) +provideIntegrationTestingDependencies() +enableJavaPerformanceTesting() + +description = "Apache Beam :: SDKs :: Java :: ML :: Inference :: Gemini" +ext.summary = "Gemini model handler for remote inference" + +dependencies { + implementation project(":sdks:java:ml:inference:remote") + implementation "com.google.genai:google-genai:1.59.0" + + + testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") + testImplementation project(path: ":sdks:java:core", configuration: "shadow") + testImplementation library.java.slf4j_api + testRuntimeOnly library.java.slf4j_simple + testImplementation library.java.junit + testImplementation library.java.mockito_core +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java new file mode 100644 index 000000000000..5d74c21347cd --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java @@ -0,0 +1,33 @@ +/* + * 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.ml.inference.gemini; + +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; + +/** Wrapper for image (byte array) responses from Gemini. */ +public class GeminiImageResponse implements BaseResponse { + public final byte[] imageBytes; + + public GeminiImageResponse(byte[] imageBytes) { + this.imageBytes = imageBytes; + } + + public byte[] getImageBytes() { + return imageBytes; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java new file mode 100644 index 000000000000..5ba009b7194e --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java @@ -0,0 +1,67 @@ +/* + * 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.ml.inference.gemini; + +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponse; +import com.google.genai.types.GenerateImagesConfig; +import com.google.genai.types.GenerateImagesResponse; +import java.util.ArrayList; +import java.util.List; + +/** Common inference functions for Gemini. */ +public class GeminiInferenceFunctions { + + /** Generates content from string prompts using the standard generateContent API. */ + public static GeminiRequestFunction + generateFromString() { + return (modelName, batch, client) -> { + List results = new ArrayList<>(); + for (GeminiStringInput input : batch) { + GenerateContentResponse response = + client.models.generateContent( + modelName, input.getText(), GenerateContentConfig.builder().build()); + String text = response.text(); + results.add(new GeminiStringResponse(text != null ? text : "")); + } + return results; + }; + } + + /** Generates images from string prompts using the generateImages API. */ + public static GeminiRequestFunction + generateImageFromString() { + return (modelName, batch, client) -> { + List results = new ArrayList<>(); + for (GeminiStringInput input : batch) { + GenerateImagesResponse response = + client.models.generateImages( + modelName, input.getText(), GenerateImagesConfig.builder().build()); + // Retrieve the base64 string or bytes from the first generated image + List images = response.images(); + if (images != null && !images.isEmpty()) { + byte[] imageBytes = images.get(0).imageBytes().orElse(new byte[0]); + results.add(new GeminiImageResponse(imageBytes)); + } else { + results.add(new GeminiImageResponse(new byte[0])); + } + } + return results; + }; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java new file mode 100644 index 000000000000..ecdbf807816f --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java @@ -0,0 +1,86 @@ +/* + * 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.ml.inference.gemini; + +import com.google.genai.Client; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; +import org.apache.beam.sdk.ml.inference.remote.BaseModelHandler; +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; + +/** + * Model handler for Google Gemini API inference requests. + * + *

This handler manages communication with Google's Gemini API, including client initialization, + * request formatting, and response parsing. It allows executing a custom {@link + * GeminiRequestFunction} against a batch of inputs. + */ +@SuppressWarnings("nullness") +public class GeminiModelHandler + implements BaseModelHandler, InputT, OutputT> { + + private transient Client client; + private GeminiModelParameters modelParameters; + + @Override + public void createClient(GeminiModelParameters parameters) { + if (parameters == null) { + throw new NullPointerException("GeminiModelParameters must not be null"); + } + this.modelParameters = parameters; + + // Configure client based on vertex or API key + if (parameters.getApiKey() != null) { + if (parameters.getProject() != null || parameters.getLocation() != null) { + throw new IllegalArgumentException("Project and location must be null if API key is set"); + } + this.client = Client.builder().apiKey(parameters.getApiKey()).build(); + } else { + Client.Builder builder = Client.builder(); + if (parameters.getProject() != null && parameters.getLocation() != null) { + builder.vertexAI(true).project(parameters.getProject()).location(parameters.getLocation()); + } else if (parameters.getProject() != null || parameters.getLocation() != null) { + throw new IllegalArgumentException( + "Project and location must both be provided if one is provided"); + } + this.client = builder.build(); + } + } + + @Override + public Iterable> request(List input) { + try { + GeminiRequestFunction requestFn = modelParameters.getRequestFn(); + List responses = requestFn.apply(modelParameters.getModelName(), input, client); + + if (responses.size() != input.size()) { + throw new IllegalStateException("Number of responses must match number of inputs"); + } + + List> results = new ArrayList<>(); + for (int i = 0; i < input.size(); i++) { + results.add(PredictionResult.create(input.get(i), responses.get(i))); + } + return results; + } catch (Exception e) { + throw new RuntimeException("Error during Gemini inference request", e); + } + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java new file mode 100644 index 000000000000..6ac4f578a24e --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java @@ -0,0 +1,60 @@ +/* + * 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.ml.inference.gemini; + +import com.google.auto.value.AutoValue; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; +import org.apache.beam.sdk.ml.inference.remote.BaseModelParameters; +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; +import org.checkerframework.checker.nullness.qual.Nullable; + +@AutoValue +public abstract class GeminiModelParameters + implements BaseModelParameters { + + public abstract @Nullable String getApiKey(); + + public abstract @Nullable String getProject(); + + public abstract @Nullable String getLocation(); + + public abstract String getModelName(); + + public abstract GeminiRequestFunction getRequestFn(); + + public static + Builder builder() { + return new AutoValue_GeminiModelParameters.Builder(); + } + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setApiKey(String apiKey); + + public abstract Builder setProject(String project); + + public abstract Builder setLocation(String location); + + public abstract Builder setModelName(String modelName); + + public abstract Builder setRequestFn( + GeminiRequestFunction requestFn); + + public abstract GeminiModelParameters build(); + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java new file mode 100644 index 000000000000..15da8b360c96 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java @@ -0,0 +1,31 @@ +/* + * 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.ml.inference.gemini; + +import com.google.genai.Client; +import java.io.Serializable; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; + +/** Functional interface for custom request functions to the Gemini API. */ +@FunctionalInterface +public interface GeminiRequestFunction + extends Serializable { + List apply(String modelName, List batch, Client client) throws Exception; +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java new file mode 100644 index 000000000000..8b0cc7c26892 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java @@ -0,0 +1,33 @@ +/* + * 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.ml.inference.gemini; + +import org.apache.beam.sdk.ml.inference.remote.BaseInput; + +/** Wrapper for string inputs to Gemini. */ +public class GeminiStringInput implements BaseInput { + public final String text; + + public GeminiStringInput(String text) { + this.text = text; + } + + public String getText() { + return text; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java new file mode 100644 index 000000000000..94bb4d7168dc --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java @@ -0,0 +1,33 @@ +/* + * 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.ml.inference.gemini; + +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; + +/** Wrapper for string responses from Gemini. */ +public class GeminiStringResponse implements BaseResponse { + public final String text; + + public GeminiStringResponse(String text) { + this.text = text; + } + + public String getText() { + return text; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java new file mode 100644 index 000000000000..9dee9da9cc74 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Gemini model handler for remote inference. */ +package org.apache.beam.sdk.ml.inference.gemini; diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java new file mode 100644 index 000000000000..c1188ad6d04a --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java @@ -0,0 +1,163 @@ +/* + * 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.ml.inference.gemini; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeNotNull; +import static org.junit.Assume.assumeTrue; + +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; +import org.apache.beam.sdk.ml.inference.remote.RemoteInference; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; + +/** + * Execute Gemini model handler integration test. + * + * 

+ * ./gradlew :sdks:java:ml:inference:gemini:integrationTest \
+ *   --tests org.apache.beam.sdk.ml.inference.gemini.GeminiModelHandlerIT \
+ *   --info
+ *
+ */ +public class GeminiModelHandlerIT { + + public static class GeminiStringContentHandler + extends GeminiModelHandler {} + + public static class GeminiStringImageHandler + extends GeminiModelHandler {} + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private String apiKey; + private static final String API_KEY_ENV = "GEMINI_API_KEY"; + private static final String DEFAULT_MODEL = "gemini-3.5-flash"; + + @Before + public void setUp() { + // Get API key + apiKey = System.getenv(API_KEY_ENV); + + // Skip tests if API key is not provided + assumeNotNull( + "Gemini API key not found. Set " + + API_KEY_ENV + + " environment variable to run integration tests.", + apiKey); + assumeTrue( + "Gemini API key is empty. Set " + + API_KEY_ENV + + " environment variable to run integration tests.", + !apiKey.trim().isEmpty()); + } + + @Test + public void testSentimentAnalysisWithSingleInput() { + String input = + "Classify the sentiment of the following text as either positive or negative: This product is absolutely amazing! I love it!"; + + PCollection inputs = + pipeline.apply("CreateSingleInput", Create.of(new GeminiStringInput(input))); + + PCollection>> results = + inputs.apply( + "SentimentInference", + RemoteInference.invoke() + .handler(GeminiStringContentHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName(DEFAULT_MODEL) + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build())); + + // Verify results + PAssert.that(results) + .satisfies( + batches -> { + int count = 0; + for (Iterable> batch : + batches) { + for (PredictionResult result : batch) { + count++; + assertNotNull("Input should not be null", result.getInput()); + assertNotNull("Output should not be null", result.getOutput()); + + String sentiment = result.getOutput().getText().toLowerCase(); + assertTrue( + "Sentiment should be positive or negative, got: " + sentiment, + sentiment.contains("positive") || sentiment.contains("negative")); + } + } + assertEquals("Should have exactly 1 result", 1, count); + return null; + }); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void testGenerateImageFromString() { + String input = "A beautiful sunset over the mountains, digital art style."; + + PCollection inputs = + pipeline.apply("CreateImageInput", Create.of(new GeminiStringInput(input))); + + PCollection>> results = + inputs.apply( + "ImageInference", + RemoteInference.invoke() + .handler(GeminiStringImageHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName("imagen-4.0-generate-001") + .setRequestFn(GeminiInferenceFunctions.generateImageFromString()) + .build())); + + // Verify results + PAssert.that(results) + .satisfies( + batches -> { + int count = 0; + for (Iterable> batch : + batches) { + for (PredictionResult result : batch) { + count++; + assertNotNull("Input should not be null", result.getInput()); + assertNotNull("Output should not be null", result.getOutput()); + assertTrue( + "Output should have generated images", + result.getOutput().getImageBytes().length > 0); + } + } + assertEquals("Should have exactly 1 result", 1, count); + return null; + }); + + pipeline.run().waitUntilFinish(); + } +} diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java new file mode 100644 index 000000000000..125f1e9c3e92 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java @@ -0,0 +1,133 @@ +/* + * 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.ml.inference.gemini; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class GeminiModelHandlerTest { + + @Test + public void testAllParamsSet() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setProject("test-project") + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexLocationParam() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setProject("test-project") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexProjectParam() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testNullParameters() { + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(NullPointerException.class, () -> handler.createClient(null)); + } + + @Test + public void testRequest() throws Exception { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setModelName("gemini-model-123") + .setRequestFn( + (modelName, batch, client) -> + Arrays.asList( + new GeminiStringResponse("response1"), + new GeminiStringResponse("response2"))) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + handler.createClient(parameters); + + List input = + Arrays.asList(new GeminiStringInput("input1"), new GeminiStringInput("input2")); + Iterable> results = + handler.request(input); + + int count = 0; + for (PredictionResult result : results) { + if (count == 0) { + assertEquals("input1", result.getInput().getText()); + assertEquals("response1", result.getOutput().getText()); + } else { + assertEquals("input2", result.getInput().getText()); + assertEquals("response2", result.getOutput().getText()); + } + count++; + } + assertEquals(2, count); + } + + @Test + public void testRequestMismatchedResponseSize() throws Exception { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setModelName("gemini-model-123") + .setRequestFn( + (modelName, batch, client) -> Arrays.asList(new GeminiStringResponse("response1"))) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + handler.createClient(parameters); + + List input = + Arrays.asList(new GeminiStringInput("input1"), new GeminiStringInput("input2")); + assertThrows(RuntimeException.class, () -> handler.request(input)); + } +} diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java new file mode 100644 index 000000000000..0d43870b880d --- /dev/null +++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java @@ -0,0 +1,74 @@ +/* + * 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.ml.inference.remote; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.StructuredCoder; + +/** A {@link Coder} for {@link PredictionResult}. */ +public class PredictionResultCoder + extends StructuredCoder> { + + private final Coder inputCoder; + private final Coder outputCoder; + + private PredictionResultCoder(Coder inputCoder, Coder outputCoder) { + this.inputCoder = inputCoder; + this.outputCoder = outputCoder; + } + + public static PredictionResultCoder of( + Coder inputCoder, Coder outputCoder) { + return new PredictionResultCoder<>(inputCoder, outputCoder); + } + + @Override + public void encode(PredictionResult value, OutputStream outStream) + throws CoderException, IOException { + if (value == null) { + throw new CoderException("cannot encode a null PredictionResult"); + } + inputCoder.encode(value.getInput(), outStream); + outputCoder.encode(value.getOutput(), outStream); + } + + @Override + public PredictionResult decode(InputStream inStream) + throws CoderException, IOException { + InputT input = inputCoder.decode(inStream); + OutputT output = outputCoder.decode(inStream); + return PredictionResult.create(input, output); + } + + @Override + public List> getCoderArguments() { + return Arrays.asList(inputCoder, outputCoder); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + verifyDeterministic(this, "Input coder must be deterministic", inputCoder); + verifyDeterministic(this, "Output coder must be deterministic", outputCoder); + } +} diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java index c571911a484e..a998f4d0f077 100644 --- a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java +++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java @@ -21,6 +21,7 @@ import com.google.auto.value.AutoValue; import java.util.List; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.io.components.throttling.ReactiveThrottler; import org.apache.beam.sdk.transforms.BatchElements; import org.apache.beam.sdk.transforms.DoFn; @@ -88,6 +89,8 @@ public abstract static class Invoke outputCoder(); + abstract Builder builder(); @AutoValue.Builder @@ -108,6 +111,8 @@ abstract Builder setHandler( abstract Builder setOverloadRatio(Double overloadRatio); + abstract Builder setOutputCoder(Coder outputCoder); + abstract Invoke build(); } @@ -163,6 +168,14 @@ public Invoke withOverloadRatio(double overloadRatio) { return builder().setOverloadRatio(overloadRatio).build(); } + /** + * Configures the coder for the output of the model. If not provided, it will fallback to using + * standard Java serialization for the output element. + */ + public Invoke withOutputCoder(Coder outputCoder) { + return builder().setOutputCoder(outputCoder).build(); + } + @Override public PCollection>> expand( PCollection input) { @@ -177,9 +190,24 @@ public PCollection>> expand( batchedInput = input.apply("BatchElements", BatchElements.withDefaults()); } - return batchedInput - // Pass the list to the inference function - .apply("RemoteInference", ParDo.of(new RemoteInferenceFn(this))); + PCollection>> result = + batchedInput + // Pass the list to the inference function + .apply("RemoteInference", ParDo.of(new RemoteInferenceFn(this))); + + Coder outCoder = outputCoder(); + if (outCoder != null) { + result.setCoder( + org.apache.beam.sdk.coders.IterableCoder.of( + PredictionResultCoder.of(input.getCoder(), outCoder))); + } else { + result.setCoder( + (org.apache.beam.sdk.coders.Coder) + org.apache.beam.sdk.coders.IterableCoder.of( + org.apache.beam.sdk.coders.SerializableCoder.of(PredictionResult.class))); + } + + return result; } /** diff --git a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java index 64691aa1fedd..5fe8ea1cce19 100644 --- a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java +++ b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java @@ -45,7 +45,7 @@ public class RemoteInferenceTest { @Rule public final transient TestPipeline pipeline = TestPipeline.create(); // Test input class - public static class TestInput implements BaseInput { + public static class TestInput implements BaseInput, java.io.Serializable { private final String value; private TestInput(String value) { @@ -84,7 +84,7 @@ public String toString() { } // Test output class - public static class TestOutput implements BaseResponse { + public static class TestOutput implements BaseResponse, java.io.Serializable { private final String result; private TestOutput(String result) { diff --git a/settings.gradle.kts b/settings.gradle.kts index 44e93029b54c..8a116eeb3944 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -285,6 +285,7 @@ include(":sdks:java:maven-archetypes:gcp-bom-examples") include(":sdks:java:maven-archetypes:starter") include(":sdks:java:ml:inference:remote") include(":sdks:java:ml:inference:openai") +include(":sdks:java:ml:inference:gemini") include(":sdks:java:testing:nexmark") include(":sdks:java:testing:expansion-service") include(":sdks:java:testing:jpms-tests") From 55e1ecbb138ccf8a8ce423f68b618db22bc0b23d Mon Sep 17 00:00:00 2001 From: KRITI MITTAL <218429678+Kriti-dev07@users.noreply.github.com> Date: Wed, 22 Jul 2026 01:36:50 +0530 Subject: [PATCH 18/29] Add metric for Kafka offset commit failures (#38889) * Add metric for Kafka offset commit failures * Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java index ac6650c354d4..281a5b9a47c6 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java @@ -25,6 +25,8 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.MapElements; @@ -62,6 +64,9 @@ public class KafkaCommitOffset static class CommitOffsetDoFn extends DoFn, Void> { private static final Logger LOG = LoggerFactory.getLogger(CommitOffsetDoFn.class); + private static final Counter COMMIT_FAILURES = + Metrics.counter(CommitOffsetDoFn.class, "commit-failures"); + private final Map consumerConfig; private final SerializableFunction, Consumer> consumerFactoryFn; @@ -85,6 +90,7 @@ public void processElement(@Element KV element) { element.getKey().getTopicPartition(), new OffsetAndMetadata(element.getValue() + 1))); } catch (Exception e) { + COMMIT_FAILURES.inc(); // TODO: consider retrying. LOG.warn("Getting exception when committing offset: {}", e.getMessage()); } From 51b2ad6bc084a65cd5af24d267b9c5c863d1aa52 Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Wed, 22 Jul 2026 03:24:01 -0500 Subject: [PATCH 19/29] Add drain states to PipelineResult (#39020) * Add drain states to PipelineResult * Fix draining job lookup for Dataflow updates --- CHANGES.md | 1 + .../flink/FlinkDetachedRunnerResult.java | 4 +-- .../runners/flink/FlinkRunnerResultTest.java | 12 +++---- .../runners/dataflow/DataflowPipelineJob.java | 1 + .../beam/runners/dataflow/DataflowRunner.java | 3 +- .../runners/dataflow/util/MonitoringUtil.java | 6 ++-- .../dataflow/DataflowPipelineJobTest.java | 2 +- .../runners/dataflow/DataflowRunnerTest.java | 26 ++++++++++++-- .../dataflow/util/MonitoringUtilTest.java | 4 +-- .../runners/jobsubmission/JobInvocation.java | 15 ++++++-- .../jobsubmission/JobInvocationTest.java | 36 +++++++++++++++++++ .../portability/JobServicePipelineResult.java | 7 ++-- .../portability/PortableRunnerTest.java | 21 +++++++++++ .../org/apache/beam/sdk/PipelineResult.java | 6 ++++ .../beam/sdk/nexmark/NexmarkLauncher.java | 2 ++ 15 files changed, 123 insertions(+), 23 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 991c341460a0..87fdaa8d0216 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -75,6 +75,7 @@ ## Breaking Changes * (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)). +* (Java) Added `DRAINING` and `DRAINED` states to `PipelineResult`, including runner state mappings and Dataflow update handling ([#39020](https://github.com/apache/beam/issues/39020)). ## Deprecations diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java index b26e865526dd..d7c8912ded01 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java @@ -120,11 +120,11 @@ public synchronized State drain() throws IOException { private State getDrainState(CompletableFuture drainFuture) throws IOException { if (!drainFuture.isDone()) { - return State.RUNNING; + return State.DRAINING; } try { drainFuture.get(); - return State.DONE; + return State.DRAINED; } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IOException("Failed to drain Flink job", e); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java index 908d940f5efb..09ceede9d586 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java @@ -66,18 +66,18 @@ public void testDrainDoneResultDoesNotThrowAnException() throws Exception { } @Test - public void testDetachedDrainReturnsRunningThenDone() throws Exception { + public void testDetachedDrainReturnsDrainingThenDrained() throws Exception { JobClient jobClient = mock(JobClient.class); CompletableFuture drainFuture = new CompletableFuture<>(); when(jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT)) .thenReturn(drainFuture); FlinkDetachedRunnerResult result = new FlinkDetachedRunnerResult(jobClient, 1); - assertThat(result.drain(), is(PipelineResult.State.RUNNING)); - assertThat(result.getState(), is(PipelineResult.State.RUNNING)); + assertThat(result.drain(), is(PipelineResult.State.DRAINING)); + assertThat(result.getState(), is(PipelineResult.State.DRAINING)); drainFuture.complete("savepoint"); - assertThat(result.getState(), is(PipelineResult.State.DONE)); + assertThat(result.getState(), is(PipelineResult.State.DRAINED)); verify(jobClient).stopWithSavepoint(true, null, SavepointFormatType.DEFAULT); } @@ -132,11 +132,11 @@ public void testDetachedDrainRetriesAfterFailure() throws Exception { result.drain(); fail("Expected IOException"); } catch (IOException expected) { - assertThat(result.drain(), is(PipelineResult.State.RUNNING)); + assertThat(result.drain(), is(PipelineResult.State.DRAINING)); } retryDrainFuture.complete("savepoint"); - assertThat(result.getState(), is(PipelineResult.State.DONE)); + assertThat(result.getState(), is(PipelineResult.State.DRAINED)); verify(jobClient, times(2)).stopWithSavepoint(true, null, SavepointFormatType.DEFAULT); } } diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java index 0d7e5eaf68d9..278dace6ef48 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java @@ -365,6 +365,7 @@ private void logTerminalState(State state) { switch (state) { case DONE: case CANCELLED: + case DRAINED: LOG.info("Job {} finished with status {}.", getJobId(), state); break; case UPDATED: diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java index 011f60f4fd14..1b03cc10351e 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java @@ -2535,8 +2535,9 @@ private String getJobIdFromName(String jobName) { listResult = dataflowClient.listJobs(token); token = listResult.getNextPageToken(); for (Job job : listResult.getJobs()) { + State state = MonitoringUtil.toState(job.getCurrentState()); if (job.getName().equals(jobName) - && MonitoringUtil.toState(job.getCurrentState()).equals(State.RUNNING)) { + && (state.equals(State.RUNNING) || state.equals(State.DRAINING))) { return job.getId(); } } diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java index d117cf786129..0e25bbd265ae 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java @@ -221,17 +221,19 @@ public static State toState(@Nullable String stateName) { return State.CANCELLED; case "JOB_STATE_UPDATED": return State.UPDATED; + case "JOB_STATE_DRAINING": + return State.DRAINING; + case "JOB_STATE_DRAINED": + return State.DRAINED; case "JOB_STATE_RUNNING": case "JOB_STATE_PENDING": // Job has not yet started; closest mapping is RUNNING - case "JOB_STATE_DRAINING": // Job is still active; the closest mapping is RUNNING case "JOB_STATE_CANCELLING": // Job is still active; the closest mapping is RUNNING case "JOB_STATE_PAUSING": // Job is still active; the closest mapping is RUNNING case "JOB_STATE_RESOURCE_CLEANING_UP": // Job is still active; the closest mapping is RUNNING return State.RUNNING; case "JOB_STATE_DONE": - case "JOB_STATE_DRAINED": // Job has successfully terminated; closest mapping is DONE return State.DONE; default: LOG.warn( diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java index 4b088eb41a7d..be7ad2e6e110 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java @@ -417,7 +417,7 @@ public void testDrainUnterminatedJobThatSucceeds() throws IOException { DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); - assertEquals(State.RUNNING, job.drain()); + assertEquals(State.DRAINING, job.drain()); Job content = new Job(); content.setProjectId(PROJECT_ID); content.setId(JOB_ID); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java index 073b30f928dc..6c32dd749d0b 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java @@ -285,6 +285,11 @@ private static Pipeline buildDataflowPipelineWithLargeGraph(DataflowPipelineOpti } static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) throws IOException { + return buildMockDataflow(mockJobs, "JOB_STATE_RUNNING"); + } + + static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs, String currentState) + throws IOException { Dataflow mockDataflowClient = mock(Dataflow.class); Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); Dataflow.Projects.Locations mockLocations = mock(Dataflow.Projects.Locations.class); @@ -308,7 +313,7 @@ static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) thr new Job() .setName("oldjobname") .setId("oldJobId") - .setCurrentState("JOB_STATE_RUNNING")))); + .setCurrentState(currentState)))); Job resultJob = new Job(); resultJob.setId("newid"); @@ -375,6 +380,10 @@ static GcsUtil buildMockGcsUtil() throws IOException { } private DataflowPipelineOptions buildPipelineOptions() throws IOException { + return buildPipelineOptions("JOB_STATE_RUNNING"); + } + + private DataflowPipelineOptions buildPipelineOptions(String currentState) throws IOException { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); options.setRunner(DataflowRunner.class); options.setProject(PROJECT_ID); @@ -382,7 +391,7 @@ private DataflowPipelineOptions buildPipelineOptions() throws IOException { options.setRegion(REGION_ID); // Set FILES_PROPERTY to empty to prevent a default value calculated from classpath. options.setFilesToStage(new ArrayList<>()); - options.setDataflowClient(buildMockDataflow(mockJobs)); + options.setDataflowClient(buildMockDataflow(mockJobs, currentState)); options.setGcsUtil(mockGcsUtil); options.setGcpCredential(new TestCredential()); @@ -793,6 +802,19 @@ public void testUpdate() throws IOException { assertValidJob(jobCaptor.getValue()); } + @Test + public void testUpdateDrainingJob() throws IOException { + DataflowPipelineOptions options = buildPipelineOptions("JOB_STATE_DRAINING"); + options.setUpdate(true); + options.setJobName("oldJobName"); + Pipeline p = buildDataflowPipeline(options); + p.run(); + + ArgumentCaptor jobCaptor = ArgumentCaptor.forClass(Job.class); + Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture()); + assertEquals("oldJobId", jobCaptor.getValue().getReplaceJobId()); + } + @Test public void testUploadGraph() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java index 5f76b6750ffa..35af1be4cdfc 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java @@ -100,9 +100,9 @@ public void testToStateNormal() { // Non-trivially mapped cases assertEquals(State.STOPPED, MonitoringUtil.toState("JOB_STATE_PAUSED")); - assertEquals(State.RUNNING, MonitoringUtil.toState("JOB_STATE_DRAINING")); + assertEquals(State.DRAINING, MonitoringUtil.toState("JOB_STATE_DRAINING")); assertEquals(State.RUNNING, MonitoringUtil.toState("JOB_STATE_PAUSING")); - assertEquals(State.DONE, MonitoringUtil.toState("JOB_STATE_DRAINED")); + assertEquals(State.DRAINED, MonitoringUtil.toState("JOB_STATE_DRAINED")); } @Test diff --git a/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java b/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java index 9da3bf38e575..3bd0a1311b04 100644 --- a/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java +++ b/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java @@ -116,6 +116,12 @@ public void onSuccess(PortablePipelineResult pipelineResult) { case RUNNING: setState(JobState.Enum.RUNNING); break; + case DRAINING: + setState(JobState.Enum.DRAINING); + break; + case DRAINED: + setState(JobState.Enum.DRAINED); + break; case CANCELLED: setState(JobState.Enum.CANCELLED); break; @@ -169,9 +175,12 @@ public synchronized void cancel() { new FutureCallback() { @Override public void onSuccess(PortablePipelineResult pipelineResult) { - // Do not cancel when we are already done. - if (pipelineResult != null - && pipelineResult.getState() != PipelineResult.State.DONE) { + // Do not cancel when the runner has already successfully finished. + if (pipelineResult != null) { + PipelineResult.State state = pipelineResult.getState(); + if (state == PipelineResult.State.DONE || state == PipelineResult.State.DRAINED) { + return; + } try { pipelineResult.cancel(); setState(JobState.Enum.CANCELLED); diff --git a/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java b/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java index 31380ace43c3..7100e57717be 100644 --- a/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java +++ b/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java @@ -80,6 +80,28 @@ public void testStateAfterCompletion() throws Exception { awaitJobState(jobInvocation, JobApi.JobState.Enum.DONE); } + @Test(timeout = 10_000) + public void testStateAfterDrainCompleted() throws Exception { + jobInvocation.start(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.RUNNING)); + + TestPipelineResult pipelineResult = new TestPipelineResult(PipelineResult.State.DRAINED); + runner.setResult(pipelineResult); + + awaitJobState(jobInvocation, JobApi.JobState.Enum.DRAINED); + } + + @Test(timeout = 10_000) + public void testStateAfterDrainStarted() throws Exception { + jobInvocation.start(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.RUNNING)); + + TestPipelineResult pipelineResult = new TestPipelineResult(PipelineResult.State.DRAINING); + runner.setResult(pipelineResult); + + awaitJobState(jobInvocation, JobApi.JobState.Enum.DRAINING); + } + @Test(timeout = 10_000) public void testStateAfterCompletionWithoutResult() throws Exception { jobInvocation.start(); @@ -128,6 +150,20 @@ public void testNoCancellationWhenDone() throws Exception { assertThat(pipelineResult.cancelLatch.getCount(), is(1L)); } + @Test(timeout = 10_000) + public void testNoCancellationWhenDrained() throws Exception { + jobInvocation.start(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.RUNNING)); + + TestPipelineResult pipelineResult = new TestPipelineResult(PipelineResult.State.DRAINED); + runner.setResult(pipelineResult); + awaitJobState(jobInvocation, JobApi.JobState.Enum.DRAINED); + + jobInvocation.cancel(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.DRAINED)); + assertThat(pipelineResult.cancelLatch.getCount(), is(1L)); + } + @Test(timeout = 10_000) public void testReturnsMetricsFromJobInvocationAfterSuccess() throws Exception { JobApi.MetricResults expectedMonitoringInfos = JobApi.MetricResults.newBuilder().build(); diff --git a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java index 0f3d3b6d1910..7c72a1039c0e 100644 --- a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java +++ b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java @@ -160,7 +160,7 @@ private void waitForTerminalState() { } private void propagateErrors() { - if (terminalState != State.DONE) { + if (terminalState != State.DONE && terminalState != State.DRAINED) { JobMessagesRequest messageStreamRequest = JobMessagesRequest.newBuilder().setJobIdBytes(jobId).build(); Iterator messageStreamIterator = @@ -196,10 +196,9 @@ private static State getJavaState(JobApi.JobState.Enum protoState) { case UPDATED: return State.UPDATED; case DRAINING: - // TODO: Determine the correct mappings for the states below. - return State.UNKNOWN; + return State.DRAINING; case DRAINED: - return State.UNKNOWN; + return State.DRAINED; case STARTING: return State.RUNNING; case CANCELLING: diff --git a/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java b/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java index 68f3f6eae396..068d70b0d250 100644 --- a/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java +++ b/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java @@ -109,6 +109,27 @@ public void stagesAndRunsJob() throws Exception { assertThat(state, is(State.DONE)); } + @Test + public void mapsDrainingJobState() throws Exception { + createJobServer(JobState.Enum.DRAINING, JobApi.MetricResults.getDefaultInstance()); + PortableRunner runner = PortableRunner.create(options, ManagedChannelFactory.createInProcess()); + PipelineResult result = runner.run(p); + + try { + assertThat(result.getState(), is(State.DRAINING)); + } finally { + ((AutoCloseable) result).close(); + } + } + + @Test + public void mapsDrainedJobState() throws Exception { + createJobServer(JobState.Enum.DRAINED, JobApi.MetricResults.getDefaultInstance()); + PortableRunner runner = PortableRunner.create(options, ManagedChannelFactory.createInProcess()); + State state = runner.run(p).waitUntilFinish(); + assertThat(state, is(State.DRAINED)); + } + @Test public void extractsMetrics() throws Exception { JobApi.MetricResults metricResults = generateMetricResults(); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java index 46cca7833e52..eff031157e12 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java @@ -107,6 +107,12 @@ enum State { /** The job has been updated. */ UPDATED(true, true), + /** The job is draining: no longer accepting new input while finishing in-flight work. */ + DRAINING(false, false), + + /** The job has finished draining. */ + DRAINED(true, false), + /** The job state reported by a runner cannot be interpreted by the SDK. */ UNRECOGNIZED(false, false); diff --git a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java index 7e4e5da0d853..aaf63705ad2a 100644 --- a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java +++ b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java @@ -515,9 +515,11 @@ private void invokeBuilderForPublishOnlyPipeline(PipelineBuilder case UNRECOGNIZED: case STOPPED: case RUNNING: + case DRAINING: // Keep going. break; case DONE: + case DRAINED: // All done. running = false; break; From ee44102d01bb2769ac9e0e9648424300cffbf8dd Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:35:47 -0400 Subject: [PATCH 20/29] Bump github.com/aws/aws-sdk-go-v2/credentials in /sdks (#39425) Bumps [github.com/aws/aws-sdk-go-v2/credentials](https://github.com/aws/aws-sdk-go-v2) from 1.19.29 to 1.19.30. - [Release notes](https://github.com/aws/aws-sdk-go-v2/releases) - [Commits](https://github.com/aws/aws-sdk-go-v2/compare/credentials/v1.19.29...credentials/v1.19.30) --- updated-dependencies: - dependency-name: github.com/aws/aws-sdk-go-v2/credentials dependency-version: 1.19.30 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- sdks/go.mod | 22 +++++++++++----------- sdks/go.sum | 44 ++++++++++++++++++++++---------------------- 2 files changed, 33 insertions(+), 33 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index c438f2b351d2..752003b3552f 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -32,9 +32,9 @@ require ( cloud.google.com/go/pubsub v1.51.0 cloud.google.com/go/spanner v1.93.0 cloud.google.com/go/storage v1.63.1 - github.com/aws/aws-sdk-go-v2 v1.42.1 + github.com/aws/aws-sdk-go-v2 v1.43.0 github.com/aws/aws-sdk-go-v2/config v1.32.30 - github.com/aws/aws-sdk-go-v2/credentials v1.19.29 + github.com/aws/aws-sdk-go-v2/credentials v1.19.30 github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34 github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 github.com/aws/smithy-go v1.27.4 @@ -91,7 +91,7 @@ require ( github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.57.0 // indirect github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op // indirect github.com/apache/arrow/go/v15 v15.0.2 // indirect - github.com/aws/aws-sdk-go-v2/service/signin v1.4.1 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.5.0 // indirect github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect github.com/containerd/log v0.1.0 // indirect @@ -148,17 +148,17 @@ require ( github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40 // indirect github.com/apache/thrift v0.23.0 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect - github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 // indirect github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 // indirect - github.com/aws/aws-sdk-go-v2/service/sso v1.32.1 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 // indirect - github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.45.0 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect diff --git a/sdks/go.sum b/sdks/go.sum index ba65418653ef..2db1288b176b 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -196,8 +196,8 @@ github.com/aws/aws-sdk-go v1.37.0/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zK github.com/aws/aws-sdk-go v1.43.31/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo= github.com/aws/aws-sdk-go-v2 v1.16.2/go.mod h1:ytwTPBG6fXTZLxxeeCCWj2/EMYp/xDUgX+OET6TLNNU= github.com/aws/aws-sdk-go-v2 v1.23.0/go.mod h1:i1XDttT4rnf6vxc9AuskLc6s7XBee8rlLilKlc03uAA= -github.com/aws/aws-sdk-go-v2 v1.42.1 h1:9eOTgu1z/dVtYpNZ3/8/XbbaX0x/BqE3HUzAzs6K0ek= -github.com/aws/aws-sdk-go-v2 v1.42.1/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= +github.com/aws/aws-sdk-go-v2 v1.43.0 h1:fharf/WhbRAVZ1du0QL7roNFxZ6T/sWr+4Ni617bwSI= +github.com/aws/aws-sdk-go-v2 v1.43.0/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.1/go.mod h1:n8Bs1ElDD2wJ9kCRTczA83gYbBmjSwZp3umc6zF4EeM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.5.1/go.mod h1:t8PYl/6LzdAqsU4/9tz28V/kU+asFePvpOMkdul0gEQ= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHzbkHfPzy4ACPcEjVG0x87KOwtpqGY= @@ -208,29 +208,29 @@ github.com/aws/aws-sdk-go-v2/config v1.32.30 h1:XwsEzpTJfQYJbFicz/QMLwAZdyeNVVoO github.com/aws/aws-sdk-go-v2/config v1.32.30/go.mod h1:Ud32SuMc+/9BGxfpSVld7HrE2o05JwKmXY4M3jOQNZU= github.com/aws/aws-sdk-go-v2/credentials v1.11.2/go.mod h1:j8YsY9TXTm31k4eFhspiQicfXPLZ0gYXA50i4gxPE8g= github.com/aws/aws-sdk-go-v2/credentials v1.16.2/go.mod h1:sDdvGhXrSVT5yzBDR7qXz+rhbpiMpUYfF3vJ01QSdrc= -github.com/aws/aws-sdk-go-v2/credentials v1.19.29 h1:WHZGssHH887cO0ox07SIQZsFx3MKD4ps6w0xUEmnKYQ= -github.com/aws/aws-sdk-go-v2/credentials v1.19.29/go.mod h1:Mhl0xR6zjguiuj00XRx2wMx22sAltk7oya39sT7fdg8= +github.com/aws/aws-sdk-go-v2/credentials v1.19.30 h1:TTCvvzFU6gXa4iJecNG/0F/B0oYTiazoRECr2XyLHrY= +github.com/aws/aws-sdk-go-v2/credentials v1.19.30/go.mod h1:jKxAp2AEncnliinzpgOSZDFv6+VjvWhjw/AtbfsWT9U= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.3/go.mod h1:uk1vhHHERfSVCUnqSqz8O48LBYDSC+k6brng09jcMOk= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.14.4/go.mod h1:t4i+yGHMCcUNIX1x7YVYa6bH/Do7civ5I6cG/6PMfyA= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQu0EtP23pP553Cf9KgsDV4= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30/go.mod h1:/3AOgy4K17Dm4ucMZVC/MJkzy5kmfKUcINRHZyo0koQ= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 h1:kfVL5wAunCJycL6MOQ6aNh6PlAYEymflcjuKmrWUA0o= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31/go.mod h1:nWfRNDAppujCQgOUd43lKT4yeLv9z3nJ3bw1G3BgQKo= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.3/go.mod h1:0dHuD2HZZSiwfJSy1FO5bX1hQ1TxVV1QXXjpn3XUE44= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.14.0/go.mod h1:UcgIwJ9KHquYxs6Q5skC9qXjhYMK+JASDYcXQ4X7JZE= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34 h1:Pn7OsMwBLbkZ6OnCxWHAjf0L/22H8cnhxZC0uPwtMtg= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34/go.mod h1:eToXR/Gk1uqpn04eSmdgVXwfS0WvH8aG4eBFr8ygbpU= github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.9/go.mod h1:AnVH5pvai0pAF4lXRq0bmhbes1u9R8wTE+g+183bZNM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.3/go.mod h1:7sGSz1JCKHWWBHq98m6sMtWQikmYPpxjqOydDemiVoM= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 h1:xM/Is9cKMHa8Jj8zkvWhvrFkZsXJV9E+BB4g0HW0duQ= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30/go.mod h1:WueJeNDZvK1fMYEWJIkcivBfEzUkTpBhzlrUKKY8EuA= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 h1:Z8F3hfCY33IGpJjFAnv0wvtv1FIKj1GHmRDEYqy64tw= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31/go.mod h1:aVyUoytEyOViR6jhq6jula0xkc5NfBE2hgeF6BvOrao= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.3/go.mod h1:ssOhaLpRlh88H3UmEcsBoVKq309quMvm3Ds8e9d4eJM= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.5.3/go.mod h1:ify42Rb7nKeDDPkFjKn7q1bPscVPu/+gmHH8d2c+anU= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 h1:jn46zC9LdsVR/ZpMIJqMqb8hHv31BlLx3ulVqNspUOk= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30/go.mod h1:1hTMsAgbdS/AtUi4bw8+gUuh1pceo+eXRLfpSuSQj3M= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 h1:hyOxUyXdh3AyjE93gBgsfziJag9ACwcs+ZpDBLzi8mw= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31/go.mod h1:OERqI9k0draSLB8O8woxY3q25ZWTELRK4RRoLMuMZFo= github.com/aws/aws-sdk-go-v2/internal/ini v1.3.10/go.mod h1:8DcYQcz0+ZJaSxANlHIsbbi6S+zMwjwdDqwW3r9AzaE= github.com/aws/aws-sdk-go-v2/internal/ini v1.7.1/go.mod h1:6fQQgfuGmw8Al/3M2IgIllycxV7ZW7WCdVSqfBeUiCY= github.com/aws/aws-sdk-go-v2/internal/v4a v1.2.3/go.mod h1:5yzAuE9i2RkVAttBl8yxZgQr5OCq4D5yDnG7j9x2L0U= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 h1:3GUprIsfmGcC5SACIyB0e7E0BM1O1b3Erl5CePYIAeQ= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31/go.mod h1:7PuV1yl5e2xnUbm+RqvVg5i2iBM8EyijZNoI9wsOoOc= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 h1:0MrUL35H/Y4kdFfItoR5jCgtDQ4Z/8LudAoIHRfA4hE= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32/go.mod h1:2tNZkuWz54arj8mHVf+8Y7cKkcD8Wr/fBpENgEXpjLc= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.1/go.mod h1:GeUru+8VzrTXV/83XyMJ80KpH8xO89VPoUileyNQ+tc= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.10.1/go.mod h1:l9ymW25HOqymeU2m1gbUQ3rUIsTwKs8gYHXkqDQUhiI= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur/BiHK6SKPjoBIXSE/hJ6g6JGRLuxQy1jGjlN4= @@ -241,8 +241,8 @@ github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 h1:9Fjh6fi/U5JESt github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23/go.mod h1:iMoT2f1tClxrWAAnKCXjZQ6LOmfLrMG14wmnWpM+F14= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.3/go.mod h1:wlY6SVjuwvh3TVRpTqdy4I1JpBFLX4UGeKZdWntaocw= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.10.3/go.mod h1:Owv1I59vaghv1Ax8zz8ELY8DN7/Y0rGS+WWAmjgi950= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 h1:/Z5jmNrKsSD7EmDjzAPsm/3L9IuOkzaynklJZ1qX7S4= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30/go.mod h1:lEzEZnOosE7zi8Z6royW1cFJTD9fpab4Ul1SBrllewk= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 h1:w2SIhW92DZPFrSL4ksVCr8IYff5OZwIcxg8+95tzvAI= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31/go.mod h1:wAhpCQbkov+IcvjozJbd2xRCoZybUEHNkcFunssNACg= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.3/go.mod h1:Bm/v2IaN6rZ+Op7zX+bOUMdL4fsrYZiD0dsjLhNKwZc= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.16.3/go.mod h1:KZgs2ny8HsxRIRbDwgvJcHHBZPOzQr/+NtGwnP+w2ec= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 h1:uao4A3QZ5UmB326V6KF+qRpv9Tjz7IlnlnTbbANntlU= @@ -253,22 +253,22 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.43.0/go.mod h1:NXRKkiRF+erX2hnybnVU66 github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 h1:5C00eQYpTrgQXnp6V3P6P7zPElna3AXvlukbANE6nJI= github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2/go.mod h1:zdmCoFO/dSI7GlrwsPqFJI+WlFnSU4Tc8TJnlXrM1Do= github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.15.4/go.mod h1:PJc8s+lxyU8rrre0/4a0pn2wgwiDvOEzoOjcJUBr67o= -github.com/aws/aws-sdk-go-v2/service/signin v1.4.1 h1:V7ZZ300WPXGjvkyore5DGe0ljVPOxCXie/thWdtSBXE= -github.com/aws/aws-sdk-go-v2/service/signin v1.4.1/go.mod h1:mxC0nT/C8wMMS97DemZPzvUZxvIt+2Iq+eS3JdFZGgg= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.0 h1:OHH5iTQvVGmfHjX/5Q+vFuA/Rf2x6/95aJ/75QCQSm4= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.0/go.mod h1:mCF3AK9PpL49oOrhniUXWAfhVBVQ/XbytoE5eccZUIs= github.com/aws/aws-sdk-go-v2/service/sns v1.17.4/go.mod h1:kElt+uCcXxcqFyc+bQqZPFD9DME/eC6oHBXvFzQ9Bcw= github.com/aws/aws-sdk-go-v2/service/sqs v1.18.3/go.mod h1:skmQo0UPvsjsuYYSYMVmrPc1HWCbHUJyrCEp+ZaLzqM= github.com/aws/aws-sdk-go-v2/service/ssm v1.24.1/go.mod h1:NR/xoKjdbRJ+qx0pMR4mI+N/H1I1ynHwXnO6FowXJc0= github.com/aws/aws-sdk-go-v2/service/sso v1.11.3/go.mod h1:7UQ/e69kU7LDPtY40OyoHYgRmgfGM4mgsLYtcObdveU= github.com/aws/aws-sdk-go-v2/service/sso v1.17.2/go.mod h1:/pE21vno3q1h4bbhUOEi+6Zu/aT26UK2WKkDXd+TssQ= -github.com/aws/aws-sdk-go-v2/service/sso v1.32.1 h1:gYFYh4iLLcAOJRLNPY2aD2g9DIhKn4eof8UkIrr1rTk= -github.com/aws/aws-sdk-go-v2/service/sso v1.32.1/go.mod h1:u8af9Nqkmqnr96f7v9nHqzZT9XBwbXEkTiqT4ROuJSE= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 h1:CaJyYhxBE0M/HJX/YvSaSmQlsI91VHB0lKU8LtLxL3A= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.0/go.mod h1:+e6BMRMPjBQoCw/WovYR9GLy2IU0z4Q77smOB1DraSg= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.20.0/go.mod h1:dWqm5G767qwKPuayKfzm4rjzFmVjiBFbOJrpSPnAMDs= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 h1:arjT9Cm3/WYbGmD5TUZHk4UQn4Lle1fUNZs5FC6CtF0= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 h1:tC323YV77QdafeBr6LUhLDTsboyuyHLNRwAyCP44kGU= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0/go.mod h1:SfLK1sgviHmbI+MozR9iDwDjL4cdCVZtahsjoR+z7wg= github.com/aws/aws-sdk-go-v2/service/sts v1.16.3/go.mod h1:bfBj0iVmsUyUg4weDB4NxktD9rDGeKSVWnjTnwbx9b8= github.com/aws/aws-sdk-go-v2/service/sts v1.25.3/go.mod h1:4EqRHDCKP78hq3zOnmFXu5k0j4bXbRFfCh/zQ6KnEfQ= -github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 h1:RvfHDg+xvAeZ+5741vUEjpOVtYSIm93W2zhx10Xtydw= -github.com/aws/aws-sdk-go-v2/service/sts v1.44.1/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.0 h1:Pd6PNlp4t8PTXxqzstICl52Wsy78vpjFZ7PRUj44mJc= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.0/go.mod h1:rmQ0TnHzuLPmabgjPcsywhsSOmaBDgzR4zvDxSPsGdg= github.com/aws/smithy-go v1.11.2/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM= github.com/aws/smithy-go v1.17.0/go.mod h1:NukqUGpCZIILqqiV0NIjeFh24kd/FAa4beRb6nbIUPE= github.com/aws/smithy-go v1.27.4 h1:JQcphmBN4f0q/sPqXqROIItRNV/hy10cgu7CsFy616M= From 0dd47882822fb58cd8ba2b0612031c99351cbeac Mon Sep 17 00:00:00 2001 From: janaom <83917694+janaom@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:57:05 +0300 Subject: [PATCH 21/29] Interview with Raj Katakam from Intuit Credit Karma (#39412) * Create beam-summit-2026-interview-with-raj-katakam.md Interview with Raj Katakam * Update authors.yml Added my info to authors.yml * Update beam-summit-2026-interview-with-raj-katakam.md * Update website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md --------- Co-authored-by: Danny McCormick --- ...-summit-2026-interview-with-raj-katakam.md | 88 +++++++++++++++++++ website/www/site/data/authors.yml | 3 + 2 files changed, 91 insertions(+) create mode 100644 website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md diff --git a/website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md b/website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md new file mode 100644 index 000000000000..6d4e649c7247 --- /dev/null +++ b/website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md @@ -0,0 +1,88 @@ +--- +title: "Beam Summit 2026 | Interview with Raj Katakam, Intuit Credit Karma" +date: 2026-07-22 15:00:00 -0500 +authors: + - janaom +categories: + - blog +--- + + +# Beam Summit 2026 | Interview with Raj Katakam, Intuit Credit Karma + +Apache Beam and Google Cloud Dataflow are usually known for ETL and batch/streaming processing. [Intuit Credit Karma](https://www.creditkarma.com/) said, why stop there, and built an entire ML platform on top of them. + +Here is my interview with **Raj Katakam**, Staff Machine Learning Engineer at Intuit Credit Karma. We talked about how they built a unified ML platform on Apache Beam and Dataflow to serve 140 million users, why being active in the Apache Beam community matters to him, and what advice he would give anyone just getting started with Beam and ML. + +-------------- + +**Jana: Intuit Credit Karma’s tech stack looks like a dream! You are on Google Cloud and using services like Dataflow, BigQuery, and Managed Airflow (previously, Cloud Composer). What were the main reasons for choosing this stack, and what trade-offs did you have to consider when selecting GCP and Apache Beam/Dataflow?** + +**Raj**: The selection of Google Cloud, utilizing BigQuery, Dataflow, and Managed Airflow, was a strategic decision to standardize the ML lifecycle and enable high-velocity development at scale. By leveraging these technologies, we established a unified architectural standard that allows for elastic scalability and reliability across petabyte-scale workloads. This approach ensures that the entire ML lifecycle, from initial experimentation to global production execution, is governed by a consistent framework, providing the necessary infrastructure to manage complex data movements and distributed computing requirements without sacrificing architectural integrity. + +Our goal was to create a unified abstraction layer that could bridge various data processing tools like Spark and BigQuery with ML libraries such as scikit-learn and TensorFlow. GCP provided the deep, native integrations we required to smoothly transition from local prototypes to massive, distributed cloud execution. + +Talking about trade-offs, the biggest hurdle was simplifying the inherent complexity of a multi-framework environment. Even with powerful connectors for GCS and BigQuery, we still had to develop Vega (Intuit Credit Karma’s internal ML platform) to shield our data scientists from manual orchestration and resource management, allowing them to remain focused on a clean, unified Python API. + +**Jana: Dataflow is widely known for ETL and batch/streaming processing, but you took it a step further by using it for ML. You built Vega, a full ML explainability platform on Apache Beam and Dataflow. What motivated that choice, and what advantages did these technologies bring to the project?** + +**Raj**: We utilized Dataflow as the distributed execution engine to power the broader Vega ecosystem, which serves as our comprehensive ML platform. It is important to emphasize that explainability is a foundational pillar of this ecosystem, integrated directly into the core architecture rather than treated as a standalone task. By building on Apache Beam, we ensure that every model, whether focused on feature engineering, scoring, or interpretability, benefits from a unified processing model. This integration allows us to maintain rigorous standards for model transparency and governance as a native component of our large-scale distributed ML operations. + +The main motivation was to eliminate the experimentation-to-production gap. Traditional platforms forced data scientists to manually rewrite code for production frameworks, but by using Dataflow and Apache Beam as the core engine, we were able to support both batch and streaming processing alongside ML scoring within a single workflow. + +The primary advantage that came out of this was the unified Python API. It allowed data scientists to build complex, end-to-end data and ML pipelines that transition seamlessly from local development on sampled data to cloud-scale production execution, without architectural changes or manual re-engineering. + +**Jana: Looking back at building Vega, what were the biggest challenges you faced, and what would you do differently today?** + +**Raj**: The primary challenges centered on managing the inherent complexities of distributed systems and maintaining architectural coherence across a diverse ML lifecycle. Operating at the scale of over 140 million members requires solving for high-concurrency data access, managing complex dependency graphs in distributed environments, and ensuring that performance remains consistent as workloads transition from sampled data to massive production clusters. Balancing the need for developer flexibility with the strict latency and reliability constraints of a globally distributed platform represents the most significant architectural hurdle in modern ML engineering. + +Reflecting on the evolution of Vega, the biggest challenge was navigating the trade-offs of building critical infrastructure while scaling to support massive business growth with limited resources. We prioritized speed-to-market and immediate utility to enable the business, which was necessary but created technical debt in terms of documentation and onboarding. If I were starting today, I would prioritize earlier investments in developer tooling and standardized templates. We are now shifting our focus from that initial fast-and-lean build phase to long-term sustainability, refining the developer experience, automating maintenance even further, and broadening support for modern ML paradigms. This ensures that the platform remains a scalable, durable asset as it evolves to meet the next generation of business needs. + +**Jana: Are there any features or improvements you would love to see in Apache Beam or Dataflow in the future?** + +**Raj**: While Vega has successfully leveraged Beam for our core workflows, there are a few areas where future improvements could make a real difference. One would be enhanced data engineering primitives. We would value deeper native support for large-scale data manipulation, specifically generalized tiling operations for distributed datasets and built-in capabilities for saving and managing partial processing states. This would significantly reduce the complexity of checkpointing and state recovery in our pipelines. + +The other area is real-time inference. While we have strong batch capabilities, we need more streamlined primitives in Beam to better integrate our data pipelines with real-time serving infrastructure. Reducing the complexity of the hand-off between processing and prediction would be a major win for our personalization engines. + +**Jana: You are a Staff Machine Learning Engineer at Intuit Credit Karma — could you tell us a bit about your background and journey to this role? And for readers who are interested in pursuing a career in ML, what advice would you give them?** + +**Raj**: As a Staff Machine Learning Engineer, my career has been defined by the practice of platform-scale engineering. I have focused on architecting ecosystems that provide high-velocity developer experiences, allowing engineering teams to move from conceptual design to production-scale deployment within a unified framework. By treating the ML lifecycle as a first-class engineering problem, we build systems that automate the complexities of infrastructure management, enabling data scientists to focus on model innovation while the platform handles the rigorous demands of distributed execution and governance. + +My advice to anyone pursuing a career in ML would be to focus on first-principles thinking. Don’t just learn a specific library, understand the underlying data infrastructure, including data movement, latency constraints, and caching. The most impactful engineers I have worked with, including those I have collaborated with at Google and beyond, are the ones who can reason across domains, from low-level feature serving latency to high-level platform architecture. + +**Jana: You are an active member of the Apache Beam community. This is your third Beam Summit. How important do you think community involvement is for engineers, and how has it shaped your career?** + +**Raj**: Community involvement is critical. Participating in initiatives like the Apache Beam Summit has allowed me to stay at the frontier of distributed computing and directly influence the evolution of the tools we depend on. It transforms engineering from a siloed task into a collaborative effort; seeing how other engineers solve (or struggle with) similar problems helps refine your own architectural decisions and often provides the reference architectures that we would otherwise have taken years to develop independently. + +**Jana: If you could give one piece of advice to someone just starting with Apache Beam, what would it be?** + +**Raj**: Focus on the abstraction. Do not treat Beam as just another processing framework; treat it as the “universal language” for your data pipeline. Start by leveraging the high-level transforms and built-in IO connectors rather than trying to optimize low-level custom code immediately. Understanding how to structure your pipeline to take advantage of Beam’s windowing and triggering capabilities is what will eventually separate a good engineer from a great one when you hit the scale of millions of users. + +**Jana: Thank you so much, Raj, for sharing these insights!** + +**Raj**: Thank you for having me, Jana, this was a great conversation! + +------------ + +🌐 To learn more about **Intuit Credit Karma**, visit their [website](https://www.creditkarma.com/) or check out their [LinkedIn](https://www.linkedin.com/company/intuitcreditkarma/). You can also connect with **Raj** directly on [LinkedIn](https://www.linkedin.com/in/rajkiran2190). + +📌 Raj and Pallav Anand presented at Beam Summit 2026. Check out their talk, [‘Beyond the Black Box: How Intuit Credit Karma Runs ML Explainability for 140M Members with Beam’](https://beamsummit.org/sessions/2026/beyond-the-black-box-how-intuit-credit-karma-runs-ml-explainability-for-140m-members-with-beam/), and explore the full program at [beamsummit.org](https://beamsummit.org/sessions/2026/). + +📺 Session recording will be available on the [Apache Beam YouTube channel](https://www.youtube.com/@ApacheBeamYT). + +📰 Missed the Beam Summit 2026? Read my recap: [Beam Summit 2026: Apache Beam Just Got Even More Interesting 🐝](https://medium.com/google-cloud/beam-summit-2026-apache-beam-just-got-even-more-interesting-7c6c9967ffb9) + +– [Jana Polianskaja](https://www.linkedin.com/in/jana-polianskaja/) diff --git a/website/www/site/data/authors.yml b/website/www/site/data/authors.yml index b74a6456d609..d98f2281bddc 100644 --- a/website/www/site/data/authors.yml +++ b/website/www/site/data/authors.yml @@ -87,6 +87,9 @@ jamesmalone: name: James Malone email: jamesmalone@apache.org twitter: chimerasaurus +janaom: + name: Jana Polianskaja + email: jana.polianskaja@gmail.com jesseanderson: name: Jesse Anderson twitter: jessetanderson From 22a595161680ed170ebc196deb30f8f7adbbd2e3 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Wed, 22 Jul 2026 07:57:50 -0400 Subject: [PATCH 22/29] Document job management and interactive API protos (#39294) Add documentation comments to previously undocumented messages and fields in the Beam job management and interactive API protos: beam_artifact_api.proto: ArtifactRetrievalService, ArtifactStagingService, ArtifactRequestWrapper, ProxyManifest, PutArtifactMetadata, and more. beam_expansion_api.proto: ExpansionRequest/Response, DiscoverSchemaTransformRequest/Response, SchemaTransformConfig. beam_job_api.proto: JobService, JobInfo, JobStateEvent, JobMessage, GetJobMetricsRequest/Response, MetricResults, PipelineOptionType. beam_interactive_api.proto: TestStreamService, EventsRequest. --- .../interactive/v1/beam_interactive_api.proto | 4 +++ .../job_management/v1/beam_artifact_api.proto | 19 +++++++++-- .../v1/beam_expansion_api.proto | 9 ++++++ .../job_management/v1/beam_job_api.proto | 32 +++++++++++++++++++ 4 files changed, 62 insertions(+), 2 deletions(-) diff --git a/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto b/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto index 6351b2472579..faae80261402 100644 --- a/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto +++ b/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto @@ -50,11 +50,15 @@ message TestStreamFileRecord { org.apache.beam.model.pipeline.v1.TestStreamPayload.Event recorded_event = 1; } +// A gRPC service that serves cached TestStream events to a TestStream source. +// Used by Interactive Beam to replay recorded pipeline elements from a streaming cache. service TestStreamService { // A TestStream will request for events using this RPC. rpc Events(EventsRequest) returns (stream org.apache.beam.model.pipeline.v1.TestStreamPayload.Event) {} } +// A request for TestStream events from a streaming cache. +// The client specifies which PCollection output tags to read events for. message EventsRequest { // The set of PCollections to read from. These are the PTransform outputs // local names. These are a subset of the TestStream's outputs. This allows diff --git a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto index 35cd4f0c73d1..a64352c9289e 100644 --- a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto +++ b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto @@ -78,18 +78,23 @@ message ResolveArtifactsResponse { // A request to get an artifact. message GetArtifactRequest { + // (Required) The artifact to retrieve, typically a resolved artifact from ResolveArtifacts. org.apache.beam.model.pipeline.v1.ArtifactInformation artifact = 1; } // Part of a response to getting an artifact. message GetArtifactResponse { + // A chunk of the artifact's binary content. bytes data = 1; } // Wraps an ArtifactRetrievalService request for use in ReverseArtifactRetrievalService. +// The server sends these requests to the client during the reverse staging flow. message ArtifactRequestWrapper { oneof request { + // Request to resolve a set of artifacts into concrete artifact references. ResolveArtifactsRequest resolve_artifact = 1000; + // Request to retrieve the bytes of a specific artifact. GetArtifactRequest get_artifact = 1001; } } @@ -151,16 +156,22 @@ message ArtifactMetadata { // A collection of artifacts. message Manifest { + // (Required) The list of artifacts staged for a job. repeated ArtifactMetadata artifact = 1; } // A manifest with location information. message ProxyManifest { + // (Required) The manifest of staged artifacts. Manifest manifest = 1; + // A location (name + URI) at which a staged artifact can be retrieved. message Location { - string name = 1; - string uri = 2; + // (Required) The name of the artifact, matching the name in the Manifest. + string name = 1; + // (Required) The URI at which the artifact content can be retrieved. + string uri = 2; } + // (Required) The list of artifact locations for retrieval. repeated Location location = 2; } @@ -173,6 +184,7 @@ message GetManifestRequest { // A response containing a job manifest. message GetManifestResponse { + // (Required) The manifest of staged artifacts for the job. Manifest manifest = 1; } @@ -187,9 +199,11 @@ message LegacyGetArtifactRequest { // Part of an artifact. message ArtifactChunk { + // A chunk of the artifact's binary content. bytes data = 1; } +// Metadata for an artifact being staged, sent as the first message in a PutArtifact stream. message PutArtifactMetadata { // (Required) A token for artifact staging session. This token can be obtained // from PrepareJob request in JobService @@ -211,6 +225,7 @@ message PutArtifactRequest { } } +// A response to staging an artifact. An empty response indicating the artifact was staged successfully. message PutArtifactResponse { } diff --git a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto index a4736f8b4938..0e3b1adbffd6 100644 --- a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto +++ b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto @@ -33,6 +33,8 @@ import "google/protobuf/struct.proto"; import "org/apache/beam/model/pipeline/v1/beam_runner_api.proto"; import "org/apache/beam/model/pipeline/v1/schema.proto"; +// A request to expand a single PTransform in a remote SDK. The expansion service +// resolves the transform into its constituent subtransforms and outputs. message ExpansionRequest { // Set of components needed to interpret the transform, or which // may be useful for its expansion. This includes the input @@ -64,6 +66,8 @@ message ExpansionRequest { google.protobuf.Struct pipeline_options = 6; } +// The result of expanding a PTransform, containing the expanded transform and +// all newly created components. message ExpansionResponse { // Set of components needed to execute the expanded transform, // including the (original) inputs, outputs, and subtransforms. @@ -82,9 +86,13 @@ message ExpansionResponse { string error = 10; } +// A request to discover all SchemaTransformProviders registered with the expansion service. +// This is used by cross-language pipeline construction to enumerate available transforms. message DiscoverSchemaTransformRequest { } +// Configuration metadata for a SchemaTransform, describing its inputs, outputs, +// and configuration schema. Used to enable cross-language transform discovery. message SchemaTransformConfig { // Config schema of the SchemaTransform org.apache.beam.model.pipeline.v1.Schema config_schema = 1; @@ -102,6 +110,7 @@ message SchemaTransformConfig { string description = 4; } +// A response containing all discovered SchemaTransform configurations. message DiscoverSchemaTransformResponse { // A mapping from SchemaTransform ID to schema transform config of discovered // SchemaTransforms diff --git a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto index 8e196d5bedd3..353eaec57994 100644 --- a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto +++ b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto @@ -112,6 +112,7 @@ message RunJobRequest { } +// The response from submitting a job for execution. message RunJobResponse { string job_id = 1; // (required) The ID for the executing job } @@ -144,6 +145,7 @@ message DrainJobResponse { } // A subset of info provided by ProvisionApi.ProvisionInfo +// Represents metadata about a submitted job. message JobInfo { string job_id = 1; // (required) string job_name = 2; // (required) @@ -155,6 +157,7 @@ message JobInfo { // Throws error GRPC_STATUS_UNAVAILABLE if server is down message GetJobsRequest { } +// The response from requesting a list of all invoked jobs. message GetJobsResponse { repeated JobInfo job_info = 1; // (required) } @@ -168,6 +171,7 @@ message GetJobStateRequest { } +// An event representing a job state transition. Emitted by GetState and GetStateStream. message JobStateEvent { JobState.Enum state = 1; // (required) google.protobuf.Timestamp timestamp = 2; // (required) @@ -182,6 +186,7 @@ message GetJobPipelineRequest { } +// The response from requesting a job's pipeline representation. message GetJobPipelineResponse { org.apache.beam.model.pipeline.v1.Pipeline pipeline = 1; // (required) } @@ -195,25 +200,41 @@ message JobMessagesRequest { string job_id = 1; // (required) } +// A single log message or diagnostic event from a running job. message JobMessage { + // (Required) A unique identifier for this message. string message_id = 1; + // (Required) The time at which the message was emitted, as a string representation. string time = 2; + // (Required) The importance level of this message. MessageImportance importance = 3; + // (Required) The text content of the message. string message_text = 4; + // Importance levels for job messages, ordered from least to most severe. enum MessageImportance { + // The importance was not specified. MESSAGE_IMPORTANCE_UNSPECIFIED = 0; + // Debug-level messages, typically very verbose. JOB_MESSAGE_DEBUG = 1; + // Detailed informational messages. JOB_MESSAGE_DETAILED = 2; + // Basic informational messages. JOB_MESSAGE_BASIC = 3; + // Warning messages indicating potential issues. JOB_MESSAGE_WARNING = 4; + // Error messages indicating failures. JOB_MESSAGE_ERROR = 5; } } +// A streaming response from GetMessageStream, containing either a job message +// or a job state change event. message JobMessagesResponse { oneof response { + // A log message or diagnostic event from the job. JobMessage message_response = 1; + // A job state transition event. JobStateEvent state_response = 2; } } @@ -270,17 +291,22 @@ message JobState { } +// A request to fetch metrics for a given job. message GetJobMetricsRequest { string job_id = 1; // (required) } +// A response containing metrics for a given job. message GetJobMetricsResponse { + // (Required) The metrics results containing both attempted and committed values. MetricResults metrics = 1; } // All metrics for a given job. Runners may support one or the other or both. message MetricResults { + // Metrics reflecting the result of attempted (non-committed) computations. repeated org.apache.beam.model.pipeline.v1.MonitoringInfo attempted = 1; + // Metrics reflecting the result of committed computations. repeated org.apache.beam.model.pipeline.v1.MonitoringInfo committed = 2; } @@ -296,12 +322,17 @@ message DescribePipelineOptionsRequest { // Types mirror those of JSON, since that's how pipeline options are serialized. message PipelineOptionType { enum Enum { + // A string value. STRING = 0; + // A boolean (true/false) value. BOOLEAN = 1; // whole numbers, see https://json-schema.org/understanding-json-schema/reference/numeric.html INTEGER = 2; + // A floating-point number. NUMBER = 3; + // An array of values. ARRAY = 4; + // A nested object. OBJECT = 5; }; } @@ -324,6 +355,7 @@ message PipelineOptionDescriptor { string group = 5; } +// The response from describing pipeline options, containing a list of option descriptors. message DescribePipelineOptionsResponse { // List of pipeline option descriptors. repeated PipelineOptionDescriptor options = 1; From d2d7ef88eac2f37ef5f0403d3a9ab2a4906577e7 Mon Sep 17 00:00:00 2001 From: Abdelrahman Ibrahim Date: Wed, 22 Jul 2026 15:18:21 +0300 Subject: [PATCH 23/29] Prefer binary wheels when installing Python container dependencies (#39402) * Speed up Python container builds with prefer-binary and registry cache * Address review * Drop docker registry cache --------- Co-authored-by: Abdelrahman Ibrahim --- sdks/python/container/Dockerfile | 3 ++- sdks/python/container/run_generate_requirements.sh | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/sdks/python/container/Dockerfile b/sdks/python/container/Dockerfile index 7f5aa3604fe9..a1e17edb82d1 100644 --- a/sdks/python/container/Dockerfile +++ b/sdks/python/container/Dockerfile @@ -53,7 +53,8 @@ RUN \ # Install required packages for Beam Python SDK and common dependencies used by users. # use --no-deps to ensure the list includes all transitive dependencies. - pip install --no-deps -r /tmp/base_image_requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu && \ + # use --prefer-binary to avoid compiling wheels from source when prebuilt wheels exist. + pip install --prefer-binary --no-deps -r /tmp/base_image_requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu && \ rm -rf /tmp/base_image_requirements.txt && \ python -c "import nltk; nltk.download('stopwords')" && \ rm /root/nltk_data/corpora/stopwords.zip && \ diff --git a/sdks/python/container/run_generate_requirements.sh b/sdks/python/container/run_generate_requirements.sh index 4c7eea0b7398..ed1f57faac4e 100755 --- a/sdks/python/container/run_generate_requirements.sh +++ b/sdks/python/container/run_generate_requirements.sh @@ -93,8 +93,8 @@ fi # Force torch dependencies to be pulled from the PyTorch CPU wheel # repository so that they don't include GPU dependencies with # non-compliant licenses -pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir "$SDK_TARBALL""$EXTRAS" $INDEX_URL_OPTION -pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir -r "$PWD"/sdks/python/container/base_image_requirements_manual.txt +pip install --prefer-binary ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir "$SDK_TARBALL""$EXTRAS" $INDEX_URL_OPTION +pip install --prefer-binary ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir -r "$PWD"/sdks/python/container/base_image_requirements_manual.txt pip uninstall -y apache-beam echo "Checking for broken dependencies:" From 3f3f18a8ce0073cdf036f541f48eb8f3bd930426 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:23:28 -0400 Subject: [PATCH 24/29] Bump github.com/aws/aws-sdk-go-v2/feature/s3/manager in /sdks (#39421) Bumps [github.com/aws/aws-sdk-go-v2/feature/s3/manager](https://github.com/aws/aws-sdk-go-v2) from 1.22.34 to 1.22.35. - [Release notes](https://github.com/aws/aws-sdk-go-v2/releases) - [Commits](https://github.com/aws/aws-sdk-go-v2/compare/feature/s3/manager/v1.22.34...feature/s3/manager/v1.22.35) --- updated-dependencies: - dependency-name: github.com/aws/aws-sdk-go-v2/feature/s3/manager dependency-version: 1.22.35 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Derrick Williams --- sdks/go.mod | 10 +++++----- sdks/go.sum | 20 ++++++++++---------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index 752003b3552f..02dccbc101b1 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -33,10 +33,10 @@ require ( cloud.google.com/go/spanner v1.93.0 cloud.google.com/go/storage v1.63.1 github.com/aws/aws-sdk-go-v2 v1.43.0 - github.com/aws/aws-sdk-go-v2/config v1.32.30 + github.com/aws/aws-sdk-go-v2/config v1.32.31 github.com/aws/aws-sdk-go-v2/credentials v1.19.30 - github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34 - github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 + github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.35 + github.com/aws/aws-sdk-go-v2/service/s3 v1.106.0 github.com/aws/smithy-go v1.27.4 github.com/docker/go-connections v0.7.0 // indirect github.com/dustin/go-humanize v1.0.1 @@ -153,9 +153,9 @@ require ( github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 // indirect github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.24 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.32 // indirect github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.45.0 // indirect diff --git a/sdks/go.sum b/sdks/go.sum index 2db1288b176b..9d3f1375ec10 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -204,8 +204,8 @@ github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHz github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14/go.mod h1:zwM6veDkhGgQFqkBy+uT28AAYpLu+uFMlPl+rCg/73E= github.com/aws/aws-sdk-go-v2/config v1.15.3/go.mod h1:9YL3v07Xc/ohTsxFXzan9ZpFpdTOFl4X65BAKYaz8jg= github.com/aws/aws-sdk-go-v2/config v1.25.3/go.mod h1:tAByZy03nH5jcq0vZmkcVoo6tRzRHEwSFx3QW4NmDw8= -github.com/aws/aws-sdk-go-v2/config v1.32.30 h1:XwsEzpTJfQYJbFicz/QMLwAZdyeNVVoOEkbF7R3gPJk= -github.com/aws/aws-sdk-go-v2/config v1.32.30/go.mod h1:Ud32SuMc+/9BGxfpSVld7HrE2o05JwKmXY4M3jOQNZU= +github.com/aws/aws-sdk-go-v2/config v1.32.31 h1:n4nY9O3QKoHIkL85EX+V8RcMFtOhlpTFhGArg915PXk= +github.com/aws/aws-sdk-go-v2/config v1.32.31/go.mod h1:PN0NYDCCoOpGGsZ2+elDUidmHfQBPyYzN2GCgl8HEBs= github.com/aws/aws-sdk-go-v2/credentials v1.11.2/go.mod h1:j8YsY9TXTm31k4eFhspiQicfXPLZ0gYXA50i4gxPE8g= github.com/aws/aws-sdk-go-v2/credentials v1.16.2/go.mod h1:sDdvGhXrSVT5yzBDR7qXz+rhbpiMpUYfF3vJ01QSdrc= github.com/aws/aws-sdk-go-v2/credentials v1.19.30 h1:TTCvvzFU6gXa4iJecNG/0F/B0oYTiazoRECr2XyLHrY= @@ -216,8 +216,8 @@ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 h1:kfVL5wAunCJycL6MOQ6aNh github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31/go.mod h1:nWfRNDAppujCQgOUd43lKT4yeLv9z3nJ3bw1G3BgQKo= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.3/go.mod h1:0dHuD2HZZSiwfJSy1FO5bX1hQ1TxVV1QXXjpn3XUE44= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.14.0/go.mod h1:UcgIwJ9KHquYxs6Q5skC9qXjhYMK+JASDYcXQ4X7JZE= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34 h1:Pn7OsMwBLbkZ6OnCxWHAjf0L/22H8cnhxZC0uPwtMtg= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34/go.mod h1:eToXR/Gk1uqpn04eSmdgVXwfS0WvH8aG4eBFr8ygbpU= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.35 h1:TwCjUC1rnFKTtfqEpQY9ClYPFpGpUaouODrdGPB5b3Y= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.35/go.mod h1:V0zqtP3iJk9zu86GxuQGN09RYyQB+3mjcPiyfLC6wlg= github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.9/go.mod h1:AnVH5pvai0pAF4lXRq0bmhbes1u9R8wTE+g+183bZNM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.3/go.mod h1:7sGSz1JCKHWWBHq98m6sMtWQikmYPpxjqOydDemiVoM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 h1:Z8F3hfCY33IGpJjFAnv0wvtv1FIKj1GHmRDEYqy64tw= @@ -237,21 +237,21 @@ github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13/go.mod h1:ITg9em2KbJx1s0y4aqRX5OYWG6HBZ5TVR//OdpEZ2CQ= github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.3/go.mod h1:Seb8KNmD6kVTjwRjVEgOT5hPin6sq+v4C2ycJQDwuH8= github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.2.3/go.mod h1:R+/S1O4TYpcktbVwddeOYg+uwUfLhADP2S/x4QwsCTM= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 h1:9Fjh6fi/U5JEStVZijmaMpUwE/gvBJj7x2B/PjbO9To= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23/go.mod h1:iMoT2f1tClxrWAAnKCXjZQ6LOmfLrMG14wmnWpM+F14= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.24 h1:mdPwDQPqxlw9Sc62Nt15yjEcARaDbPXkjRYtXsUripo= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.24/go.mod h1:ls5ytnwLTcQaUu32fMYXFI3MjpKuTwL840PAm9iqyEg= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.3/go.mod h1:wlY6SVjuwvh3TVRpTqdy4I1JpBFLX4UGeKZdWntaocw= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.10.3/go.mod h1:Owv1I59vaghv1Ax8zz8ELY8DN7/Y0rGS+WWAmjgi950= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 h1:w2SIhW92DZPFrSL4ksVCr8IYff5OZwIcxg8+95tzvAI= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31/go.mod h1:wAhpCQbkov+IcvjozJbd2xRCoZybUEHNkcFunssNACg= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.3/go.mod h1:Bm/v2IaN6rZ+Op7zX+bOUMdL4fsrYZiD0dsjLhNKwZc= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.16.3/go.mod h1:KZgs2ny8HsxRIRbDwgvJcHHBZPOzQr/+NtGwnP+w2ec= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 h1:uao4A3QZ5UmB326V6KF+qRpv9Tjz7IlnlnTbbANntlU= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31/go.mod h1:I/1+z0VwL1GhQyLgkoHDlygpUZ+iTAwOQ/NsftiUL2I= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.32 h1:jWXtZdCnhXa9sGFixRaU2AxT4DIVse9HS4E2f+/KwV0= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.32/go.mod h1:9JS1UpfVvyD/ZPX8GsKb/Pq8scEM+7GP5fqh9SwH7po= github.com/aws/aws-sdk-go-v2/service/kms v1.16.3/go.mod h1:QuiHPBqlOFCi4LqdSskYYAWpQlx3PKmohy+rE2F+o5g= github.com/aws/aws-sdk-go-v2/service/s3 v1.26.3/go.mod h1:g1qvDuRsJY+XghsV6zg00Z4KJ7DtFFCx8fJD2a491Ak= github.com/aws/aws-sdk-go-v2/service/s3 v1.43.0/go.mod h1:NXRKkiRF+erX2hnybnVU660cYT5/KChRD4iUgJ97cI8= -github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 h1:5C00eQYpTrgQXnp6V3P6P7zPElna3AXvlukbANE6nJI= -github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2/go.mod h1:zdmCoFO/dSI7GlrwsPqFJI+WlFnSU4Tc8TJnlXrM1Do= +github.com/aws/aws-sdk-go-v2/service/s3 v1.106.0 h1:7QZWVJZWzHivHWIa+5TELLaBBkbuoj0GPwQtMlJ0sqk= +github.com/aws/aws-sdk-go-v2/service/s3 v1.106.0/go.mod h1:fcvq5L7dK+5cQFicEJwpI6e6Wn8NY2i6yT5wRLYVc7s= github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.15.4/go.mod h1:PJc8s+lxyU8rrre0/4a0pn2wgwiDvOEzoOjcJUBr67o= github.com/aws/aws-sdk-go-v2/service/signin v1.5.0 h1:OHH5iTQvVGmfHjX/5Q+vFuA/Rf2x6/95aJ/75QCQSm4= github.com/aws/aws-sdk-go-v2/service/signin v1.5.0/go.mod h1:mCF3AK9PpL49oOrhniUXWAfhVBVQ/XbytoE5eccZUIs= From aa8869c633d23598337977556393414bba51049f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:23:56 -0400 Subject: [PATCH 25/29] Bump cloud.google.com/go/storage from 1.63.1 to 1.64.0 in /sdks (#39422) Bumps [cloud.google.com/go/storage](https://github.com/googleapis/google-cloud-go) from 1.63.1 to 1.64.0. - [Release notes](https://github.com/googleapis/google-cloud-go/releases) - [Changelog](https://github.com/googleapis/google-cloud-go/blob/main/CHANGES.md) - [Commits](https://github.com/googleapis/google-cloud-go/compare/storage/v1.63.1...compute/v1.64.0) --- updated-dependencies: - dependency-name: cloud.google.com/go/storage dependency-version: 1.64.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Derrick Williams --- sdks/go.mod | 2 +- sdks/go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index 02dccbc101b1..12bd9e9ba2f4 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -31,7 +31,7 @@ require ( cloud.google.com/go/profiler v0.6.0 cloud.google.com/go/pubsub v1.51.0 cloud.google.com/go/spanner v1.93.0 - cloud.google.com/go/storage v1.63.1 + cloud.google.com/go/storage v1.64.0 github.com/aws/aws-sdk-go-v2 v1.43.0 github.com/aws/aws-sdk-go-v2/config v1.32.31 github.com/aws/aws-sdk-go-v2/credentials v1.19.30 diff --git a/sdks/go.sum b/sdks/go.sum index 9d3f1375ec10..81b86614757b 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -99,8 +99,8 @@ cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RX cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0= cloud.google.com/go/storage v1.12.0/go.mod h1:fFLk2dp2oAhDz8QFKwqrjdJvxSp/W2g7nillojlL5Ho= cloud.google.com/go/storage v1.21.0/go.mod h1:XmRlxkgPjlBONznT2dDUU/5XlpU2OjMnKuqnZI01LAA= -cloud.google.com/go/storage v1.63.1 h1:CYXILV9G4CH0C18IQ9+V0h4XiqD2LhKnMLO0o7uJWNs= -cloud.google.com/go/storage v1.63.1/go.mod h1:lWyAtwvDZHdL3k68WVKbESP6bmWaV23ZJJ/JEVw/ZaQ= +cloud.google.com/go/storage v1.64.0 h1:KLpxI/oX9LxeRsNqn877d2WyeT3ryiEwnGt8pwcSPZg= +cloud.google.com/go/storage v1.64.0/go.mod h1:lWyAtwvDZHdL3k68WVKbESP6bmWaV23ZJJ/JEVw/ZaQ= cloud.google.com/go/trace v1.0.0/go.mod h1:4iErSByzxkyHWzzlAj63/Gmjz0NH1ASqhJguHpGcr6A= cloud.google.com/go/trace v1.2.0/go.mod h1:Wc8y/uYyOhPy12KEnXG9XGrvfMz5F5SrYecQlbW1rwM= cloud.google.com/go/trace v1.16.0 h1:GmQovzFc5F0CNfl0VLgL64aoTtu7xsM0YajW2GlG9+E= From 583601b13bf592ccb5a3ba9d910239e493f36ac8 Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:07:34 -0400 Subject: [PATCH 26/29] Fix pyrefly check bad-context-manager (#39418) --- .../python/apache_beam/runners/interactive/recording_manager.py | 2 +- sdks/python/pyproject.toml | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/runners/interactive/recording_manager.py b/sdks/python/apache_beam/runners/interactive/recording_manager.py index c19b60b64fd2..cabcca558dca 100644 --- a/sdks/python/apache_beam/runners/interactive/recording_manager.py +++ b/sdks/python/apache_beam/runners/interactive/recording_manager.py @@ -105,7 +105,7 @@ def _cancel_clicked(self, b): self.cancel() def update_display(self, msg: str, progress: Optional[float] = None): - if not IS_IPYTHON: + if not IS_IPYTHON or self._output_widget is None: print(f'AsyncCompute: {msg}') return diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index 52f713466246..5d78db863366 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -213,7 +213,6 @@ invalid-inheritance = "ignore" not-iterable = "ignore" unexpected-keyword = "ignore" bad-specialization = "ignore" -bad-context-manager = "ignore" invalid-yield = "ignore" bad-argument-count = "ignore" bad-typed-dict-key = "ignore" From aaad9b4b8e69a272b6e6a6fa4df516b93bdd1506 Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:08:37 -0400 Subject: [PATCH 27/29] [Experimental] Use zstd compression in Docker builds (#39409) * [Experimental] Use zstd compression in Docker builds * fix load errors --- ...it_Python_ValidatesContainer_Dataflow.json | 2 +- .../beam_PreCommit_Flink_Container.json | 2 +- .../beam/gradle/BeamDockerPlugin.groovy | 24 +++++++++++++++---- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json index 60eb7272a936..327fe1987c50 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 4 + "modification": 5 } \ No newline at end of file diff --git a/.github/trigger_files/beam_PreCommit_Flink_Container.json b/.github/trigger_files/beam_PreCommit_Flink_Container.json index bbdc3a3910ef..62ae7886c573 100644 --- a/.github/trigger_files/beam_PreCommit_Flink_Container.json +++ b/.github/trigger_files/beam_PreCommit_Flink_Container.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 3 + "modification": 4 } diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy index 6963f96d7313..574e34118905 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy @@ -61,6 +61,7 @@ class BeamDockerPlugin implements Plugin { boolean push = false String builder = null String target = null + String compression = 'zstd' File resolvedDockerfile = null File resolvedDockerComposeTemplate = null @@ -72,6 +73,9 @@ class BeamDockerPlugin implements Plugin { DockerExtension(Project project) { this.project = project this.copySpec = project.copySpec() + if (project.hasProperty('docker-compression')) { + this.compression = project.property('docker-compression') + } } void resolvePathsAndValidate() { @@ -235,13 +239,23 @@ class BeamDockerPlugin implements Plugin { if (!ext.platform.isEmpty()) { buildCommandLine.addAll('--platform', String.join(',', ext.platform)) } - if (ext.load) { - buildCommandLine.add '--load' + if (ext.load && ext.push) { + throw new Exception("cannot combine 'push' and 'load' options") } - if (ext.push) { - buildCommandLine.add '--push' + if (ext.compression != null && !ext.compression.isEmpty()) { + if (ext.push) { + buildCommandLine.add "--output=type=registry,compression=${ext.compression},force-compression=true,oci-mediatypes=true" + } else if (ext.load) { + buildCommandLine.add '--load' + } else { + buildCommandLine.add "--output=type=image,compression=${ext.compression},force-compression=true,oci-mediatypes=true" + } + } else { if (ext.load) { - throw new Exception("cannot combine 'push' and 'load' options") + buildCommandLine.add '--load' + } + if (ext.push) { + buildCommandLine.add '--push' } } if (ext.builder != null) { From 7ebc023bb3eff613302e5e55891c6aea7f730932 Mon Sep 17 00:00:00 2001 From: Abdelrahman Ibrahim Date: Wed, 22 Jul 2026 18:37:54 +0300 Subject: [PATCH 28/29] Isolate Python 3.14 SDK snapshot publishes from the main matrix --- .../beam_Publish_Beam_SDK_Snapshots.yml | 66 ++++++++++++++++++- 1 file changed, 64 insertions(+), 2 deletions(-) diff --git a/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml b/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml index c37ddf7e4f3d..a958ff6eeac2 100644 --- a/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml +++ b/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml @@ -62,12 +62,10 @@ jobs: - "python:container:py311:docker" - "python:container:py312:docker" - "python:container:py313:docker" - - "python:container:py314:docker" - "python:container:distroless:py310:docker" - "python:container:distroless:py311:docker" - "python:container:distroless:py312:docker" - "python:container:distroless:py313:docker" - - "python:container:distroless:py314:docker" - "python:container:ml:py310:docker" - "python:container:ml:py311:docker" - "python:container:ml:py312:docker" @@ -127,3 +125,67 @@ jobs: -Pcontainer-architecture-list=arm64,amd64 \ -Ppush-containers \ -Pdocker-pull-licenses + + beam_Publish_Beam_SDK_Snapshots_py314: + needs: beam_Publish_Beam_SDK_Snapshots + if: | + always() && + !cancelled() && + (github.event_name == 'workflow_dispatch' || + (github.event_name == 'schedule' && github.repository == 'apache/beam')) + runs-on: [self-hosted, ubuntu-24.04, main] + timeout-minutes: 300 + name: ${{ matrix.job_name }} (${{ matrix.container_task }}) + strategy: + fail-fast: false + matrix: + job_name: ["beam_Publish_Beam_SDK_Snapshots"] + job_phrase: ["N/A"] + container_task: + - "python:container:py314:docker" + - "python:container:distroless:py314:docker" + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + - name: Free Disk Space (Ubuntu) + uses: jlumbroso/free-disk-space@v1.3.1 + - name: Setup repository + uses: ./.github/actions/setup-action + with: + comment_phrase: ${{ matrix.job_phrase }} + github_token: ${{ secrets.GITHUB_TOKEN }} + github_job: ${{ matrix.job_name }} (${{ matrix.container_task }}) + - name: Find Beam Version + run: | + BEAM_VERSION_LINE=$(cat gradle.properties | grep "sdk_version") + echo "BEAM_VERSION=${BEAM_VERSION_LINE#*sdk_version=}" >> $GITHUB_ENV + - name: Set latest tag only on master branch + if: github.ref == 'refs/heads/master' + run: echo "LATEST_TAG=,latest" >> $GITHUB_ENV + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c + - name: Authenticate on GCP + uses: google-github-actions/auth@7c6bc770dae815cd3e89ee6cdf493a5fab2cc093 + with: + service_account: ${{ secrets.GCP_SA_EMAIL }} + credentials_json: ${{ secrets.GCP_SA_KEY }} + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db + - name: GCloud Docker credential helper + run: | + gcloud auth configure-docker ${{ env.docker_registry }} + - name: Setup Python environment + uses: ./.github/actions/setup-environment-action + with: + python-version: default + - name: run Publish Beam SDK Snapshots script + uses: ./.github/actions/gradle-command-self-hosted-action + with: + gradle-command: :sdks:${{ matrix.container_task }} + arguments: | + -Pdocker-repository-root=gcr.io/apache-beam-testing/beam-sdk \ + -Pdocker-tag-list=${{ github.sha }},${BEAM_VERSION}${LATEST_TAG} \ + -Pcontainer-architecture-list=arm64,amd64 \ + -Ppush-containers \ + -Pdocker-pull-licenses From 797ab19e551a900c026b061f5586423cd3589a3f Mon Sep 17 00:00:00 2001 From: Abdelrahman Ibrahim Date: Wed, 22 Jul 2026 18:53:40 +0300 Subject: [PATCH 29/29] Fix CodeQL env injection --- .../beam_Publish_Beam_SDK_Snapshots.yml | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml b/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml index a958ff6eeac2..6bf05e99bd23 100644 --- a/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml +++ b/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml @@ -88,8 +88,16 @@ jobs: # This is needed to run pipelines that use the default environment at HEAD, for example, when a # pipeline uses an expansion service built from HEAD. run: | - BEAM_VERSION_LINE=$(cat gradle.properties | grep "sdk_version") - echo "BEAM_VERSION=${BEAM_VERSION_LINE#*sdk_version=}" >> $GITHUB_ENV + BEAM_VERSION=$(grep -E '^sdk_version=' gradle.properties | cut -d= -f2-) + if [[ ! "${BEAM_VERSION}" =~ ^[0-9A-Za-z._-]+$ ]]; then + echo "Invalid sdk_version in gradle.properties: ${BEAM_VERSION}" + exit 1 + fi + { + echo "BEAM_VERSION<> "$GITHUB_ENV" - name: Set latest tag only on master branch if: github.ref == 'refs/heads/master' run: echo "LATEST_TAG=,latest" >> $GITHUB_ENV @@ -158,8 +166,16 @@ jobs: github_job: ${{ matrix.job_name }} (${{ matrix.container_task }}) - name: Find Beam Version run: | - BEAM_VERSION_LINE=$(cat gradle.properties | grep "sdk_version") - echo "BEAM_VERSION=${BEAM_VERSION_LINE#*sdk_version=}" >> $GITHUB_ENV + BEAM_VERSION=$(grep -E '^sdk_version=' gradle.properties | cut -d= -f2-) + if [[ ! "${BEAM_VERSION}" =~ ^[0-9A-Za-z._-]+$ ]]; then + echo "Invalid sdk_version in gradle.properties: ${BEAM_VERSION}" + exit 1 + fi + { + echo "BEAM_VERSION<> "$GITHUB_ENV" - name: Set latest tag only on master branch if: github.ref == 'refs/heads/master' run: echo "LATEST_TAG=,latest" >> $GITHUB_ENV