diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java index dc55756241e9..5b65d3a00119 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java @@ -85,7 +85,7 @@ public static Mutation createUnsafe() { @BetaApi public static Mutation fromProtoUnsafe(List protos) { Mutation mutation = new Mutation(true); - mutation.mutations.addAll(protos); + mutation.addAllFromProto(protos); return mutation; } @@ -98,7 +98,7 @@ public static Mutation fromProtoUnsafe(List pro @BetaApi public static Mutation fromProtoUnsafe(Iterable protos) { Mutation mutation = new Mutation(true); - mutation.mutations.addAll(protos); + mutation.addAllFromProto(protos); return mutation; } @@ -114,7 +114,7 @@ public static Mutation fromProtoUnsafe(Iterable */ static Mutation fromProto(List protos) { Mutation mutation = new Mutation(false); - mutation.mutations.addAll(protos); + mutation.addAllFromProto(protos); return mutation; } @@ -333,10 +333,22 @@ private void addMutation(com.google.bigtable.v2.Mutation mutation) { byteSize + mutation.getSerializedSize() <= MAX_BYTE_SIZE, "Byte size of mutations is too large"); + countTowardsLimits(mutation); + + mutations.add(mutation); + } + + private void countTowardsLimits(com.google.bigtable.v2.Mutation mutation) { numMutations++; byteSize += mutation.getSerializedSize(); + } - mutations.add(mutation); + private void addAllFromProto(Iterable protos) { + // One traversal: protos may be a non-repeatable Iterable. + for (com.google.bigtable.v2.Mutation proto : protos) { + mutations.add(proto); + countTowardsLimits(proto); + } } private static ByteString wrapByteString(String str) { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/MutationTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/MutationTest.java index 3ba1de67011e..4a19effbd625 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/MutationTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/MutationTest.java @@ -23,6 +23,7 @@ import com.google.bigtable.v2.Mutation.DeleteFromRow; import com.google.bigtable.v2.Mutation.MergeToCell; import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange; +import com.google.common.collect.ImmutableList; import com.google.common.primitives.Longs; import com.google.protobuf.ByteString; import java.io.ByteArrayInputStream; @@ -30,6 +31,7 @@ import java.io.IOException; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; +import java.util.Iterator; import java.util.List; import org.junit.Before; import org.junit.Test; @@ -261,6 +263,105 @@ public void tooLargeRequest() { assertThat(actualError).isInstanceOf(IllegalStateException.class); } + @Test + public void tooManyMutationsCountsWrappedProtosTest() { + Mutation mutation = + Mutation.fromProtoUnsafe( + ImmutableList.of( + com.google.bigtable.v2.Mutation.newBuilder() + .setDeleteFromRow(com.google.bigtable.v2.Mutation.DeleteFromRow.newBuilder()) + .build())); + + for (int i = 0; i < Mutation.MAX_MUTATIONS - 1; i++) { + mutation.setCell("f", "", ""); + } + + Exception actualError = null; + try { + mutation.setCell("f", "", ""); + } catch (Exception e) { + actualError = e; + } + + assertThat(actualError).isInstanceOf(IllegalStateException.class); + assertThat(mutation.getMutations()).hasSize(Mutation.MAX_MUTATIONS); + } + + @Test + public void tooLargeRequestCountsWrappedProtosTest() { + Mutation mutation = + Mutation.fromProtoUnsafe( + ImmutableList.of( + com.google.bigtable.v2.Mutation.newBuilder() + .setSetCell( + com.google.bigtable.v2.Mutation.SetCell.newBuilder() + .setFamilyName("f") + .setValue(ByteString.copyFrom(new byte[Mutation.MAX_BYTE_SIZE / 2]))) + .build(), + com.google.bigtable.v2.Mutation.newBuilder() + .setSetCell( + com.google.bigtable.v2.Mutation.SetCell.newBuilder() + .setFamilyName("f") + .setValue(ByteString.copyFrom(new byte[Mutation.MAX_BYTE_SIZE / 2]))) + .build())); + + Exception actualError = null; + try { + mutation.setCell("f", "", ""); + } catch (Exception e) { + actualError = e; + } + + assertThat(actualError).isInstanceOf(IllegalStateException.class); + } + + @Test + public void fromProtoUnsafeTraversesTheIterableOnceTest() { + Iterable oneShot = + oneShot( + com.google.bigtable.v2.Mutation.newBuilder() + .setDeleteFromRow(com.google.bigtable.v2.Mutation.DeleteFromRow.newBuilder()) + .build(), + com.google.bigtable.v2.Mutation.newBuilder() + .setDeleteFromRow(com.google.bigtable.v2.Mutation.DeleteFromRow.newBuilder()) + .build()); + + Mutation mutation = Mutation.fromProtoUnsafe(oneShot); + + assertThat(mutation.getMutations()).hasSize(2); + + for (int i = 0; i < Mutation.MAX_MUTATIONS - 2; i++) { + mutation.setCell("f", "", ""); + } + + Exception actualError = null; + try { + mutation.setCell("f", "", ""); + } catch (Exception e) { + actualError = e; + } + + assertThat(actualError).isInstanceOf(IllegalStateException.class); + } + + /** An Iterable that refuses a second traversal. */ + private static Iterable oneShot( + com.google.bigtable.v2.Mutation... protos) { + List source = ImmutableList.copyOf(protos); + return new Iterable() { + private boolean consumed; + + @Override + public Iterator iterator() { + if (consumed) { + throw new IllegalStateException("Iterable traversed more than once"); + } + consumed = true; + return source.iterator(); + } + }; + } + @Test public void testWithLongValue() { Mutation mutation =