+ *
+ * @see ServerStreamingCallable For call styles.
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public ServerStream generateInitialChangeStreamPartitions(String tableId) {
+ return generateInitialChangeStreamPartitionsCallable().call(tableId);
+ }
+
+ /**
+ * Convenience method for asynchronously streaming the partitions of a table.
+ *
+ *
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public void generateInitialChangeStreamPartitionsAsync(
+ String tableId, ResponseObserver observer) {
+ generateInitialChangeStreamPartitionsCallable().call(tableId, observer);
+ }
+
+ /**
+ * Streams back the results of the query. The returned callable object allows for customization of
+ * api invocation.
+ *
+ *
+ *
+ * @see ServerStreamingCallable For call styles.
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public ServerStreamingCallable
+ generateInitialChangeStreamPartitionsCallable() {
+ return stub.generateInitialChangeStreamPartitionsCallable();
+ }
+
+ /**
+ * Convenience method for synchronously streaming the results of a {@link ReadChangeStreamQuery}.
+ * The returned ServerStream instance is not threadsafe, it can only be used from single thread.
+ *
+ *
+ *
+ * @see ServerStreamingCallable For call styles.
+ * @see ReadChangeStreamQuery For query options.
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public ServerStream readChangeStream(ReadChangeStreamQuery query) {
+ return readChangeStreamCallable().call(query);
+ }
+
+ /**
+ * Convenience method for asynchronously streaming the results of a {@link ReadChangeStreamQuery}.
+ *
+ *
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public void readChangeStreamAsync(
+ ReadChangeStreamQuery query, ResponseObserver observer) {
+ readChangeStreamCallable().call(query, observer);
+ }
+
+ /**
+ * Streams back the results of the query. The returned callable object allows for customization of
+ * api invocation.
+ *
+ *
+ *
+ * @see ServerStreamingCallable For call styles.
+ * @see ReadChangeStreamQuery For query options.
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public ServerStreamingCallable
+ readChangeStreamCallable() {
+ return stub.readChangeStreamCallable();
+ }
+
/** Close the clients and releases all associated resources. */
@Override
public void close() {
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamContinuationToken.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamContinuationToken.java
new file mode 100644
index 0000000000..f619d9dae0
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamContinuationToken.java
@@ -0,0 +1,88 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
+import java.io.Serializable;
+import javax.annotation.Nonnull;
+
+/** A simple wrapper for {@link StreamContinuationToken}. */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class ChangeStreamContinuationToken implements Serializable {
+ private static final long serialVersionUID = 524679926247095L;
+
+ private static ChangeStreamContinuationToken create(@Nonnull StreamContinuationToken tokenProto) {
+ return new AutoValue_ChangeStreamContinuationToken(tokenProto);
+ }
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public static ChangeStreamContinuationToken create(
+ @Nonnull ByteStringRange byteStringRange, @Nonnull String token) {
+ return create(
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(byteStringRange.getStart())
+ .setEndKeyOpen(byteStringRange.getEnd())
+ .build())
+ .build())
+ .setToken(token)
+ .build());
+ }
+
+ /** Wraps the protobuf {@link StreamContinuationToken}. */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ static ChangeStreamContinuationToken fromProto(
+ @Nonnull StreamContinuationToken streamContinuationToken) {
+ return create(streamContinuationToken);
+ }
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public static ChangeStreamContinuationToken fromByteString(ByteString byteString)
+ throws InvalidProtocolBufferException {
+ return create(StreamContinuationToken.newBuilder().mergeFrom(byteString).build());
+ }
+
+ @Nonnull
+ public abstract StreamContinuationToken getTokenProto();
+
+ /**
+ * Get the partition of the current continuation token, represented by a {@link ByteStringRange}.
+ */
+ public ByteStringRange getPartition() {
+ return ByteStringRange.create(
+ getTokenProto().getPartition().getRowRange().getStartKeyClosed(),
+ getTokenProto().getPartition().getRowRange().getEndKeyOpen());
+ }
+
+ public String getToken() {
+ return getTokenProto().getToken();
+ }
+
+ public ByteString toByteString() {
+ return getTokenProto().toByteString();
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamMutation.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamMutation.java
new file mode 100644
index 0000000000..42ef300b9d
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamMutation.java
@@ -0,0 +1,231 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
+import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMerger;
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.ByteString;
+import java.io.Serializable;
+import javax.annotation.Nonnull;
+
+/**
+ * A ChangeStreamMutation represents a list of mods(represented by List<{@link Entry}>) targeted at
+ * a single row, which is concatenated by {@link ChangeStreamRecordMerger}. It represents a logical
+ * row mutation and can be converted to the original write request(i.e. {@link RowMutation} or
+ * {@link RowMutationEntry}.
+ *
+ *
A ChangeStreamMutation can be constructed in two ways, depending on whether it's a user
+ * initiated mutation or a Garbage Collection mutation. Either way, the caller should explicitly set
+ * `token` and `estimatedLowWatermark` before build(), otherwise it'll raise an error.
+ *
+ *
+ */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class ChangeStreamMutation implements ChangeStreamRecord, Serializable {
+ private static final long serialVersionUID = 8419520253162024218L;
+
+ public enum MutationType {
+ USER,
+ GARBAGE_COLLECTION
+ }
+
+ /**
+ * Creates a new instance of a user initiated mutation. It returns a builder instead of a
+ * ChangeStreamMutation because `token` and `loWatermark` must be set later when we finish
+ * building the logical mutation.
+ */
+ static Builder createUserMutation(
+ @Nonnull ByteString rowKey,
+ @Nonnull String sourceClusterId,
+ long commitTimestamp,
+ int tieBreaker) {
+ return builder()
+ .setRowKey(rowKey)
+ .setType(MutationType.USER)
+ .setSourceClusterId(sourceClusterId)
+ .setCommitTimestamp(commitTimestamp)
+ .setTieBreaker(tieBreaker);
+ }
+
+ /**
+ * Creates a new instance of a GC mutation. It returns a builder instead of a ChangeStreamMutation
+ * because `token` and `loWatermark` must be set later when we finish building the logical
+ * mutation.
+ */
+ static Builder createGcMutation(
+ @Nonnull ByteString rowKey, long commitTimestamp, int tieBreaker) {
+ return builder()
+ .setRowKey(rowKey)
+ .setType(MutationType.GARBAGE_COLLECTION)
+ .setSourceClusterId("")
+ .setCommitTimestamp(commitTimestamp)
+ .setTieBreaker(tieBreaker);
+ }
+
+ /** Get the row key of the current mutation. */
+ @Nonnull
+ public abstract ByteString getRowKey();
+
+ /** Get the type of the current mutation. */
+ @Nonnull
+ public abstract MutationType getType();
+
+ /** Get the source cluster id of the current mutation. */
+ @Nonnull
+ public abstract String getSourceClusterId();
+
+ /** Get the commit timestamp of the current mutation. */
+ public abstract long getCommitTimestamp();
+
+ /**
+ * Get the tie breaker of the current mutation. This is used to resolve conflicts when multiple
+ * mutations are applied to different clusters at the same time.
+ */
+ public abstract int getTieBreaker();
+
+ /** Get the token of the current mutation, which can be used to resume the changestream. */
+ @Nonnull
+ public abstract String getToken();
+
+ /** Get the low watermark of the current mutation. */
+ public abstract long getEstimatedLowWatermark();
+
+ /** Get the list of mods of the current mutation. */
+ @Nonnull
+ public abstract ImmutableList getEntries();
+
+ /** Returns a new builder for this class. */
+ static Builder builder() {
+ return new AutoValue_ChangeStreamMutation.Builder();
+ }
+
+ /** Helper class to create a ChangeStreamMutation. */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setRowKey(@Nonnull ByteString rowKey);
+
+ abstract Builder setType(@Nonnull MutationType type);
+
+ abstract Builder setSourceClusterId(@Nonnull String sourceClusterId);
+
+ abstract Builder setCommitTimestamp(long commitTimestamp);
+
+ abstract Builder setTieBreaker(int tieBreaker);
+
+ abstract ImmutableList.Builder entriesBuilder();
+
+ abstract Builder setToken(@Nonnull String token);
+
+ abstract Builder setEstimatedLowWatermark(long estimatedLowWatermark);
+
+ Builder setCell(
+ @Nonnull String familyName,
+ @Nonnull ByteString qualifier,
+ long timestamp,
+ @Nonnull ByteString value) {
+ this.entriesBuilder().add(SetCell.create(familyName, qualifier, timestamp, value));
+ return this;
+ }
+
+ Builder deleteCells(
+ @Nonnull String familyName,
+ @Nonnull ByteString qualifier,
+ @Nonnull TimestampRange timestampRange) {
+ this.entriesBuilder().add(DeleteCells.create(familyName, qualifier, timestampRange));
+ return this;
+ }
+
+ Builder deleteFamily(@Nonnull String familyName) {
+ this.entriesBuilder().add(DeleteFamily.create(familyName));
+ return this;
+ }
+
+ abstract ChangeStreamMutation build();
+ }
+
+ public RowMutation toRowMutation(@Nonnull String tableId) {
+ RowMutation rowMutation = RowMutation.create(tableId, getRowKey());
+ for (Entry entry : getEntries()) {
+ if (entry instanceof DeleteFamily) {
+ rowMutation.deleteFamily(((DeleteFamily) entry).getFamilyName());
+ } else if (entry instanceof DeleteCells) {
+ DeleteCells deleteCells = (DeleteCells) entry;
+ rowMutation.deleteCells(
+ deleteCells.getFamilyName(),
+ deleteCells.getQualifier(),
+ deleteCells.getTimestampRange());
+ } else if (entry instanceof SetCell) {
+ SetCell setCell = (SetCell) entry;
+ rowMutation.setCell(
+ setCell.getFamilyName(),
+ setCell.getQualifier(),
+ setCell.getTimestamp(),
+ setCell.getValue());
+ } else {
+ throw new IllegalArgumentException("Unexpected Entry type.");
+ }
+ }
+ return rowMutation;
+ }
+
+ public RowMutationEntry toRowMutationEntry() {
+ RowMutationEntry rowMutationEntry = RowMutationEntry.create(getRowKey());
+ for (Entry entry : getEntries()) {
+ if (entry instanceof DeleteFamily) {
+ rowMutationEntry.deleteFamily(((DeleteFamily) entry).getFamilyName());
+ } else if (entry instanceof DeleteCells) {
+ DeleteCells deleteCells = (DeleteCells) entry;
+ rowMutationEntry.deleteCells(
+ deleteCells.getFamilyName(),
+ deleteCells.getQualifier(),
+ deleteCells.getTimestampRange());
+ } else if (entry instanceof SetCell) {
+ SetCell setCell = (SetCell) entry;
+ rowMutationEntry.setCell(
+ setCell.getFamilyName(),
+ setCell.getQualifier(),
+ setCell.getTimestamp(),
+ setCell.getValue());
+ } else {
+ throw new IllegalArgumentException("Unexpected Entry type.");
+ }
+ }
+ return rowMutationEntry;
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecord.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecord.java
new file mode 100644
index 0000000000..edf0c1a26e
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecord.java
@@ -0,0 +1,25 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+
+/**
+ * Default representation of a change stream record, which can be a Heartbeat, a CloseStream, or a
+ * logical mutation.
+ */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+public interface ChangeStreamRecord {}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordAdapter.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordAdapter.java
new file mode 100644
index 0000000000..f94a3b4c3c
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordAdapter.java
@@ -0,0 +1,172 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
+import com.google.protobuf.ByteString;
+import javax.annotation.Nonnull;
+
+/**
+ * An extension point that allows end users to plug in a custom implementation of logical change
+ * stream records. This is useful in cases where the user would like to apply advanced client side
+ * filtering(for example, only keep DeleteFamily in the mutations). This adapter acts like a factory
+ * for a SAX style change stream record builder.
+ */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+public interface ChangeStreamRecordAdapter {
+ /** Creates a new instance of a {@link ChangeStreamRecordBuilder}. */
+ ChangeStreamRecordBuilder createChangeStreamRecordBuilder();
+
+ /** Checks if the given change stream record is a Heartbeat. */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ boolean isHeartbeat(ChangeStreamRecordT record);
+
+ /**
+ * Get the token from the given Heartbeat record. If the given record is not a Heartbeat, it will
+ * throw an Exception.
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ String getTokenFromHeartbeat(ChangeStreamRecordT heartbeatRecord);
+
+ /** Checks if the given change stream record is a ChangeStreamMutation. */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ boolean isChangeStreamMutation(ChangeStreamRecordT record);
+
+ /**
+ * Get the token from the given ChangeStreamMutation record. If the given record is not a
+ * ChangeStreamMutation, it will throw an Exception.
+ */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ String getTokenFromChangeStreamMutation(ChangeStreamRecordT record);
+
+ /**
+ * A SAX style change stream record factory. It is responsible for creating one of the three types
+ * of change stream record: heartbeat, close stream, and a change stream mutation.
+ *
+ *
State management is handled external to the implementation of this class:
+ *
+ *
+ * Case 1: Heartbeat
+ *
Exactly 1 {@code onHeartbeat}.
+ *
+ *
+ *
+ * Case 2: CloseStream
+ *
Exactly 1 {@code onCloseStream}.
+ *
+ *
+ *
+ * Case 3: ChangeStreamMutation. A change stream mutation consists of one or more mods, where
+ * the SetCells might be chunked. There are 3 different types of mods that a ReadChangeStream
+ * response can have:
+ *
SetCell -> Exactly 1 {@code startCell}, At least 1 {@code CellValue}, Exactly 1 {@code
+ * finishCell}.
+ *
+ *
+ *
The whole flow of constructing a ChangeStreamMutation is:
+ *
+ *
+ *
Exactly 1 {@code startUserMutation} or {@code startGcMutation}.
+ *
At least 1 DeleteFamily/DeleteCell/SetCell mods.
+ *
Exactly 1 {@code finishChangeStreamMutation}.
+ *
+ *
+ *
Note: For a non-chunked SetCell, only 1 {@code CellValue} will be called. For a chunked
+ * SetCell, more than 1 {@code CellValue}s will be called.
+ *
+ *
Note: DeleteRow's won't appear in data changes since they'll be converted to multiple
+ * DeleteFamily's.
+ */
+ interface ChangeStreamRecordBuilder {
+ /**
+ * Called to create a heartbeat. This will be called at most once. If called, the current change
+ * stream record must not include any data changes or close stream messages.
+ */
+ ChangeStreamRecordT onHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat);
+
+ /**
+ * Called to create a close stream message. This will be called at most once. If called, the
+ * current change stream record must not include any data changes or heartbeats.
+ */
+ ChangeStreamRecordT onCloseStream(ReadChangeStreamResponse.CloseStream closeStream);
+
+ /**
+ * Called to start a new user initiated ChangeStreamMutation. This will be called at most once.
+ * If called, the current change stream record must not include any close stream message or
+ * heartbeat.
+ */
+ void startUserMutation(
+ @Nonnull ByteString rowKey,
+ @Nonnull String sourceClusterId,
+ long commitTimestamp,
+ int tieBreaker);
+
+ /**
+ * Called to start a new Garbage Collection ChangeStreamMutation. This will be called at most
+ * once. If called, the current change stream record must not include any close stream message
+ * or heartbeat.
+ */
+ void startGcMutation(@Nonnull ByteString rowKey, long commitTimestamp, int tieBreaker);
+
+ /** Called to add a DeleteFamily mod. */
+ void deleteFamily(@Nonnull String familyName);
+
+ /** Called to add a DeleteCell mod. */
+ void deleteCells(
+ @Nonnull String familyName,
+ @Nonnull ByteString qualifier,
+ @Nonnull TimestampRange timestampRange);
+
+ /**
+ * Called to start a SetCell.
+ *
+ *
+ * In case of a non-chunked cell, the following order is guaranteed:
+ *
Exactly 1 {@code startCell}.
+ *
Exactly 1 {@code cellValue}.
+ *
Exactly 1 {@code finishCell}.
+ *
+ *
+ *
+ * In case of a chunked cell, the following order is guaranteed:
+ *
Exactly 1 {@code startCell}.
+ *
At least 2 {@code cellValue}.
+ *
Exactly 1 {@code finishCell}.
+ *
+ */
+ void startCell(String family, ByteString qualifier, long timestampMicros);
+
+ /**
+ * Called once per non-chunked cell, or at least twice per chunked cell to concatenate the cell
+ * value.
+ */
+ void cellValue(ByteString value);
+
+ /** Called once per cell to signal the end of the value (unless reset). */
+ void finishCell();
+
+ /** Called once per stream record to signal that all mods have been processed (unless reset). */
+ ChangeStreamRecordT finishChangeStreamMutation(
+ @Nonnull String token, long estimatedLowWatermark);
+
+ /** Called when the current in progress change stream record should be dropped */
+ void reset();
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/CloseStream.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/CloseStream.java
new file mode 100644
index 0000000000..4760e511e9
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/CloseStream.java
@@ -0,0 +1,58 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.common.collect.ImmutableList;
+import com.google.rpc.Status;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nonnull;
+
+/**
+ * A simple wrapper for {@link ReadChangeStreamResponse.CloseStream}. This message is received when
+ * the stream reading is finished(i.e. read past the stream end time), or an error has occurred.
+ */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class CloseStream implements ChangeStreamRecord, Serializable {
+ private static final long serialVersionUID = 7316215828353608505L;
+
+ private static CloseStream create(
+ Status status, List changeStreamContinuationTokens) {
+ return new AutoValue_CloseStream(status, changeStreamContinuationTokens);
+ }
+
+ /** Wraps the protobuf {@link ReadChangeStreamResponse.CloseStream}. */
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public static CloseStream fromProto(@Nonnull ReadChangeStreamResponse.CloseStream closeStream) {
+ return create(
+ closeStream.getStatus(),
+ closeStream.getContinuationTokensList().stream()
+ .map(ChangeStreamContinuationToken::fromProto)
+ .collect(ImmutableList.toImmutableList()));
+ }
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ @Nonnull
+ public abstract Status getStatus();
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ @Nonnull
+ public abstract List getChangeStreamContinuationTokens();
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapter.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapter.java
new file mode 100644
index 0000000000..79dec5b17f
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapter.java
@@ -0,0 +1,176 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
+import com.google.common.base.Preconditions;
+import com.google.protobuf.ByteString;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Default implementation of a {@link ChangeStreamRecordAdapter} that uses {@link
+ * ChangeStreamRecord}s to represent change stream records.
+ */
+@InternalApi
+public class DefaultChangeStreamRecordAdapter
+ implements ChangeStreamRecordAdapter {
+
+ /** {@inheritDoc} */
+ @Override
+ public ChangeStreamRecordBuilder createChangeStreamRecordBuilder() {
+ return new DefaultChangeStreamRecordBuilder();
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean isHeartbeat(ChangeStreamRecord record) {
+ return record instanceof Heartbeat;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public String getTokenFromHeartbeat(ChangeStreamRecord record) {
+ Preconditions.checkArgument(isHeartbeat(record), "record is not a Heartbeat.");
+ return ((Heartbeat) record).getChangeStreamContinuationToken().getToken();
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public boolean isChangeStreamMutation(ChangeStreamRecord record) {
+ return record instanceof ChangeStreamMutation;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public String getTokenFromChangeStreamMutation(ChangeStreamRecord record) {
+ Preconditions.checkArgument(
+ isChangeStreamMutation(record), "record is not a ChangeStreamMutation.");
+ return ((ChangeStreamMutation) record).getToken();
+ }
+
+ /** {@inheritDoc} */
+ static class DefaultChangeStreamRecordBuilder
+ implements ChangeStreamRecordBuilder {
+ private ChangeStreamMutation.Builder changeStreamMutationBuilder = null;
+
+ // For the current SetCell.
+ @Nullable private String family;
+ @Nullable private ByteString qualifier;
+ private long timestampMicros;
+ @Nullable private ByteString value;
+
+ public DefaultChangeStreamRecordBuilder() {
+ reset();
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public ChangeStreamRecord onHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ Preconditions.checkState(
+ this.changeStreamMutationBuilder == null,
+ "Can not create a Heartbeat when there is an existing ChangeStreamMutation being built.");
+ return Heartbeat.fromProto(heartbeat);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public ChangeStreamRecord onCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ Preconditions.checkState(
+ this.changeStreamMutationBuilder == null,
+ "Can not create a CloseStream when there is an existing ChangeStreamMutation being built.");
+ return CloseStream.fromProto(closeStream);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void startUserMutation(
+ @Nonnull ByteString rowKey,
+ @Nonnull String sourceClusterId,
+ long commitTimestamp,
+ int tieBreaker) {
+ this.changeStreamMutationBuilder =
+ ChangeStreamMutation.createUserMutation(
+ rowKey, sourceClusterId, commitTimestamp, tieBreaker);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void startGcMutation(@Nonnull ByteString rowKey, long commitTimestamp, int tieBreaker) {
+ this.changeStreamMutationBuilder =
+ ChangeStreamMutation.createGcMutation(rowKey, commitTimestamp, tieBreaker);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void deleteFamily(@Nonnull String familyName) {
+ this.changeStreamMutationBuilder.deleteFamily(familyName);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void deleteCells(
+ @Nonnull String familyName,
+ @Nonnull ByteString qualifier,
+ @Nonnull TimestampRange timestampRange) {
+ this.changeStreamMutationBuilder.deleteCells(familyName, qualifier, timestampRange);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void startCell(String family, ByteString qualifier, long timestampMicros) {
+ this.family = family;
+ this.qualifier = qualifier;
+ this.timestampMicros = timestampMicros;
+ this.value = ByteString.EMPTY;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void cellValue(ByteString value) {
+ this.value = this.value.concat(value);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void finishCell() {
+ this.changeStreamMutationBuilder.setCell(
+ this.family, this.qualifier, this.timestampMicros, this.value);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public ChangeStreamRecord finishChangeStreamMutation(
+ @Nonnull String token, long estimatedLowWatermark) {
+ this.changeStreamMutationBuilder.setToken(token);
+ this.changeStreamMutationBuilder.setEstimatedLowWatermark(estimatedLowWatermark);
+ return this.changeStreamMutationBuilder.build();
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void reset() {
+ changeStreamMutationBuilder = null;
+
+ family = null;
+ qualifier = null;
+ timestampMicros = 0;
+ value = null;
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DeleteCells.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DeleteCells.java
new file mode 100644
index 0000000000..26fcdd1083
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DeleteCells.java
@@ -0,0 +1,49 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
+import com.google.protobuf.ByteString;
+import java.io.Serializable;
+import javax.annotation.Nonnull;
+
+/** Representation of a DeleteCells mod in a data change. */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class DeleteCells implements Entry, Serializable {
+ private static final long serialVersionUID = 851772158721462017L;
+
+ public static DeleteCells create(
+ @Nonnull String familyName,
+ @Nonnull ByteString qualifier,
+ @Nonnull TimestampRange timestampRange) {
+ return new AutoValue_DeleteCells(familyName, qualifier, timestampRange);
+ }
+
+ /** Get the column family of the current DeleteCells. */
+ @Nonnull
+ public abstract String getFamilyName();
+
+ /** Get the column qualifier of the current DeleteCells. */
+ @Nonnull
+ public abstract ByteString getQualifier();
+
+ /** Get the timestamp range of the current DeleteCells. */
+ @Nonnull
+ public abstract TimestampRange getTimestampRange();
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DeleteFamily.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DeleteFamily.java
new file mode 100644
index 0000000000..367811c386
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DeleteFamily.java
@@ -0,0 +1,36 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import javax.annotation.Nonnull;
+
+/** Representation of a DeleteFamily mod in a data change. */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class DeleteFamily implements Entry, Serializable {
+ private static final long serialVersionUID = 81806775917145615L;
+
+ public static DeleteFamily create(@Nonnull String familyName) {
+ return new AutoValue_DeleteFamily(familyName);
+ }
+
+ /** Get the column family of the current DeleteFamily. */
+ @Nonnull
+ public abstract String getFamilyName();
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Entry.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Entry.java
new file mode 100644
index 0000000000..44abf53d5f
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Entry.java
@@ -0,0 +1,26 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+
+/**
+ * Default representation of a mod in a data change, which can be a {@link DeleteFamily}, a {@link
+ * DeleteCells}, or a {@link SetCell} This class will be used by {@link ChangeStreamMutation} to
+ * represent a list of mods in a logical change stream mutation.
+ */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+public interface Entry {}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Heartbeat.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Heartbeat.java
new file mode 100644
index 0000000000..2e2b40b327
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Heartbeat.java
@@ -0,0 +1,49 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.protobuf.Timestamp;
+import java.io.Serializable;
+import javax.annotation.Nonnull;
+
+/** A simple wrapper for {@link ReadChangeStreamResponse.Heartbeat}. */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class Heartbeat implements ChangeStreamRecord, Serializable {
+ private static final long serialVersionUID = 7316215828353608504L;
+
+ private static Heartbeat create(
+ ChangeStreamContinuationToken changeStreamContinuationToken,
+ Timestamp estimatedLowWatermark) {
+ return new AutoValue_Heartbeat(changeStreamContinuationToken, estimatedLowWatermark);
+ }
+
+ /** Wraps the protobuf {@link ReadChangeStreamResponse.Heartbeat}. */
+ static Heartbeat fromProto(@Nonnull ReadChangeStreamResponse.Heartbeat heartbeat) {
+ return create(
+ ChangeStreamContinuationToken.fromProto(heartbeat.getContinuationToken()),
+ heartbeat.getEstimatedLowWatermark());
+ }
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public abstract ChangeStreamContinuationToken getChangeStreamContinuationToken();
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public abstract Timestamp getEstimatedLowWatermark();
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java
index 4d7a10ab2a..a3cdff5912 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java
@@ -15,10 +15,13 @@
*/
package com.google.cloud.bigtable.data.v2.models;
+import com.google.api.core.InternalApi;
import com.google.api.core.InternalExtensionOnly;
+import com.google.bigtable.v2.RowRange;
import com.google.common.base.Objects;
import com.google.common.base.Preconditions;
import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
@@ -395,6 +398,22 @@ private void writeObject(ObjectOutputStream output) throws IOException {
output.defaultWriteObject();
}
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public static ByteString serializeToByteString(ByteStringRange byteStringRange) {
+ return RowRange.newBuilder()
+ .setStartKeyClosed(byteStringRange.getStart())
+ .setEndKeyOpen(byteStringRange.getEnd())
+ .build()
+ .toByteString();
+ }
+
+ @InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+ public static ByteStringRange toByteStringRange(ByteString byteString)
+ throws InvalidProtocolBufferException {
+ RowRange rowRange = RowRange.newBuilder().mergeFrom(byteString).build();
+ return ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen());
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) {
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ReadChangeStreamQuery.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ReadChangeStreamQuery.java
new file mode 100644
index 0000000000..e6bfd8c431
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ReadChangeStreamQuery.java
@@ -0,0 +1,271 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationTokens;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Objects;
+import com.google.common.base.Preconditions;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Duration;
+import com.google.protobuf.util.Timestamps;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.util.List;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/** A simple wrapper to construct a query for the ReadChangeStream RPC. */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+public final class ReadChangeStreamQuery implements Serializable, Cloneable {
+ private static final long serialVersionUID = 948588515749969176L;
+
+ private final String tableId;
+ private transient ReadChangeStreamRequest.Builder builder = ReadChangeStreamRequest.newBuilder();
+
+ /**
+ * Constructs a new ReadChangeStreamQuery object for the specified table id. The table id will be
+ * combined with the instance name specified in the {@link
+ * com.google.cloud.bigtable.data.v2.BigtableDataSettings}.
+ */
+ public static ReadChangeStreamQuery create(String tableId) {
+ return new ReadChangeStreamQuery(tableId);
+ }
+
+ private ReadChangeStreamQuery(String tableId) {
+ this.tableId = tableId;
+ }
+
+ private void readObject(ObjectInputStream input) throws IOException, ClassNotFoundException {
+ input.defaultReadObject();
+ builder = ReadChangeStreamRequest.newBuilder().mergeFrom(input);
+ }
+
+ private void writeObject(ObjectOutputStream output) throws IOException {
+ output.defaultWriteObject();
+ builder.build().writeTo(output);
+ }
+
+ /**
+ * Adds a partition.
+ *
+ * @param rowRange Represents the partition in the form [startKey, endKey). startKey can be null
+ * to represent negative infinity. endKey can be null to represent positive infinity.
+ */
+ public ReadChangeStreamQuery streamPartition(@Nonnull RowRange rowRange) {
+ builder.setPartition(StreamPartition.newBuilder().setRowRange(rowRange).build());
+ return this;
+ }
+
+ /**
+ * Adds a partition.
+ *
+ * @param start The beginning of the range (inclusive). Can be null to represent negative
+ * infinity.
+ * @param end The end of the range (exclusive). Can be null to represent positive infinity.
+ */
+ public ReadChangeStreamQuery streamPartition(String start, String end) {
+ return streamPartition(wrapKey(start), wrapKey(end));
+ }
+
+ /**
+ * Adds a partition.
+ *
+ * @param start The beginning of the range (inclusive). Can be null to represent negative
+ * infinity.
+ * @param end The end of the range (exclusive). Can be null to represent positive infinity.
+ */
+ public ReadChangeStreamQuery streamPartition(
+ @Nullable ByteString start, @Nullable ByteString end) {
+ RowRange.Builder rangeBuilder = RowRange.newBuilder();
+ if (start != null) {
+ rangeBuilder.setStartKeyClosed(start);
+ }
+ if (end != null) {
+ rangeBuilder.setEndKeyOpen(end);
+ }
+ return streamPartition(rangeBuilder.build());
+ }
+
+ /** Adds a partition. */
+ public ReadChangeStreamQuery streamPartition(ByteStringRange range) {
+ RowRange.Builder rangeBuilder = RowRange.newBuilder();
+
+ switch (range.getStartBound()) {
+ case OPEN:
+ throw new IllegalStateException("Start bound should be closed.");
+ case CLOSED:
+ rangeBuilder.setStartKeyClosed(range.getStart());
+ break;
+ case UNBOUNDED:
+ rangeBuilder.clearStartKey();
+ break;
+ default:
+ throw new IllegalStateException("Unknown start bound: " + range.getStartBound());
+ }
+
+ switch (range.getEndBound()) {
+ case OPEN:
+ rangeBuilder.setEndKeyOpen(range.getEnd());
+ break;
+ case CLOSED:
+ throw new IllegalStateException("End bound should be open.");
+ case UNBOUNDED:
+ rangeBuilder.clearEndKey();
+ break;
+ default:
+ throw new IllegalStateException("Unknown end bound: " + range.getEndBound());
+ }
+
+ return streamPartition(rangeBuilder.build());
+ }
+
+ /** Sets the startTime(Nanosecond) to read the change stream. */
+ public ReadChangeStreamQuery startTime(long value) {
+ Preconditions.checkState(
+ !builder.hasContinuationTokens(),
+ "startTime and continuationTokens can't be specified together");
+ builder.setStartTime(Timestamps.fromNanos(value));
+ return this;
+ }
+
+ /** Sets the endTime(Nanosecond) to read the change stream. */
+ public ReadChangeStreamQuery endTime(long value) {
+ builder.setEndTime(Timestamps.fromNanos(value));
+ return this;
+ }
+
+ /** Sets the stream continuation tokens to read the change stream. */
+ public ReadChangeStreamQuery continuationTokens(
+ List changeStreamContinuationTokens) {
+ Preconditions.checkState(
+ !builder.hasStartTime(), "startTime and continuationTokens can't be specified together");
+ StreamContinuationTokens.Builder streamContinuationTokensBuilder =
+ StreamContinuationTokens.newBuilder();
+ for (ChangeStreamContinuationToken changeStreamContinuationToken :
+ changeStreamContinuationTokens) {
+ streamContinuationTokensBuilder.addTokens(changeStreamContinuationToken.getTokenProto());
+ }
+ builder.setContinuationTokens(streamContinuationTokensBuilder);
+ return this;
+ }
+
+ /** Sets the heartbeat duration for the change stream. */
+ public ReadChangeStreamQuery heartbeatDuration(java.time.Duration duration) {
+ builder.setHeartbeatDuration(
+ Duration.newBuilder()
+ .setSeconds(duration.getSeconds())
+ .setNanos(duration.getNano())
+ .build());
+ return this;
+ }
+
+ /**
+ * Creates the request protobuf. This method is considered an internal implementation detail and
+ * not meant to be used by applications.
+ */
+ @InternalApi("Used in Changestream beam pipeline.")
+ public ReadChangeStreamRequest toProto(RequestContext requestContext) {
+ String tableName =
+ NameUtil.formatTableName(
+ requestContext.getProjectId(), requestContext.getInstanceId(), tableId);
+
+ return builder
+ .setTableName(tableName)
+ .setAppProfileId(requestContext.getAppProfileId())
+ .build();
+ }
+
+ /**
+ * Wraps the protobuf {@link ReadChangeStreamRequest}.
+ *
+ *
WARNING: Please note that the project id & instance id in the table name will be overwritten
+ * by the configuration in the BigtableDataClient.
+ */
+ public static ReadChangeStreamQuery fromProto(@Nonnull ReadChangeStreamRequest request) {
+ ReadChangeStreamQuery query =
+ new ReadChangeStreamQuery(NameUtil.extractTableIdFromTableName(request.getTableName()));
+ query.builder = request.toBuilder();
+
+ return query;
+ }
+
+ @Override
+ protected ReadChangeStreamQuery clone() {
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.create(tableId);
+ query.builder = this.builder.clone();
+ return query;
+ }
+
+ @Nullable
+ private static ByteString wrapKey(@Nullable String key) {
+ if (key == null) {
+ return null;
+ }
+ return ByteString.copyFromUtf8(key);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ ReadChangeStreamQuery query = (ReadChangeStreamQuery) o;
+ return Objects.equal(tableId, query.tableId)
+ && Objects.equal(builder.getPartition(), query.builder.getPartition())
+ && Objects.equal(builder.getStartTime(), query.builder.getStartTime())
+ && Objects.equal(builder.getEndTime(), query.builder.getEndTime())
+ && Objects.equal(builder.getContinuationTokens(), query.builder.getContinuationTokens())
+ && Objects.equal(builder.getHeartbeatDuration(), query.builder.getHeartbeatDuration());
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hashCode(
+ tableId,
+ builder.getPartition(),
+ builder.getStartTime(),
+ builder.getEndTime(),
+ builder.getContinuationTokens(),
+ builder.getHeartbeatDuration());
+ }
+
+ @Override
+ public String toString() {
+ ReadChangeStreamRequest request = builder.build();
+
+ return MoreObjects.toStringHelper(this)
+ .add("tableId", tableId)
+ .add("partition", request.getPartition())
+ .add("startTime", request.getStartTime())
+ .add("endTime", request.getEndTime())
+ .add("continuationTokens", request.getContinuationTokens())
+ .add("heartbeatDuration", request.getHeartbeatDuration())
+ .toString();
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SetCell.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SetCell.java
new file mode 100644
index 0000000000..92f9b6d386
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SetCell.java
@@ -0,0 +1,56 @@
+/*
+ * 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.models;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMerger;
+import com.google.protobuf.ByteString;
+import java.io.Serializable;
+import javax.annotation.Nonnull;
+
+/**
+ * Representation of a SetCell mod in a data change, whose value is concatenated by {@link
+ * ChangeStreamRecordMerger} in case of SetCell value chunking.
+ */
+@InternalApi("Intended for use by the BigtableIO in apache/beam only.")
+@AutoValue
+public abstract class SetCell implements Entry, Serializable {
+ private static final long serialVersionUID = 77123872266724154L;
+
+ public static SetCell create(
+ @Nonnull String familyName,
+ @Nonnull ByteString qualifier,
+ long timestamp,
+ @Nonnull ByteString value) {
+ return new AutoValue_SetCell(familyName, qualifier, timestamp, value);
+ }
+
+ /** Get the column family of the current SetCell. */
+ @Nonnull
+ public abstract String getFamilyName();
+
+ /** Get the column qualifier of the current SetCell. */
+ @Nonnull
+ public abstract ByteString getQualifier();
+
+ /** Get the timestamp of the current SetCell. */
+ public abstract long getTimestamp();
+
+ /** Get the value of the current SetCell. */
+ @Nonnull
+ public abstract ByteString getValue();
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallable.java
index afc517bbc3..7ea1f90b38 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallable.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallable.java
@@ -26,28 +26,29 @@
/**
* This callable converts the "Received rst stream" exception into a retryable {@link ApiException}.
*/
-final class ConvertExceptionCallable
- extends ServerStreamingCallable {
+final class ConvertExceptionCallable
+ extends ServerStreamingCallable {
- private final ServerStreamingCallable innerCallable;
+ private final ServerStreamingCallable innerCallable;
- public ConvertExceptionCallable(ServerStreamingCallable innerCallable) {
+ public ConvertExceptionCallable(ServerStreamingCallable innerCallable) {
this.innerCallable = innerCallable;
}
@Override
public void call(
- ReadRowsRequest request, ResponseObserver responseObserver, ApiCallContext context) {
- ReadRowsConvertExceptionResponseObserver observer =
- new ReadRowsConvertExceptionResponseObserver<>(responseObserver);
+ RequestT request, ResponseObserver responseObserver, ApiCallContext context) {
+ ConvertExceptionResponseObserver observer =
+ new ConvertExceptionResponseObserver<>(responseObserver);
innerCallable.call(request, observer, context);
}
- private class ReadRowsConvertExceptionResponseObserver extends SafeResponseObserver {
+ private class ConvertExceptionResponseObserver
+ extends SafeResponseObserver {
- private final ResponseObserver outerObserver;
+ private final ResponseObserver outerObserver;
- ReadRowsConvertExceptionResponseObserver(ResponseObserver outerObserver) {
+ ConvertExceptionResponseObserver(ResponseObserver outerObserver) {
super(outerObserver);
this.outerObserver = outerObserver;
}
@@ -58,7 +59,7 @@ protected void onStartImpl(StreamController controller) {
}
@Override
- protected void onResponseImpl(RowT response) {
+ protected void onResponseImpl(ResponseT response) {
outerObserver.onResponse(response);
}
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 e8cec34e84..2b50224957 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
@@ -47,31 +47,46 @@
import com.google.bigtable.v2.BigtableGrpc;
import com.google.bigtable.v2.CheckAndMutateRowRequest;
import com.google.bigtable.v2.CheckAndMutateRowResponse;
+import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsRequest;
+import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse;
import com.google.bigtable.v2.MutateRowRequest;
import com.google.bigtable.v2.MutateRowResponse;
import com.google.bigtable.v2.MutateRowsRequest;
import com.google.bigtable.v2.MutateRowsResponse;
import com.google.bigtable.v2.PingAndWarmRequest;
import com.google.bigtable.v2.PingAndWarmResponse;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
import com.google.bigtable.v2.ReadModifyWriteRowRequest;
import com.google.bigtable.v2.ReadModifyWriteRowResponse;
import com.google.bigtable.v2.ReadRowsRequest;
import com.google.bigtable.v2.ReadRowsResponse;
+import com.google.bigtable.v2.RowRange;
import com.google.bigtable.v2.SampleRowKeysRequest;
import com.google.bigtable.v2.SampleRowKeysResponse;
import com.google.cloud.bigtable.Version;
import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience;
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamMutation;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter;
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
+import com.google.cloud.bigtable.data.v2.models.DefaultChangeStreamRecordAdapter;
import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter;
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.cloud.bigtable.data.v2.models.ReadChangeStreamQuery;
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.google.cloud.bigtable.data.v2.models.RowAdapter;
import com.google.cloud.bigtable.data.v2.models.RowMutation;
import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
+import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMergingCallable;
+import com.google.cloud.bigtable.data.v2.stub.changestream.GenerateInitialChangeStreamPartitionsUserCallable;
+import com.google.cloud.bigtable.data.v2.stub.changestream.ReadChangeStreamResumptionStrategy;
+import com.google.cloud.bigtable.data.v2.stub.changestream.ReadChangeStreamUserCallable;
import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerStreamingCallable;
import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerUnaryCallable;
import com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsTracerFactory;
@@ -146,6 +161,12 @@ public class EnhancedBigtableStub implements AutoCloseable {
private final UnaryCallable readModifyWriteRowCallable;
private final UnaryCallable pingAndWarmCallable;
+ private final ServerStreamingCallable
+ generateInitialChangeStreamPartitionsCallable;
+
+ private final ServerStreamingCallable
+ readChangeStreamCallable;
+
public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
throws IOException {
settings = finalizeSettings(settings, Tags.getTagger(), Stats.getStatsRecorder());
@@ -287,6 +308,10 @@ public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext
bulkMutateRowsCallable = createBulkMutateRowsCallable();
checkAndMutateRowCallable = createCheckAndMutateRowCallable();
readModifyWriteRowCallable = createReadModifyWriteRowCallable();
+ generateInitialChangeStreamPartitionsCallable =
+ createGenerateInitialChangeStreamPartitionsCallable();
+ readChangeStreamCallable =
+ createReadChangeStreamCallable(new DefaultChangeStreamRecordAdapter());
pingAndWarmCallable = createPingAndWarmCallable();
}
@@ -815,6 +840,166 @@ public Map extract(ReadModifyWriteRowRequest request) {
methodName, new ReadModifyWriteRowCallable(retrying, requestContext));
}
+ /**
+ * Creates a callable chain to handle streaming GenerateInitialChangeStreamPartitions RPCs. The
+ * chain will:
+ *
+ *
+ *
Convert a String format tableId into a {@link
+ * GenerateInitialChangeStreamPartitionsRequest} and dispatch the RPC.
+ *
Upon receiving the response stream, it will convert the {@link
+ * com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse}s into {@link
+ * RowRange}.
+ *
+ */
+ private ServerStreamingCallable
+ createGenerateInitialChangeStreamPartitionsCallable() {
+ ServerStreamingCallable<
+ GenerateInitialChangeStreamPartitionsRequest,
+ GenerateInitialChangeStreamPartitionsResponse>
+ base =
+ GrpcRawCallableFactory.createServerStreamingCallable(
+ GrpcCallSettings
+ .
+ newBuilder()
+ .setMethodDescriptor(
+ BigtableGrpc.getGenerateInitialChangeStreamPartitionsMethod())
+ .setParamsExtractor(
+ new RequestParamsExtractor() {
+ @Override
+ public Map extract(
+ GenerateInitialChangeStreamPartitionsRequest
+ generateInitialChangeStreamPartitionsRequest) {
+ return ImmutableMap.of(
+ "table_name",
+ generateInitialChangeStreamPartitionsRequest.getTableName(),
+ "app_profile_id",
+ generateInitialChangeStreamPartitionsRequest.getAppProfileId());
+ }
+ })
+ .build(),
+ settings.generateInitialChangeStreamPartitionsSettings().getRetryableCodes());
+
+ ServerStreamingCallable userCallable =
+ new GenerateInitialChangeStreamPartitionsUserCallable(base, requestContext);
+
+ ServerStreamingCallable withStatsHeaders =
+ new StatsHeadersServerStreamingCallable<>(userCallable);
+
+ // Sometimes GenerateInitialChangeStreamPartitions connections are disconnected via an RST
+ // frame. This error is transient and should be treated similar to UNAVAILABLE. However, this
+ // exception has an INTERNAL error code which by default is not retryable. Convert the exception
+ // so it can be retried in the client.
+ ServerStreamingCallable convertException =
+ new ConvertExceptionCallable<>(withStatsHeaders);
+
+ // Copy idle timeout settings for watchdog.
+ ServerStreamingCallSettings innerSettings =
+ ServerStreamingCallSettings.newBuilder()
+ .setRetryableCodes(
+ settings.generateInitialChangeStreamPartitionsSettings().getRetryableCodes())
+ .setRetrySettings(
+ settings.generateInitialChangeStreamPartitionsSettings().getRetrySettings())
+ .setIdleTimeout(
+ settings.generateInitialChangeStreamPartitionsSettings().getIdleTimeout())
+ .build();
+
+ ServerStreamingCallable watched =
+ Callables.watched(convertException, innerSettings, clientContext);
+
+ ServerStreamingCallable withBigtableTracer =
+ new BigtableTracerStreamingCallable<>(watched);
+
+ ServerStreamingCallable retrying =
+ Callables.retrying(withBigtableTracer, innerSettings, clientContext);
+
+ SpanName span = getSpanName("GenerateInitialChangeStreamPartitions");
+ ServerStreamingCallable traced =
+ new TracedServerStreamingCallable<>(retrying, clientContext.getTracerFactory(), span);
+
+ return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+ }
+
+ /**
+ * Creates a callable chain to handle streaming ReadChangeStream RPCs. The chain will:
+ *
+ *
+ *
Convert a {@link ReadChangeStreamQuery} into a {@link ReadChangeStreamRequest} and
+ * dispatch the RPC.
+ *
Upon receiving the response stream, it will produce a stream of ChangeStreamRecordT. In
+ * case of mutations, it will merge the {@link ReadChangeStreamResponse.DataChange}s into
+ * {@link ChangeStreamMutation}. The actual change stream record implementation can be
+ * configured by the {@code changeStreamRecordAdapter} parameter.
+ *
Retry/resume on failure.
+ *
Add tracing & metrics.
+ *
+ */
+ public
+ ServerStreamingCallable
+ createReadChangeStreamCallable(
+ ChangeStreamRecordAdapter changeStreamRecordAdapter) {
+ ServerStreamingCallable base =
+ GrpcRawCallableFactory.createServerStreamingCallable(
+ GrpcCallSettings.newBuilder()
+ .setMethodDescriptor(BigtableGrpc.getReadChangeStreamMethod())
+ .setParamsExtractor(
+ new RequestParamsExtractor() {
+ @Override
+ public Map extract(
+ ReadChangeStreamRequest readChangeStreamRequest) {
+ return ImmutableMap.of(
+ "table_name", readChangeStreamRequest.getTableName(),
+ "app_profile_id", readChangeStreamRequest.getAppProfileId());
+ }
+ })
+ .build(),
+ settings.readChangeStreamSettings().getRetryableCodes());
+
+ ServerStreamingCallable withStatsHeaders =
+ new StatsHeadersServerStreamingCallable<>(base);
+
+ // Sometimes ReadChangeStream connections are disconnected via an RST frame. This error is
+ // transient and should be treated similar to UNAVAILABLE. However, this exception has an
+ // INTERNAL error code which by default is not retryable. Convert the exception it can be
+ // retried in the client.
+ ServerStreamingCallable convertException =
+ new ConvertExceptionCallable<>(withStatsHeaders);
+
+ ServerStreamingCallable merging =
+ new ChangeStreamRecordMergingCallable<>(convertException, changeStreamRecordAdapter);
+
+ // Copy idle timeout settings for watchdog.
+ ServerStreamingCallSettings innerSettings =
+ ServerStreamingCallSettings.newBuilder()
+ .setResumptionStrategy(
+ new ReadChangeStreamResumptionStrategy<>(changeStreamRecordAdapter))
+ .setRetryableCodes(settings.readChangeStreamSettings().getRetryableCodes())
+ .setRetrySettings(settings.readChangeStreamSettings().getRetrySettings())
+ .setIdleTimeout(settings.readChangeStreamSettings().getIdleTimeout())
+ .build();
+
+ ServerStreamingCallable watched =
+ Callables.watched(merging, innerSettings, clientContext);
+
+ ServerStreamingCallable withBigtableTracer =
+ new BigtableTracerStreamingCallable<>(watched);
+
+ ServerStreamingCallable readChangeStreamCallable =
+ Callables.retrying(withBigtableTracer, innerSettings, clientContext);
+
+ ServerStreamingCallable
+ readChangeStreamUserCallable =
+ new ReadChangeStreamUserCallable<>(readChangeStreamCallable, requestContext);
+
+ SpanName span = getSpanName("ReadChangeStream");
+ ServerStreamingCallable traced =
+ new TracedServerStreamingCallable<>(
+ readChangeStreamUserCallable, clientContext.getTracerFactory(), span);
+
+ return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+ }
+
/**
* Wraps a callable chain in a user presentable callable that will inject the default call context
* and trace the call.
@@ -891,6 +1076,18 @@ public UnaryCallable readModifyWriteRowCallable() {
return readModifyWriteRowCallable;
}
+ /** Returns a streaming generate initial change stream partitions callable */
+ public ServerStreamingCallable
+ generateInitialChangeStreamPartitionsCallable() {
+ return generateInitialChangeStreamPartitionsCallable;
+ }
+
+ /** Returns a streaming read change stream callable. */
+ public ServerStreamingCallable
+ readChangeStreamCallable() {
+ return readChangeStreamCallable;
+ }
+
UnaryCallable pingAndWarmCallable() {
return pingAndWarmCallable;
}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
index c78bdafbf3..b6dd063cb6 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
@@ -34,9 +34,12 @@
import com.google.auth.Credentials;
import com.google.bigtable.v2.PingAndWarmRequest;
import com.google.cloud.bigtable.Version;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord;
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.cloud.bigtable.data.v2.models.ReadChangeStreamQuery;
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.google.cloud.bigtable.data.v2.models.RowMutation;
@@ -140,6 +143,42 @@ public class EnhancedBigtableStubSettings extends StubSettings GENERATE_INITIAL_CHANGE_STREAM_PARTITIONS_RETRY_CODES =
+ ImmutableSet.builder().addAll(IDEMPOTENT_RETRY_CODES).add(Code.ABORTED).build();
+
+ private static final RetrySettings GENERATE_INITIAL_CHANGE_STREAM_PARTITIONS_RETRY_SETTINGS =
+ RetrySettings.newBuilder()
+ .setInitialRetryDelay(Duration.ofMillis(10))
+ .setRetryDelayMultiplier(2.0)
+ .setMaxRetryDelay(Duration.ofMinutes(1))
+ .setMaxAttempts(10)
+ .setJittered(true)
+ .setInitialRpcTimeout(Duration.ofMinutes(1))
+ .setRpcTimeoutMultiplier(2.0)
+ .setMaxRpcTimeout(Duration.ofMinutes(10))
+ .setTotalTimeout(Duration.ofMinutes(60))
+ .build();
+
+ // Allow retrying ABORTED statuses. These will be returned by the server when the client is
+ // too slow to read the change stream records. This makes sense for the java client because
+ // retries happen after the mutation merging logic. Which means that the retry will not be
+ // invoked until the current buffered change stream mutations are consumed.
+ private static final Set READ_CHANGE_STREAM_RETRY_CODES =
+ ImmutableSet.builder().addAll(IDEMPOTENT_RETRY_CODES).add(Code.ABORTED).build();
+
+ private static final RetrySettings READ_CHANGE_STREAM_RETRY_SETTINGS =
+ RetrySettings.newBuilder()
+ .setInitialRetryDelay(Duration.ofMillis(10))
+ .setRetryDelayMultiplier(2.0)
+ .setMaxRetryDelay(Duration.ofMinutes(1))
+ .setMaxAttempts(10)
+ .setJittered(true)
+ .setInitialRpcTimeout(Duration.ofMinutes(5))
+ .setRpcTimeoutMultiplier(2.0)
+ .setMaxRpcTimeout(Duration.ofMinutes(5))
+ .setTotalTimeout(Duration.ofHours(12))
+ .build();
+
/**
* Scopes that are equivalent to JWT's audience.
*
@@ -176,6 +215,10 @@ public class EnhancedBigtableStubSettings extends StubSettings checkAndMutateRowSettings;
private final UnaryCallSettings readModifyWriteRowSettings;
+ private final ServerStreamingCallSettings
+ generateInitialChangeStreamPartitionsSettings;
+ private final ServerStreamingCallSettings
+ readChangeStreamSettings;
private final UnaryCallSettings pingAndWarmSettings;
private EnhancedBigtableStubSettings(Builder builder) {
@@ -212,6 +255,9 @@ private EnhancedBigtableStubSettings(Builder builder) {
bulkReadRowsSettings = builder.bulkReadRowsSettings.build();
checkAndMutateRowSettings = builder.checkAndMutateRowSettings.build();
readModifyWriteRowSettings = builder.readModifyWriteRowSettings.build();
+ generateInitialChangeStreamPartitionsSettings =
+ builder.generateInitialChangeStreamPartitionsSettings.build();
+ readChangeStreamSettings = builder.readChangeStreamSettings.build();
pingAndWarmSettings = builder.pingAndWarmSettings.build();
}
@@ -503,6 +549,16 @@ public UnaryCallSettings readModifyWriteRowSettings() {
return readModifyWriteRowSettings;
}
+ public ServerStreamingCallSettings
+ generateInitialChangeStreamPartitionsSettings() {
+ return generateInitialChangeStreamPartitionsSettings;
+ }
+
+ public ServerStreamingCallSettings
+ readChangeStreamSettings() {
+ return readChangeStreamSettings;
+ }
+
/**
* Returns the object with the settings used for calls to PingAndWarm.
*
@@ -536,6 +592,10 @@ public static class Builder extends StubSettings.Builder
checkAndMutateRowSettings;
private final UnaryCallSettings.Builder readModifyWriteRowSettings;
+ private final ServerStreamingCallSettings.Builder
+ generateInitialChangeStreamPartitionsSettings;
+ private final ServerStreamingCallSettings.Builder
+ readChangeStreamSettings;
private final UnaryCallSettings.Builder pingAndWarmSettings;
/**
@@ -649,6 +709,18 @@ private Builder() {
readModifyWriteRowSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
copyRetrySettings(baseDefaults.readModifyWriteRowSettings(), readModifyWriteRowSettings);
+ generateInitialChangeStreamPartitionsSettings = ServerStreamingCallSettings.newBuilder();
+ generateInitialChangeStreamPartitionsSettings
+ .setRetryableCodes(GENERATE_INITIAL_CHANGE_STREAM_PARTITIONS_RETRY_CODES)
+ .setRetrySettings(GENERATE_INITIAL_CHANGE_STREAM_PARTITIONS_RETRY_SETTINGS)
+ .setIdleTimeout(Duration.ofMinutes(5));
+
+ readChangeStreamSettings = ServerStreamingCallSettings.newBuilder();
+ readChangeStreamSettings
+ .setRetryableCodes(READ_CHANGE_STREAM_RETRY_CODES)
+ .setRetrySettings(READ_CHANGE_STREAM_RETRY_SETTINGS)
+ .setIdleTimeout(Duration.ofMinutes(5));
+
pingAndWarmSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
pingAndWarmSettings.setRetrySettings(
RetrySettings.newBuilder()
@@ -677,6 +749,9 @@ private Builder(EnhancedBigtableStubSettings settings) {
bulkReadRowsSettings = settings.bulkReadRowsSettings.toBuilder();
checkAndMutateRowSettings = settings.checkAndMutateRowSettings.toBuilder();
readModifyWriteRowSettings = settings.readModifyWriteRowSettings.toBuilder();
+ generateInitialChangeStreamPartitionsSettings =
+ settings.generateInitialChangeStreamPartitionsSettings.toBuilder();
+ readChangeStreamSettings = settings.readChangeStreamSettings.toBuilder();
pingAndWarmSettings = settings.pingAndWarmSettings.toBuilder();
}
//
@@ -851,6 +926,20 @@ public UnaryCallSettings.Builder readModifyWriteRowSett
return readModifyWriteRowSettings;
}
+ /** Returns the builder for the settings used for calls to ReadChangeStream. */
+ public ServerStreamingCallSettings.Builder
+ readChangeStreamSettings() {
+ return readChangeStreamSettings;
+ }
+
+ /**
+ * Returns the builder for the settings used for calls to GenerateInitialChangeStreamPartitions.
+ */
+ public ServerStreamingCallSettings.Builder
+ generateInitialChangeStreamPartitionsSettings() {
+ return generateInitialChangeStreamPartitionsSettings;
+ }
+
/** Returns the builder with the settings used for calls to PingAndWarm. */
public UnaryCallSettings.Builder pingAndWarmSettings() {
return pingAndWarmSettings;
@@ -903,6 +992,10 @@ public String toString() {
.add("bulkReadRowsSettings", bulkReadRowsSettings)
.add("checkAndMutateRowSettings", checkAndMutateRowSettings)
.add("readModifyWriteRowSettings", readModifyWriteRowSettings)
+ .add(
+ "generateInitialChangeStreamPartitionsSettings",
+ generateInitialChangeStreamPartitionsSettings)
+ .add("readChangeStreamSettings", readChangeStreamSettings)
.add("pingAndWarmSettings", pingAndWarmSettings)
.add("parent", super.toString())
.toString();
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMerger.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMerger.java
new file mode 100644
index 0000000000..30c6eb94b6
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMerger.java
@@ -0,0 +1,118 @@
+/*
+ * 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.changestream;
+
+import com.google.api.core.InternalApi;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter;
+import com.google.cloud.bigtable.gaxx.reframing.Reframer;
+import com.google.cloud.bigtable.gaxx.reframing.ReframingResponseObserver;
+import com.google.common.base.Preconditions;
+import java.util.ArrayDeque;
+import java.util.Queue;
+
+/**
+ * An implementation of a {@link Reframer} that feeds the change stream record merging {@link
+ * ChangeStreamStateMachine}.
+ *
+ *
{@link ReframingResponseObserver} pushes {@link ReadChangeStreamResponse}s into this class and
+ * pops a change stream record containing one of the following: 1) Heartbeat. 2) CloseStream. 3)
+ * ChangeStreamMutation(a representation of a fully merged logical mutation).
+ *
+ *
This class is considered an internal implementation detail and not meant to be used by
+ * applications.
+ *
+ *
Package-private for internal use.
+ *
+ * @see ReframingResponseObserver for more details
+ */
+@InternalApi
+public class ChangeStreamRecordMerger
+ implements Reframer {
+ private final ChangeStreamStateMachine changeStreamStateMachine;
+ private final Queue changeStreamRecord;
+
+ public ChangeStreamRecordMerger(
+ ChangeStreamRecordAdapter.ChangeStreamRecordBuilder
+ changeStreamRecordBuilder) {
+ changeStreamStateMachine = new ChangeStreamStateMachine<>(changeStreamRecordBuilder);
+ changeStreamRecord = new ArrayDeque<>();
+ }
+
+ @Override
+ public void push(ReadChangeStreamResponse response) {
+ switch (response.getStreamRecordCase()) {
+ case HEARTBEAT:
+ changeStreamStateMachine.handleHeartbeat(response.getHeartbeat());
+ break;
+ case CLOSE_STREAM:
+ changeStreamStateMachine.handleCloseStream(response.getCloseStream());
+ break;
+ case DATA_CHANGE:
+ changeStreamStateMachine.handleDataChange(response.getDataChange());
+ break;
+ case STREAMRECORD_NOT_SET:
+ throw new IllegalStateException("Illegal stream record.");
+ }
+ if (changeStreamStateMachine.hasCompleteChangeStreamRecord()) {
+ changeStreamRecord.add(changeStreamStateMachine.consumeChangeStreamRecord());
+ }
+ }
+
+ @Override
+ public boolean hasFullFrame() {
+ return !changeStreamRecord.isEmpty();
+ }
+
+ @Override
+ public boolean hasPartialFrame() {
+ // Check if buffer in this class contains data. If an assembled is still not available, then
+ // that means `buffer` has been fully consumed. The last place to check is the
+ // ChangeStreamStateMachine buffer, to see if it's holding on to an incomplete change
+ // stream record.
+ return hasFullFrame() || changeStreamStateMachine.isChangeStreamRecordInProgress();
+ }
+
+ @Override
+ public ChangeStreamRecordT pop() {
+ return Preconditions.checkNotNull(
+ changeStreamRecord.poll(),
+ "ChangeStreamRecordMerger.pop() called when there are no change stream records.");
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallable.java
new file mode 100644
index 0000000000..5c6c07451b
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallable.java
@@ -0,0 +1,63 @@
+/*
+ * 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.changestream;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter;
+import com.google.cloud.bigtable.gaxx.reframing.ReframingResponseObserver;
+
+/**
+ * A ServerStreamingCallable that consumes {@link ReadChangeStreamResponse}s and produces change
+ * stream records.
+ *
+ *
This class delegates all the work to gax's {@link ReframingResponseObserver} and the logic to
+ * {@link ChangeStreamRecordMerger}.
+ *
+ *
This class is considered an internal implementation detail and not meant to be used by
+ * applications.
+ */
+@InternalApi
+public class ChangeStreamRecordMergingCallable
+ extends ServerStreamingCallable {
+ private final ServerStreamingCallable inner;
+ private final ChangeStreamRecordAdapter changeStreamRecordAdapter;
+
+ public ChangeStreamRecordMergingCallable(
+ ServerStreamingCallable inner,
+ ChangeStreamRecordAdapter changeStreamRecordAdapter) {
+ this.inner = inner;
+ this.changeStreamRecordAdapter = changeStreamRecordAdapter;
+ }
+
+ @Override
+ public void call(
+ ReadChangeStreamRequest request,
+ ResponseObserver responseObserver,
+ ApiCallContext context) {
+ ChangeStreamRecordAdapter.ChangeStreamRecordBuilder
+ changeStreamRecordBuilder = changeStreamRecordAdapter.createChangeStreamRecordBuilder();
+ ChangeStreamRecordMerger merger =
+ new ChangeStreamRecordMerger<>(changeStreamRecordBuilder);
+ ReframingResponseObserver innerObserver =
+ new ReframingResponseObserver<>(responseObserver, merger);
+ inner.call(request, innerObserver, context);
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamStateMachine.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamStateMachine.java
new file mode 100644
index 0000000000..5190276368
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamStateMachine.java
@@ -0,0 +1,629 @@
+/*
+ * 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.changestream;
+
+import com.google.bigtable.v2.Mutation;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.bigtable.v2.ReadChangeStreamResponse.DataChange.Type;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter.ChangeStreamRecordBuilder;
+import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
+import com.google.common.base.Preconditions;
+import com.google.protobuf.util.Timestamps;
+
+/**
+ * A state machine to produce change stream records from a stream of {@link
+ * ReadChangeStreamResponse}. A change stream record can be a Heartbeat, a CloseStream or a
+ * ChangeStreamMutation.
+ *
+ *
There could be two types of chunking for a ChangeStreamMutation:
+ *
+ *
+ *
Non-SetCell chunking. For example, a ChangeStreamMutation has two mods, DeleteFamily and
+ * DeleteColumn. DeleteFamily is sent in the first {@link ReadChangeStreamResponse} and
+ * DeleteColumn is sent in the second {@link ReadChangeStreamResponse}.
+ *
{@link ReadChangeStreamResponse.MutationChunk} has a chunked {@link
+ * com.google.bigtable.v2.Mutation.SetCell} mutation. For example, a logical mutation has one
+ * big {@link Mutation.SetCell} mutation which is chunked into two {@link
+ * ReadChangeStreamResponse}s. The first {@link ReadChangeStreamResponse.DataChange} has the
+ * first half of the cell value, and the second {@link ReadChangeStreamResponse.DataChange}
+ * has the second half.
+ *
+ *
+ * This state machine handles both types of chunking.
+ *
+ *
Building of the actual change stream record object is delegated to a {@link
+ * ChangeStreamRecordBuilder}. This class is not thread safe.
+ *
+ *
{@link ReadChangeStreamResponse.DataChange}s, that must be merged to a
+ * ChangeStreamMutation.
+ *
ChangeStreamRecord consumption events that reset the state machine for the next change
+ * stream record.
+ *
+ *
+ *
The outputs are:
+ *
+ *
+ *
Heartbeat records.
+ *
CloseStream records.
+ *
ChangeStreamMutation records.
+ *
+ *
+ *
Expected Usage:
+ *
+ *
{@code
+ * ChangeStreamStateMachine changeStreamStateMachine = new ChangeStreamStateMachine<>(myChangeStreamRecordAdapter);
+ * while(responseIterator.hasNext()) {
+ * ReadChangeStreamResponse response = responseIterator.next();
+ * switch (response.getStreamRecordCase()) {
+ * case HEARTBEAT:
+ * changeStreamStateMachine.handleHeartbeat(response.getHeartbeat());
+ * break;
+ * case CLOSE_STREAM:
+ * changeStreamStateMachine.handleCloseStream(response.getCloseStream());
+ * break;
+ * case DATA_CHANGE:
+ * changeStreamStateMachine.handleDataChange(response.getDataChange());
+ * break;
+ * case STREAMRECORD_NOT_SET:
+ * throw new IllegalStateException("Illegal stream record.");
+ * }
+ * if (changeStreamStateMachine.hasCompleteChangeStreamRecord()) {
+ * MyChangeStreamRecord = changeStreamStateMachine.consumeChangeStreamRecord();
+ * // do something with the change stream record.
+ * }
+ * }
+ * }
+ *
+ *
Package-private for internal use.
+ *
+ * @param The type of row the adapter will build
+ */
+final class ChangeStreamStateMachine {
+ private final ChangeStreamRecordBuilder builder;
+ private State currentState;
+ // debug stats
+ private int numHeartbeats = 0;
+ private int numCloseStreams = 0;
+ private int numDataChanges = 0;
+ private int numNonCellMods = 0;
+ private int numCellChunks = 0; // 1 for non-chunked cell.
+ /**
+ * Expected total size of a chunked SetCell value, given by the {@link
+ * ReadChangeStreamResponse.MutationChunk.ChunkInfo}. This value should be the same for all chunks
+ * of a SetCell.
+ */
+ private int expectedTotalSizeOfChunkedSetCell = 0;
+
+ private int actualTotalSizeOfChunkedSetCell = 0;
+ private ChangeStreamRecordT completeChangeStreamRecord;
+
+ /**
+ * Initialize a new state machine that's ready for a new change stream record.
+ *
+ * @param builder The builder that will build the final change stream record.
+ */
+ ChangeStreamStateMachine(ChangeStreamRecordBuilder builder) {
+ this.builder = builder;
+ reset();
+ }
+
+ /**
+ * Handle heartbeat events from the server.
+ *
+ *
+ */
+ void handleCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ try {
+ numCloseStreams++;
+ currentState = currentState.handleCloseStream(closeStream);
+ } catch (RuntimeException e) {
+ currentState = ERROR;
+ throw e;
+ }
+ }
+
+ /**
+ * Feeds a new dataChange into the state machine. If the dataChange is invalid, the state machine
+ * will throw an exception and should not be used for further input.
+ *
+ *
+ *
+ * @param dataChange The new chunk to process.
+ * @throws ChangeStreamStateMachine.InvalidInputException When the chunk is not applicable to the
+ * current state.
+ * @throws IllegalStateException When the internal state is inconsistent
+ */
+ void handleDataChange(ReadChangeStreamResponse.DataChange dataChange) {
+ try {
+ numDataChanges++;
+ currentState = currentState.handleMod(dataChange, 0);
+ } catch (RuntimeException e) {
+ currentState = ERROR;
+ throw e;
+ }
+ }
+
+ /**
+ * Returns the completed change stream record and transitions to {@link
+ * ChangeStreamStateMachine#AWAITING_NEW_STREAM_RECORD}.
+ *
+ * @return The completed change stream record.
+ * @throws IllegalStateException If the last dataChange did not complete a change stream record.
+ */
+ ChangeStreamRecordT consumeChangeStreamRecord() {
+ Preconditions.checkState(
+ completeChangeStreamRecord != null, "No change stream record to consume.");
+ Preconditions.checkState(
+ currentState == AWAITING_STREAM_RECORD_CONSUME,
+ "Change stream record is not ready to consume: " + currentState);
+ ChangeStreamRecordT changeStreamRecord = completeChangeStreamRecord;
+ reset();
+ return changeStreamRecord;
+ }
+
+ /** Checks if there is a complete change stream record to be consumed. */
+ boolean hasCompleteChangeStreamRecord() {
+ return completeChangeStreamRecord != null && currentState == AWAITING_STREAM_RECORD_CONSUME;
+ }
+ /**
+ * Checks if the state machine is in the middle of processing a change stream record.
+ *
+ * @return True If there is a change stream record in progress.
+ */
+ boolean isChangeStreamRecordInProgress() {
+ return currentState != AWAITING_NEW_STREAM_RECORD;
+ }
+
+ private void reset() {
+ currentState = AWAITING_NEW_STREAM_RECORD;
+ numHeartbeats = 0;
+ numCloseStreams = 0;
+ numDataChanges = 0;
+ numNonCellMods = 0;
+ numCellChunks = 0;
+ expectedTotalSizeOfChunkedSetCell = 0;
+ actualTotalSizeOfChunkedSetCell = 0;
+ completeChangeStreamRecord = null;
+
+ builder.reset();
+ }
+
+ /**
+ * Base class for all the state machine's internal states.
+ *
+ *
Each state can consume 3 events: Heartbeat, CloseStream and a Mod. By default, the default
+ * implementation will just throw an IllegalStateException unless the subclass adds explicit
+ * handling for these events.
+ */
+ abstract static class State {
+ /**
+ * Accepts a Heartbeat by the server. And completes the current change stream record.
+ *
+ * @throws IllegalStateException If the subclass can't handle heartbeat events.
+ */
+ ChangeStreamStateMachine.State handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ throw new IllegalStateException();
+ }
+
+ /**
+ * Accepts a CloseStream by the server. And completes the current change stream record.
+ *
+ * @throws IllegalStateException If the subclass can't handle CloseStream events.
+ */
+ ChangeStreamStateMachine.State handleCloseStream(
+ ReadChangeStreamResponse.CloseStream closeStream) {
+ throw new IllegalStateException();
+ }
+
+ /**
+ * Accepts a new mod and transitions to the next state. A mod could be a DeleteFamily, a
+ * DeleteColumn, or a SetCell.
+ *
+ * @param dataChange The DataChange that holds the new mod to process.
+ * @param index The index of the mod in the DataChange.
+ * @return The next state.
+ * @throws IllegalStateException If the subclass can't handle the mod.
+ * @throws ChangeStreamStateMachine.InvalidInputException If the subclass determines that this
+ * dataChange is invalid.
+ */
+ ChangeStreamStateMachine.State handleMod(
+ ReadChangeStreamResponse.DataChange dataChange, int index) {
+ throw new IllegalStateException();
+ }
+ }
+
+ /**
+ * The default state when the state machine is awaiting a ReadChangeStream response to start a new
+ * change stream record. It will notify the builder of the new change stream record and transits
+ * to one of the following states:
+ *
+ *
+ *
{@link ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME}, in case of a Heartbeat
+ * or a CloseStream.
+ *
Same as {@link ChangeStreamStateMachine#AWAITING_NEW_MOD}, depending on the DataChange.
+ *
+ */
+ private final State AWAITING_NEW_STREAM_RECORD =
+ new State() {
+ @Override
+ State handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ validate(
+ completeChangeStreamRecord == null,
+ "AWAITING_NEW_STREAM_RECORD: Existing ChangeStreamRecord not consumed yet.");
+ completeChangeStreamRecord = builder.onHeartbeat(heartbeat);
+ return AWAITING_STREAM_RECORD_CONSUME;
+ }
+
+ @Override
+ State handleCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ validate(
+ completeChangeStreamRecord == null,
+ "AWAITING_NEW_STREAM_RECORD: Existing ChangeStreamRecord not consumed yet.");
+ completeChangeStreamRecord = builder.onCloseStream(closeStream);
+ return AWAITING_STREAM_RECORD_CONSUME;
+ }
+
+ @Override
+ State handleMod(ReadChangeStreamResponse.DataChange dataChange, int index) {
+ validate(
+ completeChangeStreamRecord == null,
+ "AWAITING_NEW_STREAM_RECORD: Existing ChangeStreamRecord not consumed yet.");
+ validate(
+ !dataChange.getRowKey().isEmpty(),
+ "AWAITING_NEW_STREAM_RECORD: First data change missing rowKey.");
+ validate(
+ dataChange.hasCommitTimestamp(),
+ "AWAITING_NEW_STREAM_RECORD: First data change missing commit timestamp.");
+ validate(
+ index == 0,
+ "AWAITING_NEW_STREAM_RECORD: First data change should start with the first mod.");
+ validate(
+ dataChange.getChunksCount() > 0,
+ "AWAITING_NEW_STREAM_RECORD: First data change missing mods.");
+ if (dataChange.getType() == Type.GARBAGE_COLLECTION) {
+ validate(
+ dataChange.getSourceClusterId().isEmpty(),
+ "AWAITING_NEW_STREAM_RECORD: GC mutation shouldn't have source cluster id.");
+ builder.startGcMutation(
+ dataChange.getRowKey(),
+ Timestamps.toNanos(dataChange.getCommitTimestamp()),
+ dataChange.getTiebreaker());
+ } else if (dataChange.getType() == Type.USER) {
+ validate(
+ !dataChange.getSourceClusterId().isEmpty(),
+ "AWAITING_NEW_STREAM_RECORD: User initiated data change missing source cluster id.");
+ builder.startUserMutation(
+ dataChange.getRowKey(),
+ dataChange.getSourceClusterId(),
+ Timestamps.toNanos(dataChange.getCommitTimestamp()),
+ dataChange.getTiebreaker());
+ } else {
+ validate(false, "AWAITING_NEW_STREAM_RECORD: Unexpected type: " + dataChange.getType());
+ }
+ return AWAITING_NEW_MOD.handleMod(dataChange, index);
+ }
+ };
+
+ /**
+ * A state to handle the next Mod.
+ *
+ *
+ *
Valid exit states:
+ *
{@link ChangeStreamStateMachine#AWAITING_NEW_MOD}. Current mod is added, and we have more
+ * mods to expect.
+ *
{@link ChangeStreamStateMachine#AWAITING_CELL_VALUE}. Current mod is the first chunk of a
+ * chunked SetCell.
+ *
{@link ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME}. Current mod is the last
+ * mod of the current logical mutation.
+ *
+ */
+ private final State AWAITING_NEW_MOD =
+ new State() {
+ @Override
+ State handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ throw new IllegalStateException(
+ "AWAITING_NEW_MOD: Can't handle a Heartbeat in the middle of building a ChangeStreamMutation.");
+ }
+
+ @Override
+ State handleCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ throw new IllegalStateException(
+ "AWAITING_NEW_MOD: Can't handle a CloseStream in the middle of building a ChangeStreamMutation.");
+ }
+
+ @Override
+ State handleMod(ReadChangeStreamResponse.DataChange dataChange, int index) {
+ validate(
+ 0 <= index && index <= dataChange.getChunksCount() - 1,
+ "AWAITING_NEW_MOD: Index out of bound.");
+ ReadChangeStreamResponse.MutationChunk chunk = dataChange.getChunks(index);
+ Mutation mod = chunk.getMutation();
+ // Case 1: SetCell
+ if (mod.hasSetCell()) {
+ // Start the Cell and delegate to AWAITING_CELL_VALUE to add the cell value.
+ Mutation.SetCell setCell = chunk.getMutation().getSetCell();
+ if (chunk.hasChunkInfo()) {
+ // If it has chunk info, it must be the first chunk of a chunked SetCell.
+ validate(
+ chunk.getChunkInfo().getChunkedValueOffset() == 0,
+ "AWAITING_NEW_MOD: First chunk of a chunked cell must start with offset==0.");
+ validate(
+ chunk.getChunkInfo().getChunkedValueSize() > 0,
+ "AWAITING_NEW_MOD: First chunk of a chunked cell must have a positive chunked value size.");
+ expectedTotalSizeOfChunkedSetCell = chunk.getChunkInfo().getChunkedValueSize();
+ actualTotalSizeOfChunkedSetCell = 0;
+ }
+ builder.startCell(
+ setCell.getFamilyName(),
+ setCell.getColumnQualifier(),
+ setCell.getTimestampMicros());
+ return AWAITING_CELL_VALUE.handleMod(dataChange, index);
+ }
+ // Case 2: DeleteFamily
+ if (mod.hasDeleteFromFamily()) {
+ numNonCellMods++;
+ builder.deleteFamily(mod.getDeleteFromFamily().getFamilyName());
+ return checkAndFinishMutationIfNeeded(dataChange, index + 1);
+ }
+ // Case 3: DeleteCell
+ if (mod.hasDeleteFromColumn()) {
+ numNonCellMods++;
+ builder.deleteCells(
+ mod.getDeleteFromColumn().getFamilyName(),
+ mod.getDeleteFromColumn().getColumnQualifier(),
+ TimestampRange.create(
+ mod.getDeleteFromColumn().getTimeRange().getStartTimestampMicros(),
+ mod.getDeleteFromColumn().getTimeRange().getEndTimestampMicros()));
+ return checkAndFinishMutationIfNeeded(dataChange, index + 1);
+ }
+ throw new IllegalStateException("AWAITING_NEW_MOD: Unexpected mod type");
+ }
+ };
+
+ /**
+ * A state that represents a cell's value continuation.
+ *
+ *
+ *
Valid exit states:
+ *
{@link ChangeStreamStateMachine#AWAITING_NEW_MOD}. Current chunked SetCell is added, and
+ * we have more mods to expect.
+ *
{@link ChangeStreamStateMachine#AWAITING_CELL_VALUE}. Current chunked SetCell has more
+ * cell values to expect.
+ *
{@link ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME}. Current chunked SetCell
+ * is the last mod of the current logical mutation.
+ *
+ */
+ private final State AWAITING_CELL_VALUE =
+ new State() {
+ @Override
+ State handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ throw new IllegalStateException(
+ "AWAITING_CELL_VALUE: Can't handle a Heartbeat in the middle of building a SetCell.");
+ }
+
+ @Override
+ State handleCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ throw new IllegalStateException(
+ "AWAITING_CELL_VALUE: Can't handle a CloseStream in the middle of building a SetCell.");
+ }
+
+ @Override
+ State handleMod(ReadChangeStreamResponse.DataChange dataChange, int index) {
+ validate(
+ 0 <= index && index <= dataChange.getChunksCount() - 1,
+ "AWAITING_CELL_VALUE: Index out of bound.");
+ ReadChangeStreamResponse.MutationChunk chunk = dataChange.getChunks(index);
+ validate(
+ chunk.getMutation().hasSetCell(),
+ "AWAITING_CELL_VALUE: Current mod is not a SetCell.");
+ Mutation.SetCell setCell = chunk.getMutation().getSetCell();
+ numCellChunks++;
+ builder.cellValue(setCell.getValue());
+ // Case 1: Current SetCell is chunked. For example: [ReadChangeStreamResponse1:
+ // {DeleteColumn, DeleteFamily, SetCell_1}, ReadChangeStreamResponse2: {SetCell_2,
+ // DeleteFamily}].
+ if (chunk.hasChunkInfo()) {
+ validate(
+ chunk.getChunkInfo().getChunkedValueSize() > 0,
+ "AWAITING_CELL_VALUE: Chunked value size must be positive.");
+ validate(
+ chunk.getChunkInfo().getChunkedValueSize() == expectedTotalSizeOfChunkedSetCell,
+ "AWAITING_CELL_VALUE: Chunked value size must be the same for all chunks.");
+ actualTotalSizeOfChunkedSetCell += setCell.getValue().size();
+ // If it's the last chunk of the chunked SetCell, finish the cell.
+ if (chunk.getChunkInfo().getLastChunk()) {
+ builder.finishCell();
+ validate(
+ actualTotalSizeOfChunkedSetCell == expectedTotalSizeOfChunkedSetCell,
+ "Chunked value size in ChunkInfo doesn't match the actual total size. "
+ + "Expected total size: "
+ + expectedTotalSizeOfChunkedSetCell
+ + "; actual total size: "
+ + actualTotalSizeOfChunkedSetCell);
+ return checkAndFinishMutationIfNeeded(dataChange, index + 1);
+ } else {
+ // If this is not the last chunk of a chunked SetCell, then this must be the last mod
+ // of the current response, and we're expecting the rest of the chunked cells in the
+ // following ReadChangeStream response.
+ validate(
+ index == dataChange.getChunksCount() - 1,
+ "AWAITING_CELL_VALUE: Current mod is a chunked SetCell "
+ + "but not the last chunk, but it's not the last mod of the current response.");
+ return AWAITING_CELL_VALUE;
+ }
+ }
+ // Case 2: Current SetCell is not chunked.
+ builder.finishCell();
+ return checkAndFinishMutationIfNeeded(dataChange, index + 1);
+ }
+ };
+
+ /**
+ * A state that represents a completed change stream record. It prevents new change stream records
+ * from being read until the current one has been consumed. The caller is supposed to consume the
+ * change stream record by calling {@link ChangeStreamStateMachine#consumeChangeStreamRecord()}
+ * which will reset the state to {@link ChangeStreamStateMachine#AWAITING_NEW_STREAM_RECORD}.
+ */
+ private final State AWAITING_STREAM_RECORD_CONSUME =
+ new State() {
+ @Override
+ State handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ throw new IllegalStateException(
+ "AWAITING_STREAM_RECORD_CONSUME: Skipping completed change stream record.");
+ }
+
+ @Override
+ State handleCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ throw new IllegalStateException(
+ "AWAITING_STREAM_RECORD_CONSUME: Skipping completed change stream record.");
+ }
+
+ @Override
+ State handleMod(ReadChangeStreamResponse.DataChange dataChange, int index) {
+ throw new IllegalStateException(
+ "AWAITING_STREAM_RECORD_CONSUME: Skipping completed change stream record.");
+ }
+ };
+
+ /**
+ * A state that represents a broken state of the state machine. Any method called on this state
+ * will get an exception.
+ */
+ private final State ERROR =
+ new State() {
+ @Override
+ State handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) {
+ throw new IllegalStateException("ERROR: Failed to handle Heartbeat.");
+ }
+
+ @Override
+ State handleCloseStream(ReadChangeStreamResponse.CloseStream closeStream) {
+ throw new IllegalStateException("ERROR: Failed to handle CloseStream.");
+ }
+
+ @Override
+ State handleMod(ReadChangeStreamResponse.DataChange dataChange, int index) {
+ throw new IllegalStateException("ERROR: Failed to handle DataChange.");
+ }
+ };
+
+ /**
+ * Check if we should continue handling mods in the current DataChange or wrap up. There are 3
+ * cases:
+ *
+ *
+ *
1) index < dataChange.getChunksCount() -> continue to handle the next mod.
+ *
2_1) index == dataChange.getChunksCount() && dataChange.done == true -> current change
+ * stream mutation is complete. Wrap it up and return {@link
+ * ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME}.
+ *
2_2) index == dataChange.getChunksCount() && dataChange.done != true -> current change
+ * stream mutation isn't complete. Return {@link ChangeStreamStateMachine#AWAITING_NEW_MOD}
+ * to wait for more mods in the next ReadChangeStreamResponse.
+ *
+ */
+ private State checkAndFinishMutationIfNeeded(
+ ReadChangeStreamResponse.DataChange dataChange, int index) {
+ validate(
+ 0 <= index && index <= dataChange.getChunksCount(),
+ "checkAndFinishMutationIfNeeded: index out of bound.");
+ // Case 1): Handle the next mod.
+ if (index < dataChange.getChunksCount()) {
+ return AWAITING_NEW_MOD.handleMod(dataChange, index);
+ }
+ // If we reach here, it means that all the mods in this DataChange have been handled. We should
+ // finish up the logical mutation or wait for more mods in the next ReadChangeStreamResponse,
+ // depending on whether the current response is the last response for the logical mutation.
+ if (dataChange.getDone()) {
+ // Case 2_1): Current change stream mutation is complete.
+ validate(!dataChange.getToken().isEmpty(), "Last data change missing token");
+ validate(dataChange.hasEstimatedLowWatermark(), "Last data change missing lowWatermark");
+ completeChangeStreamRecord =
+ builder.finishChangeStreamMutation(
+ dataChange.getToken(), Timestamps.toNanos(dataChange.getEstimatedLowWatermark()));
+ return AWAITING_STREAM_RECORD_CONSUME;
+ }
+ // Case 2_2): The current DataChange itself is chunked, so wait for the next
+ // ReadChangeStreamResponse. Note that we should wait for the new mods instead
+ // of for the new change stream record since the current record hasn't finished yet.
+ return AWAITING_NEW_MOD;
+ }
+
+ private void validate(boolean condition, String message) {
+ if (!condition) {
+ throw new ChangeStreamStateMachine.InvalidInputException(
+ message
+ + ". numHeartbeats: "
+ + numHeartbeats
+ + ", numCloseStreams: "
+ + numCloseStreams
+ + ", numDataChanges: "
+ + numDataChanges
+ + ", numNonCellMods: "
+ + numNonCellMods
+ + ", numCellChunks: "
+ + numCellChunks
+ + ", expectedTotalSizeOfChunkedSetCell: "
+ + expectedTotalSizeOfChunkedSetCell
+ + ", actualTotalSizeOfChunkedSetCell: "
+ + actualTotalSizeOfChunkedSetCell);
+ }
+ }
+
+ static class InvalidInputException extends RuntimeException {
+ InvalidInputException(String message) {
+ super(message);
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/GenerateInitialChangeStreamPartitionsUserCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/GenerateInitialChangeStreamPartitionsUserCallable.java
new file mode 100644
index 0000000000..ce07018c52
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/GenerateInitialChangeStreamPartitionsUserCallable.java
@@ -0,0 +1,98 @@
+/*
+ * Copyright 2018 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.changestream;
+
+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.StreamController;
+import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsRequest;
+import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+
+/**
+ * Simple wrapper for GenerateInitialChangeStreamPartitions to wrap the request and response
+ * protobufs.
+ */
+public class GenerateInitialChangeStreamPartitionsUserCallable
+ extends ServerStreamingCallable {
+ private final RequestContext requestContext;
+ private final ServerStreamingCallable<
+ GenerateInitialChangeStreamPartitionsRequest,
+ GenerateInitialChangeStreamPartitionsResponse>
+ inner;
+
+ public GenerateInitialChangeStreamPartitionsUserCallable(
+ ServerStreamingCallable<
+ GenerateInitialChangeStreamPartitionsRequest,
+ GenerateInitialChangeStreamPartitionsResponse>
+ inner,
+ RequestContext requestContext) {
+ this.requestContext = requestContext;
+ this.inner = inner;
+ }
+
+ @Override
+ public void call(
+ String tableId, ResponseObserver responseObserver, ApiCallContext context) {
+ String tableName =
+ NameUtil.formatTableName(
+ requestContext.getProjectId(), requestContext.getInstanceId(), tableId);
+ GenerateInitialChangeStreamPartitionsRequest request =
+ GenerateInitialChangeStreamPartitionsRequest.newBuilder()
+ .setTableName(tableName)
+ .setAppProfileId(requestContext.getAppProfileId())
+ .build();
+
+ inner.call(request, new ConvertPartitionToRangeObserver(responseObserver), context);
+ }
+
+ private static class ConvertPartitionToRangeObserver
+ implements ResponseObserver {
+
+ private final ResponseObserver outerObserver;
+
+ ConvertPartitionToRangeObserver(ResponseObserver observer) {
+ this.outerObserver = observer;
+ }
+
+ @Override
+ public void onStart(final StreamController controller) {
+ outerObserver.onStart(controller);
+ }
+
+ @Override
+ public void onResponse(GenerateInitialChangeStreamPartitionsResponse response) {
+ ByteStringRange byteStringRange =
+ ByteStringRange.create(
+ response.getPartition().getRowRange().getStartKeyClosed(),
+ response.getPartition().getRowRange().getEndKeyOpen());
+ outerObserver.onResponse(byteStringRange);
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ outerObserver.onError(t);
+ }
+
+ @Override
+ public void onComplete() {
+ outerObserver.onComplete();
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamResumptionStrategy.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamResumptionStrategy.java
new file mode 100644
index 0000000000..660466db95
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamResumptionStrategy.java
@@ -0,0 +1,100 @@
+/*
+ * 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.changestream;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.retrying.StreamResumptionStrategy;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamRequest.Builder;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamContinuationTokens;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter;
+
+/**
+ * An implementation of a {@link StreamResumptionStrategy} for change stream records. This class
+ * tracks the continuation token and upon retry can build a request to resume the stream from where
+ * it left off.
+ *
+ *
This class is considered an internal implementation detail and not meant to be used by
+ * applications.
+ */
+@InternalApi
+public class ReadChangeStreamResumptionStrategy
+ implements StreamResumptionStrategy {
+ private final ChangeStreamRecordAdapter changeStreamRecordAdapter;
+ private String token = null;
+
+ public ReadChangeStreamResumptionStrategy(
+ ChangeStreamRecordAdapter changeStreamRecordAdapter) {
+ this.changeStreamRecordAdapter = changeStreamRecordAdapter;
+ }
+
+ @Override
+ public boolean canResume() {
+ return true;
+ }
+
+ @Override
+ public StreamResumptionStrategy createNew() {
+ return new ReadChangeStreamResumptionStrategy<>(changeStreamRecordAdapter);
+ }
+
+ @Override
+ public ChangeStreamRecordT processResponse(ChangeStreamRecordT response) {
+ // Update the token from a Heartbeat or a ChangeStreamMutation.
+ // We don't worry about resumption after CloseStream, since the server
+ // will return an OK status right after sending a CloseStream.
+ if (changeStreamRecordAdapter.isHeartbeat(response)) {
+ this.token = changeStreamRecordAdapter.getTokenFromHeartbeat(response);
+ } else if (changeStreamRecordAdapter.isChangeStreamMutation(response)) {
+ this.token = changeStreamRecordAdapter.getTokenFromChangeStreamMutation(response);
+ }
+ return response;
+ }
+
+ /**
+ * {@inheritDoc}
+ *
+ *
Given a request, this implementation will narrow that request to include data changes that
+ * come after {@link #token}.
+ */
+ @Override
+ public ReadChangeStreamRequest getResumeRequest(ReadChangeStreamRequest originalRequest) {
+ // A null token means that we have not successfully read a Heartbeat nor a ChangeStreamMutation,
+ // so start from the beginning.
+ if (this.token == null) {
+ return originalRequest;
+ }
+
+ Builder builder = originalRequest.toBuilder();
+ // We need to clear the start_from and use the updated continuation_tokens
+ // to resume the request.
+ // The partition should always be the same as the one from the original request,
+ // otherwise we would receive a CloseStream with different
+ // partitions(which indicates tablet split/merge events).
+ builder.clearStartFrom();
+ builder.setContinuationTokens(
+ StreamContinuationTokens.newBuilder()
+ .addTokens(
+ StreamContinuationToken.newBuilder()
+ .setPartition(originalRequest.getPartition())
+ .setToken(this.token)
+ .build())
+ .build());
+
+ return builder.build();
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamUserCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamUserCallable.java
new file mode 100644
index 0000000000..0c78199ccd
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamUserCallable.java
@@ -0,0 +1,51 @@
+/*
+ * 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.changestream;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.ReadChangeStreamQuery;
+
+/**
+ * A ServerStreamingCallable that converts a {@link ReadChangeStreamQuery} to a {@link
+ * ReadChangeStreamRequest}.
+ */
+@InternalApi("Used in Changestream beam pipeline.")
+public class ReadChangeStreamUserCallable
+ extends ServerStreamingCallable {
+ private final ServerStreamingCallable inner;
+ private final RequestContext requestContext;
+
+ public ReadChangeStreamUserCallable(
+ ServerStreamingCallable inner,
+ RequestContext requestContext) {
+ this.inner = inner;
+ this.requestContext = requestContext;
+ }
+
+ @Override
+ public void call(
+ ReadChangeStreamQuery request,
+ ResponseObserver responseObserver,
+ ApiCallContext context) {
+ ReadChangeStreamRequest innerRequest = request.toProto(requestContext);
+ inner.call(innerRequest, responseObserver, context);
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java
index 34c9a29d71..f4f23085a2 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java
@@ -25,11 +25,14 @@
import com.google.api.gax.rpc.ServerStreamingCallable;
import com.google.api.gax.rpc.UnaryCallable;
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord;
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
import com.google.cloud.bigtable.data.v2.models.Filters.Filter;
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
import com.google.cloud.bigtable.data.v2.models.Mutation;
import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.cloud.bigtable.data.v2.models.ReadChangeStreamQuery;
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.google.cloud.bigtable.data.v2.models.RowCell;
@@ -79,6 +82,14 @@ public class BigtableDataClientTests {
@Mock private Batcher mockBulkMutationBatcher;
@Mock private Batcher mockBulkReadRowsBatcher;
+ @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+ private ServerStreamingCallable
+ mockGenerateInitialChangeStreamPartitionsCallable;
+
+ @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+ private ServerStreamingCallable
+ mockReadChangeStreamCallable;
+
private BigtableDataClient bigtableDataClient;
@Before
@@ -153,6 +164,21 @@ public void proxyReadRowsCallableTest() {
assertThat(bigtableDataClient.readRowsCallable()).isSameInstanceAs(mockReadRowsCallable);
}
+ @Test
+ public void proxyGenerateInitialChangeStreamPartitionsCallableTest() {
+ Mockito.when(mockStub.generateInitialChangeStreamPartitionsCallable())
+ .thenReturn(mockGenerateInitialChangeStreamPartitionsCallable);
+ assertThat(bigtableDataClient.generateInitialChangeStreamPartitionsCallable())
+ .isSameInstanceAs(mockGenerateInitialChangeStreamPartitionsCallable);
+ }
+
+ @Test
+ public void proxyReadChangeStreamCallableTest() {
+ Mockito.when(mockStub.readChangeStreamCallable()).thenReturn(mockReadChangeStreamCallable);
+ assertThat(bigtableDataClient.readChangeStreamCallable())
+ .isSameInstanceAs(mockReadChangeStreamCallable);
+ }
+
@Test
public void proxyReadRowAsyncTest() {
Mockito.when(mockStub.readRowCallable()).thenReturn(mockReadRowCallable);
@@ -300,6 +326,51 @@ public void proxyReadRowsAsyncTest() {
Mockito.verify(mockReadRowsCallable).call(query, mockObserver);
}
+ @Test
+ public void proxyGenerateInitialChangeStreamPartitionsSyncTest() {
+ Mockito.when(mockStub.generateInitialChangeStreamPartitionsCallable())
+ .thenReturn(mockGenerateInitialChangeStreamPartitionsCallable);
+
+ bigtableDataClient.generateInitialChangeStreamPartitions("fake-table");
+
+ Mockito.verify(mockGenerateInitialChangeStreamPartitionsCallable).call("fake-table");
+ }
+
+ @Test
+ public void proxyGenerateInitialChangeStreamPartitionsAsyncTest() {
+ Mockito.when(mockStub.generateInitialChangeStreamPartitionsCallable())
+ .thenReturn(mockGenerateInitialChangeStreamPartitionsCallable);
+
+ @SuppressWarnings("unchecked")
+ ResponseObserver mockObserver = Mockito.mock(ResponseObserver.class);
+ bigtableDataClient.generateInitialChangeStreamPartitionsAsync("fake-table", mockObserver);
+
+ Mockito.verify(mockGenerateInitialChangeStreamPartitionsCallable)
+ .call("fake-table", mockObserver);
+ }
+
+ @Test
+ public void proxyReadChangeStreamSyncTest() {
+ Mockito.when(mockStub.readChangeStreamCallable()).thenReturn(mockReadChangeStreamCallable);
+
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.create("fake-table");
+ bigtableDataClient.readChangeStream(query);
+
+ Mockito.verify(mockReadChangeStreamCallable).call(query);
+ }
+
+ @Test
+ public void proxyReadChangeStreamAsyncTest() {
+ Mockito.when(mockStub.readChangeStreamCallable()).thenReturn(mockReadChangeStreamCallable);
+
+ @SuppressWarnings("unchecked")
+ ResponseObserver mockObserver = Mockito.mock(ResponseObserver.class);
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.create("fake-table");
+ bigtableDataClient.readChangeStreamAsync(query, mockObserver);
+
+ Mockito.verify(mockReadChangeStreamCallable).call(query, mockObserver);
+ }
+
@Test
public void proxySampleRowKeysCallableTest() {
Mockito.when(mockStub.sampleRowKeysCallable()).thenReturn(mockSampleRowKeysCallable);
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamContinuationTokenTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamContinuationTokenTest.java
new file mode 100644
index 0000000000..7e15ad5bbb
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamContinuationTokenTest.java
@@ -0,0 +1,97 @@
+/*
+ * 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.models;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ChangeStreamContinuationTokenTest {
+
+ private final String TOKEN = "token";
+
+ private ByteStringRange createFakeByteStringRange() {
+ return ByteStringRange.create("a", "b");
+ }
+
+ private RowRange rowRangeFromPartition(ByteStringRange partition) {
+ return RowRange.newBuilder()
+ .setStartKeyClosed(partition.getStart())
+ .setEndKeyOpen(partition.getEnd())
+ .build();
+ }
+
+ @Test
+ public void basicTest() throws Exception {
+ ByteStringRange byteStringRange = createFakeByteStringRange();
+ ChangeStreamContinuationToken changeStreamContinuationToken =
+ ChangeStreamContinuationToken.create(byteStringRange, TOKEN);
+ assertThat(changeStreamContinuationToken.getPartition()).isEqualTo(byteStringRange);
+ assertThat(changeStreamContinuationToken.getToken()).isEqualTo(TOKEN);
+
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(changeStreamContinuationToken);
+ oos.close();
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+ ChangeStreamContinuationToken actual = (ChangeStreamContinuationToken) ois.readObject();
+ assertThat(actual).isEqualTo(changeStreamContinuationToken);
+ }
+
+ @Test
+ public void fromProtoTest() {
+ ByteStringRange byteStringRange = createFakeByteStringRange();
+ StreamContinuationToken proto =
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(rowRangeFromPartition(byteStringRange))
+ .build())
+ .setToken(TOKEN)
+ .build();
+ ChangeStreamContinuationToken changeStreamContinuationToken =
+ ChangeStreamContinuationToken.fromProto(proto);
+ assertThat(changeStreamContinuationToken.getPartition()).isEqualTo(byteStringRange);
+ assertThat(changeStreamContinuationToken.getToken()).isEqualTo(TOKEN);
+ assertThat(changeStreamContinuationToken)
+ .isEqualTo(
+ ChangeStreamContinuationToken.fromProto(changeStreamContinuationToken.getTokenProto()));
+ }
+
+ @Test
+ public void toByteStringTest() throws Exception {
+ ByteStringRange byteStringRange = createFakeByteStringRange();
+ ChangeStreamContinuationToken changeStreamContinuationToken =
+ ChangeStreamContinuationToken.create(byteStringRange, TOKEN);
+ assertThat(changeStreamContinuationToken.getPartition()).isEqualTo(byteStringRange);
+ assertThat(changeStreamContinuationToken.getToken()).isEqualTo(TOKEN);
+ assertThat(changeStreamContinuationToken)
+ .isEqualTo(
+ ChangeStreamContinuationToken.fromByteString(
+ changeStreamContinuationToken.toByteString()));
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamMutationTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamMutationTest.java
new file mode 100644
index 0000000000..04285bcc5f
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamMutationTest.java
@@ -0,0 +1,261 @@
+/*
+ * 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.models;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.MutateRowRequest;
+import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.common.primitives.Longs;
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ChangeStreamMutationTest {
+ private static final String PROJECT_ID = "fake-project";
+ private static final String INSTANCE_ID = "fake-instance";
+ private static final String TABLE_ID = "fake-table";
+ private static final String APP_PROFILE_ID = "fake-profile";
+ private static final RequestContext REQUEST_CONTEXT =
+ RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID);
+ private static final long FAKE_COMMIT_TIMESTAMP = 1000L;
+ private static final long FAKE_LOW_WATERMARK = 2000L;
+
+ @Test
+ public void userInitiatedMutationTest() throws IOException, ClassNotFoundException {
+ // Create a user initiated logical mutation.
+ ChangeStreamMutation changeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000,
+ ByteString.copyFromUtf8("fake-value"))
+ .deleteFamily("fake-family")
+ .deleteCells(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Test the getters.
+ assertThat(changeStreamMutation.getRowKey()).isEqualTo(ByteString.copyFromUtf8("key"));
+ assertThat(changeStreamMutation.getType()).isEqualTo(ChangeStreamMutation.MutationType.USER);
+ assertThat(changeStreamMutation.getSourceClusterId()).isEqualTo("fake-source-cluster-id");
+ assertThat(changeStreamMutation.getCommitTimestamp()).isEqualTo(FAKE_COMMIT_TIMESTAMP);
+ assertThat(changeStreamMutation.getTieBreaker()).isEqualTo(0);
+ assertThat(changeStreamMutation.getToken()).isEqualTo("fake-token");
+ assertThat(changeStreamMutation.getEstimatedLowWatermark()).isEqualTo(FAKE_LOW_WATERMARK);
+
+ // Test serialization.
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(changeStreamMutation);
+ oos.close();
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+ ChangeStreamMutation actual = (ChangeStreamMutation) ois.readObject();
+ assertThat(actual).isEqualTo(changeStreamMutation);
+ }
+
+ @Test
+ public void gcMutationTest() throws IOException, ClassNotFoundException {
+ // Create a GC mutation.
+ ChangeStreamMutation changeStreamMutation =
+ ChangeStreamMutation.createGcMutation(
+ ByteString.copyFromUtf8("key"), FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000,
+ ByteString.copyFromUtf8("fake-value"))
+ .deleteFamily("fake-family")
+ .deleteCells(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Test the getters.
+ assertThat(changeStreamMutation.getRowKey()).isEqualTo(ByteString.copyFromUtf8("key"));
+ assertThat(changeStreamMutation.getType())
+ .isEqualTo(ChangeStreamMutation.MutationType.GARBAGE_COLLECTION);
+ Assert.assertTrue(changeStreamMutation.getSourceClusterId().isEmpty());
+ assertThat(changeStreamMutation.getCommitTimestamp()).isEqualTo(FAKE_COMMIT_TIMESTAMP);
+ assertThat(changeStreamMutation.getTieBreaker()).isEqualTo(0);
+ assertThat(changeStreamMutation.getToken()).isEqualTo("fake-token");
+ assertThat(changeStreamMutation.getEstimatedLowWatermark()).isEqualTo(FAKE_LOW_WATERMARK);
+
+ // Test serialization.
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(changeStreamMutation);
+ oos.close();
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+ ChangeStreamMutation actual = (ChangeStreamMutation) ois.readObject();
+ assertThat(actual).isEqualTo(changeStreamMutation);
+ }
+
+ @Test
+ public void toRowMutationTest() {
+ ChangeStreamMutation changeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000,
+ ByteString.copyFromUtf8("fake-value"))
+ .deleteFamily("fake-family")
+ .deleteCells(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Convert it to a rowMutation and construct a MutateRowRequest.
+ RowMutation rowMutation = changeStreamMutation.toRowMutation(TABLE_ID);
+ MutateRowRequest mutateRowRequest = rowMutation.toProto(REQUEST_CONTEXT);
+ String tableName =
+ NameUtil.formatTableName(
+ REQUEST_CONTEXT.getProjectId(), REQUEST_CONTEXT.getInstanceId(), TABLE_ID);
+ assertThat(mutateRowRequest.getTableName()).isEqualTo(tableName);
+ assertThat(mutateRowRequest.getMutationsList()).hasSize(3);
+ assertThat(mutateRowRequest.getMutations(0).getSetCell().getValue())
+ .isEqualTo(ByteString.copyFromUtf8("fake-value"));
+ assertThat(mutateRowRequest.getMutations(1).getDeleteFromFamily().getFamilyName())
+ .isEqualTo("fake-family");
+ assertThat(mutateRowRequest.getMutations(2).getDeleteFromColumn().getFamilyName())
+ .isEqualTo("fake-family");
+ assertThat(mutateRowRequest.getMutations(2).getDeleteFromColumn().getColumnQualifier())
+ .isEqualTo(ByteString.copyFromUtf8("fake-qualifier"));
+ }
+
+ @Test
+ public void toRowMutationWithoutTokenShouldFailTest() {
+ ChangeStreamMutation.Builder builder =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK);
+ Assert.assertThrows(IllegalStateException.class, builder::build);
+ }
+
+ @Test
+ public void toRowMutationWithoutLowWatermarkShouldFailTest() {
+ ChangeStreamMutation.Builder builder =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setToken("fake-token");
+ Assert.assertThrows(IllegalStateException.class, builder::build);
+ }
+
+ @Test
+ public void toRowMutationEntryTest() {
+ ChangeStreamMutation changeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000,
+ ByteString.copyFromUtf8("fake-value"))
+ .deleteFamily("fake-family")
+ .deleteCells(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Convert it to a rowMutationEntry and construct a MutateRowRequest.
+ RowMutationEntry rowMutationEntry = changeStreamMutation.toRowMutationEntry();
+ MutateRowsRequest.Entry mutateRowsRequestEntry = rowMutationEntry.toProto();
+ assertThat(mutateRowsRequestEntry.getRowKey()).isEqualTo(ByteString.copyFromUtf8("key"));
+ assertThat(mutateRowsRequestEntry.getMutationsList()).hasSize(3);
+ assertThat(mutateRowsRequestEntry.getMutations(0).getSetCell().getValue())
+ .isEqualTo(ByteString.copyFromUtf8("fake-value"));
+ assertThat(mutateRowsRequestEntry.getMutations(1).getDeleteFromFamily().getFamilyName())
+ .isEqualTo("fake-family");
+ assertThat(mutateRowsRequestEntry.getMutations(2).getDeleteFromColumn().getFamilyName())
+ .isEqualTo("fake-family");
+ assertThat(mutateRowsRequestEntry.getMutations(2).getDeleteFromColumn().getColumnQualifier())
+ .isEqualTo(ByteString.copyFromUtf8("fake-qualifier"));
+ }
+
+ @Test
+ public void toRowMutationEntryWithoutTokenShouldFailTest() {
+ ChangeStreamMutation.Builder builder =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK);
+ Assert.assertThrows(IllegalStateException.class, builder::build);
+ }
+
+ @Test
+ public void toRowMutationEntryWithoutLowWatermarkShouldFailTest() {
+ ChangeStreamMutation.Builder builder =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setToken("fake-token");
+ Assert.assertThrows(IllegalStateException.class, builder::build);
+ }
+
+ @Test
+ public void testWithLongValue() {
+ ChangeStreamMutation changeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000L,
+ ByteString.copyFrom(Longs.toByteArray(1L)))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ RowMutation rowMutation = changeStreamMutation.toRowMutation(TABLE_ID);
+ MutateRowRequest mutateRowRequest = rowMutation.toProto(REQUEST_CONTEXT);
+ String tableName =
+ NameUtil.formatTableName(
+ REQUEST_CONTEXT.getProjectId(), REQUEST_CONTEXT.getInstanceId(), TABLE_ID);
+ assertThat(mutateRowRequest.getTableName()).isEqualTo(tableName);
+ assertThat(mutateRowRequest.getMutationsList()).hasSize(1);
+ assertThat(mutateRowRequest.getMutations(0).getSetCell().getValue())
+ .isEqualTo(ByteString.copyFromUtf8("\000\000\000\000\000\000\000\001"));
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordTest.java
new file mode 100644
index 0000000000..2637352bd8
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordTest.java
@@ -0,0 +1,171 @@
+/*
+ * 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.models;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import com.google.rpc.Status;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ChangeStreamRecordTest {
+
+ @Test
+ public void heartbeatSerializationTest() throws IOException, ClassNotFoundException {
+ ReadChangeStreamResponse.Heartbeat heartbeatProto =
+ ReadChangeStreamResponse.Heartbeat.newBuilder()
+ .setEstimatedLowWatermark(
+ com.google.protobuf.Timestamp.newBuilder().setSeconds(1000).build())
+ .setContinuationToken(
+ StreamContinuationToken.newBuilder().setToken("random-token").build())
+ .build();
+ Heartbeat heartbeat = Heartbeat.fromProto(heartbeatProto);
+
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(heartbeat);
+ oos.close();
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+ Heartbeat actual = (Heartbeat) ois.readObject();
+ assertThat(actual).isEqualTo(heartbeat);
+ }
+
+ @Test
+ public void closeStreamSerializationTest() throws IOException, ClassNotFoundException {
+ Status status = Status.newBuilder().setCode(0).build();
+ RowRange rowRange1 =
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8(""))
+ .setEndKeyOpen(ByteString.copyFromUtf8("apple"))
+ .build();
+ String token1 = "close-stream-token-1";
+ RowRange rowRange2 =
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("apple"))
+ .setEndKeyOpen(ByteString.copyFromUtf8(""))
+ .build();
+ String token2 = "close-stream-token-2";
+ ReadChangeStreamResponse.CloseStream closeStreamProto =
+ ReadChangeStreamResponse.CloseStream.newBuilder()
+ .addContinuationTokens(
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange1).build())
+ .setToken(token1)
+ .build())
+ .addContinuationTokens(
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
+ .setToken(token2)
+ .build())
+ .setStatus(status)
+ .build();
+ CloseStream closeStream = CloseStream.fromProto(closeStreamProto);
+
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(closeStream);
+ oos.close();
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+ CloseStream actual = (CloseStream) ois.readObject();
+ assertThat(actual.getChangeStreamContinuationTokens())
+ .isEqualTo(closeStream.getChangeStreamContinuationTokens());
+ assertThat(actual.getStatus()).isEqualTo(closeStream.getStatus());
+ }
+
+ @Test
+ public void heartbeatTest() {
+ Timestamp lowWatermark = Timestamp.newBuilder().setSeconds(1000).build();
+ RowRange rowRange =
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("apple"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("banana"))
+ .build();
+ String token = "heartbeat-token";
+ ReadChangeStreamResponse.Heartbeat heartbeatProto =
+ ReadChangeStreamResponse.Heartbeat.newBuilder()
+ .setEstimatedLowWatermark(lowWatermark)
+ .setContinuationToken(
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange).build())
+ .setToken(token)
+ .build())
+ .build();
+ Heartbeat actualHeartbeat = Heartbeat.fromProto(heartbeatProto);
+
+ assertThat(actualHeartbeat.getEstimatedLowWatermark()).isEqualTo(lowWatermark);
+ assertThat(actualHeartbeat.getChangeStreamContinuationToken().getPartition())
+ .isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
+ assertThat(actualHeartbeat.getChangeStreamContinuationToken().getToken()).isEqualTo(token);
+ }
+
+ @Test
+ public void closeStreamTest() {
+ Status status = Status.newBuilder().setCode(0).build();
+ RowRange rowRange1 =
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8(""))
+ .setEndKeyOpen(ByteString.copyFromUtf8("apple"))
+ .build();
+ String token1 = "close-stream-token-1";
+ RowRange rowRange2 =
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("apple"))
+ .setEndKeyOpen(ByteString.copyFromUtf8(""))
+ .build();
+ String token2 = "close-stream-token-2";
+ ReadChangeStreamResponse.CloseStream closeStreamProto =
+ ReadChangeStreamResponse.CloseStream.newBuilder()
+ .addContinuationTokens(
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange1).build())
+ .setToken(token1)
+ .build())
+ .addContinuationTokens(
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange2).build())
+ .setToken(token2)
+ .build())
+ .setStatus(status)
+ .build();
+ CloseStream actualCloseStream = CloseStream.fromProto(closeStreamProto);
+
+ assertThat(status).isEqualTo(actualCloseStream.getStatus());
+ assertThat(actualCloseStream.getChangeStreamContinuationTokens().get(0).getPartition())
+ .isEqualTo(
+ ByteStringRange.create(rowRange1.getStartKeyClosed(), rowRange1.getEndKeyOpen()));
+ assertThat(token1)
+ .isEqualTo(actualCloseStream.getChangeStreamContinuationTokens().get(0).getToken());
+ assertThat(actualCloseStream.getChangeStreamContinuationTokens().get(1).getPartition())
+ .isEqualTo(
+ ByteStringRange.create(rowRange2.getStartKeyClosed(), rowRange2.getEndKeyOpen()));
+ assertThat(token2)
+ .isEqualTo(actualCloseStream.getChangeStreamContinuationTokens().get(1).getToken());
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapterTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapterTest.java
new file mode 100644
index 0000000000..2af99577d6
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapterTest.java
@@ -0,0 +1,449 @@
+/*
+ * 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.models;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.Mutation;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.TimestampRange;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter.ChangeStreamRecordBuilder;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.util.Timestamps;
+import com.google.rpc.Status;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DefaultChangeStreamRecordAdapterTest {
+
+ private final DefaultChangeStreamRecordAdapter adapter = new DefaultChangeStreamRecordAdapter();
+ private ChangeStreamRecordBuilder changeStreamRecordBuilder;
+ private static final long FAKE_COMMIT_TIMESTAMP = 1000L;
+ private static final long FAKE_LOW_WATERMARK = 2000L;
+
+ @Rule public ExpectedException expect = ExpectedException.none();
+
+ @Before
+ public void setUp() {
+ changeStreamRecordBuilder = adapter.createChangeStreamRecordBuilder();
+ }
+
+ @Test
+ public void isHeartbeatTest() {
+ ChangeStreamRecord heartbeatRecord =
+ Heartbeat.fromProto(ReadChangeStreamResponse.Heartbeat.getDefaultInstance());
+ ChangeStreamRecord closeStreamRecord =
+ CloseStream.fromProto(ReadChangeStreamResponse.CloseStream.getDefaultInstance());
+ ChangeStreamRecord changeStreamMutationRecord =
+ ChangeStreamMutation.createGcMutation(
+ ByteString.copyFromUtf8("key"), FAKE_COMMIT_TIMESTAMP, 0)
+ .setToken("token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+ Assert.assertTrue(adapter.isHeartbeat(heartbeatRecord));
+ Assert.assertFalse(adapter.isHeartbeat(closeStreamRecord));
+ Assert.assertFalse(adapter.isHeartbeat(changeStreamMutationRecord));
+ }
+
+ @Test
+ public void getTokenFromHeartbeatTest() {
+ ChangeStreamRecord heartbeatRecord =
+ Heartbeat.fromProto(
+ ReadChangeStreamResponse.Heartbeat.newBuilder()
+ .setEstimatedLowWatermark(Timestamps.fromNanos(FAKE_LOW_WATERMARK))
+ .setContinuationToken(
+ StreamContinuationToken.newBuilder().setToken("heartbeat-token").build())
+ .build());
+ Assert.assertEquals(adapter.getTokenFromHeartbeat(heartbeatRecord), "heartbeat-token");
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void getTokenFromHeartbeatInvalidTypeTest() {
+ ChangeStreamRecord closeStreamRecord =
+ CloseStream.fromProto(ReadChangeStreamResponse.CloseStream.getDefaultInstance());
+ adapter.getTokenFromHeartbeat(closeStreamRecord);
+ expect.expectMessage("record is not a Heartbeat.");
+ }
+
+ @Test
+ public void isChangeStreamMutationTest() {
+ ChangeStreamRecord heartbeatRecord =
+ Heartbeat.fromProto(ReadChangeStreamResponse.Heartbeat.getDefaultInstance());
+ ChangeStreamRecord closeStreamRecord =
+ CloseStream.fromProto(ReadChangeStreamResponse.CloseStream.getDefaultInstance());
+ ChangeStreamRecord changeStreamMutationRecord =
+ ChangeStreamMutation.createGcMutation(
+ ByteString.copyFromUtf8("key"), FAKE_COMMIT_TIMESTAMP, 0)
+ .setToken("token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+ Assert.assertFalse(adapter.isChangeStreamMutation(heartbeatRecord));
+ Assert.assertFalse(adapter.isChangeStreamMutation(closeStreamRecord));
+ Assert.assertTrue(adapter.isChangeStreamMutation(changeStreamMutationRecord));
+ }
+
+ @Test
+ public void getTokenFromChangeStreamMutationTest() {
+ ChangeStreamRecord changeStreamMutationRecord =
+ ChangeStreamMutation.createGcMutation(
+ ByteString.copyFromUtf8("key"), FAKE_COMMIT_TIMESTAMP, 0)
+ .setToken("change-stream-mutation-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+ Assert.assertEquals(
+ adapter.getTokenFromChangeStreamMutation(changeStreamMutationRecord),
+ "change-stream-mutation-token");
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void getTokenFromChangeStreamMutationInvalidTypeTest() {
+ ChangeStreamRecord closeStreamRecord =
+ CloseStream.fromProto(ReadChangeStreamResponse.CloseStream.getDefaultInstance());
+ adapter.getTokenFromChangeStreamMutation(closeStreamRecord);
+ expect.expectMessage("record is not a ChangeStreamMutation.");
+ }
+
+ @Test
+ public void heartbeatTest() {
+ ReadChangeStreamResponse.Heartbeat expectedHeartbeat =
+ ReadChangeStreamResponse.Heartbeat.newBuilder()
+ .setEstimatedLowWatermark(Timestamps.fromNanos(FAKE_LOW_WATERMARK))
+ .setContinuationToken(
+ StreamContinuationToken.newBuilder().setToken("random-token").build())
+ .build();
+ assertThat(changeStreamRecordBuilder.onHeartbeat(expectedHeartbeat))
+ .isEqualTo(Heartbeat.fromProto(expectedHeartbeat));
+ // Call again.
+ assertThat(changeStreamRecordBuilder.onHeartbeat(expectedHeartbeat))
+ .isEqualTo(Heartbeat.fromProto(expectedHeartbeat));
+ }
+
+ @Test
+ public void closeStreamTest() {
+ ReadChangeStreamResponse.CloseStream expectedCloseStream =
+ ReadChangeStreamResponse.CloseStream.newBuilder()
+ .addContinuationTokens(
+ StreamContinuationToken.newBuilder().setToken("random-token").build())
+ .setStatus(Status.newBuilder().setCode(0).build())
+ .build();
+ assertThat(changeStreamRecordBuilder.onCloseStream(expectedCloseStream))
+ .isEqualTo(CloseStream.fromProto(expectedCloseStream));
+ // Call again.
+ assertThat(changeStreamRecordBuilder.onCloseStream(expectedCloseStream))
+ .isEqualTo(CloseStream.fromProto(expectedCloseStream));
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void createHeartbeatWithExistingMutationShouldFailTest() {
+ changeStreamRecordBuilder.startGcMutation(
+ ByteString.copyFromUtf8("key"), FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.onHeartbeat(ReadChangeStreamResponse.Heartbeat.getDefaultInstance());
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void createCloseStreamWithExistingMutationShouldFailTest() {
+ changeStreamRecordBuilder.startGcMutation(
+ ByteString.copyFromUtf8("key"), FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.onCloseStream(
+ ReadChangeStreamResponse.CloseStream.getDefaultInstance());
+ }
+
+ @Test
+ public void singleDeleteFamilyTest() {
+ // Suppose this is the mod we get from the ReadChangeStreamResponse.
+ Mutation.DeleteFromFamily deleteFromFamily =
+ Mutation.DeleteFromFamily.newBuilder().setFamilyName("fake-family").build();
+
+ // Expected logical mutation in the change stream record.
+ ChangeStreamMutation expectedChangeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Create the ChangeStreamMutation through the ChangeStreamRecordBuilder.
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.deleteFamily(deleteFromFamily.getFamilyName());
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ // Call again.
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ }
+
+ @Test
+ public void singleDeleteCellTest() {
+ // Suppose this is the mod we get from the ReadChangeStreamResponse.
+ Mutation.DeleteFromColumn deleteFromColumn =
+ Mutation.DeleteFromColumn.newBuilder()
+ .setFamilyName("fake-family")
+ .setColumnQualifier(ByteString.copyFromUtf8("fake-qualifier"))
+ .setTimeRange(
+ TimestampRange.newBuilder()
+ .setStartTimestampMicros(1000L)
+ .setEndTimestampMicros(2000L)
+ .build())
+ .build();
+
+ // Expected logical mutation in the change stream record.
+ ChangeStreamMutation expectedChangeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteCells(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Create the ChangeStreamMutation through the ChangeStreamRecordBuilder.
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.deleteCells(
+ deleteFromColumn.getFamilyName(),
+ deleteFromColumn.getColumnQualifier(),
+ Range.TimestampRange.create(
+ deleteFromColumn.getTimeRange().getStartTimestampMicros(),
+ deleteFromColumn.getTimeRange().getEndTimestampMicros()));
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ // Call again.
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ }
+
+ @Test
+ public void singleNonChunkedCellTest() {
+ // Expected logical mutation in the change stream record.
+ ChangeStreamMutation expectedChangeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 100L,
+ ByteString.copyFromUtf8("fake-value"))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Create the ChangeStreamMutation through the ChangeStreamRecordBuilder.
+ // Suppose the SetCell is not chunked and the state machine calls `cellValue()` once.
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.startCell(
+ "fake-family", ByteString.copyFromUtf8("fake-qualifier"), 100L);
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("fake-value"));
+ changeStreamRecordBuilder.finishCell();
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ // Call again.
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ }
+
+ @Test
+ public void singleChunkedCellTest() {
+ // Expected logical mutation in the change stream record.
+ ChangeStreamMutation expectedChangeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 100L,
+ ByteString.copyFromUtf8("fake-value1-value2"))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ // Create the ChangeStreamMutation through the ChangeStreamRecordBuilder.
+ // Suppose the SetCell is chunked into two pieces and the state machine calls `cellValue()`
+ // twice.
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.startCell(
+ "fake-family", ByteString.copyFromUtf8("fake-qualifier"), 100L);
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("fake-value1"));
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("-value2"));
+ changeStreamRecordBuilder.finishCell();
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ // Call again.
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ }
+
+ @Test
+ public void multipleChunkedCellsTest() {
+ // Expected logical mutation in the change stream record.
+ ChangeStreamMutation.Builder expectedChangeStreamMutationBuilder =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ for (int i = 0; i < 10; ++i) {
+ expectedChangeStreamMutationBuilder.setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 100L,
+ ByteString.copyFromUtf8(i + "-fake-value1-value2-value3"));
+ }
+ expectedChangeStreamMutationBuilder
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK);
+
+ // Create the ChangeStreamMutation through the ChangeStreamRecordBuilder.
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ for (int i = 0; i < 10; ++i) {
+ changeStreamRecordBuilder.startCell(
+ "fake-family", ByteString.copyFromUtf8("fake-qualifier"), 100L);
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8(i + "-fake-value1"));
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("-value2"));
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("-value3"));
+ changeStreamRecordBuilder.finishCell();
+ }
+ // Check that they're the same.
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutationBuilder.build());
+ // Call again.
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutationBuilder.build());
+ }
+
+ @Test
+ public void multipleDifferentModsTest() {
+ // Expected logical mutation in the change stream record, which contains one DeleteFromFamily,
+ // one non-chunked cell, and one chunked cell.
+ ChangeStreamMutation.Builder expectedChangeStreamMutationBuilder =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 100L,
+ ByteString.copyFromUtf8("non-chunked-value"))
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 100L,
+ ByteString.copyFromUtf8("chunked-value"))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK);
+
+ // Create the ChangeStreamMutation through the ChangeStreamRecordBuilder.
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.deleteFamily("fake-family");
+ // Add non-chunked cell.
+ changeStreamRecordBuilder.startCell(
+ "fake-family", ByteString.copyFromUtf8("fake-qualifier"), 100L);
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("non-chunked-value"));
+ changeStreamRecordBuilder.finishCell();
+ // Add chunked cell.
+ changeStreamRecordBuilder.startCell(
+ "fake-family", ByteString.copyFromUtf8("fake-qualifier"), 100L);
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("chunked"));
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("-value"));
+ changeStreamRecordBuilder.finishCell();
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutationBuilder.build());
+ }
+
+ @Test
+ public void resetTest() {
+ // Build a Heartbeat.
+ ReadChangeStreamResponse.Heartbeat expectedHeartbeat =
+ ReadChangeStreamResponse.Heartbeat.getDefaultInstance();
+ assertThat(changeStreamRecordBuilder.onHeartbeat(expectedHeartbeat))
+ .isEqualTo(Heartbeat.fromProto(expectedHeartbeat));
+
+ // Reset and build a CloseStream.
+ changeStreamRecordBuilder.reset();
+ ReadChangeStreamResponse.CloseStream expectedCloseStream =
+ ReadChangeStreamResponse.CloseStream.getDefaultInstance();
+ assertThat(changeStreamRecordBuilder.onCloseStream(expectedCloseStream))
+ .isEqualTo(CloseStream.fromProto(expectedCloseStream));
+
+ // Reset and build a DeleteFamily.
+ changeStreamRecordBuilder.reset();
+ Mutation deleteFromFamily =
+ Mutation.newBuilder()
+ .setDeleteFromFamily(
+ Mutation.DeleteFromFamily.newBuilder().setFamilyName("fake-family").build())
+ .build();
+ ChangeStreamMutation expectedChangeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .deleteFamily("fake-family")
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.deleteFamily(deleteFromFamily.getDeleteFromFamily().getFamilyName());
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+
+ // Reset a build a cell.
+ changeStreamRecordBuilder.reset();
+ expectedChangeStreamMutation =
+ ChangeStreamMutation.createUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0)
+ .setCell(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 100L,
+ ByteString.copyFromUtf8("fake-value1-value2"))
+ .setToken("fake-token")
+ .setEstimatedLowWatermark(FAKE_LOW_WATERMARK)
+ .build();
+
+ changeStreamRecordBuilder.startUserMutation(
+ ByteString.copyFromUtf8("key"), "fake-source-cluster-id", FAKE_COMMIT_TIMESTAMP, 0);
+ changeStreamRecordBuilder.startCell(
+ "fake-family", ByteString.copyFromUtf8("fake-qualifier"), 100L);
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("fake-value1"));
+ changeStreamRecordBuilder.cellValue(ByteString.copyFromUtf8("-value2"));
+ changeStreamRecordBuilder.finishCell();
+ assertThat(
+ changeStreamRecordBuilder.finishChangeStreamMutation("fake-token", FAKE_LOW_WATERMARK))
+ .isEqualTo(expectedChangeStreamMutation);
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/EntryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/EntryTest.java
new file mode 100644
index 0000000000..748df81af6
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/EntryTest.java
@@ -0,0 +1,101 @@
+/*
+ * 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.models;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class EntryTest {
+ private void validateSerializationRoundTrip(Object obj)
+ throws IOException, ClassNotFoundException {
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(obj);
+ oos.close();
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+ assertThat(ois.readObject()).isEqualTo(obj);
+ }
+
+ @Test
+ public void serializationTest() throws IOException, ClassNotFoundException {
+ // DeleteFamily
+ Entry deleteFamilyEntry = DeleteFamily.create("fake-family");
+ validateSerializationRoundTrip(deleteFamilyEntry);
+
+ // DeleteCell
+ Entry deleteCellsEntry =
+ DeleteCells.create(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L));
+ validateSerializationRoundTrip(deleteCellsEntry);
+
+ // SetCell
+ Entry setCellEntry =
+ SetCell.create(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000,
+ ByteString.copyFromUtf8("fake-value"));
+ validateSerializationRoundTrip(setCellEntry);
+ }
+
+ @Test
+ public void deleteFamilyTest() {
+ Entry deleteFamilyEntry = DeleteFamily.create("fake-family");
+ DeleteFamily deleteFamily = (DeleteFamily) deleteFamilyEntry;
+ assertThat("fake-family").isEqualTo(deleteFamily.getFamilyName());
+ }
+
+ @Test
+ public void deleteCellsTest() {
+ Entry deleteCellEntry =
+ DeleteCells.create(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ Range.TimestampRange.create(1000L, 2000L));
+ DeleteCells deleteCells = (DeleteCells) deleteCellEntry;
+ assertThat("fake-family").isEqualTo(deleteCells.getFamilyName());
+ assertThat(ByteString.copyFromUtf8("fake-qualifier")).isEqualTo(deleteCells.getQualifier());
+ assertThat(Range.TimestampRange.create(1000L, 2000L))
+ .isEqualTo(deleteCells.getTimestampRange());
+ }
+
+ @Test
+ public void setSellTest() {
+ Entry setCellEntry =
+ SetCell.create(
+ "fake-family",
+ ByteString.copyFromUtf8("fake-qualifier"),
+ 1000,
+ ByteString.copyFromUtf8("fake-value"));
+ SetCell setCell = (SetCell) setCellEntry;
+ assertThat("fake-family").isEqualTo(setCell.getFamilyName());
+ assertThat(ByteString.copyFromUtf8("fake-qualifier")).isEqualTo(setCell.getQualifier());
+ assertThat(1000).isEqualTo(setCell.getTimestamp());
+ assertThat(ByteString.copyFromUtf8("fake-value")).isEqualTo(setCell.getValue());
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RangeTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RangeTest.java
index eebdba5811..6f1061f8dc 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RangeTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RangeTest.java
@@ -21,6 +21,7 @@
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
@@ -306,4 +307,14 @@ public void byteStringSerializationTest() throws IOException, ClassNotFoundExcep
ByteStringRange actual = (ByteStringRange) ois.readObject();
assertThat(actual).isEqualTo(expected);
}
+
+ @Test
+ public void byteStringRangeToByteStringTest() throws InvalidProtocolBufferException {
+ ByteStringRange expected = ByteStringRange.create("a", "z");
+
+ ByteString serialized = ByteStringRange.serializeToByteString(expected);
+ ByteStringRange deserialized = ByteStringRange.toByteStringRange(serialized);
+
+ assertThat(expected).isEqualTo(deserialized);
+ }
}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ReadChangeStreamQueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ReadChangeStreamQueryTest.java
new file mode 100644
index 0000000000..cf042e736c
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/ReadChangeStreamQueryTest.java
@@ -0,0 +1,360 @@
+/*
+ * 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.models;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamRequest.Builder;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamContinuationTokens;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Duration;
+import com.google.protobuf.util.Timestamps;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.Collections;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ReadChangeStreamQueryTest {
+ private static final String PROJECT_ID = "fake-project";
+ private static final String INSTANCE_ID = "fake-instance";
+ private static final String TABLE_ID = "fake-table";
+ private static final String APP_PROFILE_ID = "fake-profile-id";
+ private RequestContext requestContext;
+ private static final long FAKE_START_TIME = 1000L;
+ private static final long FAKE_END_TIME = 2000L;
+
+ @Rule public ExpectedException expect = ExpectedException.none();
+
+ @Before
+ public void setUp() {
+ requestContext = RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID);
+ }
+
+ @Test
+ public void requestContextTest() {
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.create(TABLE_ID);
+
+ ReadChangeStreamRequest proto = query.toProto(requestContext);
+ assertThat(proto).isEqualTo(expectedProtoBuilder().build());
+ }
+
+ @Test
+ public void streamPartitionTest() {
+ // Case 1: String.
+ ReadChangeStreamQuery query1 =
+ ReadChangeStreamQuery.create(TABLE_ID).streamPartition("simple-begin", "simple-end");
+ ReadChangeStreamRequest actualProto1 = query1.toProto(requestContext);
+ Builder expectedProto1 = expectedProtoBuilder();
+ expectedProto1.setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("simple-begin"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("simple-end"))
+ .build())
+ .build());
+ assertThat(actualProto1).isEqualTo(expectedProto1.build());
+
+ // Case 2: ByteString.
+ ReadChangeStreamQuery query2 =
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .streamPartition(
+ ByteString.copyFromUtf8("byte-begin"), ByteString.copyFromUtf8("byte-end"));
+ ReadChangeStreamRequest actualProto2 = query2.toProto(requestContext);
+ Builder expectedProto2 = expectedProtoBuilder();
+ expectedProto2.setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("byte-begin"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("byte-end"))
+ .build())
+ .build());
+ assertThat(actualProto2).isEqualTo(expectedProto2.build());
+
+ // Case 3: ByteStringRange.
+ ReadChangeStreamQuery query3 =
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .streamPartition(ByteStringRange.create("range-begin", "range-end"));
+ ReadChangeStreamRequest actualProto3 = query3.toProto(requestContext);
+ Builder expectedProto3 = expectedProtoBuilder();
+ expectedProto3.setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("range-begin"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("range-end"))
+ .build())
+ .build());
+ assertThat(actualProto3).isEqualTo(expectedProto3.build());
+ }
+
+ @Test
+ public void startTimeTest() {
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.create(TABLE_ID).startTime(FAKE_START_TIME);
+
+ Builder expectedProto =
+ expectedProtoBuilder().setStartTime(Timestamps.fromNanos(FAKE_START_TIME));
+
+ ReadChangeStreamRequest actualProto = query.toProto(requestContext);
+ assertThat(actualProto).isEqualTo(expectedProto.build());
+ }
+
+ @Test
+ public void endTimeTest() {
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.create(TABLE_ID).endTime(FAKE_END_TIME);
+
+ Builder expectedProto = expectedProtoBuilder().setEndTime(Timestamps.fromNanos(FAKE_END_TIME));
+
+ ReadChangeStreamRequest actualProto = query.toProto(requestContext);
+ assertThat(actualProto).isEqualTo(expectedProto.build());
+ }
+
+ @Test
+ public void heartbeatDurationTest() {
+ ReadChangeStreamQuery query =
+ ReadChangeStreamQuery.create(TABLE_ID).heartbeatDuration(java.time.Duration.ofSeconds(5));
+
+ Builder expectedProto =
+ expectedProtoBuilder()
+ .setHeartbeatDuration(com.google.protobuf.Duration.newBuilder().setSeconds(5).build());
+
+ ReadChangeStreamRequest actualProto = query.toProto(requestContext);
+ assertThat(actualProto).isEqualTo(expectedProto.build());
+ }
+
+ @Test
+ public void continuationTokensTest() {
+ StreamContinuationToken tokenProto =
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("start"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("end"))
+ .build())
+ .build())
+ .setToken("random-token")
+ .build();
+ ChangeStreamContinuationToken token = ChangeStreamContinuationToken.fromProto(tokenProto);
+ ReadChangeStreamQuery query =
+ ReadChangeStreamQuery.create(TABLE_ID).continuationTokens(Collections.singletonList(token));
+
+ Builder expectedProto =
+ expectedProtoBuilder()
+ .setContinuationTokens(
+ StreamContinuationTokens.newBuilder().addTokens(tokenProto).build());
+
+ ReadChangeStreamRequest actualProto = query.toProto(requestContext);
+ assertThat(actualProto).isEqualTo(expectedProto.build());
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void createWithStartTimeAndContinuationTokensTest() {
+ StreamContinuationToken tokenProto =
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("start"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("end"))
+ .build())
+ .build())
+ .setToken("random-token")
+ .build();
+ ChangeStreamContinuationToken token = ChangeStreamContinuationToken.fromProto(tokenProto);
+ ReadChangeStreamQuery query =
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .startTime(FAKE_START_TIME)
+ .continuationTokens(Collections.singletonList(token));
+ expect.expect(IllegalArgumentException.class);
+ expect.expectMessage("startTime and continuationTokens can't be specified together");
+ }
+
+ @Test
+ public void serializationTest() throws IOException, ClassNotFoundException {
+ StreamContinuationToken tokenProto =
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("start"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("end"))
+ .build())
+ .build())
+ .setToken("random-token")
+ .build();
+ ChangeStreamContinuationToken token = ChangeStreamContinuationToken.fromProto(tokenProto);
+ ReadChangeStreamQuery expected =
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .streamPartition("simple-begin", "simple-end")
+ .continuationTokens(Collections.singletonList(token))
+ .endTime(FAKE_END_TIME)
+ .heartbeatDuration(java.time.Duration.ofSeconds(5));
+
+ ByteArrayOutputStream bos = new ByteArrayOutputStream();
+ ObjectOutputStream oos = new ObjectOutputStream(bos);
+ oos.writeObject(expected);
+ oos.close();
+
+ ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+
+ ReadChangeStreamQuery actual = (ReadChangeStreamQuery) ois.readObject();
+ assertThat(actual.toProto(requestContext)).isEqualTo(expected.toProto(requestContext));
+ }
+
+ private static ReadChangeStreamRequest.Builder expectedProtoBuilder() {
+ return ReadChangeStreamRequest.newBuilder()
+ .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+ .setAppProfileId(APP_PROFILE_ID);
+ }
+
+ @Test
+ public void testFromProto() {
+ StreamContinuationToken token =
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8(""))
+ .setEndKeyOpen(ByteString.copyFromUtf8(""))
+ .build())
+ .build())
+ .setToken("random-token")
+ .build();
+ ReadChangeStreamRequest request =
+ ReadChangeStreamRequest.newBuilder()
+ .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+ .setAppProfileId(APP_PROFILE_ID)
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8(""))
+ .setEndKeyClosed(ByteString.copyFromUtf8(""))
+ .build()))
+ .setContinuationTokens(StreamContinuationTokens.newBuilder().addTokens(token).build())
+ .setEndTime(Timestamps.fromNanos(FAKE_END_TIME))
+ .setHeartbeatDuration(Duration.newBuilder().setSeconds(5).build())
+ .build();
+ ReadChangeStreamQuery query = ReadChangeStreamQuery.fromProto(request);
+ assertThat(query.toProto(requestContext)).isEqualTo(request);
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testFromProtoWithEmptyTableId() {
+ ReadChangeStreamQuery.fromProto(ReadChangeStreamRequest.getDefaultInstance());
+
+ expect.expect(IllegalArgumentException.class);
+ expect.expectMessage("Invalid table name:");
+ }
+
+ @Test
+ public void testEquality() {
+ ReadChangeStreamQuery request =
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .streamPartition("simple-begin", "simple-end")
+ .startTime(FAKE_START_TIME)
+ .endTime(FAKE_END_TIME)
+ .heartbeatDuration(java.time.Duration.ofSeconds(5));
+
+ // ReadChangeStreamQuery#toProto should not change the ReadChangeStreamQuery instance state
+ request.toProto(requestContext);
+ assertThat(request)
+ .isEqualTo(
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .streamPartition("simple-begin", "simple-end")
+ .startTime(FAKE_START_TIME)
+ .endTime(FAKE_END_TIME)
+ .heartbeatDuration(java.time.Duration.ofSeconds(5)));
+
+ assertThat(ReadChangeStreamQuery.create(TABLE_ID).streamPartition("begin-1", "end-1"))
+ .isNotEqualTo(ReadChangeStreamQuery.create(TABLE_ID).streamPartition("begin-2", "end-1"));
+ assertThat(ReadChangeStreamQuery.create(TABLE_ID).startTime(FAKE_START_TIME))
+ .isNotEqualTo(ReadChangeStreamQuery.create(TABLE_ID).startTime(1001L));
+ assertThat(ReadChangeStreamQuery.create(TABLE_ID).endTime(FAKE_END_TIME))
+ .isNotEqualTo(ReadChangeStreamQuery.create(TABLE_ID).endTime(1001L));
+ assertThat(
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .heartbeatDuration(java.time.Duration.ofSeconds(5)))
+ .isNotEqualTo(
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .heartbeatDuration(java.time.Duration.ofSeconds(6)));
+ }
+
+ @Test
+ public void testClone() {
+ StreamContinuationToken tokenProto =
+ StreamContinuationToken.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("start"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("end"))
+ .build())
+ .build())
+ .setToken("random-token")
+ .build();
+ ChangeStreamContinuationToken token = ChangeStreamContinuationToken.fromProto(tokenProto);
+ ReadChangeStreamQuery query =
+ ReadChangeStreamQuery.create(TABLE_ID)
+ .streamPartition("begin", "end")
+ .continuationTokens(Collections.singletonList(token))
+ .endTime(FAKE_END_TIME)
+ .heartbeatDuration(java.time.Duration.ofSeconds(5));
+ ReadChangeStreamRequest request =
+ ReadChangeStreamRequest.newBuilder()
+ .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+ .setAppProfileId(APP_PROFILE_ID)
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8("begin"))
+ .setEndKeyOpen(ByteString.copyFromUtf8("end"))
+ .build()))
+ .setContinuationTokens(
+ StreamContinuationTokens.newBuilder().addTokens(tokenProto).build())
+ .setEndTime(Timestamps.fromNanos(FAKE_END_TIME))
+ .setHeartbeatDuration(Duration.newBuilder().setSeconds(5).build())
+ .build();
+
+ ReadChangeStreamQuery clonedReq = query.clone();
+ assertThat(clonedReq).isEqualTo(query);
+ assertThat(clonedReq.toProto(requestContext)).isEqualTo(request);
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallableTest.java
new file mode 100644
index 0000000000..534d341914
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ConvertExceptionCallableTest.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2018 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;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.ApiException;
+import com.google.api.gax.rpc.InternalException;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ConvertExceptionCallableTest {
+
+ @Test
+ public void rstStreamExceptionConvertedToRetryableTest() {
+ ApiException originalException =
+ new InternalException(
+ new StatusRuntimeException(
+ Status.INTERNAL.withDescription(
+ "INTERNAL: HTTP/2 error code: INTERNAL_ERROR\nReceived Rst Stream")),
+ GrpcStatusCode.of(Status.Code.INTERNAL),
+ false);
+ assertFalse(originalException.isRetryable());
+ SettableExceptionCallable settableExceptionCallable =
+ new SettableExceptionCallable<>(originalException);
+ ConvertExceptionCallable convertStreamExceptionCallable =
+ new ConvertExceptionCallable<>(settableExceptionCallable);
+
+ Throwable actualError = null;
+ try {
+ convertStreamExceptionCallable.all().call("fake-request");
+ } catch (Throwable t) {
+ actualError = t;
+ }
+ assert actualError instanceof InternalException;
+ InternalException actualException = (InternalException) actualError;
+ assertTrue(actualException.isRetryable());
+ }
+
+ private static final class SettableExceptionCallable
+ extends ServerStreamingCallable {
+ private final Throwable throwable;
+
+ public SettableExceptionCallable(Throwable throwable) {
+ this.throwable = throwable;
+ }
+
+ @Override
+ public void call(
+ RequestT request, ResponseObserver responseObserver, ApiCallContext context) {
+ responseObserver.onError(throwable);
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java
index 466355f892..a754421ad9 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java
@@ -645,6 +645,83 @@ public void checkAndMutateRowSettingsAreNotLostTest() {
.isEqualTo(retrySettings);
}
+ @Test
+ public void generateInitialChangeStreamPartitionsSettingsAreNotLostTest() {
+ String dummyProjectId = "my-project";
+ String dummyInstanceId = "my-instance";
+
+ EnhancedBigtableStubSettings.Builder builder =
+ EnhancedBigtableStubSettings.newBuilder()
+ .setProjectId(dummyProjectId)
+ .setInstanceId(dummyInstanceId)
+ .setRefreshingChannel(false);
+
+ RetrySettings retrySettings = RetrySettings.newBuilder().build();
+ builder
+ .generateInitialChangeStreamPartitionsSettings()
+ .setRetryableCodes(Code.ABORTED, Code.DEADLINE_EXCEEDED)
+ .setRetrySettings(retrySettings)
+ .build();
+
+ assertThat(builder.generateInitialChangeStreamPartitionsSettings().getRetryableCodes())
+ .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+ assertThat(builder.generateInitialChangeStreamPartitionsSettings().getRetrySettings())
+ .isEqualTo(retrySettings);
+
+ assertThat(builder.build().generateInitialChangeStreamPartitionsSettings().getRetryableCodes())
+ .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+ assertThat(builder.build().generateInitialChangeStreamPartitionsSettings().getRetrySettings())
+ .isEqualTo(retrySettings);
+
+ assertThat(
+ builder
+ .build()
+ .toBuilder()
+ .generateInitialChangeStreamPartitionsSettings()
+ .getRetryableCodes())
+ .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+ assertThat(
+ builder
+ .build()
+ .toBuilder()
+ .generateInitialChangeStreamPartitionsSettings()
+ .getRetrySettings())
+ .isEqualTo(retrySettings);
+ }
+
+ @Test
+ public void readChangeStreamSettingsAreNotLostTest() {
+ String dummyProjectId = "my-project";
+ String dummyInstanceId = "my-instance";
+
+ EnhancedBigtableStubSettings.Builder builder =
+ EnhancedBigtableStubSettings.newBuilder()
+ .setProjectId(dummyProjectId)
+ .setInstanceId(dummyInstanceId)
+ .setRefreshingChannel(false);
+
+ RetrySettings retrySettings = RetrySettings.newBuilder().build();
+ builder
+ .readChangeStreamSettings()
+ .setRetryableCodes(Code.ABORTED, Code.DEADLINE_EXCEEDED)
+ .setRetrySettings(retrySettings)
+ .build();
+
+ assertThat(builder.readChangeStreamSettings().getRetryableCodes())
+ .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+ assertThat(builder.readChangeStreamSettings().getRetrySettings()).isEqualTo(retrySettings);
+
+ assertThat(builder.build().readChangeStreamSettings().getRetryableCodes())
+ .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+ assertThat(builder.build().readChangeStreamSettings().getRetrySettings())
+ .isEqualTo(retrySettings);
+
+ assertThat(builder.build().toBuilder().readChangeStreamSettings().getRetryableCodes())
+ .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED);
+ assertThat(builder.build().toBuilder().readChangeStreamSettings().getRetrySettings())
+ .isEqualTo(retrySettings);
+ }
+
@Test
public void checkAndMutateRowSettingsAreSane() {
UnaryCallSettings.Builder builder =
@@ -719,6 +796,8 @@ public void isRefreshingChannelFalseValueTest() {
"bulkReadRowsSettings",
"checkAndMutateRowSettings",
"readModifyWriteRowSettings",
+ "generateInitialChangeStreamPartitionsSettings",
+ "readChangeStreamSettings",
"pingAndWarmSettings",
};
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallableTest.java
new file mode 100644
index 0000000000..17849c9250
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamRecordMergingCallableTest.java
@@ -0,0 +1,127 @@
+/*
+ * 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.changestream;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamContinuationToken;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord;
+import com.google.cloud.bigtable.data.v2.models.CloseStream;
+import com.google.cloud.bigtable.data.v2.models.DefaultChangeStreamRecordAdapter;
+import com.google.cloud.bigtable.data.v2.models.Heartbeat;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import com.google.rpc.Status;
+import java.util.Collections;
+import java.util.List;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Additional tests in addition to {@link ReadChangeStreamMergingAcceptanceTest}.
+ *
+ *
All the ChangeStreamMutation tests are in {@link ReadChangeStreamMergingAcceptanceTest}.
+ */
+@RunWith(JUnit4.class)
+public class ChangeStreamRecordMergingCallableTest {
+
+ @Test
+ public void heartbeatTest() {
+ RowRange rowRange = RowRange.newBuilder().getDefaultInstanceForType();
+ ReadChangeStreamResponse.Heartbeat heartbeatProto =
+ ReadChangeStreamResponse.Heartbeat.newBuilder()
+ .setEstimatedLowWatermark(Timestamp.newBuilder().setSeconds(1000).build())
+ .setContinuationToken(
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange))
+ .setToken("random-token")
+ .build())
+ .build();
+ ReadChangeStreamResponse response =
+ ReadChangeStreamResponse.newBuilder().setHeartbeat(heartbeatProto).build();
+ FakeStreamingApi.ServerStreamingStashCallable
+ inner = new ServerStreamingStashCallable<>(Collections.singletonList(response));
+
+ ChangeStreamRecordMergingCallable mergingCallable =
+ new ChangeStreamRecordMergingCallable<>(inner, new DefaultChangeStreamRecordAdapter());
+ List results =
+ mergingCallable.all().call(ReadChangeStreamRequest.getDefaultInstance());
+
+ // Validate the result.
+ assertThat(results.size()).isEqualTo(1);
+ ChangeStreamRecord record = results.get(0);
+ Assert.assertTrue(record instanceof Heartbeat);
+ Heartbeat heartbeat = (Heartbeat) record;
+ assertThat(heartbeat.getChangeStreamContinuationToken().getPartition())
+ .isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
+ assertThat(heartbeat.getChangeStreamContinuationToken().getToken())
+ .isEqualTo(heartbeatProto.getContinuationToken().getToken());
+ assertThat(heartbeat.getEstimatedLowWatermark())
+ .isEqualTo(heartbeatProto.getEstimatedLowWatermark());
+ }
+
+ @Test
+ public void closeStreamTest() {
+ RowRange rowRange =
+ RowRange.newBuilder()
+ .setStartKeyClosed(ByteString.copyFromUtf8(""))
+ .setEndKeyOpen(ByteString.copyFromUtf8(""))
+ .build();
+ StreamContinuationToken streamContinuationToken =
+ StreamContinuationToken.newBuilder()
+ .setPartition(StreamPartition.newBuilder().setRowRange(rowRange).build())
+ .setToken("random-token")
+ .build();
+ ReadChangeStreamResponse.CloseStream closeStreamProto =
+ ReadChangeStreamResponse.CloseStream.newBuilder()
+ .addContinuationTokens(streamContinuationToken)
+ .setStatus(Status.newBuilder().setCode(0).build())
+ .build();
+ ReadChangeStreamResponse response =
+ ReadChangeStreamResponse.newBuilder().setCloseStream(closeStreamProto).build();
+ FakeStreamingApi.ServerStreamingStashCallable
+ inner = new ServerStreamingStashCallable<>(Collections.singletonList(response));
+
+ ChangeStreamRecordMergingCallable mergingCallable =
+ new ChangeStreamRecordMergingCallable<>(inner, new DefaultChangeStreamRecordAdapter());
+ List results =
+ mergingCallable.all().call(ReadChangeStreamRequest.getDefaultInstance());
+
+ // Validate the result.
+ assertThat(results.size()).isEqualTo(1);
+ ChangeStreamRecord record = results.get(0);
+ Assert.assertTrue(record instanceof CloseStream);
+ CloseStream closeStream = (CloseStream) record;
+ assertThat(closeStream.getStatus()).isEqualTo(closeStreamProto.getStatus());
+ assertThat(closeStream.getChangeStreamContinuationTokens().size()).isEqualTo(1);
+ ChangeStreamContinuationToken changeStreamContinuationToken =
+ closeStream.getChangeStreamContinuationTokens().get(0);
+ assertThat(changeStreamContinuationToken.getPartition())
+ .isEqualTo(ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen()));
+ assertThat(changeStreamContinuationToken.getToken())
+ .isEqualTo(streamContinuationToken.getToken());
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamStateMachineTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamStateMachineTest.java
new file mode 100644
index 0000000000..d86df91c35
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ChangeStreamStateMachineTest.java
@@ -0,0 +1,61 @@
+/*
+ * 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.changestream;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord;
+import com.google.cloud.bigtable.data.v2.models.DefaultChangeStreamRecordAdapter;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ChangeStreamStateMachineTest {
+ ChangeStreamStateMachine changeStreamStateMachine;
+
+ @Before
+ public void setUp() throws Exception {
+ changeStreamStateMachine =
+ new ChangeStreamStateMachine<>(
+ new DefaultChangeStreamRecordAdapter().createChangeStreamRecordBuilder());
+ }
+
+ @Test
+ public void testErrorHandlingStats() {
+ ReadChangeStreamResponse.DataChange dataChange =
+ ReadChangeStreamResponse.DataChange.newBuilder().build();
+
+ ChangeStreamStateMachine.InvalidInputException actualError = null;
+ try {
+ changeStreamStateMachine.handleDataChange(dataChange);
+ } catch (ChangeStreamStateMachine.InvalidInputException e) {
+ actualError = e;
+ }
+
+ assertThat(actualError)
+ .hasMessageThat()
+ .containsMatch("AWAITING_NEW_STREAM_RECORD: First data change missing rowKey");
+ assertThat(actualError).hasMessageThat().contains("numHeartbeats: 0");
+ assertThat(actualError).hasMessageThat().contains("numCloseStreams: 0");
+ assertThat(actualError).hasMessageThat().contains("numDataChanges: 1");
+ assertThat(actualError).hasMessageThat().contains("numNonCellMods: 0");
+ assertThat(actualError).hasMessageThat().contains("numCellChunks: 0");
+ assertThat(actualError).hasMessageThat().contains("actualTotalSizeOfChunkedSetCell: 0");
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/GenerateInitialChangeStreamPartitionsUserCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/GenerateInitialChangeStreamPartitionsUserCallableTest.java
new file mode 100644
index 0000000000..885b1c6355
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/GenerateInitialChangeStreamPartitionsUserCallableTest.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright 2018 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.changestream;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsRequest;
+import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.cloud.bigtable.data.v2.internal.NameUtil;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
+import com.google.common.collect.Lists;
+import com.google.common.truth.Truth;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class GenerateInitialChangeStreamPartitionsUserCallableTest {
+ private final RequestContext requestContext =
+ RequestContext.create("my-project", "my-instance", "my-profile");
+
+ @Test
+ public void requestIsCorrect() {
+ FakeStreamingApi.ServerStreamingStashCallable<
+ GenerateInitialChangeStreamPartitionsRequest,
+ GenerateInitialChangeStreamPartitionsResponse>
+ inner = new FakeStreamingApi.ServerStreamingStashCallable<>(Lists.newArrayList());
+ GenerateInitialChangeStreamPartitionsUserCallable
+ generateInitialChangeStreamPartitionsUserCallable =
+ new GenerateInitialChangeStreamPartitionsUserCallable(inner, requestContext);
+
+ generateInitialChangeStreamPartitionsUserCallable.all().call("my-table");
+ assertThat(inner.getActualRequest())
+ .isEqualTo(
+ GenerateInitialChangeStreamPartitionsRequest.newBuilder()
+ .setTableName(
+ NameUtil.formatTableName(
+ requestContext.getProjectId(), requestContext.getInstanceId(), "my-table"))
+ .setAppProfileId(requestContext.getAppProfileId())
+ .build());
+ }
+
+ @Test
+ public void responseIsConverted() {
+ FakeStreamingApi.ServerStreamingStashCallable<
+ GenerateInitialChangeStreamPartitionsRequest,
+ GenerateInitialChangeStreamPartitionsResponse>
+ inner =
+ new FakeStreamingApi.ServerStreamingStashCallable<>(
+ Lists.newArrayList(
+ GenerateInitialChangeStreamPartitionsResponse.newBuilder()
+ .setPartition(
+ StreamPartition.newBuilder()
+ .setRowRange(RowRange.newBuilder().getDefaultInstanceForType())
+ .build())
+ .build()));
+ GenerateInitialChangeStreamPartitionsUserCallable
+ generateInitialChangeStreamPartitionsUserCallable =
+ new GenerateInitialChangeStreamPartitionsUserCallable(inner, requestContext);
+
+ List results =
+ generateInitialChangeStreamPartitionsUserCallable.all().call("my-table");
+ Truth.assertThat(results).containsExactly(ByteStringRange.create("", ""));
+ }
+}
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamMergingAcceptanceTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamMergingAcceptanceTest.java
new file mode 100644
index 0000000000..b745d7ef2d
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamMergingAcceptanceTest.java
@@ -0,0 +1,284 @@
+/*
+ * 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.changestream;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import com.google.api.client.util.Lists;
+import com.google.api.gax.rpc.ServerStream;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.bigtable.v2.Mutation;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamResponse;
+import com.google.bigtable.v2.ReadChangeStreamResponse.DataChange.Type;
+import com.google.bigtable.v2.RowRange;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamPartition;
+import com.google.bigtable.v2.TimestampRange;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamContinuationToken;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamMutation;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord;
+import com.google.cloud.bigtable.data.v2.models.CloseStream;
+import com.google.cloud.bigtable.data.v2.models.DefaultChangeStreamRecordAdapter;
+import com.google.cloud.bigtable.data.v2.models.DeleteCells;
+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.SetCell;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
+import com.google.cloud.conformance.bigtable.v2.ChangeStreamTestDefinition.ChangeStreamTestFile;
+import com.google.cloud.conformance.bigtable.v2.ChangeStreamTestDefinition.ReadChangeStreamTest;
+import com.google.common.base.CaseFormat;
+import com.google.protobuf.util.JsonFormat;
+import com.google.protobuf.util.Timestamps;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+/**
+ * Parses and runs the acceptance tests for read change stream. Currently, this test is only used by
+ * the JAVA library. If in the future we need cross-language support, we should move the test proto
+ * to https://github.com/googleapis/conformance-tests/tree/main/bigtable/v2/proto/google/cloud/conformance/bigtable/v2
+ * and the test data to https://github.com/googleapis/conformance-tests/blob/main/bigtable/v2/changestream.json
+ */
+@RunWith(Parameterized.class)
+public class ReadChangeStreamMergingAcceptanceTest {
+ // Location: `google-cloud-bigtable/src/test/resources/changestream.json`
+ private static final String TEST_DATA_JSON_RESOURCE = "changestream.json";
+
+ private final ReadChangeStreamTest testCase;
+
+ /**
+ * @param testData The serialized test data representing the test case.
+ * @param junitName Not used by the test, but used by the parameterized test runner as the name of
+ * the test.
+ */
+ public ReadChangeStreamMergingAcceptanceTest(
+ ReadChangeStreamTest testData, @SuppressWarnings("unused") String junitName) {
+ this.testCase = testData;
+ }
+
+ // Each tuple consists of [testData: ReadChangeStreamTest, junitName: String]
+ @Parameterized.Parameters(name = "{1}")
+ public static Collection