Examples with FixedHeaderProvider com.google.api.gax.rpc.FixedHeaderProvider used on opensource projects

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);
    }
}