Search in sources :

Example 1 with RetryHttpRequestInitializer

use of org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer in project beam by apache.

the class BigQueryServicesImpl method newBigQueryClient.

/**
 * Returns a BigQuery client builder using the specified {@link BigQueryOptions}.
 */
private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
    // Do not log 404. It clutters the output and is possibly even required by the
    // caller.
    RetryHttpRequestInitializer httpRequestInitializer = new RetryHttpRequestInitializer(ImmutableList.of(404));
    httpRequestInitializer.setCustomErrors(createBigQueryClientCustomErrors());
    httpRequestInitializer.setWriteTimeout(options.getHTTPWriteTimeout());
    ImmutableList.Builder<HttpRequestInitializer> initBuilder = ImmutableList.builder();
    Credentials credential = options.getGcpCredential();
    initBuilder.add(credential == null ? new NullCredentialInitializer() : new HttpCredentialsAdapter(credential));
    initBuilder.add(new LatencyRecordingHttpRequestInitializer(API_METRIC_LABEL));
    initBuilder.add(httpRequestInitializer);
    HttpRequestInitializer chainInitializer = new ChainingHttpRequestInitializer(Iterables.toArray(initBuilder.build(), HttpRequestInitializer.class));
    return new Bigquery.Builder(Transport.getTransport(), Transport.getJsonFactory(), chainInitializer).setApplicationName(options.getAppName()).setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
Also used : RetryHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer) HttpCredentialsAdapter(com.google.auth.http.HttpCredentialsAdapter) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) LatencyRecordingHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.LatencyRecordingHttpRequestInitializer) ChainingHttpRequestInitializer(com.google.cloud.hadoop.util.ChainingHttpRequestInitializer) RetryHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer) LatencyRecordingHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.LatencyRecordingHttpRequestInitializer) HttpRequestInitializer(com.google.api.client.http.HttpRequestInitializer) NullCredentialInitializer(org.apache.beam.sdk.extensions.gcp.auth.NullCredentialInitializer) Credentials(com.google.auth.Credentials) ChainingHttpRequestInitializer(com.google.cloud.hadoop.util.ChainingHttpRequestInitializer)

Example 2 with RetryHttpRequestInitializer

use of org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer in project beam by apache.

the class V1TestUtil method getDatastore.

/**
 * Build a new datastore client.
 */
static Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) {
    Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential();
    HttpRequestInitializer initializer;
    if (credential != null) {
        initializer = new ChainingHttpRequestInitializer(new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer());
    } else {
        initializer = new RetryHttpRequestInitializer();
    }
    DatastoreOptions.Builder builder = new DatastoreOptions.Builder().projectId(projectId).initializer(initializer);
    return DatastoreFactory.get().create(builder.build());
}
Also used : RetryHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer) GcpOptions(org.apache.beam.sdk.extensions.gcp.options.GcpOptions) HttpCredentialsAdapter(com.google.auth.http.HttpCredentialsAdapter) DatastoreOptions(com.google.datastore.v1.client.DatastoreOptions) HttpRequestInitializer(com.google.api.client.http.HttpRequestInitializer) RetryHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer) ChainingHttpRequestInitializer(com.google.cloud.hadoop.util.ChainingHttpRequestInitializer) Credentials(com.google.auth.Credentials) ChainingHttpRequestInitializer(com.google.cloud.hadoop.util.ChainingHttpRequestInitializer)

Example 3 with RetryHttpRequestInitializer

use of org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer in project beam by apache.

the class BigQueryServicesImplTest method setUp.

@Before
public void setUp() {
    MockitoAnnotations.initMocks(this);
    // Set up the MockHttpRequest for future inspection.
    request = new MockLowLevelHttpRequest() {

        int index = 0;

        @Override
        public LowLevelHttpResponse execute() throws IOException {
            Verify.verify(index < responses.length, "The number of HttpRequest invocation exceeded the number of prepared mock requests. Index: %d", index);
            return responses[index++];
        }
    };
    // A mock transport that lets us mock the API responses.
    MockHttpTransport transport = new MockHttpTransport.Builder().setLowLevelHttpRequest(request).build();
    // A sample BigQuery API client that uses default JsonFactory and RetryHttpInitializer.
    bigquery = new Bigquery.Builder(transport, Transport.getJsonFactory(), new RetryHttpRequestInitializer()).build();
    // Setup the ProcessWideContainer for testing metrics are set.
    MetricsContainerImpl container = new MetricsContainerImpl(null);
    MetricsEnvironment.setProcessWideContainer(container);
    MetricsEnvironment.setCurrentContainer(container);
}
Also used : RetryHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer) MetricsContainerImpl(org.apache.beam.runners.core.metrics.MetricsContainerImpl) MockHttpTransport(com.google.api.client.testing.http.MockHttpTransport) LowLevelHttpResponse(com.google.api.client.http.LowLevelHttpResponse) IOException(java.io.IOException) MockLowLevelHttpRequest(com.google.api.client.testing.http.MockLowLevelHttpRequest) Before(org.junit.Before)

Aggregations

RetryHttpRequestInitializer (org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer)3 HttpRequestInitializer (com.google.api.client.http.HttpRequestInitializer)2 Credentials (com.google.auth.Credentials)2 HttpCredentialsAdapter (com.google.auth.http.HttpCredentialsAdapter)2 ChainingHttpRequestInitializer (com.google.cloud.hadoop.util.ChainingHttpRequestInitializer)2 LowLevelHttpResponse (com.google.api.client.http.LowLevelHttpResponse)1 MockHttpTransport (com.google.api.client.testing.http.MockHttpTransport)1 MockLowLevelHttpRequest (com.google.api.client.testing.http.MockLowLevelHttpRequest)1 DatastoreOptions (com.google.datastore.v1.client.DatastoreOptions)1 IOException (java.io.IOException)1 MetricsContainerImpl (org.apache.beam.runners.core.metrics.MetricsContainerImpl)1 NullCredentialInitializer (org.apache.beam.sdk.extensions.gcp.auth.NullCredentialInitializer)1 GcpOptions (org.apache.beam.sdk.extensions.gcp.options.GcpOptions)1 LatencyRecordingHttpRequestInitializer (org.apache.beam.sdk.extensions.gcp.util.LatencyRecordingHttpRequestInitializer)1 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)1 Before (org.junit.Before)1