From 7b56f2cea3203c3abbfd25dd64b9b08198c67c9d Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Wed, 27 Jul 2022 11:55:11 -0400 Subject: [PATCH 1/4] fix: add a ReadFirstRow callable to set future in onComplete --- .../data/v2/stub/EnhancedBigtableStub.java | 13 ++- .../stub/readrows/ReadFirstRowCallable.java | 80 +++++++++++++++++++ .../readrows/ReadFirstRowCallableTest.java | 74 +++++++++++++++++ 3 files changed, 164 insertions(+), 3 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 301ecd66b5df..5a48631aa2fc 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -85,6 +85,7 @@ import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor; import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable; +import com.google.cloud.bigtable.data.v2.stub.readrows.ReadFirstRowCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable; @@ -366,10 +367,16 @@ public UnaryCallable createReadRowCallable(RowAdapter .build(), rowAdapter); - UnaryCallable readRowCallable = - new ReadRowsUserCallable<>(readRowsCallable, requestContext).first(); + ServerStreamingCallable readRowCallable = + new ReadRowsUserCallable<>(readRowsCallable, requestContext); + + ServerStreamingCallable traced = + new TracedServerStreamingCallable<>( + readRowCallable, clientContext.getTracerFactory(), getSpanName("ReadRow")); + + ReadFirstRowCallable firstRow = new ReadFirstRowCallable<>(traced); - return createUserFacingUnaryCallable("ReadRow", readRowCallable); + return firstRow.withDefaultCallContext(clientContext.getDefaultCallContext()); } /** diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java new file mode 100644 index 000000000000..c9c242a52fe0 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java @@ -0,0 +1,80 @@ +/* + * Copyright 2022 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.data.v2.stub.readrows; + +import com.google.api.core.ApiFuture; +import com.google.api.core.InternalApi; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.api.gax.rpc.StateCheckingResponseObserver; +import com.google.api.gax.rpc.StreamController; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.models.Query; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Enhancement for `readRowsCallable().first()` to gracefully limit the row count and only set the + * future after the call is completed so the trailers can be populated correctly. + */ +@InternalApi("For internal use only") +public class ReadFirstRowCallable extends UnaryCallable { + + private final ServerStreamingCallable inner; + + public ReadFirstRowCallable(ServerStreamingCallable inner) { + this.inner = inner; + } + + @Override + public ApiFuture futureCall(Query query, ApiCallContext context) { + ReadFirstRowResponseObserver observer = new ReadFirstRowResponseObserver<>(); + this.inner.call(query.limit(1), observer, context); + return observer.getFuture(); + } + + private class ReadFirstRowResponseObserver extends StateCheckingResponseObserver { + private StreamController innerController; + private AtomicReference newItem = new AtomicReference<>(); + private SettableApiFuture settableFuture = SettableApiFuture.create(); + + @Override + protected void onStartImpl(StreamController streamController) { + this.innerController = streamController; + innerController.disableAutoInboundFlowControl(); + innerController.request(1); + } + + @Override + protected void onResponseImpl(RowT response) { + newItem.compareAndSet(null, response); + } + + @Override + protected void onErrorImpl(Throwable throwable) { + settableFuture.setException(throwable); + } + + @Override + protected void onCompleteImpl() { + settableFuture.set(newItem.get()); + } + + protected ApiFuture getFuture() { + return settableFuture; + } + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java new file mode 100644 index 000000000000..78bcb5d49bd0 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java @@ -0,0 +1,74 @@ +/* + * Copyright 2022 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.data.v2.stub.readrows; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.verify; + +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.cloud.bigtable.data.v2.models.Row; +import com.google.common.truth.Truth; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; +import org.mockito.quality.Strictness; + +@RunWith(JUnit4.class) +public class ReadFirstRowCallableTest { + @Rule public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.WARN); + + private static final RequestContext REQUEST_CONTEXT = + RequestContext.create("fake-project", "fake-instance", "fake-profile"); + private ServerStreamingCallable innerCallable; + private ArgumentCaptor innerQuery; + + @Before + public void setUp() { + innerCallable = Mockito.mock(ServerStreamingCallable.class); + innerQuery = ArgumentCaptor.forClass(Query.class); + } + + @Test + public void testLimitAdded() { + ReadFirstRowCallable callable = new ReadFirstRowCallable(innerCallable); + callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault()); + verify(innerCallable) + .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); + Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) + .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); + } + + @Test + public void testLimitChanged() { + ReadFirstRowCallable callable = new ReadFirstRowCallable<>(innerCallable); + callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault()); + verify(innerCallable) + .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); + Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) + .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); + } +} From 3d2eb9345ca2d3ced0b6039c2959f0a8eeea7d37 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 1 Aug 2022 14:49:59 -0400 Subject: [PATCH 2/4] use ReadRowsFirst callable instead --- .../data/v2/stub/EnhancedBigtableStub.java | 6 +- .../stub/readrows/ReadFirstRowCallable.java | 80 ------------------- .../stub/readrows/ReadRowsFirstCallable.java | 48 ++++++++++- .../stub/readrows/ReadRowsUserCallable.java | 12 --- .../readrows/ReadFirstRowCallableTest.java | 74 ----------------- .../readrows/ReadRowsFirstCallableTest.java | 30 +++---- .../readrows/ReadRowsUserCallableTest.java | 13 --- 7 files changed, 63 insertions(+), 200 deletions(-) delete mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java delete mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 5a48631aa2fc..c6591d588a49 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -85,8 +85,8 @@ import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor; import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable; -import com.google.cloud.bigtable.data.v2.stub.readrows.ReadFirstRowCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor; +import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsFirstCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable; @@ -367,14 +367,14 @@ public UnaryCallable createReadRowCallable(RowAdapter .build(), rowAdapter); - ServerStreamingCallable readRowCallable = + ReadRowsUserCallable readRowCallable = new ReadRowsUserCallable<>(readRowsCallable, requestContext); ServerStreamingCallable traced = new TracedServerStreamingCallable<>( readRowCallable, clientContext.getTracerFactory(), getSpanName("ReadRow")); - ReadFirstRowCallable firstRow = new ReadFirstRowCallable<>(traced); + ReadRowsFirstCallable firstRow = new ReadRowsFirstCallable<>(traced); return firstRow.withDefaultCallContext(clientContext.getDefaultCallContext()); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java deleted file mode 100644 index c9c242a52fe0..000000000000 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallable.java +++ /dev/null @@ -1,80 +0,0 @@ -/* - * Copyright 2022 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.data.v2.stub.readrows; - -import com.google.api.core.ApiFuture; -import com.google.api.core.InternalApi; -import com.google.api.core.SettableApiFuture; -import com.google.api.gax.rpc.ApiCallContext; -import com.google.api.gax.rpc.ServerStreamingCallable; -import com.google.api.gax.rpc.StateCheckingResponseObserver; -import com.google.api.gax.rpc.StreamController; -import com.google.api.gax.rpc.UnaryCallable; -import com.google.cloud.bigtable.data.v2.models.Query; -import java.util.concurrent.atomic.AtomicReference; - -/** - * Enhancement for `readRowsCallable().first()` to gracefully limit the row count and only set the - * future after the call is completed so the trailers can be populated correctly. - */ -@InternalApi("For internal use only") -public class ReadFirstRowCallable extends UnaryCallable { - - private final ServerStreamingCallable inner; - - public ReadFirstRowCallable(ServerStreamingCallable inner) { - this.inner = inner; - } - - @Override - public ApiFuture futureCall(Query query, ApiCallContext context) { - ReadFirstRowResponseObserver observer = new ReadFirstRowResponseObserver<>(); - this.inner.call(query.limit(1), observer, context); - return observer.getFuture(); - } - - private class ReadFirstRowResponseObserver extends StateCheckingResponseObserver { - private StreamController innerController; - private AtomicReference newItem = new AtomicReference<>(); - private SettableApiFuture settableFuture = SettableApiFuture.create(); - - @Override - protected void onStartImpl(StreamController streamController) { - this.innerController = streamController; - innerController.disableAutoInboundFlowControl(); - innerController.request(1); - } - - @Override - protected void onResponseImpl(RowT response) { - newItem.compareAndSet(null, response); - } - - @Override - protected void onErrorImpl(Throwable throwable) { - settableFuture.setException(throwable); - } - - @Override - protected void onCompleteImpl() { - settableFuture.set(newItem.get()); - } - - protected ApiFuture getFuture() { - return settableFuture; - } - } -} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java index c2584e0c9366..6a0c641bca68 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java @@ -16,23 +16,63 @@ package com.google.cloud.bigtable.data.v2.stub.readrows; import com.google.api.core.ApiFuture; +import com.google.api.core.InternalApi; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.api.gax.rpc.StateCheckingResponseObserver; +import com.google.api.gax.rpc.StreamController; import com.google.api.gax.rpc.UnaryCallable; import com.google.cloud.bigtable.data.v2.models.Query; +import java.util.concurrent.atomic.AtomicReference; /** * Enhancement for `readRowsCallable().first()` to gracefully limit the row count instead of * cancelling the RPC */ -class ReadRowsFirstCallable extends UnaryCallable { - private final UnaryCallable inner; +@InternalApi +public class ReadRowsFirstCallable extends UnaryCallable { - ReadRowsFirstCallable(UnaryCallable inner) { + private final ServerStreamingCallable inner; + + public ReadRowsFirstCallable(ServerStreamingCallable inner) { this.inner = inner; } @Override public ApiFuture futureCall(Query query, ApiCallContext context) { - return inner.futureCall(query.limit(1), context); + ReadRowsFirstResponseObserver observer = new ReadRowsFirstResponseObserver<>(); + this.inner.call(query.limit(1), observer, context); + return observer.getFuture(); + } + + private class ReadRowsFirstResponseObserver extends StateCheckingResponseObserver { + private StreamController innerController; + private AtomicReference newItem = new AtomicReference<>(); + private SettableApiFuture settableFuture = SettableApiFuture.create(); + + @Override + protected void onStartImpl(StreamController streamController) { + this.innerController = streamController; + } + + @Override + protected void onResponseImpl(RowT response) { + newItem.compareAndSet(null, response); + } + + @Override + protected void onErrorImpl(Throwable throwable) { + settableFuture.setException(throwable); + } + + @Override + protected void onCompleteImpl() { + settableFuture.set(newItem.get()); + } + + protected ApiFuture getFuture() { + return settableFuture; + } } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallable.java index 94ced791c50a..3f1db6d0d84e 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallable.java @@ -19,7 +19,6 @@ import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; -import com.google.api.gax.rpc.UnaryCallable; import com.google.bigtable.v2.ReadRowsRequest; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.Query; @@ -34,14 +33,11 @@ public class ReadRowsUserCallable extends ServerStreamingCallable { private final ServerStreamingCallable inner; private final RequestContext requestContext; - private final ReadRowsFirstCallable firstCallable; public ReadRowsUserCallable( ServerStreamingCallable inner, RequestContext requestContext) { this.inner = inner; this.requestContext = requestContext; - - this.firstCallable = new ReadRowsFirstCallable<>(super.first()); } @Override @@ -49,12 +45,4 @@ public void call(Query request, ResponseObserver responseObserver, ApiCall ReadRowsRequest innerRequest = request.toProto(requestContext); inner.call(innerRequest, responseObserver, context); } - - // Optimization: since the server supports row limits, override the first callable. - // This way unnecessary data doesn't need to be buffered and the number of CANCELLED request - // statuses is minimized - @Override - public UnaryCallable first() { - return firstCallable; - } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java deleted file mode 100644 index 78bcb5d49bd0..000000000000 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadFirstRowCallableTest.java +++ /dev/null @@ -1,74 +0,0 @@ -/* - * Copyright 2022 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.data.v2.stub.readrows; - -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.verify; - -import com.google.api.gax.grpc.GrpcCallContext; -import com.google.api.gax.rpc.ApiCallContext; -import com.google.api.gax.rpc.ResponseObserver; -import com.google.api.gax.rpc.ServerStreamingCallable; -import com.google.cloud.bigtable.data.v2.internal.RequestContext; -import com.google.cloud.bigtable.data.v2.models.Query; -import com.google.cloud.bigtable.data.v2.models.Row; -import com.google.common.truth.Truth; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.ArgumentCaptor; -import org.mockito.Mockito; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; -import org.mockito.quality.Strictness; - -@RunWith(JUnit4.class) -public class ReadFirstRowCallableTest { - @Rule public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.WARN); - - private static final RequestContext REQUEST_CONTEXT = - RequestContext.create("fake-project", "fake-instance", "fake-profile"); - private ServerStreamingCallable innerCallable; - private ArgumentCaptor innerQuery; - - @Before - public void setUp() { - innerCallable = Mockito.mock(ServerStreamingCallable.class); - innerQuery = ArgumentCaptor.forClass(Query.class); - } - - @Test - public void testLimitAdded() { - ReadFirstRowCallable callable = new ReadFirstRowCallable(innerCallable); - callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault()); - verify(innerCallable) - .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); - Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) - .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); - } - - @Test - public void testLimitChanged() { - ReadFirstRowCallable callable = new ReadFirstRowCallable<>(innerCallable); - callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault()); - verify(innerCallable) - .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); - Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) - .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); - } -} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java index edc20e6081e6..07cf3478c1f8 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java @@ -15,8 +15,13 @@ */ package com.google.cloud.bigtable.data.v2.stub.readrows; -import com.google.api.core.SettableApiFuture; -import com.google.api.gax.rpc.UnaryCallable; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.verify; + +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.Row; @@ -38,26 +43,23 @@ public class ReadRowsFirstCallableTest { private static final RequestContext REQUEST_CONTEXT = RequestContext.create("fake-project", "fake-instance", "fake-profile"); - private UnaryCallable innerCallable; + + private ServerStreamingCallable innerCallable; private ArgumentCaptor innerQuery; - private SettableApiFuture innerResult; @SuppressWarnings("unchecked") @Before public void setUp() { - innerCallable = Mockito.mock(UnaryCallable.class); + innerCallable = Mockito.mock(ServerStreamingCallable.class); innerQuery = ArgumentCaptor.forClass(Query.class); - innerResult = SettableApiFuture.create(); - Mockito.when(innerCallable.futureCall(innerQuery.capture(), Mockito.any())) - .thenReturn(innerResult); } @Test public void testLimitAdded() { ReadRowsFirstCallable callable = new ReadRowsFirstCallable<>(innerCallable); - innerResult.set(null); - callable.call(Query.create("fake-table")); - + callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault()); + verify(innerCallable) + .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); } @@ -65,9 +67,9 @@ public void testLimitAdded() { @Test public void testLimitChanged() { ReadRowsFirstCallable callable = new ReadRowsFirstCallable<>(innerCallable); - innerResult.set(null); - callable.call(Query.create("fake-table").limit(1_000)); - + callable.futureCall(Query.create("fake-table").limit(10), GrpcCallContext.createDefault()); + verify(innerCallable) + .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallableTest.java index d1e03df2c660..b518a554153b 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallableTest.java @@ -41,17 +41,4 @@ public void testRequestConverted() { Truth.assertThat(innerCallable.getActualRequest()).isEqualTo(query.toProto(REQUEST_CONTEXT)); } - - @Test - public void testFirstIsLimited() { - ServerStreamingStashCallable innerCallable = - new ServerStreamingStashCallable<>(); - ReadRowsUserCallable callable = new ReadRowsUserCallable<>(innerCallable, REQUEST_CONTEXT); - Query query = Query.create("fake-table"); - - callable.first().call(query); - - Truth.assertThat(innerCallable.getActualRequest()) - .isEqualTo(query.limit(1).toProto(REQUEST_CONTEXT)); - } } From 1b0ca9ec055743e1ab7e4a5c648a1a61bc908a86 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 2 Aug 2022 10:19:18 -0400 Subject: [PATCH 3/4] don't use atomic --- .../data/v2/stub/readrows/ReadRowsFirstCallable.java | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java index 6a0c641bca68..2ef26605b478 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java @@ -24,7 +24,6 @@ import com.google.api.gax.rpc.StreamController; import com.google.api.gax.rpc.UnaryCallable; import com.google.cloud.bigtable.data.v2.models.Query; -import java.util.concurrent.atomic.AtomicReference; /** * Enhancement for `readRowsCallable().first()` to gracefully limit the row count instead of @@ -48,7 +47,7 @@ public ApiFuture futureCall(Query query, ApiCallContext context) { private class ReadRowsFirstResponseObserver extends StateCheckingResponseObserver { private StreamController innerController; - private AtomicReference newItem = new AtomicReference<>(); + private RowT firstRow; private SettableApiFuture settableFuture = SettableApiFuture.create(); @Override @@ -58,7 +57,9 @@ protected void onStartImpl(StreamController streamController) { @Override protected void onResponseImpl(RowT response) { - newItem.compareAndSet(null, response); + if (firstRow == null) { + this.firstRow = response; + } } @Override @@ -68,7 +69,7 @@ protected void onErrorImpl(Throwable throwable) { @Override protected void onCompleteImpl() { - settableFuture.set(newItem.get()); + settableFuture.set(firstRow); } protected ApiFuture getFuture() { From d3f7108bd860587bfb6e224e497d7c2bb327d42b Mon Sep 17 00:00:00 2001 From: Owl Bot Date: Tue, 2 Aug 2022 14:24:39 +0000 Subject: [PATCH 4/4] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20po?= =?UTF-8?q?st-processor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 6348e79d25f4..8680353195ad 100644 --- a/README.md +++ b/README.md @@ -56,13 +56,13 @@ implementation 'com.google.cloud:google-cloud-bigtable' If you are using Gradle without BOM, add this to your dependencies: ```Groovy -implementation 'com.google.cloud:google-cloud-bigtable:2.10.0' +implementation 'com.google.cloud:google-cloud-bigtable:2.10.1' ``` If you are using SBT, add this to your dependencies: ```Scala -libraryDependencies += "com.google.cloud" % "google-cloud-bigtable" % "2.10.0" +libraryDependencies += "com.google.cloud" % "google-cloud-bigtable" % "2.10.1" ``` ## Authentication