Search in sources :

Example 61 with ReadRowsRequest

use of com.google.bigtable.repackaged.com.google.bigtable.v2.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);
}
Also used : ReadRowsRequest(com.google.cloud.bigquery.storage.v1.ReadRowsRequest) StatusCode(com.google.api.gax.rpc.StatusCode) ApiException(com.google.api.gax.rpc.ApiException) Test(org.junit.Test)

Example 62 with ReadRowsRequest

use of com.google.bigtable.repackaged.com.google.bigtable.v2.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());
    }
}
Also used : InvalidArgumentException(com.google.api.gax.rpc.InvalidArgumentException) ReadRowsResponse(com.google.bigtable.v2.ReadRowsResponse) StatusRuntimeException(io.grpc.StatusRuntimeException) ReadRowsRequest(com.google.bigtable.v2.ReadRowsRequest) MockStreamObserver(com.google.api.gax.grpc.testing.MockStreamObserver) ExecutionException(java.util.concurrent.ExecutionException) Test(org.junit.Test)

Example 63 with ReadRowsRequest

use of com.google.bigtable.repackaged.com.google.bigtable.v2.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));
}
Also used : ReadRowsResponse(com.google.bigtable.v2.ReadRowsResponse) ReadRowsRequest(com.google.bigtable.v2.ReadRowsRequest) MockStreamObserver(com.google.api.gax.grpc.testing.MockStreamObserver) Test(org.junit.Test)

Example 64 with ReadRowsRequest

use of com.google.bigtable.repackaged.com.google.bigtable.v2.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.v1.stub.readrows.ReadRowsRetryingCallable) ReadRowsRequest(com.google.cloud.bigquery.storage.v1.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.v1.ReadRowsResponse) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap)

Example 65 with ReadRowsRequest

use of com.google.bigtable.repackaged.com.google.bigtable.v2.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);
}
Also used : ReadRowsResponse(com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse) ReadSession(com.google.cloud.bigquery.storage.v1beta2.ReadSession) 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)10 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