diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptor.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptor.java index 87f5c88d3eb1..5ac49f491b9d 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptor.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptor.java @@ -20,6 +20,7 @@ import com.google.api.gax.batching.BatchResource; import com.google.api.gax.batching.BatchingDescriptor; import com.google.api.gax.batching.BatchingRequestBuilder; +import com.google.bigtable.v2.MutateRowsRequest; import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.MutateRowsException; import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation; @@ -101,8 +102,8 @@ public long countBytes(RowMutationEntry entry) { @Override public BatchResource createResource(RowMutationEntry element) { - long byteCount = countBytes(element); - return MutateRowsBatchResource.create(1, byteCount, element.toProto().getMutationsCount()); + MutateRowsRequest.Entry proto = element.toProto(); + return MutateRowsBatchResource.create(1, proto.getSerializedSize(), proto.getMutationsCount()); } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptorTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptorTest.java index 11658d8792bd..a3115fab7a5b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptorTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptorTest.java @@ -27,6 +27,7 @@ import com.google.api.gax.rpc.DeadlineExceededException; import com.google.api.gax.rpc.InternalException; import com.google.api.gax.rpc.UnavailableException; +import com.google.bigtable.v2.MutateRowsRequest; 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.MutateRowsException; @@ -64,6 +65,34 @@ public void countBytesTest() { assertThat(underTest.countBytes(request)).isEqualTo(bytes); } + @Test + public void createResourceTest() { + RowMutationEntry request = + RowMutationEntry.create(ROW_KEY) + .setCell(FAMILY, QUALIFIER, VALUE) + .setCell(FAMILY, QUALIFIER, "another-value") + .deleteFamily(FAMILY); + MutateRowsRequest.Entry proto = request.toProto(); + + MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor(); + BatchResource resource = underTest.createResource(request); + + assertThat(resource.getElementCount()).isEqualTo(1); + assertThat(resource.getByteCount()).isEqualTo(proto.getSerializedSize()); + assertThat(((MutateRowsBatchResource) resource).getMutationCount()) + .isEqualTo(proto.getMutationsCount()); + } + + @Test + public void createEmptyResourceTest() { + MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor(); + BatchResource resource = underTest.createEmptyResource(); + + assertThat(resource.getElementCount()).isEqualTo(0); + assertThat(resource.getByteCount()).isEqualTo(0); + assertThat(((MutateRowsBatchResource) resource).getMutationCount()).isEqualTo(0); + } + @Test public void requestBuilderTest() { MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor();