diff --git a/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java index d2147e6167da..c4792944489e 100644 --- a/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java +++ b/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java @@ -43,6 +43,10 @@ import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.RowCell; import com.google.cloud.bigtable.data.v2.models.RowMutation; +import com.google.cloud.bigtable.data.v2.models.TypedCell; +import com.google.cloud.bigtable.data.v2.models.TypedQualifier; +import com.google.cloud.bigtable.data.v2.models.TypedQuery; +import com.google.cloud.bigtable.data.v2.models.TypedRow; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; @@ -51,6 +55,7 @@ import com.google.common.base.Preconditions; import com.google.protobuf.ByteString; import com.google.protobuf.util.Durations; +import com.google.protobuf.util.Timestamps; import com.google.rpc.Code; import io.grpc.ManagedChannelBuilder; import io.grpc.Status; @@ -76,6 +81,7 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; +import javax.annotation.Nullable; /** Java implementation of the CBT test proxy. Used to test the Java CBT client. */ public class CbtTestProxy extends CloudBigtableV2TestProxyImplBase implements Closeable { @@ -90,11 +96,17 @@ static CbtClient create(BigtableDataSettings settings, BigtableDataClient dataCl return new AutoValue_CbtTestProxy_CbtClient(settings, dataClient); } + @Nullable abstract BigtableDataSettings settings(); + @Nullable abstract BigtableDataClient dataClient(); } + void registerClientForTest(String clientId, CbtClient client) { + idClientMap.put(clientId, client); + } + private static final Logger logger = Logger.getLogger(CbtTestProxy.class.getName()); private CbtTestProxy() { @@ -186,8 +198,9 @@ public synchronized void createClient( .setInstanceId(request.getInstanceId()) .setAppProfileId(request.getAppProfileId()); + Duration newTimeout = null; if (request.hasPerOperationTimeout()) { - Duration newTimeout = Duration.ofMillis(Durations.toMillis(request.getPerOperationTimeout())); + newTimeout = Duration.ofMillis(Durations.toMillis(request.getPerOperationTimeout())); settingsBuilder = overrideTimeoutSetting(newTimeout, settingsBuilder); logger.info( String.format( @@ -229,6 +242,7 @@ public synchronized void createClient( } BigtableDataSettings settings = settingsBuilder.build(); BigtableDataClient client = BigtableDataClient.create(settings); + CbtClient cbtClient = CbtClient.create(settings, client); idClientMap.put(request.getClientId(), cbtClient); } catch (IOException e) { @@ -251,7 +265,9 @@ public void closeClient( return; } - client.dataClient().close(); + if (client.dataClient() != null) { + client.dataClient().close(); + } responseObserver.onNext(CloseClientResponse.getDefaultInstance()); responseObserver.onCompleted(); @@ -269,6 +285,10 @@ public void removeClient( return; } + if (client.dataClient() != null) { + client.dataClient().close(); + } + responseObserver.onNext(RemoveClientResponse.getDefaultInstance()); responseObserver.onCompleted(); } @@ -462,6 +482,64 @@ public void readRows(ReadRowsRequest request, StreamObserver respons responseObserver.onCompleted(); } + @Override + public void typedReadRows( + TypedReadRowsRequest request, StreamObserver responseObserver) { + CbtClient client; + try { + client = getClient(request.getClientId()); + } catch (StatusException e) { + responseObserver.onError(e); + return; + } + + TypedQuery query; + try { + query = TypedQuery.fromProto(request.getRequest()); + } catch (RuntimeException e) { + responseObserver.onNext( + TypedRowsResult.newBuilder() + .setStatus( + com.google.rpc.Status.newBuilder() + .setCode(Code.INVALID_ARGUMENT.getNumber()) + .setMessage(e.getMessage()) + .build()) + .build()); + responseObserver.onCompleted(); + return; + } + + TypedRowsResult.Builder resultBuilder = TypedRowsResult.newBuilder(); + try { + ServerStream rows = client.dataClient().typedReadRows(query); + readTypedRowsInto(rows, request.getCancelAfterRows(), resultBuilder); + responseObserver.onNext( + resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build()); + } catch (ApiException e) { + responseObserver.onNext(resultBuilder.setStatus(convertStatus(e)).build()); + responseObserver.onCompleted(); + return; + } catch (StatusRuntimeException e) { + responseObserver.onNext( + resultBuilder.setStatus(StatusProto.fromThrowable(e)).build()); + responseObserver.onCompleted(); + return; + } catch (RuntimeException e) { + responseObserver.onNext( + resultBuilder + .setStatus( + com.google.rpc.Status.newBuilder() + .setCode(Code.INTERNAL.getNumber()) + .setMessage(e.getMessage()) + .build()) + .build()); + responseObserver.onCompleted(); + return; + } + + responseObserver.onCompleted(); + } + /** * Helper method to convert row from type com.google.cloud.bigtable.data.v2.models.Row to type * com.google.bigtable.v2.Row. After conversion, row cells within the same column and family are @@ -530,6 +608,107 @@ private static RowsResult.Builder convertRowsResult( return resultBuilder; } + private static com.google.bigtable.v2.Value extractQualifierProtoValue(TypedQualifier qualifier) { + try { + java.lang.reflect.Field field = TypedQualifier.class.getDeclaredField("protoValue"); + field.setAccessible(true); + return (com.google.bigtable.v2.Value) field.get(qualifier); + } catch (Exception e) { + if (qualifier.isNull()) { + return Value.getDefaultInstance(); + } + return Value.newBuilder().setRawValue(qualifier.getBytes()).build(); + } + } + + private static com.google.bigtable.v2.TypedCell extractCellProto(TypedCell cell) { + try { + java.lang.reflect.Field field = cell.getClass().getDeclaredField("cellProto"); + field.setAccessible(true); + return (com.google.bigtable.v2.TypedCell) field.get(cell); + } catch (Exception e) { + return null; + } + } + + private static com.google.bigtable.v2.TypedRow convertTypedRow(TypedRow row) { + com.google.bigtable.v2.TypedRow.Builder rowBuilder = + com.google.bigtable.v2.TypedRow.newBuilder(); + + if (row.getRowKey().isRaw()) { + rowBuilder.setRowKey( + Value.newBuilder().setRawValue(row.getRowKey().getRaw()).build()); + } else { + rowBuilder.setRowKey(row.getRowKey().toProtoValue()); + } + + Map>> grouped = + row.getCells().stream() + .collect( + Collectors.groupingBy( + TypedCell::getFamily, + LinkedHashMap::new, + Collectors.groupingBy( + TypedCell::getTypedQualifier, + LinkedHashMap::new, + Collectors.toList()))); + + for (Map.Entry>> famEntry : grouped.entrySet()) { + com.google.bigtable.v2.TypedFamily.Builder familyBuilder = + rowBuilder.addFamiliesBuilder().setFamilyName(famEntry.getKey()); + + for (Map.Entry> colEntry : famEntry.getValue().entrySet()) { + com.google.bigtable.v2.TypedColumn.Builder colBuilder = + familyBuilder.addColumnsBuilder(); + + TypedQualifier qualifier = colEntry.getKey(); + colBuilder.setQualifier(extractQualifierProtoValue(qualifier)); + + for (TypedCell cell : colEntry.getValue()) { + com.google.bigtable.v2.TypedCell cellProto = extractCellProto(cell); + if (cellProto != null) { + colBuilder.addCells(cellProto); + } else { + com.google.bigtable.v2.TypedCell.Builder cellBuilder = colBuilder.addCellsBuilder(); + cellBuilder.setTimestamp(Timestamps.fromMicros(cell.getTimestamp())); + if (!cell.isNull()) { + cellBuilder.setValue(Value.newBuilder().setRawValue(cell.getBytesValue()).build()); + } + if (cell.getLabels() != null && !cell.getLabels().isEmpty()) { + cellBuilder.addAllLabels(cell.getLabels()); + } + } + } + } + } + + return rowBuilder.build(); + } + + /** + * Helper method to convert rows from type com.google.cloud.bigtable.data.v2.models.TypedRow to + * proto type com.google.bigtable.v2.TypedRow. + * + * @param rows Logical rows in ServerStream + * @param cancelAfterRows Ignore the results after this row if set positive + * @return the converted rows in TypedRowsResult Builder + */ + private static void readTypedRowsInto( + ServerStream rows, int cancelAfterRows, TypedRowsResult.Builder resultBuilder) { + int rowCounter = 0; + for (TypedRow row : rows) { + rowCounter++; + resultBuilder.addRows(convertTypedRow(row)); + + if (cancelAfterRows > 0 && rowCounter >= cancelAfterRows) { + logger.info( + String.format( + "Canceling TypedReadRows() to respect cancel_after_rows=%d", cancelAfterRows)); + break; + } + } + } + @Override public void sampleRowKeys( SampleRowKeysRequest request, StreamObserver responseObserver) { @@ -727,7 +906,9 @@ public synchronized void close() { Iterator> it = idClientMap.entrySet().iterator(); while (it.hasNext()) { Map.Entry entry = it.next(); - entry.getValue().dataClient().close(); + if (entry.getValue().dataClient() != null) { + entry.getValue().dataClient().close(); + } it.remove(); } } diff --git a/java-bigtable/test-proxy/src/main/proto/test_proxy.proto b/java-bigtable/test-proxy/src/main/proto/test_proxy.proto index 34cf534425c2..82ea16ab1590 100644 --- a/java-bigtable/test-proxy/src/main/proto/test_proxy.proto +++ b/java-bigtable/test-proxy/src/main/proto/test_proxy.proto @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -19,8 +19,8 @@ package google.bigtable.testproxy; import "google/api/client.proto"; import "google/bigtable/v2/bigtable.proto"; import "google/bigtable/v2/data.proto"; -import "google/protobuf/duration.proto"; import "google/protobuf/descriptor.proto"; +import "google/protobuf/duration.proto"; import "google/rpc/status.proto"; option go_package = "./testproxypb"; @@ -171,6 +171,28 @@ message RowsResult { repeated google.bigtable.v2.Row rows = 2; } +// Request to test proxy service to read rows using the Typed API. +message TypedReadRowsRequest { + // The ID of the target client object. + string client_id = 1; + + // The raw Typed request to the Bigtable server. + google.bigtable.v2.TypedReadRowsRequest request = 2; + + // The streaming read can be canceled before all items are seen. + // Has no effect if non-positive. + int32 cancel_after_rows = 3; +} + +// Response from test proxy service for TypedReadRowsRequest. +message TypedRowsResult { + // The RPC status from the client binding. + google.rpc.Status status = 1; + + // The successfully parsed structured rows. + repeated google.bigtable.v2.TypedRow rows = 2; +} + // Request to test proxy service to mutate a row. message MutateRowRequest { // The ID of the target client object. @@ -336,6 +358,9 @@ service CloudBigtableV2TestProxy { // Reads rows with the client instance. rpc ReadRows(ReadRowsRequest) returns (RowsResult) {} + // Reads rows with the client instance using the Typed API. + rpc TypedReadRows(TypedReadRowsRequest) returns (TypedRowsResult) {} + // Writes a row with the client instance. rpc MutateRow(MutateRowRequest) returns (MutateRowResult) {} diff --git a/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTypedReadRowsTest.java b/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTypedReadRowsTest.java new file mode 100644 index 000000000000..6c3eb1cef39c --- /dev/null +++ b/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTypedReadRowsTest.java @@ -0,0 +1,713 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.testproxy; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.bigtable.v2.BigtableGrpc; +import com.google.cloud.bigtable.data.v2.BigtableDataClient; +import com.google.cloud.bigtable.data.v2.BigtableDataSettings; +import com.google.bigtable.v2.PartialRowResponse; +import com.google.bigtable.v2.TypedCell; +import com.google.bigtable.v2.TypedColumn; +import com.google.bigtable.v2.TypedFamily; +import com.google.bigtable.v2.TypedReadRowsResponse; +import com.google.bigtable.v2.TypedRow; +import com.google.bigtable.v2.TypedRows; +import com.google.bigtable.v2.TypedRowsBatch; +import com.google.bigtable.v2.Value; +import com.google.common.hash.Hashing; +import com.google.protobuf.ByteString; +import com.google.protobuf.Timestamp; +import com.google.rpc.Code; +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.StreamObserver; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class CbtTestProxyTypedReadRowsTest { + + private static final String CLIENT_ID = "test-client"; + private static final String TABLE_NAME = "projects/p/instances/i/tables/t"; + + private Server inProcessServer; + private ManagedChannel inProcessChannel; + private CbtTestProxy testProxy; + private MockBigtableService mockBigtableService; + + @Before + public void setUp() throws IOException { + String serverName = InProcessServerBuilder.generateName(); + mockBigtableService = new MockBigtableService(); + inProcessServer = + InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(mockBigtableService) + .build() + .start(); + + inProcessChannel = + InProcessChannelBuilder.forName(serverName).directExecutor().build(); + + testProxy = CbtTestProxy.create(); + + BigtableDataSettings.Builder settingsBuilder = + BigtableDataSettings.newBuilderForEmulator("localhost", 8080) + .setProjectId("p") + .setInstanceId("i") + .setCredentialsProvider(NoCredentialsProvider.create()); + settingsBuilder + .stubSettings() + .setTransportChannelProvider( + FixedTransportChannelProvider.create(GrpcTransportChannel.create(inProcessChannel))); + BigtableDataSettings settings = settingsBuilder.build(); + BigtableDataClient dataClient = BigtableDataClient.create(settings); + CbtTestProxy.CbtClient client = CbtTestProxy.CbtClient.create(settings, dataClient); + testProxy.registerClientForTest(CLIENT_ID, client); + } + + @After + public void tearDown() throws InterruptedException { + testProxy.close(); + inProcessChannel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS); + inProcessServer.shutdownNow().awaitTermination(5, TimeUnit.SECONDS); + } + + private static TypedRow createRow(String rowKey, String family, String col, String val) { + return TypedRow.newBuilder() + .setRowKey(Value.newBuilder().setStringValue(rowKey).build()) + .addFamilies( + TypedFamily.newBuilder() + .setFamilyName(family) + .addColumns( + TypedColumn.newBuilder() + .setQualifier(Value.newBuilder().setStringValue(col).build()) + .addCells( + TypedCell.newBuilder() + .setValue(Value.newBuilder().setStringValue(val).build()) + .setTimestamp(Timestamp.newBuilder().setSeconds(12345).build()) + .build()) + .build()) + .build()) + .build(); + } + + @Test + public void testTypedReadRows_successfulStream() throws Exception { + TypedRow row1 = createRow("rk-1", "cf1", "cq1", "val1"); + TypedRow row2 = createRow("rk-2", "cf1", "cq2", "val2"); + TypedRows batch = TypedRows.newBuilder().addRows(row1).addRows(row2).build(); + ByteString batchBytes = batch.toByteString(); + int checksum = Hashing.crc32c().hashBytes(batchBytes.toByteArray()).asInt(); + + TypedReadRowsResponse response = + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder() + .setChecksum(checksum) + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .build()) + .build()) + .build(); + + mockBigtableService.responses.add(response); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.OK_VALUE); + assertThat(result.getRowsCount()).isEqualTo(2); + assertThat(result.getRows(0).getRowKey().getStringValue()).isEqualTo("rk-1"); + assertThat(result.getRows(1).getRowKey().getStringValue()).isEqualTo("rk-2"); + } + + @Test + public void testTypedReadRows_fragmentedBatchReassembly() throws Exception { + TypedRow row = createRow("rk-frag", "cf", "col", "value"); + TypedRows batch = TypedRows.newBuilder().addRows(row).build(); + ByteString batchBytes = batch.toByteString(); + int checksum = Hashing.crc32c().hashBytes(batchBytes.toByteArray()).asInt(); + + int mid = batchBytes.size() / 2; + ByteString part1 = batchBytes.substring(0, mid); + ByteString part2 = batchBytes.substring(mid); + + TypedReadRowsResponse resp1 = + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(part1).build()) + .build()) + .build(); + + TypedReadRowsResponse resp2 = + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(part2).build()) + .setFlush(PartialRowResponse.Flush.newBuilder().setChecksum(checksum).build()) + .build()) + .build(); + + mockBigtableService.responses.add(resp1); + mockBigtableService.responses.add(resp2); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.OK_VALUE); + assertThat(result.getRowsCount()).isEqualTo(1); + assertThat(result.getRows(0).getRowKey().getStringValue()).isEqualTo("rk-frag"); + } + + @Test + public void testTypedReadRows_resetDiscardsBufferedData() throws Exception { + TypedRow row = createRow("rk-valid", "cf", "col", "val"); + TypedRows batch = TypedRows.newBuilder().addRows(row).build(); + ByteString batchBytes = batch.toByteString(); + int checksum = Hashing.crc32c().hashBytes(batchBytes.toByteArray()).asInt(); + + // 1. Partial response with junk bytes + TypedReadRowsResponse junkResp = + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch( + TypedRowsBatch.newBuilder() + .setBatchData(ByteString.copyFromUtf8("invalid_bytes")) + .build()) + .build()) + .build(); + + // 2. Response with reset=true followed by valid batch and flush + TypedReadRowsResponse validResp = + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setReset(true) + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes).build()) + .setFlush(PartialRowResponse.Flush.newBuilder().setChecksum(checksum).build()) + .build()) + .build(); + + mockBigtableService.responses.add(junkResp); + mockBigtableService.responses.add(validResp); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.OK_VALUE); + assertThat(result.getRowsCount()).isEqualTo(1); + assertThat(result.getRows(0).getRowKey().getStringValue()).isEqualTo("rk-valid"); + } + + @Test + public void testTypedReadRows_checksumMismatchReturnsDataLoss() throws Exception { + ByteString data = ByteString.copyFromUtf8("some_data"); + int badChecksum = 999999; + + TypedReadRowsResponse response = + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(data).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder().setChecksum(badChecksum).build()) + .build()) + .build(); + + mockBigtableService.responses.add(response); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.UNAVAILABLE_VALUE); + assertThat(result.getStatus().getMessage()).contains("Checksum mismatch"); + } + + @Test + public void testTypedReadRows_cancelAfterRowsStopsStream() throws Exception { + for (int i = 1; i <= 5; i++) { + TypedRow row = createRow("rk-" + i, "cf", "col", "val"); + TypedRows batch = TypedRows.newBuilder().addRows(row).build(); + ByteString batchBytes = batch.toByteString(); + int checksum = Hashing.crc32c().hashBytes(batchBytes.toByteArray()).asInt(); + + mockBigtableService.responses.add( + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch( + TypedRowsBatch.newBuilder().setBatchData(batchBytes).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder().setChecksum(checksum).build()) + .build()) + .build()); + } + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setCancelAfterRows(2) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.OK_VALUE); + assertThat(result.getRowsCount()).isEqualTo(2); + assertThat(result.getRows(0).getRowKey().getStringValue()).isEqualTo("rk-1"); + assertThat(result.getRows(1).getRowKey().getStringValue()).isEqualTo("rk-2"); + } + + @Test + public void testTypedReadRows_serverErrorPropagatedInStatus() throws Exception { + mockBigtableService.errorToThrow = + Status.NOT_FOUND.withDescription("Table does not exist").asRuntimeException(); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.NOT_FOUND_VALUE); + assertThat(result.getStatus().getMessage()).contains("Table does not exist"); + } + + @Test + public void testTypedReadRows_missingTargetReturnsInvalidArgument() throws Exception { + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest(com.google.bigtable.v2.TypedReadRowsRequest.getDefaultInstance()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.INVALID_ARGUMENT_VALUE); + } + + @Test + public void testTypedReadRows_multiBatchRunningCrc32c() throws Exception { + TypedRow row1 = createRow("rk-1", "cf", "col", "val1"); + TypedRows batch1 = TypedRows.newBuilder().addRows(row1).build(); + ByteString batchBytes1 = batch1.toByteString(); + int checksum1 = (int) Hashing.crc32c().hashBytes(batchBytes1.toByteArray()).padToLong(); + + TypedRow row2 = createRow("rk-2", "cf", "col", "val2"); + TypedRows batch2 = TypedRows.newBuilder().addRows(row2).build(); + ByteString batchBytes2 = batch2.toByteString(); + ByteString allBytes = batchBytes1.concat(batchBytes2); + int checksum2 = (int) Hashing.crc32c().hashBytes(allBytes.toByteArray()).padToLong(); + + mockBigtableService.responses.add( + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes1).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder() + .setChecksum(checksum1) + .setResumeToken(ByteString.copyFromUtf8("tok-1")) + .build()) + .build()) + .build()); + + mockBigtableService.responses.add( + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes2).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder() + .setChecksum(checksum2) + .setResumeToken(ByteString.copyFromUtf8("tok-2")) + .build()) + .build()) + .build()); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.OK_VALUE); + assertThat(result.getRowsCount()).isEqualTo(2); + assertThat(result.getRows(0).getRowKey().getStringValue()).isEqualTo("rk-1"); + assertThat(result.getRows(1).getRowKey().getStringValue()).isEqualTo("rk-2"); + } + + @Test + public void testTypedReadRows_resetRollsBackToLastResumeToken() throws Exception { + // 1. Batch 1 committed with resume token "tok-1" + TypedRow row1 = createRow("rk-committed", "cf", "col", "val1"); + TypedRows batch1 = TypedRows.newBuilder().addRows(row1).build(); + ByteString batchBytes1 = batch1.toByteString(); + int checksum1 = (int) Hashing.crc32c().hashBytes(batchBytes1.toByteArray()).padToLong(); + + // 2. Batch 2 uncommitted with no resume token + TypedRow row2 = createRow("rk-uncommitted", "cf", "col", "val2"); + TypedRows batch2 = TypedRows.newBuilder().addRows(row2).build(); + ByteString batchBytes2 = batch2.toByteString(); + int checksum2 = + (int) Hashing.crc32c().hashBytes(batchBytes1.concat(batchBytes2).toByteArray()).padToLong(); + + // 3. Reset occurs, followed by Batch 3 with resume token "tok-3" + TypedRow row3 = createRow("rk-after-reset", "cf", "col", "val3"); + TypedRows batch3 = TypedRows.newBuilder().addRows(row3).build(); + ByteString batchBytes3 = batch3.toByteString(); + int checksum3 = + (int) Hashing.crc32c().hashBytes(batchBytes1.concat(batchBytes3).toByteArray()).padToLong(); + + mockBigtableService.responses.add( + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes1).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder() + .setChecksum(checksum1) + .setResumeToken(ByteString.copyFromUtf8("tok-1")) + .build()) + .build()) + .build()); + + mockBigtableService.responses.add( + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes2).build()) + .setFlush(PartialRowResponse.Flush.newBuilder().setChecksum(checksum2).build()) + .build()) + .build()); + + mockBigtableService.responses.add( + TypedReadRowsResponse.newBuilder() + .setResponse( + PartialRowResponse.newBuilder() + .setReset(true) + .setTypedRowsBatch(TypedRowsBatch.newBuilder().setBatchData(batchBytes3).build()) + .setFlush( + PartialRowResponse.Flush.newBuilder() + .setChecksum(checksum3) + .setResumeToken(ByteString.copyFromUtf8("tok-3")) + .build()) + .build()) + .build()); + + TypedReadRowsRequest request = + TypedReadRowsRequest.newBuilder() + .setClientId(CLIENT_ID) + .setRequest( + com.google.bigtable.v2.TypedReadRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .build()) + .build(); + + CountDownLatch latch = new CountDownLatch(1); + List results = new ArrayList<>(); + testProxy.typedReadRows( + request, + new StreamObserver() { + @Override + public void onNext(TypedRowsResult value) { + results.add(value); + } + + @Override + public void onError(Throwable t) { + latch.countDown(); + } + + @Override + public void onCompleted() { + latch.countDown(); + } + }); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(results).hasSize(1); + TypedRowsResult result = results.get(0); + assertThat(result.getStatus().getCode()).isEqualTo(Code.OK_VALUE); + // row2 must have been discarded on reset, leaving only row1 and row3 + assertThat(result.getRowsCount()).isEqualTo(2); + assertThat(result.getRows(0).getRowKey().getStringValue()).isEqualTo("rk-committed"); + assertThat(result.getRows(1).getRowKey().getStringValue()).isEqualTo("rk-after-reset"); + } + + private static class MockBigtableService extends BigtableGrpc.BigtableImplBase { + final List responses = new ArrayList<>(); + RuntimeException errorToThrow = null; + final AtomicBoolean wasCancelled = new AtomicBoolean(false); + + @Override + public void typedReadRows( + com.google.bigtable.v2.TypedReadRowsRequest request, + StreamObserver responseObserver) { + if (errorToThrow != null) { + responseObserver.onError(errorToThrow); + return; + } + for (TypedReadRowsResponse resp : responses) { + responseObserver.onNext(resp); + } + responseObserver.onCompleted(); + } + } +}