diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml index 2f7631c873..5e5f8c9733 100644 --- a/google-cloud-bigtable/clirr-ignored-differences.xml +++ b/google-cloud-bigtable/clirr-ignored-differences.xml @@ -39,6 +39,11 @@ 8001 com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerFactory + + + 8001 + com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsConvertExceptionCallable + 8001 diff --git a/google-cloud-bigtable/pom.xml b/google-cloud-bigtable/pom.xml index b6f3eb4c3c..0197aa44db 100644 --- a/google-cloud-bigtable/pom.xml +++ b/google-cloud-bigtable/pom.xml @@ -36,6 +36,12 @@ batch-bigtable.googleapis.com:443 + + + 1.52.1 + 3.21.12 + ${protobuf.version} @@ -593,7 +599,63 @@ + + + kr.motd.maven + os-maven-plugin + 1.6.0 + + + + org.apache.maven.plugins + maven-enforcer-plugin + 3.1.0 + + + enforce-declared-grpc-and-proto-version + + enforce + + + + + + io.grpc:*:[${grpc.version}] + com.google.protobuf:*:[${protobuf.version}] + + + io.grpc:* + com.google.protobuf:* + + + + + + + + + org.xolstice.maven.plugins + protobuf-maven-plugin + 0.6.1 + + + + test-compile + test-compile-custom + + + + + + com.google.protobuf:protoc:${protoc.version}:exe:${os.detected.classifier} + + grpc-java + + io.grpc:protoc-gen-grpc-java:${grpc.version}:exe:${os.detected.classifier} + + + org.codehaus.mojo build-helper-maven-plugin diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index ce9a57fa7e..968ebaef26 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -30,11 +30,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; 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.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; @@ -1489,6 +1492,298 @@ public UnaryCallable readModifyWriteRowCallable() { return stub.readModifyWriteRowCallable(); } + /** + * Convenience method for synchronously streaming the partitions of a table. The returned + * ServerStream instance is not threadsafe, it can only be used from single thread. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   try {
+   *     ServerStream stream = bigtableDataClient.generateInitialChangeStreamPartitions(tableId);
+   *     int count = 0;
+   *
+   *     // Iterator style
+   *     for (ByteStringRange partition : stream) {
+   *       if (++count > 10) {
+   *         stream.cancel();
+   *         break;
+   *       }
+   *       // Do something with partition
+   *     }
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   * }
+   * }
+ * + * @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. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   bigtableDataClient.generateInitialChangeStreamPartitionsAsync(tableId, new ResponseObserver() {
+   *     StreamController controller;
+   *     int count = 0;
+   *
+   *     public void onStart(StreamController controller) {
+   *       this.controller = controller;
+   *     }
+   *     public void onResponse(ByteStringRange partition) {
+   *       if (++count > 10) {
+   *         controller.cancel();
+   *         return;
+   *       }
+   *       // Do something with partition
+   *     }
+   *     public void onError(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to read a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onComplete() {
+   *       // Handle stream completion
+   *     }
+   *   });
+   * }
+   * }
+ */ + @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. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   // Iterator style
+   *   try {
+   *     for(ByteStringRange partition : bigtableDataClient.generateInitialChangeStreamPartitionsCallable().call(tableId)) {
+   *       // Do something with partition
+   *     }
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   *
+   *   // Sync style
+   *   try {
+   *     List partitions = bigtableDataClient.generateInitialChangeStreamPartitionsCallable().all().call(tableId);
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   *
+   *   // Point look up
+   *   ApiFuture partitionFuture =
+   *     bigtableDataClient.generateInitialChangeStreamPartitionsCallable().first().futureCall(tableId);
+   *
+   *   ApiFutures.addCallback(partitionFuture, new ApiFutureCallback() {
+   *     public void onFailure(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to read a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onSuccess(RowRange result) {
+   *       System.out.println("Got partition: " + result);
+   *     }
+   *   }, MoreExecutors.directExecutor());
+   *
+   *   // etc
+   * }
+   * }
+ * + * @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. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   ReadChangeStreamQuery query = ReadChangeStreamQuery.create(tableId)
+   *          .streamPartition("START_KEY", "END_KEY")
+   *          .startTime(Timestamp.newBuilder().setSeconds(100).build());
+   *
+   *   try {
+   *     ServerStream stream = bigtableDataClient.readChangeStream(query);
+   *     int count = 0;
+   *
+   *     // Iterator style
+   *     for (ChangeStreamRecord record : stream) {
+   *       if (++count > 10) {
+   *         stream.cancel();
+   *         break;
+   *       }
+   *       // Do something with the change stream record.
+   *     }
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   * }
+   * }
+ * + * @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}. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   ReadChangeStreamQuery query = ReadChangeStreamQuery.create(tableId)
+   *          .streamPartition("START_KEY", "END_KEY")
+   *          .startTime(Timestamp.newBuilder().setSeconds(100).build());
+   *
+   *   bigtableDataClient.readChangeStreamAsync(query, new ResponseObserver() {
+   *     StreamController controller;
+   *     int count = 0;
+   *
+   *     public void onStart(StreamController controller) {
+   *       this.controller = controller;
+   *     }
+   *     public void onResponse(ChangeStreamRecord record) {
+   *       if (++count > 10) {
+   *         controller.cancel();
+   *         return;
+   *       }
+   *       // Do something with the change stream record.
+   *     }
+   *     public void onError(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to read a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onComplete() {
+   *       // Handle stream completion
+   *     }
+   *   });
+   * }
+   * }
+ */ + @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. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE]";
+   *
+   *   ReadChangeStreamQuery query = ReadChangeStreamQuery.create(tableId)
+   *          .streamPartition("START_KEY", "END_KEY")
+   *          .startTime(Timestamp.newBuilder().setSeconds(100).build());
+   *
+   *   // Iterator style
+   *   try {
+   *     for(ChangeStreamRecord record : bigtableDataClient.readChangeStreamCallable().call(query)) {
+   *       // Do something with record
+   *     }
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   *
+   *   // Sync style
+   *   try {
+   *     List records = bigtableDataClient.readChangeStreamCallable().all().call(query);
+   *   } catch (NotFoundException e) {
+   *     System.out.println("Tried to read a non-existent table");
+   *   } catch (RuntimeException e) {
+   *     e.printStackTrace();
+   *   }
+   *
+   *   // Point look up
+   *   ApiFuture recordFuture =
+   *     bigtableDataClient.readChangeStreamCallable().first().futureCall(query);
+   *
+   *   ApiFutures.addCallback(recordFuture, new ApiFutureCallback() {
+   *     public void onFailure(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to read a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onSuccess(ChangeStreamRecord result) {
+   *       System.out.println("Got record: " + result);
+   *     }
+   *   }, MoreExecutors.directExecutor());
+   *
+   *   // etc
+   * }
+   * }
+ * + * @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. + * + *

Case 1) User initiated mutation. + * + *

{@code
+ * ChangeStreamMutation.Builder builder = ChangeStreamMutation.createUserMutation(...);
+ * builder.setCell(...);
+ * builder.deleteFamily(...);
+ * builder.deleteCells(...);
+ * ChangeStreamMutation changeStreamMutation = builder.setToken(...).setEstimatedLowWatermark().build();
+ * }
+ * + * Case 2) Garbage Collection mutation. + * + *
{@code
+ * ChangeStreamMutation.Builder builder = ChangeStreamMutation.createGcMutation(...);
+ * builder.setCell(...);
+ * builder.deleteFamily(...);
+ * builder.deleteCells(...);
+ * ChangeStreamMutation changeStreamMutation = builder.setToken(...).setEstimatedLowWatermark().build();
+ * }
+ */ +@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 + *
  1. Exactly 1 {@code onHeartbeat}. + *
+ * + *
    + * Case 2: CloseStream + *
  1. 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: + *
  1. DeleteFamily -> Exactly 1 {@code deleteFamily} + *
  2. DeleteCell -> Exactly 1 {@code deleteCell} + *
  3. SetCell -> Exactly 1 {@code startCell}, At least 1 {@code CellValue}, Exactly 1 {@code + * finishCell}. + *
+ * + *

The whole flow of constructing a ChangeStreamMutation is: + * + *

    + *
  1. Exactly 1 {@code startUserMutation} or {@code startGcMutation}. + *
  2. At least 1 DeleteFamily/DeleteCell/SetCell mods. + *
  3. 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: + *
  1. Exactly 1 {@code startCell}. + *
  2. Exactly 1 {@code cellValue}. + *
  3. Exactly 1 {@code finishCell}. + *
+ * + *
    + * In case of a chunked cell, the following order is guaranteed: + *
  1. Exactly 1 {@code startCell}. + *
  2. At least 2 {@code cellValue}. + *
  3. 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). + * + *

Example usage: + * + *

{@code
+ * ChangeStreamRecordMerger changeStreamRecordMerger =
+ *     new ChangeStreamRecordMerger<>(myChangeStreamRecordAdaptor);
+ *
+ * while(responseIterator.hasNext()) {
+ *   ReadChangeStreamResponse response = responseIterator.next();
+ *
+ *   if (changeStreamRecordMerger.hasFullFrame()) {
+ *     ChangeStreamRecord changeStreamRecord = changeStreamRecordMerger.pop();
+ *     // Do something with change stream record.
+ *   } else {
+ *     changeStreamRecordMerger.push(response);
+ *   }
+ * }
+ *
+ * if (changeStreamRecordMerger.hasPartialFrame()) {
+ *   throw new RuntimeException("Incomplete stream");
+ * }
+ *
+ * }
+ * + *

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. + * + *

The inputs are: + * + *

    + *
  • {@link ReadChangeStreamResponse.Heartbeat}s. + *
  • {@link ReadChangeStreamResponse.CloseStream}s. + *
  • {@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. + * + *

+ *
Valid states: + *
{@link ChangeStreamStateMachine#AWAITING_NEW_STREAM_RECORD} + *
Resulting states: + *
{@link ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME} + *
+ */ + void handleHeartbeat(ReadChangeStreamResponse.Heartbeat heartbeat) { + try { + numHeartbeats++; + currentState = currentState.handleHeartbeat(heartbeat); + } catch (RuntimeException e) { + currentState = ERROR; + throw e; + } + } + + /** + * Handle CloseStream events from the server. + * + *
+ *
Valid states: + *
{@link ChangeStreamStateMachine#AWAITING_NEW_STREAM_RECORD} + *
Resulting states: + *
{@link ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME} + *
+ */ + 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. + * + *
+ *
Valid states: + *
{@link ChangeStreamStateMachine#AWAITING_NEW_STREAM_RECORD} + *
{@link ChangeStreamStateMachine#AWAITING_NEW_MOD} + *
{@link ChangeStreamStateMachine#AWAITING_CELL_VALUE} + *
Resulting states: + *
{@link ChangeStreamStateMachine#AWAITING_NEW_MOD} + *
{@link ChangeStreamStateMachine#AWAITING_CELL_VALUE} + *
{@link ChangeStreamStateMachine#AWAITING_STREAM_RECORD_CONSUME} + *
+ * + * @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 data() throws IOException { + ClassLoader cl = Thread.currentThread().getContextClassLoader(); + InputStream dataJson = cl.getResourceAsStream(TEST_DATA_JSON_RESOURCE); + assertWithMessage("Unable to load test definition: %s", TEST_DATA_JSON_RESOURCE) + .that(dataJson) + .isNotNull(); + + InputStreamReader reader = new InputStreamReader(dataJson); + ChangeStreamTestFile.Builder testBuilder = ChangeStreamTestFile.newBuilder(); + JsonFormat.parser().merge(reader, testBuilder); + ChangeStreamTestFile testDefinition = testBuilder.build(); + + List tests = testDefinition.getReadChangeStreamTestsList(); + ArrayList data = new ArrayList<>(tests.size()); + for (ReadChangeStreamTest test : tests) { + String junitName = + CaseFormat.LOWER_HYPHEN.to( + CaseFormat.LOWER_CAMEL, test.getDescription().replace(" ", "-")); + data.add(new Object[] {test, junitName}); + } + return data; + } + + @Test + public void test() throws Exception { + List responses = testCase.getApiResponsesList(); + + // Wrap the responses in a callable. + ServerStreamingCallable source = + new FakeStreamingApi.ServerStreamingStashCallable<>(responses); + ChangeStreamRecordMergingCallable mergingCallable = + new ChangeStreamRecordMergingCallable<>(source, new DefaultChangeStreamRecordAdapter()); + + // Invoke the callable to get the change stream records. + ServerStream stream = + mergingCallable.call(ReadChangeStreamRequest.getDefaultInstance()); + + // Transform the change stream records into ReadChangeStreamTest.Result's. + List actualResults = Lists.newArrayList(); + Exception error = null; + try { + for (ChangeStreamRecord record : stream) { + if (record instanceof Heartbeat) { + Heartbeat heartbeat = (Heartbeat) record; + ChangeStreamContinuationToken token = heartbeat.getChangeStreamContinuationToken(); + ReadChangeStreamResponse.Heartbeat heartbeatProto = + ReadChangeStreamResponse.Heartbeat.newBuilder() + .setContinuationToken( + StreamContinuationToken.newBuilder() + .setPartition( + StreamPartition.newBuilder() + .setRowRange( + RowRange.newBuilder() + .setStartKeyClosed(token.getPartition().getStart()) + .setEndKeyOpen(token.getPartition().getEnd()) + .build()) + .build()) + .setToken(heartbeat.getChangeStreamContinuationToken().getToken()) + .build()) + .setEstimatedLowWatermark(heartbeat.getEstimatedLowWatermark()) + .build(); + actualResults.add( + ReadChangeStreamTest.Result.newBuilder() + .setRecord( + ReadChangeStreamTest.TestChangeStreamRecord.newBuilder() + .setHeartbeat(heartbeatProto) + .build()) + .build()); + } else if (record instanceof CloseStream) { + CloseStream closeStream = (CloseStream) record; + ReadChangeStreamResponse.CloseStream.Builder builder = + ReadChangeStreamResponse.CloseStream.newBuilder().setStatus(closeStream.getStatus()); + for (ChangeStreamContinuationToken token : + closeStream.getChangeStreamContinuationTokens()) { + builder.addContinuationTokens( + StreamContinuationToken.newBuilder() + .setPartition( + StreamPartition.newBuilder() + .setRowRange( + RowRange.newBuilder() + .setStartKeyClosed(token.getPartition().getStart()) + .setEndKeyOpen(token.getPartition().getEnd()) + .build())) + .setToken(token.getToken()) + .build()); + } + ReadChangeStreamResponse.CloseStream closeStreamProto = builder.build(); + actualResults.add( + ReadChangeStreamTest.Result.newBuilder() + .setRecord( + ReadChangeStreamTest.TestChangeStreamRecord.newBuilder() + .setCloseStream(closeStreamProto) + .build()) + .build()); + } else if (record instanceof ChangeStreamMutation) { + ChangeStreamMutation changeStreamMutation = (ChangeStreamMutation) record; + ReadChangeStreamTest.TestChangeStreamMutation.Builder builder = + ReadChangeStreamTest.TestChangeStreamMutation.newBuilder(); + builder.setRowKey(changeStreamMutation.getRowKey()); + Type type = Type.UNRECOGNIZED; + if (changeStreamMutation.getType() == ChangeStreamMutation.MutationType.USER) { + type = Type.USER; + } else if (changeStreamMutation.getType() + == ChangeStreamMutation.MutationType.GARBAGE_COLLECTION) { + type = Type.GARBAGE_COLLECTION; + } + builder.setType(type); + if (changeStreamMutation.getSourceClusterId() != null) { + builder.setSourceClusterId(changeStreamMutation.getSourceClusterId()); + } + builder.setCommitTimestamp( + Timestamps.fromNanos(changeStreamMutation.getCommitTimestamp())); + builder.setTiebreaker(changeStreamMutation.getTieBreaker()); + builder.setToken(changeStreamMutation.getToken()); + builder.setEstimatedLowWatermark( + Timestamps.fromNanos(changeStreamMutation.getEstimatedLowWatermark())); + for (Entry entry : changeStreamMutation.getEntries()) { + if (entry instanceof DeleteFamily) { + DeleteFamily deleteFamily = (DeleteFamily) entry; + builder.addMutations( + Mutation.newBuilder() + .setDeleteFromFamily( + Mutation.DeleteFromFamily.newBuilder() + .setFamilyName(deleteFamily.getFamilyName()) + .build())); + } else if (entry instanceof DeleteCells) { + DeleteCells deleteCells = (DeleteCells) entry; + builder.addMutations( + Mutation.newBuilder() + .setDeleteFromColumn( + Mutation.DeleteFromColumn.newBuilder() + .setFamilyName(deleteCells.getFamilyName()) + .setColumnQualifier(deleteCells.getQualifier()) + .setTimeRange( + TimestampRange.newBuilder() + .setStartTimestampMicros( + deleteCells.getTimestampRange().getStart()) + .setEndTimestampMicros( + deleteCells.getTimestampRange().getEnd()) + .build()) + .build())); + } else if (entry instanceof SetCell) { + SetCell setCell = (SetCell) entry; + builder.addMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName(setCell.getFamilyName()) + .setColumnQualifier(setCell.getQualifier()) + .setTimestampMicros(setCell.getTimestamp()) + .setValue(setCell.getValue()))); + } else { + throw new IllegalStateException("Unexpected Entry type"); + } + } + actualResults.add( + ReadChangeStreamTest.Result.newBuilder() + .setRecord( + ReadChangeStreamTest.TestChangeStreamRecord.newBuilder() + .setChangeStreamMutation(builder)) + .build()); + } else { + throw new IllegalStateException("Unexpected ChangeStreamRecord type"); + } + } + } catch (Exception e) { + error = e; + } + + // Verify the results. + if (expectsError(testCase)) { + assertThat(error).isNotNull(); + } else { + if (error != null) { + throw error; + } + } + + assertThat(getNonExceptionResults(testCase)).isEqualTo(actualResults); + } + + private static boolean expectsError(ReadChangeStreamTest testCase) { + List results = testCase.getResultsList(); + return results != null && !results.isEmpty() && results.get(results.size() - 1).getError(); + } + + private static List getNonExceptionResults( + ReadChangeStreamTest testCase) { + List results = testCase.getResultsList(); + List response = new ArrayList<>(); + if (results != null) { + for (ReadChangeStreamTest.Result result : results) { + if (!result.getError()) { + response.add(result); + } + } + } + return response; + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamRetryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamRetryTest.java new file mode 100644 index 0000000000..2f4cc14327 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamRetryTest.java @@ -0,0 +1,478 @@ +/* + * 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.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.api.gax.rpc.InternalException; +import com.google.api.gax.rpc.ServerStream; +import com.google.bigtable.v2.BigtableGrpc; +import com.google.bigtable.v2.Mutation; +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.StreamContinuationTokens; +import com.google.bigtable.v2.StreamPartition; +import com.google.cloud.bigtable.data.v2.BigtableDataClient; +import com.google.cloud.bigtable.data.v2.BigtableDataSettings; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +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.Heartbeat; +import com.google.cloud.bigtable.data.v2.models.ReadChangeStreamQuery; +import com.google.common.collect.Lists; +import com.google.common.collect.Queues; +import com.google.common.truth.Truth; +import com.google.protobuf.ByteString; +import com.google.protobuf.Timestamp; +import com.google.protobuf.util.Timestamps; +import io.grpc.Status; +import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.StreamObserver; +import io.grpc.testing.GrpcServerRule; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Queue; +import javax.annotation.Nonnull; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ReadChangeStreamRetryTest { + 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 START_KEY_CLOSED = "a"; + private static final String END_KEY_OPEN = "b"; + private static final String HEARTBEAT_TOKEN = "heartbeat-token"; + private static final String CLOSE_STREAM_TOKEN = "close-stream-token"; + private static final String DATA_CHANGE_TOKEN = "data-change-token"; + private static final long REQUEST_START_TIME = 1000L; + + @Rule public GrpcServerRule serverRule = new GrpcServerRule(); + private TestBigtableService service; + private BigtableDataClient client; + + @Before + public void setUp() throws IOException { + service = new TestBigtableService(); + serverRule.getServiceRegistry().addService(service); + + BigtableDataSettings.Builder settings = + BigtableDataSettings.newBuilderForEmulator(serverRule.getServer().getPort()) + .setProjectId(PROJECT_ID) + .setInstanceId(INSTANCE_ID) + .setCredentialsProvider(NoCredentialsProvider.create()); + + settings + .stubSettings() + .setTransportChannelProvider( + FixedTransportChannelProvider.create( + GrpcTransportChannel.create(serverRule.getChannel()))) + .build(); + + client = BigtableDataClient.create(settings.build()); + } + + @After + public void tearDown() { + if (client != null) { + client.close(); + } + } + + private StreamContinuationToken createStreamContinuationToken(@Nonnull String token) { + return StreamContinuationToken.newBuilder() + .setPartition( + StreamPartition.newBuilder() + .setRowRange( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8(START_KEY_CLOSED)) + .setEndKeyOpen(ByteString.copyFromUtf8(END_KEY_OPEN)) + .build()) + .build()) + .setToken(token) + .build(); + } + + private ReadChangeStreamResponse.Heartbeat createHeartbeat( + StreamContinuationToken streamContinuationToken) { + return ReadChangeStreamResponse.Heartbeat.newBuilder() + .setContinuationToken(streamContinuationToken) + .setEstimatedLowWatermark(Timestamp.newBuilder().setSeconds(1000).build()) + .build(); + } + + private ReadChangeStreamResponse.CloseStream createCloseStream() { + return ReadChangeStreamResponse.CloseStream.newBuilder() + .addContinuationTokens(createStreamContinuationToken(CLOSE_STREAM_TOKEN)) + .setStatus(com.google.rpc.Status.newBuilder().setCode(0).build()) + .build(); + } + + private ReadChangeStreamResponse.DataChange createDataChange(boolean done) { + Mutation deleteFromFamily = + Mutation.newBuilder() + .setDeleteFromFamily( + Mutation.DeleteFromFamily.newBuilder().setFamilyName("fake-family").build()) + .build(); + ReadChangeStreamResponse.DataChange.Builder dataChangeBuilder = + ReadChangeStreamResponse.DataChange.newBuilder() + .setType(ReadChangeStreamResponse.DataChange.Type.USER) + .setSourceClusterId("fake-source-cluster-id") + .setRowKey(ByteString.copyFromUtf8("key")) + .setCommitTimestamp(Timestamp.newBuilder().setSeconds(100).build()) + .setTiebreaker(100) + .addChunks( + ReadChangeStreamResponse.MutationChunk.newBuilder().setMutation(deleteFromFamily)); + if (done) { + dataChangeBuilder.setDone(true); + dataChangeBuilder.setEstimatedLowWatermark(Timestamp.newBuilder().setSeconds(1).build()); + dataChangeBuilder.setToken(DATA_CHANGE_TOKEN); + } + return dataChangeBuilder.build(); + } + + // [{ReadChangeStreamResponse.Heartbeat}] -> [{Heartbeat}] + @Test + public void happyPathHeartbeatTest() { + ReadChangeStreamResponse heartbeatResponse = + ReadChangeStreamResponse.newBuilder() + .setHeartbeat(createHeartbeat(createStreamContinuationToken(HEARTBEAT_TOKEN))) + .build(); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(heartbeatResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof Heartbeat); + } + + // [{ReadChangeStreamResponse.CloseStream}] -> [{CloseStream}] + @Test + public void happyPathCloseStreamTest() { + ReadChangeStreamResponse closeStreamResponse = + ReadChangeStreamResponse.newBuilder().setCloseStream(createCloseStream()).build(); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(closeStreamResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof CloseStream); + } + + // [{DataChange(done==true)}] -> [{ReadChangeStreamMutation}] + @Test + public void happyPathCompleteDataChangeTest() { + // Setting `done==true` to complete the ChangeStreamMutation. + ReadChangeStreamResponse dataChangeResponse = + ReadChangeStreamResponse.newBuilder().setDataChange(createDataChange(true)).build(); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(dataChangeResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof ChangeStreamMutation); + } + + // [{UNAVAILABLE}, {ReadChangeStreamResponse.Heartbeat}] -> [{Heartbeat}] + @Test + public void singleHeartbeatImmediateRetryTest() { + ReadChangeStreamResponse heartbeatResponse = + ReadChangeStreamResponse.newBuilder() + .setHeartbeat(createHeartbeat(createStreamContinuationToken(HEARTBEAT_TOKEN))) + .build(); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWithStatus(Code.UNAVAILABLE)); + // Resume with the exact same request. + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(heartbeatResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof Heartbeat); + } + + // [{UNAVAILABLE}, {ReadChangeStreamResponse.CloseStream}] -> [{CloseStream}] + @Test + public void singleCloseStreamImmediateRetryTest() { + // CloseStream. + ReadChangeStreamResponse closeStreamResponse = + ReadChangeStreamResponse.newBuilder().setCloseStream(createCloseStream()).build(); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWithStatus(Code.UNAVAILABLE)); + // Resume with the exact same request. + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(closeStreamResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof CloseStream); + } + + // [{UNAVAILABLE}, {DataChange with done==true}] -> [{(ChangeStreamRecord) ChangeStreamMutation}] + @Test + public void singleCompleteDataChangeImmediateRetryTest() { + // DataChange + ReadChangeStreamResponse dataChangeResponse = + ReadChangeStreamResponse.newBuilder().setDataChange(createDataChange(true)).build(); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWithStatus(Code.UNAVAILABLE)); + // Resume with the exact same request. + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(dataChangeResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof ChangeStreamMutation); + } + + // [{ReadChangeStreamResponse.Heartbeat}, {UNAVAILABLE}] -> Resume with token from heartbeat. + @Test + public void errorAfterHeartbeatShouldResumeWithTokenTest() { + StreamContinuationToken streamContinuationToken = + createStreamContinuationToken(HEARTBEAT_TOKEN); + ReadChangeStreamResponse heartbeatResponse = + ReadChangeStreamResponse.newBuilder() + .setHeartbeat(createHeartbeat(streamContinuationToken)) + .build(); + service.expectations.add( + RpcExpectation.create() + .expectInitialRequest() + .respondWith(heartbeatResponse) + .respondWithStatus(Code.UNAVAILABLE)); + // Resume the request with the token from the Heartbeat. `startTime` is cleared. + // We don't care about the response here so just do expectRequest. + service.expectations.add( + RpcExpectation.create() + .expectRequest( + StreamContinuationTokens.newBuilder().addTokens(streamContinuationToken).build())); + List actualResults = getResults(); + // This is the Heartbeat we get before UNAVAILABLE. + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof Heartbeat); + } + + // [{DataChange with done==true}, {UNAVAILABLE}] -> Resume with token from DataChange. + @Test + public void errorAfterDataChangeWithDoneShouldResumeWithTokenTest() { + // DataChange + ReadChangeStreamResponse dataChangeResponse = + ReadChangeStreamResponse.newBuilder().setDataChange(createDataChange(true)).build(); + service.expectations.add( + RpcExpectation.create() + .expectInitialRequest() + .respondWith(dataChangeResponse) + .respondWithStatus(Code.UNAVAILABLE)); + // Resume the request with the token from the ChangeStreamMutation. `startTime` is cleared. + // We don't care about the response here so just do expectRequest. + service.expectations.add( + RpcExpectation.create() + .expectRequest( + StreamContinuationTokens.newBuilder() + .addTokens(createStreamContinuationToken(DATA_CHANGE_TOKEN)) + .build())); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof ChangeStreamMutation); + } + + // [{DataChange with done==false}, {UNAVAILABLE}] -> Resume with original request. + @Test + public void errorAfterDataChangeWithoutDoneShouldResumeWithTokenTest() { + // DataChange + ReadChangeStreamResponse dataChangeResponse = + ReadChangeStreamResponse.newBuilder().setDataChange(createDataChange(false)).build(); + service.expectations.add( + RpcExpectation.create() + .expectInitialRequest() + .respondWith(dataChangeResponse) + .respondWithStatus(Code.UNAVAILABLE)); + // Resume the request with the original request, because the previous DataChange didn't + // complete the ChangeStreamMutation(i.e. without `done==true`). + // We don't care about the response here so just do expectRequest. + service.expectations.add(RpcExpectation.create().expectInitialRequest()); + List actualResults = getResults(); + Truth.assertThat(actualResults).isEmpty(); + } + + // [{DataChange with done==true}, {Heartbeat}, {UNAVAILABLE}] -> Resume with token from Heartbeat. + @Test + public void shouldResumeWithLastTokenTest() { + // DataChange + ReadChangeStreamResponse dataChangeResponse = + ReadChangeStreamResponse.newBuilder().setDataChange(createDataChange(true)).build(); + // Heartbeat. + ReadChangeStreamResponse heartbeatResponse = + ReadChangeStreamResponse.newBuilder() + .setHeartbeat(createHeartbeat(createStreamContinuationToken(HEARTBEAT_TOKEN))) + .build(); + service.expectations.add( + RpcExpectation.create() + .expectInitialRequest() + .respondWith(dataChangeResponse) + .respondWith(heartbeatResponse) + .respondWithStatus(Code.UNAVAILABLE)); + // If we receive a DataChange with done==true and a Heartbeat then a retryable error, it should + // resume with the last token, which is the one from the heartbeat. + // If the original request reads with start_time, it'll be resumed with the continuation token. + // We don't care about the response here so just do expectRequest. + service.expectations.add( + RpcExpectation.create() + .expectRequest( + StreamContinuationTokens.newBuilder() + .addTokens(createStreamContinuationToken(HEARTBEAT_TOKEN)) + .build())); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(2); + Assert.assertTrue(actualResults.get(0) instanceof ChangeStreamMutation); + Assert.assertTrue(actualResults.get(1) instanceof Heartbeat); + } + + @Test + public void retryRstStreamExceptionTest() { + ApiException exception = + new InternalException( + new StatusRuntimeException( + Status.INTERNAL.withDescription( + "INTERNAL: HTTP/2 error code: INTERNAL_ERROR\nReceived Rst Stream")), + GrpcStatusCode.of(Code.INTERNAL), + false); + ReadChangeStreamResponse heartbeatResponse = + ReadChangeStreamResponse.newBuilder() + .setHeartbeat(createHeartbeat(createStreamContinuationToken(HEARTBEAT_TOKEN))) + .build(); + service.expectations.add( + RpcExpectation.create() + .expectInitialRequest() + .respondWithException(Code.INTERNAL, exception)); + service.expectations.add( + RpcExpectation.create().expectInitialRequest().respondWith(heartbeatResponse)); + List actualResults = getResults(); + assertThat(actualResults.size()).isEqualTo(1); + Assert.assertTrue(actualResults.get(0) instanceof Heartbeat); + } + + private List getResults() { + ReadChangeStreamQuery query = + ReadChangeStreamQuery.create(TABLE_ID).startTime(REQUEST_START_TIME); + // Always give it this partition. We don't care. + ServerStream actualRecords = + client.readChangeStream(query.streamPartition(START_KEY_CLOSED, END_KEY_OPEN)); + List actualValues = Lists.newArrayList(); + for (ChangeStreamRecord record : actualRecords) { + actualValues.add(record); + } + return actualValues; + } + + private static class TestBigtableService extends BigtableGrpc.BigtableImplBase { + Queue expectations = Queues.newArrayDeque(); + int i = -1; + + @Override + public void readChangeStream( + ReadChangeStreamRequest request, + StreamObserver responseObserver) { + + RpcExpectation expectedRpc = expectations.poll(); + i++; + + Truth.assertWithMessage("Unexpected request#" + i + ":" + request.toString()) + .that(expectedRpc) + .isNotNull(); + Truth.assertWithMessage("Unexpected request#" + i) + .that(request) + .isEqualTo(expectedRpc.getExpectedRequest()); + + for (ReadChangeStreamResponse response : expectedRpc.responses) { + responseObserver.onNext(response); + } + if (expectedRpc.statusCode.toStatus().isOk()) { + responseObserver.onCompleted(); + } else if (expectedRpc.exception != null) { + responseObserver.onError(expectedRpc.exception); + } else { + responseObserver.onError(expectedRpc.statusCode.toStatus().asRuntimeException()); + } + } + } + + private static class RpcExpectation { + ReadChangeStreamRequest.Builder requestBuilder; + Status.Code statusCode; + ApiException exception; + List responses; + + private RpcExpectation() { + this.requestBuilder = + ReadChangeStreamRequest.newBuilder() + .setTableName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID)) + .setPartition( + StreamPartition.newBuilder() + .setRowRange( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8(START_KEY_CLOSED)) + .setEndKeyOpen(ByteString.copyFromUtf8(END_KEY_OPEN)) + .build()) + .build()); + this.statusCode = Status.Code.OK; + this.responses = Lists.newArrayList(); + } + + static RpcExpectation create() { + return new RpcExpectation(); + } + + RpcExpectation expectInitialRequest() { + requestBuilder.setStartTime(Timestamps.fromNanos(REQUEST_START_TIME)); + return this; + } + + RpcExpectation expectRequest(StreamContinuationTokens continuationTokens) { + requestBuilder.setContinuationTokens(continuationTokens); + return this; + } + + RpcExpectation respondWithStatus(Status.Code code) { + this.statusCode = code; + return this; + } + + RpcExpectation respondWithException(Status.Code code, ApiException exception) { + this.statusCode = code; + this.exception = exception; + return this; + } + + RpcExpectation respondWith(ReadChangeStreamResponse... responses) { + Collections.addAll(this.responses, responses); + return this; + } + + ReadChangeStreamRequest getExpectedRequest() { + return requestBuilder.build(); + } + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamUserCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamUserCallableTest.java new file mode 100644 index 0000000000..2f50c7065b --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamUserCallableTest.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.stub.changestream; + +import com.google.bigtable.v2.ReadChangeStreamRequest; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord; +import com.google.cloud.bigtable.data.v2.models.ReadChangeStreamQuery; +import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; +import com.google.common.truth.Truth; +import java.time.Duration; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ReadChangeStreamUserCallableTest { + private static final RequestContext REQUEST_CONTEXT = + RequestContext.create("fake-project", "fake-instance", "fake-profile"); + + @Test + public void testRequestIsConverted() { + ServerStreamingStashCallable innerCallable = + new ServerStreamingStashCallable<>(); + ReadChangeStreamUserCallable callable = + new ReadChangeStreamUserCallable<>(innerCallable, REQUEST_CONTEXT); + ReadChangeStreamQuery query = + ReadChangeStreamQuery.create("fake-table") + .streamPartition("begin", "end") + .startTime(1000L) + .endTime(2000L) + .heartbeatDuration(Duration.ofSeconds(1)); + callable.call(query); + Truth.assertThat(innerCallable.getActualRequest()).isEqualTo(query.toProto(REQUEST_CONTEXT)); + } +} diff --git a/google-cloud-bigtable/src/test/proto/changestream_tests.proto b/google-cloud-bigtable/src/test/proto/changestream_tests.proto new file mode 100644 index 0000000000..82f0fff492 --- /dev/null +++ b/google-cloud-bigtable/src/test/proto/changestream_tests.proto @@ -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. + +syntax = "proto3"; + +package google.cloud.conformance.bigtable.v2; + +import "google/bigtable/v2/bigtable.proto"; +import "google/protobuf/timestamp.proto"; +import "google/bigtable/v2/data.proto"; + +option csharp_namespace = "Google.Cloud.Bigtable.V2.Tests.Conformance"; +option java_outer_classname = "ChangeStreamTestDefinition"; +option java_package = "com.google.cloud.conformance.bigtable.v2"; +option go_package = "google/cloud/conformance/bigtable/v2"; + +message ChangeStreamTestFile { + repeated ReadChangeStreamTest read_change_stream_tests = 1; +} + +message ReadChangeStreamTest { + + message TestChangeStreamMutation { + bytes row_key = 1; + google.bigtable.v2.ReadChangeStreamResponse.DataChange.Type type = 2; + string source_cluster_id = 3; + google.protobuf.Timestamp commit_timestamp = 4; + int64 tiebreaker = 5; + string token = 6; + google.protobuf.Timestamp estimated_low_watermark = 7; + repeated google.bigtable.v2.Mutation mutations = 8; + } + + message TestChangeStreamRecord { + oneof record { + google.bigtable.v2.ReadChangeStreamResponse.Heartbeat heartbeat = 1; + google.bigtable.v2.ReadChangeStreamResponse.CloseStream close_stream = 2; + TestChangeStreamMutation change_stream_mutation = 3; + } + } + + // Expected results of reading the change stream. + // Only the last result can be an error. + message Result { + TestChangeStreamRecord record = 1; + bool error = 2; + } + + string description = 1; + repeated google.bigtable.v2.ReadChangeStreamResponse api_responses = 2; + repeated Result results = 3; +} diff --git a/google-cloud-bigtable/src/test/resources/changestream.json b/google-cloud-bigtable/src/test/resources/changestream.json new file mode 100644 index 0000000000..9d9e2d46cc --- /dev/null +++ b/google-cloud-bigtable/src/test/resources/changestream.json @@ -0,0 +1,1379 @@ +{ + "readChangeStreamTests": [ + { + "description": "1 heartbeat", + "api_responses": [ + { + "heartbeat": { + "continuation_token": { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "" + } + }, + "token": "heartbeat-token" + }, + "estimated_low_watermark": "2022-07-01T00:00:00Z" + } + } + ], + "results": [ + { + "record" : { + "heartbeat": { + "continuation_token": { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "" + } + }, + "token": "heartbeat-token" + }, + "estimated_low_watermark": "2022-07-01T00:00:00Z" + } + }, + "error": false + } + ] + }, + { + "description": "1 CloseStream", + "api_responses": [ + { + "close_stream": { + "status": { + "code": "11", + "message": "Partition boundaries are misaligned." + }, + "continuation_tokens": [ + { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "0000000000000001" + } + }, + "token": "close-stream-token-1" + }, + { + "partition": { + "row_range": { + "start_key_closed": "0000000000000001", + "end_key_open": "0000000000000002" + } + }, + "token": "close-stream-token-2" + } + ] + } + } + ], + "results": [ + { + "record" : { + "close_stream": { + "status": { + "code": "11", + "message": "Partition boundaries are misaligned." + }, + "continuation_tokens": [ + { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "0000000000000001" + } + }, + "token": "close-stream-token-1" + }, + { + "partition": { + "row_range": { + "start_key_closed": "0000000000000001", + "end_key_open": "0000000000000002" + } + }, + "token": "close-stream-token-2" + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 heartbeat + 1 CloseStream", + "api_responses": [ + { + "heartbeat": { + "continuation_token": { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "" + } + }, + "token": "heartbeat-token" + }, + "estimated_low_watermark": "2022-07-01T00:00:00Z" + } + }, + { + "close_stream": { + "status": { + "code": "11", + "message": "Partition boundaries are misaligned." + }, + "continuation_tokens": [ + { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "0000000000000001" + } + }, + "token": "close-stream-token-1" + } + ] + } + } + ], + "results": [ + { + "record" : { + "heartbeat": { + "continuation_token": { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "" + } + }, + "token": "heartbeat-token" + }, + "estimated_low_watermark": "2022-07-01T00:00:00Z" + } + }, + "error": false + }, + { + "record" : { + "close_stream": { + "status": { + "code": "11", + "message": "Partition boundaries are misaligned." + }, + "continuation_tokens": [ + { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "0000000000000001" + } + }, + "token": "close-stream-token-1" + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 logical mutation no chunking([{DF,DC,SC}]->ChangeStreamMutation{DF,DC,SC})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + }, + { + "mutation": { + "delete_from_column" : { + "family_name": "family", + "column_qualifier": "dg==", + "time_range": { + "start_timestamp_micros": 5000, + "end_timestamp_micros": 15000 + } + } + } + }, + { + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "mutations": [ + { + "delete_from_family": { + "family_name": "family" + } + }, + { + "delete_from_column" : { + "family_name": "family", + "column_qualifier": "dg==", + "time_range": { + "start_timestamp_micros": 5000, + "end_timestamp_micros": 15000 + } + } + }, + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 incomplete logical mutation(missing `done: true`)", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ] + } + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "GC mutation no source cluster id", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "GARBAGE_COLLECTION", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "GARBAGE_COLLECTION", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "delete_from_family": { + "family_name": "family" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 chunked SetCell([{SC_chunk1(v)}, {SC_chunk2(alue-VAL)}]->ChangeStreamMutation{SC(value-VAL)})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "ChunkedValueSize mismatch for a chunked SetCell([{SC_chunk1(v)}, {SC_chunk2(alue-VAL)}]->error)", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "chunk_info": { + "chunked_value_size": 1 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "1 chunked SetCell([{SC_chunk1(v)}, {SC_chunk2(alue-VAL)}, {SC_chunk3(-VAL)}]->ChangeStreamMutation{SC(value-VAL-VAL)})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "chunk_info": { + "chunked_value_size": 13 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 13 + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 9, + "chunked_value_size": 13, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "LVZBTA==" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFMLVZBTA==" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "2 chunked SetCells([{SC1_chunk1(v)}, {SC1_chunk2(alue-VAL), SC2_chunk1(v)}, {SC2_chunk2(alue-VAL)}]->ChangeStreamMutation{SC1(value-VAL),SC2(value-VAL)})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + }, + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + }, + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 chunked SetCell + 1 unchunked SetCell([{SC1_chunk1(v)}, {SC1_chunk2(alue-VAL), SC2(value-VAL)}]->ChangeStreamMutation{SC1(value-VAL),SC2(value-VAL)})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + }, + { + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + }, + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 unchunked SetCell + 1 chunked SetCell([{SC1(v), SC2_chunk1(v)}, {SC2_chunk2(alue-VAL)}]->ChangeStreamMutation{SC1(v),SC2(value-VAL)})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + }, + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + }, + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 mod + 1 chunked SetCell + 1 mod([{DF1,SC_chunk1(v)}, {SC_chunk2(alue-VAL), DF2}]->ChangeStreamMutation{DF1,SC(value-VAL),DF2})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + }, + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + }, + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "delete_from_family": { + "family_name": "family" + } + }, + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + }, + { + "delete_from_family": { + "family_name": "family" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "1 chunked SetCell + many nonchunked mods([{SC_chunk1(v)}, {SC_chunk2(alue-VAL),DF,DC}]->ChangeStreamMutation{SC(value-VAL),DF,DC})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "chunk_info": { + "chunked_value_size": 9 + }, + "mutation": { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dg==" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "chunk_info": { + "chunked_value_offset": 1, + "chunked_value_size": 9, + "last_chunk": true + }, + "mutation": { + "set_cell": { + "value": "YWx1ZS1WQUw=" + } + } + }, + { + "mutation": { + "delete_from_column" : { + "family_name": "family", + "column_qualifier": "dg==", + "time_range": { + "start_timestamp_micros": 5000, + "end_timestamp_micros": 15000 + } + } + } + }, + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "set_cell": { + "family_name": "family", + "column_qualifier": "0000000000000000", + "timestamp_micros": 1000, + "value": "dmFsdWUtVkFM" + } + }, + { + "delete_from_column" : { + "family_name": "family", + "column_qualifier": "dg==", + "time_range": { + "start_timestamp_micros": 5000, + "end_timestamp_micros": 15000 + } + } + }, + { + "delete_from_family": { + "family_name": "family" + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "non SetCell chunking([{DF1},{DF2,DC}]->ChangeStreamMutation{DF1,DF2,DC})", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + }, + { + "mutation": { + "delete_from_column" : { + "family_name": "family", + "column_qualifier": "dg==", + "time_range": { + "start_timestamp_micros": 5000, + "end_timestamp_micros": 15000 + } + } + } + } + ], + "done": true + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "delete_from_family": { + "family_name": "family" + } + }, + { + "delete_from_family": { + "family_name": "family" + } + }, + { + "delete_from_column" : { + "family_name": "family", + "column_qualifier": "dg==", + "time_range": { + "start_timestamp_micros": 5000, + "end_timestamp_micros": 15000 + } + } + } + ] + } + }, + "error": false + } + ] + }, + { + "description": "2 logical mutations with non SetCell chunking + CloseStream([{Change1_DF1}, {Change1_DF2}, {Change2_DF3}, {Change2_DF4}, {CloseStream}]->[ChangeStreamMutation1{DF1,DF2}),ChangeStreamMutation2{DF3,DF4}),CloseStream]", + "api_responses": [ + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ], + "done": true + } + }, + { + "data_change": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ] + } + }, + { + "data_change": { + "type": "CONTINUATION", + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "chunks": [ + { + "mutation": { + "delete_from_family": { + "family_name": "family" + } + } + } + ], + "done": true + } + }, + { + "close_stream": { + "status": { + "code": "11", + "message": "Partition boundaries are misaligned." + }, + "continuation_tokens": [ + { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "0000000000000001" + } + }, + "token": "close-stream-token-1" + }, + { + "partition": { + "row_range": { + "start_key_closed": "0000000000000001", + "end_key_open": "0000000000000002" + } + }, + "token": "close-stream-token-2" + } + ] + } + } + ], + "results": [ + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "delete_from_family": { + "family_name": "family" + } + }, + { + "delete_from_family": { + "family_name": "family" + } + } + ] + } + }, + "error": false + }, + { + "record": { + "change_stream_mutation": { + "row_key": "0000000000000000", + "type": "USER", + "source_cluster_id": "source-cluster-id", + "commit_timestamp": "2022-07-01T00:00:00Z", + "tiebreaker": 100, + "token": "data-change-token", + "estimated_low_watermark": "2022-07-01T00:00:00Z", + "mutations": [ + { + "delete_from_family": { + "family_name": "family" + } + }, + { + "delete_from_family": { + "family_name": "family" + } + } + ] + } + }, + "error": false + }, + { + "record" : { + "close_stream": { + "status": { + "code": "11", + "message": "Partition boundaries are misaligned." + }, + "continuation_tokens": [ + { + "partition": { + "row_range": { + "start_key_closed": "", + "end_key_open": "0000000000000001" + } + }, + "token": "close-stream-token-1" + }, + { + "partition": { + "row_range": { + "start_key_closed": "0000000000000001", + "end_key_open": "0000000000000002" + } + }, + "token": "close-stream-token-2" + } + ] + } + }, + "error": false + } + ] + } + ] +} \ No newline at end of file