Add ReturnConsumedCapacity to BatchGetItemEnhancedRequest (#4421)

* Add ReturnConsumedCapacity to BatchGetItemEnhancedRequest * Change log for "Add ReturnConsumedCapacity to BatchGetItemEnhancedRequest" --------- Co-authored-by: Nilesh <psnilesh@amazon.com>
aws · Sep 12, 2023 · d906387 · d906387
1 parent 3abf1c3
commit d906387
Show file tree

Hide file tree

Showing 7 changed files with 365 additions and 48 deletions.
diff --git a/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-ee18d70.json b/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-ee18d70.json
@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "Amazon DynamoDB Enhanced Client",
+    "contributor": "psnilesh",
+    "description": "This commit introduces returnConsumedCapacity input to BatchGetItemEnhancedRequest that allows customers to find out exactly how much read units were consumed by the operation."
+}
diff --git a/...a/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java b/...a/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java
@@ -62,6 +62,7 @@ public BatchGetItemRequest generateRequest(DynamoDbEnhancedClientExtension exten
 
         return BatchGetItemRequest.builder()
                                   .requestItems(Collections.unmodifiableMap(requestItems))
+                                  .returnConsumedCapacity(request.returnConsumedCapacityAsString())
                                   .build();
     }
 

diff --git a/...main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequest.java b/...main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequest.java
@@ -20,10 +20,14 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
+import software.amazon.awssdk.services.dynamodb.model.BatchGetItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
 
 /**
  * Defines parameters used for the batchGetItem() operation (such as
@@ -36,9 +40,11 @@
 public final class BatchGetItemEnhancedRequest {
 
     private final List<ReadBatch> readBatches;
+    private final String returnConsumedCapacity;
 
     private BatchGetItemEnhancedRequest(Builder builder) {
         this.readBatches = getListIfExist(builder.readBatches);
+        this.returnConsumedCapacity = builder.returnConsumedCapacity;
     }
 
     /**
@@ -52,7 +58,7 @@ public static Builder builder() {
      * Returns a builder initialized with all existing values on the request object.
      */
     public Builder toBuilder() {
-        return new Builder().readBatches(readBatches);
+        return new Builder().readBatches(readBatches).returnConsumedCapacity(this.returnConsumedCapacity);
     }
 
     /**
@@ -62,6 +68,25 @@ public Collection<ReadBatch> readBatches() {
         return readBatches;
     }
 
+    /**
+     * Whether to return the capacity consumed by this operation.
+     *
+     * @see GetItemRequest#returnConsumedCapacity()
+     */
+    public ReturnConsumedCapacity returnConsumedCapacity() {
+        return ReturnConsumedCapacity.fromValue(returnConsumedCapacity);
+    }
+
+    /**
+     * Whether to return the capacity consumed by this operation.
+     * <p>
+     * Similar to {@link #returnConsumedCapacity()} but return the value as a string. This is useful in situations where the
+     * value is not defined in {@link ReturnConsumedCapacity}.
+     */
+    public String returnConsumedCapacityAsString() {
+        return returnConsumedCapacity;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -73,12 +98,15 @@ public boolean equals(Object o) {
 
         BatchGetItemEnhancedRequest that = (BatchGetItemEnhancedRequest) o;
 
-        return readBatches != null ? readBatches.equals(that.readBatches) : that.readBatches == null;
+        return Objects.equals(this.readBatches, that.readBatches) &&
+               Objects.equals(this.returnConsumedCapacity, that.returnConsumedCapacity);
     }
 
     @Override
     public int hashCode() {
-        return readBatches != null ? readBatches.hashCode() : 0;
+        int hc = readBatches != null ? readBatches.hashCode() : 0;
+        hc = 31 * hc + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
+        return hc;
     }
 
     private static List<ReadBatch> getListIfExist(List<ReadBatch> readBatches) {
@@ -91,6 +119,7 @@ private static List<ReadBatch> getListIfExist(List<ReadBatch> readBatches) {
     @NotThreadSafe
     public static final class Builder {
         private List<ReadBatch> readBatches;
+        private String returnConsumedCapacity;
 
         private Builder() {
         }
@@ -132,9 +161,30 @@ public Builder addReadBatch(ReadBatch readBatch) {
             return this;
         }
 
+        /**
+         * Whether to return the capacity consumed by this operation.
+         *
+         * @see BatchGetItemRequest.Builder#returnConsumedCapacity(ReturnConsumedCapacity)
+         * @return a builder of this type
+         */
+        public Builder returnConsumedCapacity(ReturnConsumedCapacity returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity == null ? null : returnConsumedCapacity.toString();
+            return this;
+        }
+
+        /**
+         * Whether to return the capacity consumed by this operation.
+         *
+         * @see BatchGetItemRequest.Builder#returnConsumedCapacity(String)
+         * @return a builder of this type
+         */
+        public Builder returnConsumedCapacity(String returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity;
+            return this;
+        }
+
         public BatchGetItemEnhancedRequest build() {
             return new BatchGetItemEnhancedRequest(this);
         }
     }
-
 }
diff --git a/...nced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java b/...nced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java
@@ -33,6 +33,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.internal.operations.DefaultOperationContext;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.BatchGetItemResponse;
+import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity;
 import software.amazon.awssdk.services.dynamodb.model.KeysAndAttributes;
 
 /**
@@ -111,6 +112,10 @@ public List<Key> unprocessedKeysForTable(MappedTableResource<?> mappedTable) {
                                 .collect(Collectors.toList());
     }
 
+    public List<ConsumedCapacity> consumedCapacity() {
+        return this.batchGetItemResponse.consumedCapacity();
+    }
+
     /**
      * A builder that is used to create a result object with the desired parameters.
      */