Visitar URL original
fix: add a ReadFirstRow callable to set future in onComplete (#1326) · googleapis/java-bigtable@cb539b5 · GitHub
Skip to content
This repository was archived by the owner on Aug 30, 2026. It is now read-only.

Commit cb539b5

Browse files
mutianfgcf-owl-bot[bot]
mutianf
and
gcf-owl-bot[bot]
authored
fix: add a ReadFirstRow callable to set future in onComplete (#1326)
* fix: add a ReadFirstRow callable to set future in onComplete * use ReadRowsFirst callable instead * don't use atomic * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 644aeb3 commit cb539b5

5 files changed

Lines changed: 71 additions & 46 deletions

File tree

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@
8686
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
8787
import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
8888
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
89+
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsFirstCallable;
8990
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy;
9091
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
9192
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
@@ -366,10 +367,16 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
366367
.build(),
367368
rowAdapter);
368369

369-
UnaryCallable<Query, RowT> readRowCallable =
370-
new ReadRowsUserCallable<>(readRowsCallable, requestContext).first();
370+
ReadRowsUserCallable<RowT> readRowCallable =
371+
new ReadRowsUserCallable<>(readRowsCallable, requestContext);
372+
373+
ServerStreamingCallable<Query, RowT> traced =
374+
new TracedServerStreamingCallable<>(
375+
readRowCallable, clientContext.getTracerFactory(), getSpanName("ReadRow"));
376+
377+
ReadRowsFirstCallable<RowT> firstRow = new ReadRowsFirstCallable<>(traced);
371378

372-
return createUserFacingUnaryCallable("ReadRow", readRowCallable);
379+
return firstRow.withDefaultCallContext(clientContext.getDefaultCallContext());
373380
}
374381

375382
/**

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java‎

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,23 +16,64 @@
1616
package com.google.cloud.bigtable.data.v2.stub.readrows;
1717

1818
import com.google.api.core.ApiFuture;
19+
import com.google.api.core.InternalApi;
20+
import com.google.api.core.SettableApiFuture;
1921
import com.google.api.gax.rpc.ApiCallContext;
22+
import com.google.api.gax.rpc.ServerStreamingCallable;
23+
import com.google.api.gax.rpc.StateCheckingResponseObserver;
24+
import com.google.api.gax.rpc.StreamController;
2025
import com.google.api.gax.rpc.UnaryCallable;
2126
import com.google.cloud.bigtable.data.v2.models.Query;
2227

2328
/**
2429
* Enhancement for `readRowsCallable().first()` to gracefully limit the row count instead of
2530
* cancelling the RPC
2631
*/
27-
class ReadRowsFirstCallable<RowT> extends UnaryCallable<Query, RowT> {
28-
private final UnaryCallable<Query, RowT> inner;
32+
@InternalApi
33+
public class ReadRowsFirstCallable<RowT> extends UnaryCallable<Query, RowT> {
2934

30-
ReadRowsFirstCallable(UnaryCallable<Query, RowT> inner) {
35+
private final ServerStreamingCallable<Query, RowT> inner;
36+
37+
public ReadRowsFirstCallable(ServerStreamingCallable<Query, RowT> inner) {
3138
this.inner = inner;
3239
}
3340

3441
@Override
3542
public ApiFuture<RowT> futureCall(Query query, ApiCallContext context) {
36-
return inner.futureCall(query.limit(1), context);
43+
ReadRowsFirstResponseObserver<RowT> observer = new ReadRowsFirstResponseObserver<>();
44+
this.inner.call(query.limit(1), observer, context);
45+
return observer.getFuture();
46+
}
47+
48+
private class ReadRowsFirstResponseObserver<RowT> extends StateCheckingResponseObserver<RowT> {
49+
private StreamController innerController;
50+
private RowT firstRow;
51+
private SettableApiFuture<RowT> settableFuture = SettableApiFuture.create();
52+
53+
@Override
54+
protected void onStartImpl(StreamController streamController) {
55+
this.innerController = streamController;
56+
}
57+
58+
@Override
59+
protected void onResponseImpl(RowT response) {
60+
if (firstRow == null) {
61+
this.firstRow = response;
62+
}
63+
}
64+
65+
@Override
66+
protected void onErrorImpl(Throwable throwable) {
67+
settableFuture.setException(throwable);
68+
}
69+
70+
@Override
71+
protected void onCompleteImpl() {
72+
settableFuture.set(firstRow);
73+
}
74+
75+
protected ApiFuture<RowT> getFuture() {
76+
return settableFuture;
77+
}
3778
}
3879
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallable.java‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import com.google.api.gax.rpc.ApiCallContext;
2020
import com.google.api.gax.rpc.ResponseObserver;
2121
import com.google.api.gax.rpc.ServerStreamingCallable;
22-
import com.google.api.gax.rpc.UnaryCallable;
2322
import com.google.bigtable.v2.ReadRowsRequest;
2423
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
2524
import com.google.cloud.bigtable.data.v2.models.Query;
@@ -34,27 +33,16 @@
3433
public class ReadRowsUserCallable<RowT> extends ServerStreamingCallable<Query, RowT> {
3534
private final ServerStreamingCallable<ReadRowsRequest, RowT> inner;
3635
private final RequestContext requestContext;
37-
private final ReadRowsFirstCallable<RowT> firstCallable;
3836

3937
public ReadRowsUserCallable(
4038
ServerStreamingCallable<ReadRowsRequest, RowT> inner, RequestContext requestContext) {
4139
this.inner = inner;
4240
this.requestContext = requestContext;
43-
44-
this.firstCallable = new ReadRowsFirstCallable<>(super.first());
4541
}
4642

4743
@Override
4844
public void call(Query request, ResponseObserver<RowT> responseObserver, ApiCallContext context) {
4945
ReadRowsRequest innerRequest = request.toProto(requestContext);
5046
inner.call(innerRequest, responseObserver, context);
5147
}
52-
53-
// Optimization: since the server supports row limits, override the first callable.
54-
// This way unnecessary data doesn't need to be buffered and the number of CANCELLED request
55-
// statuses is minimized
56-
@Override
57-
public UnaryCallable<Query, RowT> first() {
58-
return firstCallable;
59-
}
6048
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,13 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.stub.readrows;
1717

18-
import com.google.api.core.SettableApiFuture;
19-
import com.google.api.gax.rpc.UnaryCallable;
18+
import static org.mockito.ArgumentMatchers.any;
19+
import static org.mockito.Mockito.verify;
20+
21+
import com.google.api.gax.grpc.GrpcCallContext;
22+
import com.google.api.gax.rpc.ApiCallContext;
23+
import com.google.api.gax.rpc.ResponseObserver;
24+
import com.google.api.gax.rpc.ServerStreamingCallable;
2025
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
2126
import com.google.cloud.bigtable.data.v2.models.Query;
2227
import com.google.cloud.bigtable.data.v2.models.Row;
@@ -38,36 +43,33 @@ public class ReadRowsFirstCallableTest {
3843

3944
private static final RequestContext REQUEST_CONTEXT =
4045
RequestContext.create("fake-project", "fake-instance", "fake-profile");
41-
private UnaryCallable<Query, Row> innerCallable;
46+
47+
private ServerStreamingCallable<Query, Row> innerCallable;
4248
private ArgumentCaptor<Query> innerQuery;
43-
private SettableApiFuture<Row> innerResult;
4449

4550
@SuppressWarnings("unchecked")
4651
@Before
4752
public void setUp() {
48-
innerCallable = Mockito.mock(UnaryCallable.class);
53+
innerCallable = Mockito.mock(ServerStreamingCallable.class);
4954
innerQuery = ArgumentCaptor.forClass(Query.class);
50-
innerResult = SettableApiFuture.create();
51-
Mockito.when(innerCallable.futureCall(innerQuery.capture(), Mockito.any()))
52-
.thenReturn(innerResult);
5355
}
5456

5557
@Test
5658
public void testLimitAdded() {
5759
ReadRowsFirstCallable<Row> callable = new ReadRowsFirstCallable<>(innerCallable);
58-
innerResult.set(null);
59-
callable.call(Query.create("fake-table"));
60-
60+
callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault());
61+
verify(innerCallable)
62+
.call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class));
6163
Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT))
6264
.isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT));
6365
}
6466

6567
@Test
6668
public void testLimitChanged() {
6769
ReadRowsFirstCallable<Row> callable = new ReadRowsFirstCallable<>(innerCallable);
68-
innerResult.set(null);
69-
callable.call(Query.create("fake-table").limit(1_000));
70-
70+
callable.futureCall(Query.create("fake-table").limit(10), GrpcCallContext.createDefault());
71+
verify(innerCallable)
72+
.call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class));
7173
Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT))
7274
.isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT));
7375
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsUserCallableTest.java‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,17 +41,4 @@ public void testRequestConverted() {
4141

4242
Truth.assertThat(innerCallable.getActualRequest()).isEqualTo(query.toProto(REQUEST_CONTEXT));
4343
}
44-
45-
@Test
46-
public void testFirstIsLimited() {
47-
ServerStreamingStashCallable<ReadRowsRequest, Row> innerCallable =
48-
new ServerStreamingStashCallable<>();
49-
ReadRowsUserCallable<Row> callable = new ReadRowsUserCallable<>(innerCallable, REQUEST_CONTEXT);
50-
Query query = Query.create("fake-table");
51-
52-
callable.first().call(query);
53-
54-
Truth.assertThat(innerCallable.getActualRequest())
55-
.isEqualTo(query.limit(1).toProto(REQUEST_CONTEXT));
56-
}
5744
}

0 commit comments

Comments
 (0)