Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable FirestoreV1IT to run against a non default host #27256

Merged
merged 8 commits into from
Jul 26, 2023
Merged

Enable FirestoreV1IT to run against a non default host #27256

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c2a1ef1
Added FIRESTORE_HOST env variable to allow endpoint (i.e. host + port…
SabaSathya Jun 26, 2023
c4258cc
Fix format.
SabaSathya Jun 27, 2023
f45ecf6
Set host only if options don't have it set already.
SabaSathya Jun 29, 2023
a1f5c6c
Used FIRESTORE_EMULATOR_HOST env var for emulator host.
SabaSathya Jul 7, 2023
045e2b3
Updated to use host as argument in FirestoreV1IT
SabaSathya Jul 20, 2023
0a70e8b
Don't inject value back into FirestoreOptions
SabaSathya Jul 20, 2023
f59b99f
Updates for review comments
SabaSathya Jul 26, 2023
346aea6
Removed unused import
SabaSathya Jul 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions ...e-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.gcp.firestore;

import javax.annotation.Nonnull;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
Expand Down Expand Up @@ -57,4 +58,21 @@ public interface FirestoreOptions extends PipelineOptions {

/** Set the Firestore database ID to connect to. */
void setFirestoreDb(String firestoreDb);

/**
* A host port pair to allow connecting to a Cloud Firestore instead of the default live service.
*
* @return the string representation of a host and port pair to be used when constructing Cloud
* Firestore clients.
*/
@Nonnull
SabaSathya marked this conversation as resolved.
Show resolved Hide resolved
String getHost();

/**
* Define a host port pair to allow connecting to a Cloud Firestore instead of the default live
* service.
*
* @param host the host and port to connect to
*/
void setHost(String host);
}
14 changes: 13 additions & 1 deletion ...src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreStatefulComponentFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,10 @@
@Immutable
class FirestoreStatefulComponentFactory implements Serializable {

private static final String DEFAULT_FIRESTORE_HOST = "batch-firestore.googleapis.com:443";
private static final String FIRESTORE_HOST_ENV_VARIABLE = "FIRESTORE_HOST";
SabaSathya marked this conversation as resolved.
Show resolved Hide resolved
private static final String FIRESTORE_EMULATOR_HOST_ENV_VARIABLE = "FIRESTORE_EMULATOR_HOST";

static final FirestoreStatefulComponentFactory INSTANCE = new FirestoreStatefulComponentFactory();

private FirestoreStatefulComponentFactory() {}
Expand Down Expand Up @@ -86,7 +90,11 @@ FirestoreStub getFirestoreStub(PipelineOptions options) {

FirestoreOptions firestoreOptions = options.as(FirestoreOptions.class);
String emulatorHostPort = firestoreOptions.getEmulatorHost();
if (emulatorHostPort == null) {
SabaSathya marked this conversation as resolved.
Show resolved Hide resolved
emulatorHostPort = System.getenv(FIRESTORE_EMULATOR_HOST_ENV_VARIABLE);
}
if (emulatorHostPort != null) {
firestoreOptions.setEmulatorHost(emulatorHostPort);
SabaSathya marked this conversation as resolved.
Show resolved Hide resolved
builder
.setCredentialsProvider(FixedCredentialsProvider.create(new EmulatorCredentials()))
.setEndpoint(emulatorHostPort)
Expand All @@ -97,9 +105,13 @@ FirestoreStub getFirestoreStub(PipelineOptions options) {
.build());
} else {
GcpOptions gcpOptions = options.as(GcpOptions.class);
String host = firestoreOptions.getHost();
if (host == null) {
host = System.getenv().getOrDefault(FIRESTORE_HOST_ENV_VARIABLE, DEFAULT_FIRESTORE_HOST);
SabaSathya marked this conversation as resolved.
Show resolved Hide resolved
}
builder
.setCredentialsProvider(FixedCredentialsProvider.create(gcpOptions.getGcpCredential()))
.setEndpoint("batch-firestore.googleapis.com:443");
.setEndpoint(host);
}

ClientContext clientContext = ClientContext.create(builder.build());
Expand Down
45 changes: 38 additions & 7 deletions ...latform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/it/FirestoreTestingHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,13 @@
import com.google.api.core.ApiFunction;
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.cloud.firestore.CollectionReference;
import com.google.cloud.firestore.DocumentReference;
import com.google.cloud.firestore.Firestore;
import com.google.cloud.firestore.FirestoreOptions;
import com.google.cloud.firestore.FirestoreOptions.EmulatorCredentials;
import com.google.cloud.firestore.WriteBatch;
import com.google.cloud.firestore.WriteResult;
import com.google.cloud.firestore.spi.v1.FirestoreRpc;
Expand Down Expand Up @@ -77,6 +80,9 @@

final class FirestoreTestingHelper implements TestRule {
private static final Logger LOG = LoggerFactory.getLogger(FirestoreTestingHelper.class);
private static final String DEFAULT_FIRESTORE_HOST = "batch-firestore.googleapis.com:443";
private static final String FIRESTORE_HOST_ENV_VARIABLE = "FIRESTORE_HOST";
private static final String FIRESTORE_EMULATOR_HOST_ENV_VARIABLE = "FIRESTORE_EMULATOR_HOST";

static final String BASE_COLLECTION_ID;

Expand Down Expand Up @@ -129,13 +135,38 @@ public FirestoreTestingHelper(CleanupMode cleanupMode) {
firestoreBeamOptions =
TestPipeline.testingPipelineOptions()
.as(org.apache.beam.sdk.io.gcp.firestore.FirestoreOptions.class);
firestoreOptions =
FirestoreOptions.newBuilder()
.setCredentials(options.getGcpCredential())
.setProjectId(options.getProject())
.setDatabaseId(firestoreBeamOptions.getFirestoreDb())
.build();

String emulatorHostPort = firestoreBeamOptions.getEmulatorHost();
if (emulatorHostPort == null) {
emulatorHostPort = System.getenv(FIRESTORE_EMULATOR_HOST_ENV_VARIABLE);
}
if (emulatorHostPort != null) {
firestoreBeamOptions.setEmulatorHost(emulatorHostPort);
firestoreOptions =
FirestoreOptions.newBuilder()
.setCredentialsProvider(FixedCredentialsProvider.create(new EmulatorCredentials()))
.setProjectId(options.getProject())
.setDatabaseId(firestoreBeamOptions.getFirestoreDb())
.setHost(emulatorHostPort)
.setChannelProvider(
InstantiatingGrpcChannelProvider.newBuilder()
.setEndpoint(emulatorHostPort)
.setChannelConfigurator(c -> c.usePlaintext())
.build())
.build();
} else {
String host = firestoreBeamOptions.getHost();
if (host == null) {
host = System.getenv().getOrDefault(FIRESTORE_HOST_ENV_VARIABLE, DEFAULT_FIRESTORE_HOST);
firestoreBeamOptions.setHost(host);
}
firestoreOptions =
FirestoreOptions.newBuilder()
.setCredentials(options.getGcpCredential())
.setProjectId(options.getProject())
.setDatabaseId(firestoreBeamOptions.getFirestoreDb())
.setHost(firestoreBeamOptions.getHost())
.build();
}
fs = firestoreOptions.getService();
rpc = (FirestoreRpc) firestoreOptions.getRpc();
}
Expand Down