feature: Support of ReturnConsumedCapacity in DynamoDBEnhancedClient's BatchWrite operation

.changes/next-release/feature-DynamoDBEnhancedClient-a84a047.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "DynamoDBEnhancedClient",
+    "description": "This commit introduces ConsumedCapacity in the response of a BatchWrite response",
+    "contributor": "prateek-vats"
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchWriteItemOperation.java

@@ -67,13 +67,19 @@ public BatchWriteItemRequest generateRequest(DynamoDbEnhancedClientExtension ext
         return BatchWriteItemRequest.builder()
                                     .requestItems(
                                         Collections.unmodifiableMap(CollectionUtils.deepCopyMap(allRequestItems)))
+                                    .returnConsumedCapacity(request.returnConsumedCapacity())
+                                    .returnItemCollectionMetrics(request.returnItemCollectionMetrics())
                                     .build();
     }
 
     @Override
     public BatchWriteResult transformResponse(BatchWriteItemResponse response,
                                               DynamoDbEnhancedClientExtension extension) {
-        return BatchWriteResult.builder().unprocessedRequests(response.unprocessedItems()).build();
+        return BatchWriteResult.builder()
+                               .unprocessedRequests(response.unprocessedItems())
+                               .consumedCapacity(response.consumedCapacity())
+                               .itemCollectionMetrics(response.itemCollectionMetrics())
+                               .build();
     }
 
     @Override

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchWriteItemEnhancedRequest.java

@@ -20,26 +20,34 @@
 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.BatchWriteItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
+import software.amazon.awssdk.services.dynamodb.model.ReturnItemCollectionMetrics;
 
 /**
  * Defines parameters used for the batchWriteItem() operation (such as
  * {@link DynamoDbEnhancedClient#batchWriteItem(BatchWriteItemEnhancedRequest)}).
  * <p>
- * A request contains references to keys for delete actions and items for put actions,
- * organized into one {@link WriteBatch} object per accessed table.
+ * A request contains references to keys for delete actions and items for put actions, organized into one {@link WriteBatch}
+ * object per accessed table.
  */
 @SdkPublicApi
 @ThreadSafe
 public final class BatchWriteItemEnhancedRequest {
 
     private final List<WriteBatch> writeBatches;
+    private final String returnConsumedCapacity;
+    private final String returnItemCollectionMetrics;
 
     private BatchWriteItemEnhancedRequest(Builder builder) {
         this.writeBatches = getListIfExist(builder.writeBatches);
+        this.returnConsumedCapacity = builder.returnConsumedCapacity;
+        this.returnItemCollectionMetrics = builder.returnItemCollectionMetrics;
     }
 
     /**
@@ -50,10 +58,48 @@ public static Builder builder() {
     }
 
     /**
-     * Returns a builder initialized with all existing values on the request object.
+     * @return a builder with all existing values set
      */
     public Builder toBuilder() {
-        return new Builder().writeBatches(writeBatches);
+        return builder().writeBatches(writeBatches).returnConsumedCapacity(returnConsumedCapacity).returnItemCollectionMetrics(returnItemCollectionMetrics);
+    }
+
+    /**
+     * Whether to return the capacity consumed by this operation.
+     *
+     * @see BatchWriteItemEnhancedRequest#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;
+    }
+
+    /**
+     * Whether to return the item collection metrics.
+     *
+     * @see BatchWriteItemRequest#returnItemCollectionMetrics()
+     */
+    public ReturnItemCollectionMetrics returnItemCollectionMetrics() {
+        return ReturnItemCollectionMetrics.fromValue(returnItemCollectionMetrics);
+    }
+
+    /**
+     * Whether to return the item collection metrics.
+     * <p>
+     * Similar to {@link #returnItemCollectionMetrics()} but return the value as a string. This is useful in situations where the
+     * value is not defined in {@link ReturnItemCollectionMetrics}.
+     */
+    public String returnItemCollectionMetricsAsString() {
+        return returnItemCollectionMetrics;
     }
 
     /**
@@ -71,15 +117,14 @@ public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) {
             return false;
         }
-
         BatchWriteItemEnhancedRequest that = (BatchWriteItemEnhancedRequest) o;
-
-        return writeBatches != null ? writeBatches.equals(that.writeBatches) : that.writeBatches == null;
+        return Objects.equals(writeBatches, that.writeBatches) && Objects.equals(returnConsumedCapacity,
+                                                                                 that.returnConsumedCapacity) && Objects.equals(returnItemCollectionMetrics, that.returnItemCollectionMetrics);
     }
 
     @Override
     public int hashCode() {
-        return writeBatches != null ? writeBatches.hashCode() : 0;
+        return Objects.hash(writeBatches, returnConsumedCapacity, returnItemCollectionMetrics);
     }
 
     private static List<WriteBatch> getListIfExist(List<WriteBatch> writeBatches) {
@@ -92,10 +137,53 @@ private static List<WriteBatch> getListIfExist(List<WriteBatch> writeBatches) {
     @NotThreadSafe
     public static final class Builder {
         private List<WriteBatch> writeBatches;
+        private String returnConsumedCapacity;
+        private String returnItemCollectionMetrics;
 
         private Builder() {
         }
 
+        /**
+         * Whether to return the capacity consumed by this operation.
+         *
+         * @see Builder#returnConsumedCapacity(ReturnConsumedCapacity)
+         */
+        public Builder returnConsumedCapacity(ReturnConsumedCapacity returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity == null ? null : returnConsumedCapacity.toString();
+            return this;
+        }
+
+        /**
+         * Whether to return the capacity consumed by this operation.
+         *
+         * @see Builder#returnConsumedCapacity(String)
+         */
+        public Builder returnConsumedCapacity(String returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity;
+            return this;
+        }
+
+        /**
+         * Whether to return the item collection metrics.
+         *
+         * @see BatchWriteItemEnhancedRequest.Builder#returnItemCollectionMetrics(ReturnItemCollectionMetrics)
+         */
+        public BatchWriteItemEnhancedRequest.Builder returnItemCollectionMetrics(ReturnItemCollectionMetrics returnItemCollectionMetrics) {
+            this.returnItemCollectionMetrics = returnItemCollectionMetrics == null ? null :
+                                               returnItemCollectionMetrics.toString();
+            return this;
+        }
+
+        /**
+         * Whether to return the item collection metrics.
+         *
+         * @see BatchWriteItemEnhancedRequest.Builder#returnItemCollectionMetrics(String)
+         */
+        public BatchWriteItemEnhancedRequest.Builder returnItemCollectionMetrics(String returnItemCollectionMetrics) {
+            this.returnItemCollectionMetrics = returnItemCollectionMetrics;
+            return this;
+        }
+
         /**
          * Sets a collection of write batches to use in the batchWriteItem operation.
          *
@@ -119,8 +207,8 @@ public Builder writeBatches(WriteBatch... writeBatches) {
         }
 
         /**
-         * Adds a write batch to the collection of batches on this builder.
-         * If this is the first batch, the method creates a new list.
+         * Adds a write batch to the collection of batches on this builder. If this is the first batch, the method creates a new
+         * list.
          *
          * @param writeBatch a single write batch
          * @return a builder of this type

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchWriteResult.java

@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.annotations.SdkPublicApi;
@@ -30,29 +31,52 @@
 import software.amazon.awssdk.enhanced.dynamodb.MappedTableResource;
 import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
 import software.amazon.awssdk.enhanced.dynamodb.internal.operations.DefaultOperationContext;
+import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity;
 import software.amazon.awssdk.services.dynamodb.model.DeleteRequest;
+import software.amazon.awssdk.services.dynamodb.model.ItemCollectionMetrics;
 import software.amazon.awssdk.services.dynamodb.model.PutRequest;
 import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
 
 /**
  * Defines the result of the batchWriteItem() operation, such as
- * {@link DynamoDbEnhancedClient#batchWriteItem(BatchWriteItemEnhancedRequest)}. The result describes any unprocessed items
- * after the operation completes.
+ * {@link DynamoDbEnhancedClient#batchWriteItem(BatchWriteItemEnhancedRequest)}. The result describes any unprocessed items after
+ * the operation completes.
  * <ul>
  *     <li>Use the {@link #unprocessedPutItemsForTable(MappedTableResource)} method once for each table present in the request
  *  to get any unprocessed items from a put action on that table.</li>
  *     <li>Use the {@link #unprocessedDeleteItemsForTable(MappedTableResource)} method once for each table present in the request
  *  to get any unprocessed items from a delete action on that table.</li>
  * </ul>
- *
  */
 @SdkPublicApi
 @ThreadSafe
 public final class BatchWriteResult {
     private final Map<String, List<WriteRequest>> unprocessedRequests;
+    private final List<ConsumedCapacity> consumedCapacity;
+    private final Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics;
 
     private BatchWriteResult(Builder builder) {
         this.unprocessedRequests = Collections.unmodifiableMap(builder.unprocessedRequests);
+        this.consumedCapacity = Collections.unmodifiableList(builder.consumedCapacity);
+        this.itemCollectionMetrics = builder.itemCollectionMetrics;
+    }
+
+    /**
+     * The capacity units consumed by the {@code Write} operation.
+     *
+     * @see BatchWriteResult#consumedCapacity() for more information.
+     */
+    public List<ConsumedCapacity> consumedCapacity() {
+        return consumedCapacity;
+    }
+
+    /**
+     * The capacity units consumed by the {@code Write} operation.
+     *
+     * @see BatchWriteResult#consumedCapacity() for more information.
+     */
+    public Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics() {
+        return itemCollectionMetrics;
     }
 
     /**
@@ -63,11 +87,11 @@ public static Builder builder() {
     }
 
     /**
-     * Retrieve any unprocessed put action items belonging to the supplied table from the result .
-     * Call this method once for each table present in the batch request.
+     * Retrieve any unprocessed put action items belonging to the supplied table from the result . Call this method once for each
+     * table present in the batch request.
      *
      * @param mappedTable the table to retrieve unprocessed items for
-     * @param <T> the type of the table items
+     * @param <T>         the type of the table items
      * @return a list of items
      */
     public <T> List<T> unprocessedPutItemsForTable(MappedTableResource<T> mappedTable) {
@@ -86,9 +110,26 @@ public <T> List<T> unprocessedPutItemsForTable(MappedTableResource<T> mappedTabl
                             .collect(Collectors.toList());
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        BatchWriteResult result = (BatchWriteResult) o;
+        return Objects.equals(unprocessedRequests, result.unprocessedRequests) && Objects.equals(consumedCapacity, result.consumedCapacity) && Objects.equals(itemCollectionMetrics, result.itemCollectionMetrics);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(unprocessedRequests, consumedCapacity, itemCollectionMetrics);
+    }
+
     /**
-     * Retrieve any unprocessed delete action keys belonging to the supplied table from the result.
-     * Call this method once for each table present in the batch request.
+     * Retrieve any unprocessed delete action keys belonging to the supplied table from the result. Call this method once for each
+     * table present in the batch request.
      *
      * @param mappedTable the table to retrieve unprocessed items for.
      * @return a list of keys that were not processed as part of the batch request.
@@ -114,10 +155,41 @@ public List<Key> unprocessedDeleteItemsForTable(MappedTableResource<?> mappedTab
     @NotThreadSafe
     public static final class Builder {
         private Map<String, List<WriteRequest>> unprocessedRequests;
+        private List<ConsumedCapacity> consumedCapacity;
+        private Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics;
 
         private Builder() {
         }
 
+        /**
+         * Set the capacity units consumed by the batch write operation.
+         *
+         * <p>
+         * This is a list of ConsumedCapacity objects, one for each table in the batch write operation. The list is ordered
+         * according to the order of the request parameters.
+         *
+         * @param consumedCapacity
+         * @return a builder of this type
+         */
+        public Builder consumedCapacity(List<ConsumedCapacity> consumedCapacity) {
+            this.consumedCapacity = consumedCapacity;
+            return this;
+        }
+
+        /**
+         * Set the capacity units consumed by the batch write operation.
+         *
+         * <p>
+         * This is a Map of List of ItemCollectionMetrics objects, one for each table in the batch write operation.
+         *
+         * @param itemCollectionMetrics
+         * @return a builder of this type
+         */
+        public Builder itemCollectionMetrics(Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics) {
+            this.itemCollectionMetrics = itemCollectionMetrics;
+            return this;
+        }
+
         /**
          * Add a map of unprocessed requests to this result object.
          *