Examples with ReadRowsRequest com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest used on opensource projects

Example 61 with ReadRowsRequest

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest in project beam by apache.

the class BigQueryServicesImplTest method testReadRowsSetsRequestCountMetricOnError.

@Test
public void testReadRowsSetsRequestCountMetricOnError() throws InterruptedException, IOException {
    BigQueryServices.StorageClient client = mock(BigQueryServicesImpl.StorageClientImpl.class);
    ReadRowsRequest request = null;
    StatusCode statusCode = new StatusCode() {

        @Override
        public Code getCode() {
            return Code.INTERNAL;
        }

        @Override
        public Object getTransportCode() {
            return null;
        }
    };
    when(client.readRows(request)).thenThrow(// Mock implementation.
    new ApiException("Internal", null, statusCode, false));
    // Real implementation.
    when(client.readRows(any(), any())).thenCallRealMethod();
    thrown.expect(ApiException.class);
    thrown.expectMessage("Internal");
    client.readRows(request, "myproject:mydataset.mytable");
    verifyReadMetricWasSet("myproject", "mydataset", "mytable", "internal", 1);
}

Example 62 with ReadRowsRequest

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

the class BaseBigtableDataClientTest method readRowsExceptionTest.

@Test
public void readRowsExceptionTest() throws Exception {
    StatusRuntimeException exception = new StatusRuntimeException(io.grpc.Status.INVALID_ARGUMENT);
    mockBigtable.addException(exception);
    ReadRowsRequest request = ReadRowsRequest.newBuilder().setTableName(TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]").toString()).setAppProfileId("appProfileId704923523").setRows(RowSet.newBuilder().build()).setFilter(RowFilter.newBuilder().build()).setRowsLimit(-944199211).build();
    MockStreamObserver<ReadRowsResponse> responseObserver = new MockStreamObserver<>();
    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> callable = client.readRowsCallable();
    callable.serverStreamingCall(request, responseObserver);
    try {
        List<ReadRowsResponse> actualResponses = responseObserver.future().get();
        Assert.fail("No exception thrown");
    } catch (ExecutionException e) {
        Assert.assertTrue(e.getCause() instanceof InvalidArgumentException);
        InvalidArgumentException apiException = ((InvalidArgumentException) e.getCause());
        Assert.assertEquals(StatusCode.Code.INVALID_ARGUMENT, apiException.getStatusCode().getCode());
    }
}

Example 63 with ReadRowsRequest

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

the class BaseBigtableDataClientTest method readRowsTest.

@Test
public void readRowsTest() throws Exception {
    ReadRowsResponse expectedResponse = ReadRowsResponse.newBuilder().addAllChunks(new ArrayList<ReadRowsResponse.CellChunk>()).setLastScannedRowKey(ByteString.EMPTY).build();
    mockBigtable.addResponse(expectedResponse);
    ReadRowsRequest request = ReadRowsRequest.newBuilder().setTableName(TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]").toString()).setAppProfileId("appProfileId704923523").setRows(RowSet.newBuilder().build()).setFilter(RowFilter.newBuilder().build()).setRowsLimit(-944199211).build();
    MockStreamObserver<ReadRowsResponse> responseObserver = new MockStreamObserver<>();
    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> callable = client.readRowsCallable();
    callable.serverStreamingCall(request, responseObserver);
    List<ReadRowsResponse> actualResponses = responseObserver.future().get();
    Assert.assertEquals(1, actualResponses.size());
    Assert.assertEquals(expectedResponse, actualResponses.get(0));
}

Example 64 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());
}

Example 65 with ReadRowsRequest

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

the class ITBigQueryStorageTest method testSimpleReadAndResume.

@Test
public void testSimpleReadAndResume() {
    String table = BigQueryResource.FormatTableResource(/* projectId = */
    "bigquery-public-data", /* datasetId = */
    "samples", /* tableId = */
    "shakespeare");
    ReadSession session = client.createReadSession(/* parent = */
    parentProjectId, /* readSession = */
    ReadSession.newBuilder().setTable(table).setDataFormat(DataFormat.AVRO).build(), /* maxStreamCount = */
    1);
    assertEquals(String.format("Did not receive expected number of streams for table '%s' CreateReadSession response:%n%s", table, session.toString()), 1, session.getStreamsCount());
    // We have to read some number of rows in order to be able to resume. More details:
    long rowCount = ReadStreamToOffset(session.getStreams(0), /* rowOffset = */
    34_846);
    ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).setOffset(rowCount).build();
    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
    for (ReadRowsResponse response : stream) {
        rowCount += response.getRowCount();
    }
    // Verifies that the number of rows skipped and read equals to the total number of rows in the
    // table.
    assertEquals(164_656, rowCount);
}