Visitar URL original
Add getNewPartitions method to CloseStream for Bigtable ChangeStream · googleapis/java-bigtable@890859a · GitHub
Skip to content
This repository was archived by the owner on Aug 30, 2026. It is now read-only.

Commit 890859a

Browse files
author
Teng Zhong
committed
Add getNewPartitions method to CloseStream for Bigtable ChangeStream
1 parent 0e283bf commit 890859a

7 files changed

Lines changed: 132 additions & 6 deletions

File tree

‎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: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
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;
2223
import com.google.common.collect.ImmutableList;
2324
import java.io.Serializable;
2425
import java.util.List;
@@ -35,8 +36,10 @@ public abstract class CloseStream implements ChangeStreamRecord, Serializable {
3536

3637
private static CloseStream create(
3738
com.google.rpc.Status status,
38-
List<ChangeStreamContinuationToken> changeStreamContinuationTokens) {
39-
return new AutoValue_CloseStream(Status.fromProto(status), changeStreamContinuationTokens);
39+
List<ChangeStreamContinuationToken> changeStreamContinuationTokens,
40+
List<ByteStringRange> newPartitions) {
41+
return new AutoValue_CloseStream(
42+
Status.fromProto(status), changeStreamContinuationTokens, newPartitions);
4043
}
4144

4245
/** Wraps the protobuf {@link ReadChangeStreamResponse.CloseStream}. */
@@ -46,6 +49,13 @@ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStrea
4649
closeStream.getStatus(),
4750
closeStream.getContinuationTokensList().stream()
4851
.map(ChangeStreamContinuationToken::fromProto)
52+
.collect(ImmutableList.toImmutableList()),
53+
closeStream.getNewPartitionsList().stream()
54+
.map(
55+
newPartition ->
56+
ByteStringRange.create(
57+
newPartition.getRowRange().getStartKeyClosed(),
58+
newPartition.getRowRange().getEndKeyOpen()))
4959
.collect(ImmutableList.toImmutableList()));
5060
}
5161

@@ -56,4 +66,8 @@ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStrea
5666
@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
5767
@Nonnull
5868
public abstract List<ChangeStreamContinuationToken> getChangeStreamContinuationTokens();
69+
70+
@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
71+
@Nonnull
72+
public abstract List<ByteStringRange> getNewPartitions();
5973
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ public void closeStreamSerializationTest() throws IOException, ClassNotFoundExce
8585
.setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
8686
.setToken(token2)
8787
.build())
88+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange1))
89+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange2))
8890
.setStatus(status)
8991
.build();
9092
CloseStream closeStream = CloseStream.fromProto(closeStreamProto);
@@ -98,6 +100,7 @@ public void closeStreamSerializationTest() throws IOException, ClassNotFoundExce
98100
assertThat(actual.getChangeStreamContinuationTokens())
99101
.isEqualTo(closeStream.getChangeStreamContinuationTokens());
100102
assertThat(actual.getStatus()).isEqualTo(closeStream.getStatus());
103+
assertThat(actual.getNewPartitions()).isEqualTo(closeStream.getNewPartitions());
101104
}
102105

103106
@Test
@@ -154,6 +157,8 @@ public void closeStreamTest() {
154157
.setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
155158
.setToken(token2)
156159
.build())
160+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange1))
161+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange2))
157162
.setStatus(status)
158163
.build();
159164
CloseStream actualCloseStream = CloseStream.fromProto(closeStreamProto);
@@ -169,5 +174,11 @@ public void closeStreamTest() {
169174
ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
170175
assertThat(token2)
171176
.isEqualTo(actualCloseStream.getChangeStreamContinuationTokens().get(1).getToken());
177+
assertThat(actualCloseStream.getNewPartitions().get(0))
178+
.isEqualTo(
179+
ByteStringRange.create(rowRange1.getStartKeyClosed(), rowRange1.getEndKeyOpen()));
180+
assertThat(actualCloseStream.getNewPartitions().get(1))
181+
.isEqualTo(
182+
ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
172183
}
173184
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ public void closeStreamTest() {
102102
ReadChangeStreamResponse.CloseStream closeStreamProto =
103103
ReadChangeStreamResponse.CloseStream.newBuilder()
104104
.addContinuationTokens(streamContinuationToken)
105+
.addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange))
105106
.setStatus(Status.newBuilder().setCode(0).build())
106107
.build();
107108
ReadChangeStreamResponse response =
@@ -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: 10 additions & 0 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()
@@ -133,6 +142,7 @@ private ReadChangeStreamResponse.Heartbeat createHeartbeat(
133142
private ReadChangeStreamResponse.CloseStream createCloseStream() {
134143
return ReadChangeStreamResponse.CloseStream.newBuilder()
135144
.addContinuationTokens(createStreamContinuationToken(CLOSE_STREAM_TOKEN))
145+
.addNewPartitions(createNewPartitionForCloseStream())
136146
.setStatus(com.google.rpc.Status.newBuilder().setCode(0).build())
137147
.build();
138148
}

‎google-cloud-bigtable/src/test/resources/changestream.json‎

Lines changed: 76 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,11 +61,25 @@
6161
"partition": {
6262
"row_range": {
6363
"start_key_closed": "0000000000000001",
64-
"end_key_open": "0000000000000002"
64+
"end_key_open": "0000000000000003"
6565
}
6666
},
6767
"token": "close-stream-token-2"
6868
}
69+
],
70+
"new_partitions": [
71+
{
72+
"row_range": {
73+
"start_key_closed": "",
74+
"end_key_open": "0000000000000002"
75+
}
76+
},
77+
{
78+
"row_range": {
79+
"start_key_closed": "0000000000000002",
80+
"end_key_open": "0000000000000003"
81+
}
82+
}
6983
]
7084
}
7185
}
@@ -92,11 +106,25 @@
92106
"partition": {
93107
"row_range": {
94108
"start_key_closed": "0000000000000001",
95-
"end_key_open": "0000000000000002"
109+
"end_key_open": "0000000000000003"
96110
}
97111
},
98112
"token": "close-stream-token-2"
99113
}
114+
],
115+
"new_partitions": [
116+
{
117+
"row_range": {
118+
"start_key_closed": "",
119+
"end_key_open": "0000000000000002"
120+
}
121+
},
122+
{
123+
"row_range": {
124+
"start_key_closed": "0000000000000002",
125+
"end_key_open": "0000000000000003"
126+
}
127+
}
100128
]
101129
}
102130
},
@@ -137,6 +165,14 @@
137165
},
138166
"token": "close-stream-token-1"
139167
}
168+
],
169+
"new_partitions": [
170+
{
171+
"row_range": {
172+
"start_key_closed": "",
173+
"end_key_open": "0000000000000002"
174+
}
175+
}
140176
]
141177
}
142178
}
@@ -176,6 +212,14 @@
176212
},
177213
"token": "close-stream-token-1"
178214
}
215+
],
216+
"new_partitions": [
217+
{
218+
"row_range": {
219+
"start_key_closed": "",
220+
"end_key_open": "0000000000000002"
221+
}
222+
}
179223
]
180224
}
181225
},
@@ -1280,11 +1324,25 @@
12801324
"partition": {
12811325
"row_range": {
12821326
"start_key_closed": "0000000000000001",
1283-
"end_key_open": "0000000000000002"
1327+
"end_key_open": "0000000000000003"
12841328
}
12851329
},
12861330
"token": "close-stream-token-2"
12871331
}
1332+
],
1333+
"new_partitions": [
1334+
{
1335+
"row_range": {
1336+
"start_key_closed": "",
1337+
"end_key_open": "0000000000000002"
1338+
}
1339+
},
1340+
{
1341+
"row_range": {
1342+
"start_key_closed": "0000000000000002",
1343+
"end_key_open": "0000000000000003"
1344+
}
1345+
}
12881346
]
12891347
}
12901348
}
@@ -1363,11 +1421,25 @@
13631421
"partition": {
13641422
"row_range": {
13651423
"start_key_closed": "0000000000000001",
1366-
"end_key_open": "0000000000000002"
1424+
"end_key_open": "0000000000000003"
13671425
}
13681426
},
13691427
"token": "close-stream-token-2"
13701428
}
1429+
],
1430+
"new_partitions": [
1431+
{
1432+
"row_range": {
1433+
"start_key_closed": "",
1434+
"end_key_open": "0000000000000002"
1435+
}
1436+
},
1437+
{
1438+
"row_range": {
1439+
"start_key_closed": "0000000000000002",
1440+
"end_key_open": "0000000000000003"
1441+
}
1442+
}
13711443
]
13721444
}
13731445
},

0 commit comments

Comments
 (0)