From 4df89c704b066b9708d10ff36b8bca6af645c9ab Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Wed, 10 Jul 2024 19:05:32 +0400 Subject: [PATCH] Allow Firestore project to be configurable (#31808) --- .../beam/sdk/io/gcp/firestore/FirestoreOptions.java | 11 +++++++++++ .../beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java | 9 +++++++-- .../beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java | 9 +++++++-- .../io/gcp/firestore/it/FirestoreTestingHelper.java | 5 ++++- 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java index a292a106e51f..5adc9ef38f36 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java @@ -17,6 +17,7 @@ */ package org.apache.beam.sdk.io.gcp.firestore; +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PipelineOptions; @@ -75,4 +76,14 @@ public interface FirestoreOptions extends PipelineOptions { * @param host the host and port to connect to */ void setFirestoreHost(String host); + + /** The Firestore project ID to connect to. */ + @Description("Firestore project ID") + @Nullable + String getFirestoreProject(); + + /** + * Set the Firestore project ID, it will override the value from {@link GcpOptions#getProject()}. + */ + void setFirestoreProject(String firestoreProject); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java index b4a334b75c99..51e5efa380e8 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java @@ -635,9 +635,14 @@ public void setup() { /** {@inheritDoc} */ @Override public final void startBundle(StartBundleContext c) { - String project = c.getPipelineOptions().as(GcpOptions.class).getProject(); + String project = c.getPipelineOptions().as(FirestoreOptions.class).getFirestoreProject(); + if (project == null) { + project = c.getPipelineOptions().as(GcpOptions.class).getProject(); + } projectId = - requireNonNull(project, "project must be defined on GcpOptions of PipelineOptions"); + requireNonNull( + project, + "project must be defined on FirestoreOptions or GcpOptions of PipelineOptions"); firestoreStub = firestoreStatefulComponentFactory.getFirestoreStub(c.getPipelineOptions()); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java index 3e9e1890c1e3..09378d4f80c5 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java @@ -202,11 +202,16 @@ public void setup() { @Override public final void startBundle(StartBundleContext c) { - String project = c.getPipelineOptions().as(GcpOptions.class).getProject(); + String project = c.getPipelineOptions().as(FirestoreOptions.class).getFirestoreProject(); + if (project == null) { + project = c.getPipelineOptions().as(GcpOptions.class).getProject(); + } String databaseId = c.getPipelineOptions().as(FirestoreOptions.class).getFirestoreDb(); databaseRootName = DatabaseRootName.of( - requireNonNull(project, "project must be defined on GcpOptions of PipelineOptions"), + requireNonNull( + project, + "project must be defined on FirestoreOptions or GcpOptions of PipelineOptions"), requireNonNull( databaseId, "firestoreDb must be defined on FirestoreOptions of PipelineOptions")); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/it/FirestoreTestingHelper.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/it/FirestoreTestingHelper.java index a57dd688d4af..d3a82bf24ced 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/it/FirestoreTestingHelper.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/it/FirestoreTestingHelper.java @@ -17,6 +17,8 @@ */ package org.apache.beam.sdk.io.gcp.firestore.it; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull; + import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -132,7 +134,8 @@ public FirestoreTestingHelper(CleanupMode cleanupMode) { firestoreOptions = FirestoreOptions.newBuilder() .setCredentials(gcpOptions.getGcpCredential()) - .setProjectId(gcpOptions.getProject()) + .setProjectId( + firstNonNull(firestoreBeamOptions.getFirestoreProject(), gcpOptions.getProject())) .setDatabaseId(firestoreBeamOptions.getFirestoreDb()) .setHost(firestoreBeamOptions.getFirestoreHost()) .build();