Search in sources :

Example 1 with FixedHeaderProvider

use of com.google.api.gax.rpc.FixedHeaderProvider in project beam by apache.

the class FirestoreStatefulComponentFactory method getFirestoreStub.

/**
 * Given a {@link PipelineOptions}, return a pre-configured {@link FirestoreStub} with values set
 * based on those options.
 *
 * <p>The provided {@link PipelineOptions} is expected to provide {@link FirestoreOptions} and
 * {@link org.apache.beam.sdk.extensions.gcp.options.GcpOptions GcpOptions} for access to {@link
 * GcpOptions#getProject()}
 *
 * <p>The instance returned by this method is expected to bind to the lifecycle of a bundle.
 *
 * @param options The instance of options to read from
 * @return a new {@link FirestoreStub} pre-configured with values from the provided options
 */
FirestoreStub getFirestoreStub(PipelineOptions options) {
    try {
        FirestoreSettings.Builder builder = FirestoreSettings.newBuilder().setHeaderProvider(new FixedHeaderProvider() {

            @Override
            public Map<@NonNull String, @NonNull String> getHeaders() {
                return ImmutableMap.of("User-Agent", options.getUserAgent());
            }
        });
        RetrySettings retrySettings = RetrySettings.newBuilder().setMaxAttempts(1).build();
        builder.applyToAllUnaryMethods(b -> {
            b.setRetrySettings(retrySettings);
            return null;
        });
        FirestoreOptions firestoreOptions = options.as(FirestoreOptions.class);
        String emulatorHostPort = firestoreOptions.getEmulatorHost();
        if (emulatorHostPort != null) {
            builder.setCredentialsProvider(FixedCredentialsProvider.create(new EmulatorCredentials())).setEndpoint(emulatorHostPort).setTransportChannelProvider(InstantiatingGrpcChannelProvider.newBuilder().setEndpoint(emulatorHostPort).setChannelConfigurator(c -> c.usePlaintext()).build());
        } else {
            GcpOptions gcpOptions = options.as(GcpOptions.class);
            builder.setCredentialsProvider(FixedCredentialsProvider.create(gcpOptions.getGcpCredential())).setEndpoint("batch-firestore.googleapis.com:443");
        }
        ClientContext clientContext = ClientContext.create(builder.build());
        return GrpcFirestoreStub.create(clientContext);
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}
Also used : EmulatorCredentials(com.google.cloud.firestore.FirestoreOptions.EmulatorCredentials) FixedHeaderProvider(com.google.api.gax.rpc.FixedHeaderProvider) ClientContext(com.google.api.gax.rpc.ClientContext) RetrySettings(com.google.api.gax.retrying.RetrySettings) GcpOptions(org.apache.beam.sdk.extensions.gcp.options.GcpOptions) NonNull(org.checkerframework.checker.nullness.qual.NonNull) FirestoreSettings(com.google.cloud.firestore.v1.FirestoreSettings) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map)

Aggregations

RetrySettings (com.google.api.gax.retrying.RetrySettings)1 ClientContext (com.google.api.gax.rpc.ClientContext)1 FixedHeaderProvider (com.google.api.gax.rpc.FixedHeaderProvider)1 EmulatorCredentials (com.google.cloud.firestore.FirestoreOptions.EmulatorCredentials)1 FirestoreSettings (com.google.cloud.firestore.v1.FirestoreSettings)1 Map (java.util.Map)1 GcpOptions (org.apache.beam.sdk.extensions.gcp.options.GcpOptions)1 ImmutableMap (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap)1 NonNull (org.checkerframework.checker.nullness.qual.NonNull)1