Search in sources :

Example 16 with ReadRowsRequest

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest in project java-bigquerystorage by googleapis.

the class EnhancedBigQueryReadStub method readRowsCallable.

public ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> readRowsCallable() {
    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> innerCallable = GrpcRawCallableFactory.createServerStreamingCallable(GrpcCallSettings.<ReadRowsRequest, ReadRowsResponse>newBuilder().setMethodDescriptor(BigQueryReadGrpc.getReadRowsMethod()).setParamsExtractor(new RequestParamsExtractor<ReadRowsRequest>() {

        @Override
        public Map<String, String> extract(ReadRowsRequest request) {
            return ImmutableMap.of("read_stream", String.valueOf(request.getReadStream()));
        }
    }).build(), stubSettings.readRowsSettings().getRetryableCodes());
    ServerStreamingCallSettings<ReadRowsRequest, ReadRowsResponse> callSettings = stubSettings.readRowsSettings();
    StreamingRetryAlgorithm<Void> retryAlgorithm = new StreamingRetryAlgorithm<>(new ApiResultRetryAlgorithm<Void>(readRowsRetryAttemptListener), new ExponentialRetryAlgorithm(callSettings.getRetrySettings(), context.getClock()));
    ScheduledRetryingExecutor<Void> retryingExecutor = new ScheduledRetryingExecutor<>(retryAlgorithm, context.getExecutor());
    if (context.getStreamWatchdog() != null) {
        innerCallable = Callables.watched(innerCallable, callSettings, context);
    }
    ReadRowsRetryingCallable outerCallable = new ReadRowsRetryingCallable(context.getDefaultCallContext(), innerCallable, retryingExecutor, callSettings.getResumptionStrategy());
    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> traced = new TracedServerStreamingCallable<>(outerCallable, context.getTracerFactory(), SpanName.of(TRACING_OUTER_CLIENT_NAME, "ReadRows"));
    return traced.withDefaultCallContext(context.getDefaultCallContext());
}
Also used : TracedServerStreamingCallable(com.google.api.gax.tracing.TracedServerStreamingCallable) ReadRowsRetryingCallable(com.google.cloud.bigquery.storage.v1beta2.stub.readrows.ReadRowsRetryingCallable) ReadRowsRequest(com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest) ExponentialRetryAlgorithm(com.google.api.gax.retrying.ExponentialRetryAlgorithm) ScheduledRetryingExecutor(com.google.api.gax.retrying.ScheduledRetryingExecutor) StreamingRetryAlgorithm(com.google.api.gax.retrying.StreamingRetryAlgorithm) ReadRowsResponse(com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap)

Example 17 with ReadRowsRequest

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest in project java-bigquerystorage by googleapis.

the class ReadRowsAttemptCallable method call.

/**
 * Sends the actual RPC. The request being sent will first be transformed by the {@link
 * StreamResumptionStrategy}.
 *
 * <p>This method expects to be called by one thread at a time. Furthermore, it expects that the
 * current RPC finished before the next time it's called.
 */
@Override
public Void call() {
    Preconditions.checkState(isStarted, "Must be started first");
    ReadRowsRequest request = (++numAttempts == 1) ? initialRequest : resumptionStrategy.getResumeRequest(initialRequest);
    // Should never happen. onAttemptError will check if ResumptionStrategy can create a resume
    // request,
    // which the RetryingFuture/StreamResumptionStrategy should respect.
    Preconditions.checkState(request != null, "ResumptionStrategy returned a null request.");
    innerAttemptFuture = SettableApiFuture.create();
    seenSuccessSinceLastError = false;
    ApiCallContext attemptContext = context;
    if (!outerRetryingFuture.getAttemptSettings().getRpcTimeout().isZero()) {
        attemptContext = attemptContext.withStreamWaitTimeout(outerRetryingFuture.getAttemptSettings().getRpcTimeout());
    }
    attemptContext.getTracer().attemptStarted(outerRetryingFuture.getAttemptSettings().getOverallAttemptCount());
    innerCallable.call(request, new StateCheckingResponseObserver<ReadRowsResponse>() {

        @Override
        public void onStartImpl(StreamController controller) {
            onAttemptStart(controller);
        }

        @Override
        public void onResponseImpl(ReadRowsResponse response) {
            onAttemptResponse(response);
        }

        @Override
        public void onErrorImpl(Throwable t) {
            onAttemptError(t);
        }

        @Override
        public void onCompleteImpl() {
            onAttemptComplete();
        }
    }, attemptContext);
    outerRetryingFuture.setAttemptFuture(innerAttemptFuture);
    return null;
}
Also used : StreamController(com.google.api.gax.rpc.StreamController) ReadRowsResponse(com.google.cloud.bigquery.storage.v1.ReadRowsResponse) ReadRowsRequest(com.google.cloud.bigquery.storage.v1.ReadRowsRequest) ApiCallContext(com.google.api.gax.rpc.ApiCallContext)

Example 18 with ReadRowsRequest

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest in project java-bigquerystorage by googleapis.

the class EnhancedBigQueryReadStubSettingsTest method testReadRowsSettings.

@Test
public void testReadRowsSettings() {
    ServerStreamingCallSettings.Builder<ReadRowsRequest, ReadRowsResponse> builder = EnhancedBigQueryReadStubSettings.newBuilder().readRowsSettings();
    assertThat(builder.getRetryableCodes()).contains(Code.UNAVAILABLE);
    RetrySettings retrySettings = builder.getRetrySettings();
    assertThat(retrySettings.getInitialRetryDelay()).isEqualTo(Duration.ofMillis(100L));
    assertThat(retrySettings.getRetryDelayMultiplier()).isWithin(1e-6).of(1.3);
    assertThat(retrySettings.getMaxRetryDelay()).isEqualTo(Duration.ofMinutes(1L));
    assertThat(retrySettings.getInitialRpcTimeout()).isEqualTo(Duration.ofDays(1L));
    assertThat(retrySettings.getRpcTimeoutMultiplier()).isWithin(1e-6).of(1.0);
    assertThat(retrySettings.getMaxRpcTimeout()).isEqualTo(Duration.ofDays(1L));
    assertThat(retrySettings.getTotalTimeout()).isEqualTo(Duration.ofDays(1L));
    assertThat(builder.getIdleTimeout()).isEqualTo(Duration.ZERO);
}
Also used : ServerStreamingCallSettings(com.google.api.gax.rpc.ServerStreamingCallSettings) RetrySettings(com.google.api.gax.retrying.RetrySettings) ReadRowsResponse(com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse) ReadRowsRequest(com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest) Test(org.junit.Test)

Example 19 with ReadRowsRequest

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest in project java-bigquerystorage by googleapis.

the class ResourceHeaderTest method readRowsTest.

@Test
public void readRowsTest() {
    try {
        ReadRowsRequest request = ReadRowsRequest.newBuilder().setReadStream(TEST_STREAM_NAME).setOffset(125).build();
        client.readRowsCallable().call(request);
    } catch (UnimplementedException e) {
    // Ignore the error: none of the methods are actually implemented.
    }
    verifyHeaderSent(READ_STREAM_PATTERN);
}
Also used : ReadRowsRequest(com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest) UnimplementedException(com.google.api.gax.rpc.UnimplementedException) Test(org.junit.Test)

Example 20 with ReadRowsRequest

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest in project java-bigquerystorage by googleapis.

the class ReadRowsRetryTest method errorAtTheVeryEndTest.

@Test
public void errorAtTheVeryEndTest() {
    ReadRowsRequest request = RpcExpectation.createRequest("fake-stream", 0);
    service.expectations.add(RpcExpectation.create().expectRequest("fake-stream", 0).respondWithNumberOfRows(10).respondWithNumberOfRows(7).respondWithStatus(Code.UNAVAILABLE));
    service.expectations.add(RpcExpectation.create().expectRequest("fake-stream", 17).respondWithNumberOfRows(0));
    Assert.assertEquals(17, getRowCount(request));
}
Also used : ReadRowsRequest(com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)59 ReadRowsRequest (com.google.cloud.bigquery.storage.v1.ReadRowsRequest)31 ReadRowsRequest (com.google.bigtable.v2.ReadRowsRequest)28 ReadRowsResponse (com.google.cloud.bigquery.storage.v1.ReadRowsResponse)24 ReadSession (com.google.cloud.bigquery.storage.v1.ReadSession)17 ReadRowsRequest (com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest)16 ReadRowsResponse (com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse)11 StorageClient (org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StorageClient)10 FakeBigQueryServices (org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices)10 ReadRowsResponse (com.google.bigtable.v2.ReadRowsResponse)9 Row (com.google.cloud.bigtable.data.v2.models.Row)9 CreateReadSessionRequest (com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest)8 TableRow (com.google.api.services.bigquery.model.TableRow)7 Query (com.google.cloud.bigtable.data.v2.models.Query)7 ServerStreamingStashCallable (com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable)7 ByteString (com.google.protobuf.ByteString)7 GenericRecord (org.apache.avro.generic.GenericRecord)6 TableRowParser (org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TableRowParser)6 ReadSession (com.google.cloud.bigquery.storage.v1beta2.ReadSession)5 DefaultRowAdapter (com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter)5