From 843ae5eb8df3f69dfbe27bca417f6b1a0542c6bc Mon Sep 17 00:00:00 2001 From: Shant Stepanian Date: Wed, 2 Jul 2025 14:57:33 +0000 Subject: [PATCH] feat: Support Idempotency in MutateRow API --- .../bigtable/data/v2/models/RowMutation.java | 13 ++++ .../bigtable/data/v2/it/MutateRowIT.java | 67 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutation.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutation.java index 38d822afcbfd..470509879431 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutation.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutation.java @@ -17,6 +17,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; +import com.google.bigtable.v2.Idempotency; import com.google.bigtable.v2.MutateRowRequest; import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.MutateRowsRequest.Entry; @@ -38,6 +39,7 @@ public final class RowMutation implements MutationApi, Serializable private final TargetId targetId; private final ByteString key; private final Mutation mutation; + private ByteString idempotency_token; private RowMutation(TargetId targetId, ByteString key, Mutation mutation) { Preconditions.checkNotNull(targetId, "target id can't be null."); @@ -255,6 +257,11 @@ public RowMutation mergeToCell( return this; } + public RowMutation setIdempotency(@Nonnull ByteString token) { + this.idempotency_token = token; + return this; + } + @InternalApi public MutateRowRequest toProto(RequestContext requestContext) { MutateRowRequest.Builder builder = MutateRowRequest.newBuilder(); @@ -266,6 +273,12 @@ public MutateRowRequest toProto(RequestContext requestContext) { builder.setTableName(resourceName); } + if (idempotency_token != null) { + Idempotency.Builder idempotencyBuilder = Idempotency.newBuilder(); + idempotencyBuilder.setToken(idempotency_token); + builder.setIdempotency(idempotencyBuilder); + } + return builder .setAppProfileId(requestContext.getAppProfileId()) .setRowKey(key) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/MutateRowIT.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/MutateRowIT.java index c99000be4874..fff13abff67b 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/MutateRowIT.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/MutateRowIT.java @@ -20,6 +20,7 @@ import static com.google.cloud.bigtable.misc_utilities.AuthorizedViewTestHelper.createTestAuthorizedView; import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.TruthJUnit.assume; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; import com.google.api.gax.rpc.PermissionDeniedException; @@ -31,6 +32,8 @@ import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv; import com.google.cloud.bigtable.test_helpers.env.TestEnvRule; import com.google.protobuf.ByteString; + +import java.math.BigInteger; import java.util.UUID; import java.util.concurrent.TimeUnit; import org.junit.ClassRule; @@ -40,6 +43,16 @@ @RunWith(JUnit4.class) public class MutateRowIT { + static { + System.setProperty("bigtable.env", "cloud"); + System.setProperty("bigtable.project", "google.com:cloud-bigtable-dev"); + System.setProperty("bigtable.instance", "stepanian-agg"); + System.setProperty("bigtable.table", "my-table"); + System.setProperty("bigtable.data-endpoint", "test-bigtable.sandbox.googleapis.com:443"); + System.setProperty("bigtable.admin-endpoint", "test-bigtableadmin.sandbox.googleapis.com:443"); + +// -D -Dbigtable.data-endpoint=bigtable.googleapis.com -Dbigtable.admin-endpoint=test-bigtableadmin.sandbox.googleapis.com + } @ClassRule public static TestEnvRule testEnvRule = new TestEnvRule(); @Test @@ -152,4 +165,58 @@ public void testOnAuthorizedView() throws Exception { .getTableAdminClient() .deleteAuthorizedView(testEnvRule.env().getTableId(), testAuthorizedView.getId()); } + + @Test + public void testIdempotentAggregates() throws Exception { + String rowKey = UUID.randomUUID().toString(); + // String familyId = testEnvRule.env().getFamilyId(); + String familyId = "cfagg1"; + + // First mutation w/ idempotency token. + testEnvRule + .env() + .getDataClient() + .mutateRowAsync( + RowMutation.create(testEnvRule.env().getTableId(), rowKey) + .addToCell(familyId, "cq", 0, 10) + .setIdempotency(ByteString.copyFromUtf8("abcdefgh")) + ) + .get(1, TimeUnit.MINUTES); + + // Now replay the mutation. + testEnvRule + .env() + .getDataClient() + .mutateRowAsync( + RowMutation.create(testEnvRule.env().getTableId(), rowKey) + .addToCell(familyId, "cq", 0, 10) + .setIdempotency(ByteString.copyFromUtf8("abcdefgh")) + ) + .get(1, TimeUnit.MINUTES); + + Row row = + testEnvRule + .env() + .getDataClient() + .readRowsCallable() + .first() + .call(Query.create(testEnvRule.env().getTableId()).rowKey(rowKey)); + + System.out.println("Shant! " + row.toString()); + System.out.println("Shant! " + row.getCells().get(0)); + System.out.println("Shant! " + row.getCells().get(0).getValue().toByteArray()); + for (byte b : row.getCells().get(0).getValue().toByteArray()) { + System.out.println("SHant byte " + b); + } + System.out.println("Shant bigint " + new BigInteger(row.getCells().get(0).getValue().toByteArray())); + System.out.println("Shant bytebufsize " + row.getCells().get(0).getValue().size()); + System.out.println("Shant bytebuf " + java.nio.ByteBuffer.wrap(row.getCells().get(0).getValue().toByteArray()).getInt()); + System.out.println("Shant bytebuf " + java.nio.ByteBuffer.wrap(row.getCells().get(0).getValue().toByteArray()).order(java.nio.ByteOrder.BIG_ENDIAN).getInt()); + + // assertEquals(row.toString(), "shant"); + assertThat(row.getCells()).hasSize(1); + // Ensure that the increment is only applied once. + assertThat(new BigInteger(row.getCells().get(0).getValue().toByteArray())).isEqualTo(10); + // assertThat(java.nio.ByteBuffer.wrap(row.getCells().get(0).getValue().toByteArray()).order(java.nio.ByteOrder.BIG_ENDIAN).getInt()).isEqualTo(10); + } }