Allow Firestore project to be configurable (#31808)

apache · Jul 10, 2024 · 4df89c7 · 4df89c7
1 parent 6c829db
commit 4df89c7
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 5 deletions.
diff --git a/...e-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java b/...e-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/...-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java b/...-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/...cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java b/...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/...latform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/it/FirestoreTestingHelper.java b/...latform/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();