Skip to content

Commit 8847fed

Browse files
tengzhongertengzhonger
authored
feat: Add getNewPartitions method to CloseStream for Bigtable ChangeStream (#1655)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/java-bigtable/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> ☕️ If you write sample code, please follow the [samples format]( https://github.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
1 parent 0e283bf commit 8847fed

File tree

9 files changed

+228
-19
lines changed

9 files changed

+228
-19
lines changed

‎google-cloud-bigtable/clirr-ignored-differences.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,12 @@
100100
<method>*getStatus*</method>
101101
<to>com.google.cloud.bigtable.common.Status</to>
102102
</difference>
103+
<!-- add new method is ok because CloseStream is InternalApi -->
104+
<difference>
105+
<differenceType>7013</differenceType>
106+
<className>com/google/cloud/bigtable/data/v2/models/CloseStream</className>
107+
<method>*getNewPartitions*</method>
108+
</difference>
103109
<!-- change method return type is ok because ChangeStreamMutation is InternalApi -->
104110
<difference>
105111
<differenceType>7006</differenceType>

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/CloseStream.java

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import com.google.auto.value.AutoValue;
2020
import com.google.bigtable.v2.ReadChangeStreamResponse;
2121
import com.google.cloud.bigtable.common.Status;
22+
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
23+
import com.google.common.base.Preconditions;
2224
import com.google.common.collect.ImmutableList;
2325
import java.io.Serializable;
2426
import java.util.List;
@@ -35,8 +37,22 @@ public abstract class CloseStream implements ChangeStreamRecord, Serializable {
3537

3638
private static CloseStream create(
3739
com.google.rpc.Status status,
38-
List<ChangeStreamContinuationToken> changeStreamContinuationTokens) {
39-
return new AutoValue_CloseStream(Status.fromProto(status), changeStreamContinuationTokens);
40+
List<ChangeStreamContinuationToken> changeStreamContinuationTokens,
41+
List<ByteStringRange> newPartitions) {
42+
if (status.getCode() == 0) {
43+
Preconditions.checkState(
44+
changeStreamContinuationTokens.isEmpty(),
45+
"An OK CloseStream should not have continuation tokens.");
46+
} else {
47+
Preconditions.checkState(
48+
!changeStreamContinuationTokens.isEmpty(),
49+
"A non-OK CloseStream should have continuation token(s).");
50+
Preconditions.checkState(
51+
changeStreamContinuationTokens.size() == newPartitions.size(),
52+
"Number of continuation tokens does not match number of new partitions.");
53+
}
54+
return new AutoValue_CloseStream(
55+
Status.fromProto(status), changeStreamContinuationTokens, newPartitions);
4056
}
4157

4258
/** Wraps the protobuf {@link ReadChangeStreamResponse.CloseStream}. */
@@ -46,6 +62,13 @@ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStrea
4662
closeStream.getStatus(),
4763
closeStream.getContinuationTokensList().stream()
4864
.map(ChangeStreamContinuationToken::fromProto)
65+
.collect(ImmutableList.toImmutableList()),
66+
closeStream.getNewPartitionsList().stream()
67+
.map(
68+
newPartition ->
69+
ByteStringRange.create(
70+
newPartition.getRowRange().getStartKeyClosed(),
71+
newPartition.getRowRange().getEndKeyOpen()))
4972
.collect(ImmutableList.toImmutableList()));
5073
}
5174

@@ -56,4 +79,8 @@ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStrea
5679
@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
5780
@Nonnull
5881
public abstract List<ChangeStreamContinuationToken> getChangeStreamContinuationTokens();
82+
83+
@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
84+
@Nonnull
85+
public abstract List<ByteStringRange> getNewPartitions();
5986
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,13 @@ public StreamResumptionStrategy<ReadChangeStreamRequest, ChangeStreamRecordT> cr
5656
public ChangeStreamRecordT processResponse(ChangeStreamRecordT response) {
5757
// Update the token from a Heartbeat or a ChangeStreamMutation.
5858
// We don't worry about resumption after CloseStream, since the server
59-
// will return an OK status right after sending a CloseStream.
59+
// will close the stream with an OK status right after sending a CloseStream,
60+
// no matter what status the CloseStream.Status is:
61+
// 1) ... => CloseStream.Ok => final OK. This means the read finishes successfully.
62+
// 2) ... => CloseStream.Error => final OK. This means the client should start
63+
// a new ReadChangeStream call with the continuation tokens specified in
64+
// CloseStream.
65+
// Either case, we don't need to retry after receiving a CloseStream.
6066
if (changeStreamRecordAdapter.isHeartbeat(response)) {
6167
this.token = changeStreamRecordAdapter.getTokenFromHeartbeat(response);
6268
} else if (changeStreamRecordAdapter.isChangeStreamMutation(response)) {

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordTest.java

Lines changed: 73 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,20 @@
3030
import java.io.IOException;
3131
import java.io.ObjectInputStream;
3232
import java.io.ObjectOutputStream;
33+
import org.junit.Assert;
34+
import org.junit.Rule;
3335
import org.junit.Test;
36+
import org.junit.function.ThrowingRunnable;
37+
import org.junit.rules.ExpectedException;
3438
import org.junit.runner.RunWith;
3539
import org.junit.runners.JUnit4;
3640
import org.threeten.bp.Instant;
3741

3842
@RunWith(JUnit4.class)
3943
public class ChangeStreamRecordTest {
4044

45+
@Rule public ExpectedException expect = ExpectedException.none();
46+
4147
@Test
4248
public void heartbeatSerializationTest() throws IOException, ClassNotFoundException {
4349
ReadChangeStreamResponse.Heartbeat heartbeatProto =
@@ -60,7 +66,7 @@ public void heartbeatSerializationTest() throws IOException, ClassNotFoundExcept
6066

6167
@Test
6268
public void closeStreamSerializationTest() throws IOException, ClassNotFoundException {
63-
Status status = Status.newBuilder().setCode(0).build();
69+
Status status = Status.newBuilder().setCode(11).build();
6470
RowRange rowRange1 =
6571
RowRange.newBuilder()
6672
.setStartKeyClosed(ByteString.copyFromUtf8(""))
@@ -85,6 +91,8 @@ public void closeStreamSerializationTest() throws IOException, ClassNotFoundExce
8591
.setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
8692
.setToken(token2)
8793
.build())
94+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange1))
95+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange2))
8896
.setStatus(status)
8997
.build();
9098
CloseStream closeStream = CloseStream.fromProto(closeStreamProto);
@@ -98,6 +106,7 @@ public void closeStreamSerializationTest() throws IOException, ClassNotFoundExce
98106
assertThat(actual.getChangeStreamContinuationTokens())
99107
.isEqualTo(closeStream.getChangeStreamContinuationTokens());
100108
assertThat(actual.getStatus()).isEqualTo(closeStream.getStatus());
109+
assertThat(actual.getNewPartitions()).isEqualTo(closeStream.getNewPartitions());
101110
}
102111

103112
@Test
@@ -129,7 +138,7 @@ public void heartbeatTest() {
129138

130139
@Test
131140
public void closeStreamTest() {
132-
Status status = Status.newBuilder().setCode(0).build();
141+
Status status = Status.newBuilder().setCode(11).build();
133142
RowRange rowRange1 =
134143
RowRange.newBuilder()
135144
.setStartKeyClosed(ByteString.copyFromUtf8(""))
@@ -154,6 +163,8 @@ public void closeStreamTest() {
154163
.setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
155164
.setToken(token2)
156165
.build())
166+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange1))
167+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange2))
157168
.setStatus(status)
158169
.build();
159170
CloseStream actualCloseStream = CloseStream.fromProto(closeStreamProto);
@@ -169,5 +180,65 @@ public void closeStreamTest() {
169180
ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
170181
assertThat(token2)
171182
.isEqualTo(actualCloseStream.getChangeStreamContinuationTokens().get(1).getToken());
183+
assertThat(actualCloseStream.getNewPartitions().get(0))
184+
.isEqualTo(
185+
ByteStringRange.create(rowRange1.getStartKeyClosed(), rowRange1.getEndKeyOpen()));
186+
assertThat(actualCloseStream.getNewPartitions().get(1))
187+
.isEqualTo(
188+
ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
189+
}
190+
191+
// Tests that an OK CloseStream should not have continuation tokens.
192+
@Test(expected = IllegalStateException.class)
193+
public void closeStreamOkWithContinuationTokenShouldFail() {
194+
Status status = Status.newBuilder().setCode(0).build();
195+
RowRange rowRange =
196+
RowRange.newBuilder()
197+
.setStartKeyClosed(ByteString.copyFromUtf8(""))
198+
.setEndKeyOpen(ByteString.copyFromUtf8("apple"))
199+
.build();
200+
String token = "close-stream-token-1";
201+
ReadChangeStreamResponse.CloseStream closeStreamProto =
202+
ReadChangeStreamResponse.CloseStream.newBuilder()
203+
.addContinuationTokens(
204+
StreamContinuationToken.newBuilder()
205+
.setPartition(StreamPartition.newBuilder().setRowRange(rowRange))
206+
.setToken(token))
207+
.setStatus(status)
208+
.build();
209+
Assert.assertThrows(
210+
IllegalStateException.class, (ThrowingRunnable) CloseStream.fromProto(closeStreamProto));
211+
}
212+
213+
// Tests that a non-OK CloseStream should have continuation tokens.
214+
@Test(expected = IllegalStateException.class)
215+
public void closeStreamErrorWithoutContinuationTokenShouldFail() {
216+
Status status = Status.newBuilder().setCode(11).build();
217+
ReadChangeStreamResponse.CloseStream closeStreamProto =
218+
ReadChangeStreamResponse.CloseStream.newBuilder().setStatus(status).build();
219+
Assert.assertThrows(
220+
IllegalStateException.class, (ThrowingRunnable) CloseStream.fromProto(closeStreamProto));
221+
}
222+
223+
// Tests that the number of continuation tokens should match the number of new partitions.
224+
@Test(expected = IllegalStateException.class)
225+
public void closeStreamTokenAndNewPartitionCountMismatchedTest() {
226+
Status status = Status.newBuilder().setCode(11).build();
227+
RowRange rowRange =
228+
RowRange.newBuilder()
229+
.setStartKeyClosed(ByteString.copyFromUtf8(""))
230+
.setEndKeyOpen(ByteString.copyFromUtf8("apple"))
231+
.build();
232+
String token = "close-stream-token-1";
233+
ReadChangeStreamResponse.CloseStream closeStreamProto =
234+
ReadChangeStreamResponse.CloseStream.newBuilder()
235+
.addContinuationTokens(
236+
StreamContinuationToken.newBuilder()
237+
.setPartition(StreamPartition.newBuilder().setRowRange(rowRange))
238+
.setToken(token))
239+
.setStatus(status)
240+
.build();
241+
Assert.assertThrows(
242+
IllegalStateException.class, (ThrowingRunnable) CloseStream.fromProto(closeStreamProto));
172243
}
173244
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapterTest.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,6 @@ public void heartbeatTest() {
150150
public void closeStreamTest() {
151151
ReadChangeStreamResponse.CloseStream expectedCloseStream =
152152
ReadChangeStreamResponse.CloseStream.newBuilder()
153-
.addContinuationTokens(
154-
StreamContinuationToken.newBuilder().setToken("random-token").build())
155153
.setStatus(Status.newBuilder().setCode(0).build())
156154
.build();
157155
assertThat(changeStreamRecordBuilder.onCloseStream(expectedCloseStream))

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallableTest.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,8 @@ public void closeStreamTest() {
102102
ReadChangeStreamResponse.CloseStream closeStreamProto =
103103
ReadChangeStreamResponse.CloseStream.newBuilder()
104104
.addContinuationTokens(streamContinuationToken)
105-
.setStatus(Status.newBuilder().setCode(0).build())
105+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange))
106+
.setStatus(Status.newBuilder().setCode(11))
106107
.build();
107108
ReadChangeStreamResponse response =
108109
ReadChangeStreamResponse.newBuilder().setCloseStream(closeStreamProto).build();
@@ -127,5 +128,8 @@ public void closeStreamTest() {
127128
.isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
128129
assertThat(changeStreamContinuationToken.getToken())
129130
.isEqualTo(streamContinuationToken.getToken());
131+
assertThat(closeStream.getNewPartitions().size()).isEqualTo(1);
132+
assertThat(closeStream.getNewPartitions().get(0))
133+
.isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
130134
}
131135
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamMergingAcceptanceTest.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import com.google.cloud.bigtable.data.v2.models.DeleteFamily;
3939
import com.google.cloud.bigtable.data.v2.models.Entry;
4040
import com.google.cloud.bigtable.data.v2.models.Heartbeat;
41+
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
4142
import com.google.cloud.bigtable.data.v2.models.SetCell;
4243
import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
4344
import com.google.cloud.conformance.bigtable.v2.ChangeStreamTestDefinition.ChangeStreamTestFile;
@@ -173,6 +174,14 @@ public void test() throws Exception {
173174
.setToken(token.getToken())
174175
.build());
175176
}
177+
for (ByteStringRange newPartition : closeStream.getNewPartitions()) {
178+
builder.addNewPartitions(
179+
StreamPartition.newBuilder()
180+
.setRowRange(
181+
RowRange.newBuilder()
182+
.setStartKeyClosed(newPartition.getStart())
183+
.setEndKeyOpen(newPartition.getEnd())));
184+
}
176185
ReadChangeStreamResponse.CloseStream closeStreamProto = builder.build();
177186
actualResults.add(
178187
ReadChangeStreamTest.Result.newBuilder()

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamRetryTest.java

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,15 @@ private StreamContinuationToken createStreamContinuationToken(@Nonnull String to
122122
.build();
123123
}
124124

125+
private StreamPartition createNewPartitionForCloseStream() {
126+
return StreamPartition.newBuilder()
127+
.setRowRange(
128+
RowRange.newBuilder()
129+
.setStartKeyClosed(ByteString.copyFromUtf8(START_KEY_CLOSED))
130+
.setEndKeyOpen(ByteString.copyFromUtf8(END_KEY_OPEN)))
131+
.build();
132+
}
133+
125134
private ReadChangeStreamResponse.Heartbeat createHeartbeat(
126135
StreamContinuationToken streamContinuationToken) {
127136
return ReadChangeStreamResponse.Heartbeat.newBuilder()
@@ -130,11 +139,18 @@ private ReadChangeStreamResponse.Heartbeat createHeartbeat(
130139
.build();
131140
}
132141

133-
private ReadChangeStreamResponse.CloseStream createCloseStream() {
134-
return ReadChangeStreamResponse.CloseStream.newBuilder()
135-
.addContinuationTokens(createStreamContinuationToken(CLOSE_STREAM_TOKEN))
136-
.setStatus(com.google.rpc.Status.newBuilder().setCode(0).build())
137-
.build();
142+
private ReadChangeStreamResponse.CloseStream createCloseStream(boolean isOk) {
143+
ReadChangeStreamResponse.CloseStream.Builder builder =
144+
ReadChangeStreamResponse.CloseStream.newBuilder();
145+
if (isOk) {
146+
builder.setStatus(com.google.rpc.Status.newBuilder().setCode(0));
147+
} else {
148+
builder
149+
.setStatus(com.google.rpc.Status.newBuilder().setCode(11))
150+
.addContinuationTokens(createStreamContinuationToken(CLOSE_STREAM_TOKEN))
151+
.addNewPartitions(createNewPartitionForCloseStream());
152+
}
153+
return builder.build();
138154
}
139155

140156
private ReadChangeStreamResponse.DataChange createDataChange(boolean done) {
@@ -178,7 +194,7 @@ public void happyPathHeartbeatTest() {
178194
@Test
179195
public void happyPathCloseStreamTest() {
180196
ReadChangeStreamResponse closeStreamResponse =
181-
ReadChangeStreamResponse.newBuilder().setCloseStream(createCloseStream()).build();
197+
ReadChangeStreamResponse.newBuilder().setCloseStream(createCloseStream(true)).build();
182198
service.expectations.add(
183199
RpcExpectation.create().expectInitialRequest().respondWith(closeStreamResponse));
184200
List<ChangeStreamRecord> actualResults = getResults();
@@ -221,7 +237,7 @@ public void singleHeartbeatImmediateRetryTest() {
221237
public void singleCloseStreamImmediateRetryTest() {
222238
// CloseStream.
223239
ReadChangeStreamResponse closeStreamResponse =
224-
ReadChangeStreamResponse.newBuilder().setCloseStream(createCloseStream()).build();
240+
ReadChangeStreamResponse.newBuilder().setCloseStream(createCloseStream(false)).build();
225241
service.expectations.add(
226242
RpcExpectation.create().expectInitialRequest().respondWithStatus(Code.UNAVAILABLE));
227243
// Resume with the exact same request.

0 commit comments

Comments
 (0)