diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index cafc9e78e8b4..43340a5c8fc9 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -328,7 +328,7 @@ private UnaryCallable createPointReadCallable( ServerStreamingCallable readRowCallable = new TransformingServerStreamingCallable<>( readRowsCallable, - (query) -> query.limit(1).toProto(requestContext), + (query) -> query.toProto(requestContext).toBuilder().setRowsLimit(1).build(), Functions.identity()); BigtableUnaryOperationCallable classic = diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index f20ad1a3e5fd..977b6e1a9fa5 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -521,7 +521,9 @@ public void testCreateReadRowsCallable() throws InterruptedException { Query request = Query.create(TableId.of("table-id")).rowKey("row-key"); streamingCallable.call(request).iterator().next(); ReadRowsRequest expected = - request.toProto(RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID)); + request.toProto(RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID)).toBuilder() + .setRowsLimit(1) + .build(); assertThat(fakeDataService.popLastRequest()).isEqualTo(expected); } @@ -981,6 +983,19 @@ public void testExecuteQueryWaitTimeoutWorksWithMetadataFuture() throws IOExcept } } + @Test + public void testReadRowDoesNotMutateQuery() throws Exception { + Query query = Query.create(TABLE_ID).rowKey("row-key"); + enhancedBigtableStub.readRowCallable().call(query); + assertThat(fakeDataService.popLastRequest().getRowsLimit()).isEqualTo(1L); + assertThat(query).isEqualTo(Query.create(TABLE_ID).rowKey("row-key")); + + // Reusing the query for a multi-row scan should not inherit rows_limit = 1 + query.rowKey("row-key-2"); + enhancedBigtableStub.readRowsCallable().all().call(query); + assertThat(fakeDataService.popLastRequest().getRowsLimit()).isEqualTo(0L); + } + private static class MetadataInterceptor implements ServerInterceptor { final BlockingQueue headers = Queues.newLinkedBlockingDeque();