This is an automated email from the ASF dual-hosted git repository.
yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 4df89c704b0 Allow Firestore project to be configurable (#31808)
4df89c704b0 is described below
commit 4df89c704b066b9708d10ff36b8bca6af645c9ab
Author: Vitaly Terentyev <[email protected]>
AuthorDate: Wed Jul 10 19:05:32 2024 +0400
Allow Firestore project to be configurable (#31808)
---
.../apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java | 11 +++++++++++
.../apache/beam/sdk/io/gcp/firestore/FirestoreV1ReadFn.java | 9 +++++++--
.../apache/beam/sdk/io/gcp/firestore/FirestoreV1WriteFn.java | 9 +++++++--
.../beam/sdk/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 a292a106e51..5adc9ef38f3 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 b4a334b75c9..51e5efa380e 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 @@ final class FirestoreV1ReadFn {
/** {@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 3e9e1890c1e..09378d4f80c 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 @@ final class FirestoreV1WriteFn {
@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 a57dd688d4a..d3a82bf24ce 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 @@ final class FirestoreTestingHelper implements TestRule {
firestoreOptions =
FirestoreOptions.newBuilder()
.setCredentials(gcpOptions.getGcpCredential())
- .setProjectId(gcpOptions.getProject())
+ .setProjectId(
+ firstNonNull(firestoreBeamOptions.getFirestoreProject(),
gcpOptions.getProject()))
.setDatabaseId(firestoreBeamOptions.getFirestoreDb())
.setHost(firestoreBeamOptions.getFirestoreHost())
.build();
