Search in sources :

Example 16 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse 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.v1beta2.ReadRowsResponse) ReadRowsRequest(com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest) ApiCallContext(com.google.api.gax.rpc.ApiCallContext)

Example 17 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse 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 18 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse 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 19 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse 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 20 with ReadRowsResponse

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

the class StorageArrowSample method main.

public static void main(String... args) throws Exception {
    // Sets your Google Cloud Platform project ID.
    // String projectId = "YOUR_PROJECT_ID";
    String projectId = args[0];
    Integer snapshotMillis = null;
    if (args.length > 1) {
        snapshotMillis = Integer.parseInt(args[1]);
    }
    try (BigQueryReadClient client = BigQueryReadClient.create()) {
        String parent = String.format("projects/%s", projectId);
        // This example uses baby name data from the public datasets.
        String srcTable = String.format("projects/%s/datasets/%s/tables/%s", "bigquery-public-data", "usa_names", "usa_1910_current");
        // We specify the columns to be projected by adding them to the selected fields,
        // and set a simple filter to restrict which rows are transmitted.
        TableReadOptions options = TableReadOptions.newBuilder().addSelectedFields("name").addSelectedFields("number").addSelectedFields("state").setRowRestriction("state = \"WA\"").build();
        // Start specifying the read session we want created.
        ReadSession.Builder sessionBuilder = ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW).setReadOptions(options);
        // Optionally specify the snapshot time.  When unspecified, snapshot time is "now".
        if (snapshotMillis != null) {
            Timestamp t = Timestamp.newBuilder().setSeconds(snapshotMillis / 1000).setNanos((int) ((snapshotMillis % 1000) * 1000000)).build();
            TableModifiers modifiers = TableModifiers.newBuilder().setSnapshotTime(t).build();
            sessionBuilder.setTableModifiers(modifiers);
        }
        // Begin building the session creation request.
        CreateReadSessionRequest.Builder builder = CreateReadSessionRequest.newBuilder().setParent(parent).setReadSession(sessionBuilder).setMaxStreamCount(1);
        ReadSession session = client.createReadSession(builder.build());
        // Setup a simple reader and start a read session.
        try (SimpleRowReader reader = new SimpleRowReader(session.getArrowSchema())) {
            // Assert that there are streams available in the session.  An empty table may not have
            // data available.  If no sessions are available for an anonymous (cached) table, consider
            // writing results of a query to a named table rather than consuming cached results
            // directly.
            Preconditions.checkState(session.getStreamsCount() > 0);
            // Use the first stream to perform reading.
            String streamName = session.getStreams(0).getName();
            ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder().setReadStream(streamName).build();
            // Process each block of rows as they arrive and decode using our simple row reader.
            ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
            for (ReadRowsResponse response : stream) {
                Preconditions.checkState(response.hasArrowRecordBatch());
                reader.processRows(response.getArrowRecordBatch());
            }
        }
    }
}
Also used : ReadSession(com.google.cloud.bigquery.storage.v1.ReadSession) TableModifiers(com.google.cloud.bigquery.storage.v1.ReadSession.TableModifiers) ReadRowsRequest(com.google.cloud.bigquery.storage.v1.ReadRowsRequest) Timestamp(com.google.protobuf.Timestamp) BigQueryReadClient(com.google.cloud.bigquery.storage.v1.BigQueryReadClient) ReadRowsResponse(com.google.cloud.bigquery.storage.v1.ReadRowsResponse) TableReadOptions(com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions) CreateReadSessionRequest(com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest)

Aggregations

Test (org.junit.Test)44 ReadRowsResponse (com.google.cloud.bigquery.storage.v1.ReadRowsResponse)40 ReadRowsRequest (com.google.cloud.bigquery.storage.v1.ReadRowsRequest)25 ReadRowsResponse (com.google.bigtable.v2.ReadRowsResponse)19 ReadSession (com.google.cloud.bigquery.storage.v1.ReadSession)17 StorageClient (org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StorageClient)17 FakeBigQueryServices (org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices)17 ReadRowsRequest (com.google.bigtable.v2.ReadRowsRequest)15 TableRow (com.google.api.services.bigquery.model.TableRow)14 TableRowParser (org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TableRowParser)13 ReadRowsResponse (com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse)11 ReadRowsRequest (com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest)10 ByteString (com.google.protobuf.ByteString)10 CreateReadSessionRequest (com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest)8 GenericRecord (org.apache.avro.generic.GenericRecord)7 StatusRuntimeException (io.grpc.StatusRuntimeException)6 ReadSession (com.google.cloud.bigquery.storage.v1beta2.ReadSession)5 StreamObserver (io.grpc.stub.StreamObserver)5 Mockito.doAnswer (org.mockito.Mockito.doAnswer)5 InvocationOnMock (org.mockito.invocation.InvocationOnMock)5