feat: Add getNewPartitions method to CloseStream for Bigtable ChangeStream

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

@@ -100,6 +100,12 @@
         <method>*getStatus*</method>
         <to>com.google.cloud.bigtable.common.Status</to>
     </difference>
+    <!-- add new method is ok because CloseStream is InternalApi -->
+    <difference>
+        <differenceType>7013</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/models/CloseStream</className>
+        <method>*getNewPartitions*</method>
+    </difference>
     <!-- change method return type is ok because ChangeStreamMutation is InternalApi -->
     <difference>
         <differenceType>7006</differenceType>

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

@@ -19,6 +19,8 @@
 import com.google.auto.value.AutoValue;
 import com.google.bigtable.v2.ReadChangeStreamResponse;
 import com.google.cloud.bigtable.common.Status;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import java.io.Serializable;
 import java.util.List;
@@ -35,8 +37,22 @@ public abstract class CloseStream implements ChangeStreamRecord, Serializable {
 
   private static CloseStream create(
       com.google.rpc.Status status,
-      List<ChangeStreamContinuationToken> changeStreamContinuationTokens) {
-    return new AutoValue_CloseStream(Status.fromProto(status), changeStreamContinuationTokens);
+      List<ChangeStreamContinuationToken> changeStreamContinuationTokens,
+      List<ByteStringRange> newPartitions) {
+    if (status.getCode() == 0) {
+      Preconditions.checkState(
+          changeStreamContinuationTokens.isEmpty(),
+          "An OK CloseStream should not have continuation tokens.");
+    } else {
+      Preconditions.checkState(
+          !changeStreamContinuationTokens.isEmpty(),
+          "A non-OK CloseStream should have continuation token(s).");
+      Preconditions.checkState(
+          changeStreamContinuationTokens.size() == newPartitions.size(),
+          "Number of continuation tokens does not match number of new partitions.");
+    }
+    return new AutoValue_CloseStream(
+        Status.fromProto(status), changeStreamContinuationTokens, newPartitions);
   }
 
   /** Wraps the protobuf {@link ReadChangeStreamResponse.CloseStream}. */
@@ -46,6 +62,13 @@ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStrea
         closeStream.getStatus(),
         closeStream.getContinuationTokensList().stream()
             .map(ChangeStreamContinuationToken::fromProto)
+            .collect(ImmutableList.toImmutableList()),
+        closeStream.getNewPartitionsList().stream()
+            .map(
+                newPartition ->
+                    ByteStringRange.create(
+                        newPartition.getRowRange().getStartKeyClosed(),
+                        newPartition.getRowRange().getEndKeyOpen()))
             .collect(ImmutableList.toImmutableList()));
   }
 
@@ -56,4 +79,8 @@ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStrea
   @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
   @Nonnull
   public abstract List<ChangeStreamContinuationToken> getChangeStreamContinuationTokens();
+
+  @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+  @Nonnull
+  public abstract List<ByteStringRange> getNewPartitions();
 }

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

@@ -30,14 +30,20 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import org.junit.Assert;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.threeten.bp.Instant;
 
 @RunWith(JUnit4.class)
 public class ChangeStreamRecordTest {
 
+  @Rule public ExpectedException expect = ExpectedException.none();
+
   @Test
   public void heartbeatSerializationTest() throws IOException, ClassNotFoundException {
     ReadChangeStreamResponse.Heartbeat heartbeatProto =
@@ -60,7 +66,7 @@ public void heartbeatSerializationTest() throws IOException, ClassNotFoundExcept
 
   @Test
   public void closeStreamSerializationTest() throws IOException, ClassNotFoundException {
-    Status status = Status.newBuilder().setCode(0).build();
+    Status status = Status.newBuilder().setCode(11).build();
     RowRange rowRange1 =
         RowRange.newBuilder()
             .setStartKeyClosed(ByteString.copyFromUtf8(""))
@@ -85,6 +91,8 @@ public void closeStreamSerializationTest() throws IOException, ClassNotFoundExce
                     .setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
                     .setToken(token2)
                     .build())
+            .addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange1))
+            .addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange2))
             .setStatus(status)
             .build();
     CloseStream closeStream = CloseStream.fromProto(closeStreamProto);
@@ -98,6 +106,7 @@ public void closeStreamSerializationTest() throws IOException, ClassNotFoundExce
     assertThat(actual.getChangeStreamContinuationTokens())
         .isEqualTo(closeStream.getChangeStreamContinuationTokens());
     assertThat(actual.getStatus()).isEqualTo(closeStream.getStatus());
+    assertThat(actual.getNewPartitions()).isEqualTo(closeStream.getNewPartitions());
   }
 
   @Test
@@ -129,7 +138,7 @@ public void heartbeatTest() {
 
   @Test
   public void closeStreamTest() {
-    Status status = Status.newBuilder().setCode(0).build();
+    Status status = Status.newBuilder().setCode(11).build();
     RowRange rowRange1 =
         RowRange.newBuilder()
             .setStartKeyClosed(ByteString.copyFromUtf8(""))
@@ -154,6 +163,8 @@ public void closeStreamTest() {
                     .setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
                     .setToken(token2)
                     .build())
+            .addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange1))
+            .addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange2))
             .setStatus(status)
             .build();
     CloseStream actualCloseStream = CloseStream.fromProto(closeStreamProto);
@@ -169,5 +180,65 @@ public void closeStreamTest() {
             ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
     assertThat(token2)
         .isEqualTo(actualCloseStream.getChangeStreamContinuationTokens().get(1).getToken());
+    assertThat(actualCloseStream.getNewPartitions().get(0))
+        .isEqualTo(
+            ByteStringRange.create(rowRange1.getStartKeyClosed(), rowRange1.getEndKeyOpen()));
+    assertThat(actualCloseStream.getNewPartitions().get(1))
+        .isEqualTo(
+            ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
+  }
+
+  // Tests that an OK CloseStream should not have continuation tokens.
+  @Test(expected = IllegalStateException.class)
+  public void closeStreamOkWithContinuationTokenShouldFail() {
+    Status status = Status.newBuilder().setCode(0).build();
+    RowRange rowRange =
+        RowRange.newBuilder()
+            .setStartKeyClosed(ByteString.copyFromUtf8(""))
+            .setEndKeyOpen(ByteString.copyFromUtf8("apple"))
+            .build();
+    String token = "close-stream-token-1";
+    ReadChangeStreamResponse.CloseStream closeStreamProto =
+        ReadChangeStreamResponse.CloseStream.newBuilder()
+            .addContinuationTokens(
+                StreamContinuationToken.newBuilder()
+                    .setPartition(StreamPartition.newBuilder().setRowRange(rowRange))
+                    .setToken(token))
+            .setStatus(status)
+            .build();
+    Assert.assertThrows(
+        IllegalStateException.class, (ThrowingRunnable) CloseStream.fromProto(closeStreamProto));
+  }
+
+  // Tests that a non-OK CloseStream should have continuation tokens.
+  @Test(expected = IllegalStateException.class)
+  public void closeStreamErrorWithoutContinuationTokenShouldFail() {
+    Status status = Status.newBuilder().setCode(11).build();
+    ReadChangeStreamResponse.CloseStream closeStreamProto =
+        ReadChangeStreamResponse.CloseStream.newBuilder().setStatus(status).build();
+    Assert.assertThrows(
+        IllegalStateException.class, (ThrowingRunnable) CloseStream.fromProto(closeStreamProto));
+  }
+
+  // Tests that the number of continuation tokens should match the number of new partitions.
+  @Test(expected = IllegalStateException.class)
+  public void closeStreamTokenAndNewPartitionCountMismatchedTest() {
+    Status status = Status.newBuilder().setCode(11).build();
+    RowRange rowRange =
+        RowRange.newBuilder()
+            .setStartKeyClosed(ByteString.copyFromUtf8(""))
+            .setEndKeyOpen(ByteString.copyFromUtf8("apple"))
+            .build();
+    String token = "close-stream-token-1";
+    ReadChangeStreamResponse.CloseStream closeStreamProto =
+        ReadChangeStreamResponse.CloseStream.newBuilder()
+            .addContinuationTokens(
+                StreamContinuationToken.newBuilder()
+                    .setPartition(StreamPartition.newBuilder().setRowRange(rowRange))
+                    .setToken(token))
+            .setStatus(status)
+            .build();
+    Assert.assertThrows(
+        IllegalStateException.class, (ThrowingRunnable) CloseStream.fromProto(closeStreamProto));
   }
 }

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

@@ -102,6 +102,7 @@ public void closeStreamTest() {
     ReadChangeStreamResponse.CloseStream closeStreamProto =
         ReadChangeStreamResponse.CloseStream.newBuilder()
             .addContinuationTokens(streamContinuationToken)
+            .addNewPartitions(StreamPartition.newBuilder().setRowRange(rowRange))
             .setStatus(Status.newBuilder().setCode(0).build())
             .build();
     ReadChangeStreamResponse response =
@@ -127,5 +128,8 @@ public void closeStreamTest() {
         .isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
     assertThat(changeStreamContinuationToken.getToken())
         .isEqualTo(streamContinuationToken.getToken());
+    assertThat(closeStream.getNewPartitions().size()).isEqualTo(1);
+    assertThat(closeStream.getNewPartitions().get(0))
+        .isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
   }
 }

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

@@ -38,6 +38,7 @@
 import com.google.cloud.bigtable.data.v2.models.DeleteFamily;
 import com.google.cloud.bigtable.data.v2.models.Entry;
 import com.google.cloud.bigtable.data.v2.models.Heartbeat;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
 import com.google.cloud.bigtable.data.v2.models.SetCell;
 import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
 import com.google.cloud.conformance.bigtable.v2.ChangeStreamTestDefinition.ChangeStreamTestFile;
@@ -173,6 +174,14 @@ public void test() throws Exception {
                     .setToken(token.getToken())
                     .build());
           }
+          for (ByteStringRange newPartition : closeStream.getNewPartitions()) {
+            builder.addNewPartitions(
+                StreamPartition.newBuilder()
+                    .setRowRange(
+                        RowRange.newBuilder()
+                            .setStartKeyClosed(newPartition.getStart())
+                            .setEndKeyOpen(newPartition.getEnd())));
+          }
           ReadChangeStreamResponse.CloseStream closeStreamProto = builder.build();
           actualResults.add(
               ReadChangeStreamTest.Result.newBuilder()

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

@@ -122,6 +122,15 @@ private StreamContinuationToken createStreamContinuationToken(@Nonnull String to
         .build();
   }
 
+  private StreamPartition createNewPartitionForCloseStream() {
+    return StreamPartition.newBuilder()
+        .setRowRange(
+            RowRange.newBuilder()
+                .setStartKeyClosed(ByteString.copyFromUtf8(START_KEY_CLOSED))
+                .setEndKeyOpen(ByteString.copyFromUtf8(END_KEY_OPEN)))
+        .build();
+  }
+
   private ReadChangeStreamResponse.Heartbeat createHeartbeat(
       StreamContinuationToken streamContinuationToken) {
     return ReadChangeStreamResponse.Heartbeat.newBuilder()
@@ -133,6 +142,7 @@ private ReadChangeStreamResponse.Heartbeat createHeartbeat(
   private ReadChangeStreamResponse.CloseStream createCloseStream() {
     return ReadChangeStreamResponse.CloseStream.newBuilder()
         .addContinuationTokens(createStreamContinuationToken(CLOSE_STREAM_TOKEN))
+        .addNewPartitions(createNewPartitionForCloseStream())
         .setStatus(com.google.rpc.Status.newBuilder().setCode(0).build())
         .build();
   }