feature: Support of ReturnConsumedCapacity in DynamoDBEnhancedClient'…

…s BatchWrite operation (#5462) * Add support for ConsumedCapacity in a BatchWrite request * Added changelog * Add support for consumed capacity in transactWriteItems * Revert "Add support for consumed capacity in transactWriteItems" This reverts commit a7bd547. * Add support for consumed capacity in transactWriteItems (cherry picked from commit a7bd547) * Address review comments * Revert "Add support for consumed capacity in transactWriteItems" This reverts commit c58db98. * Address javadoc comments * Fix build --------- Co-authored-by: prateek.vats <[email protected]> Co-authored-by: Zoe Wang <[email protected]>
aws · Aug 16, 2024 · b4ecd2d · b4ecd2d
1 parent e89007e
commit b4ecd2d
Show file tree

Hide file tree

Showing 10 changed files with 668 additions and 55 deletions.
diff --git a/.changes/next-release/feature-DynamoDBEnhancedClient-a84a047.json b/.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"
+}
diff --git a/services-custom/dynamodb-enhanced/pom.xml b/services-custom/dynamodb-enhanced/pom.xml
@@ -141,6 +141,11 @@
             <artifactId>utils</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>service-test-utils</artifactId>

diff --git a/...software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchWriteItemOperation.java b/...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

diff --git a/...in/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchWriteItemEnhancedRequest.java b/...in/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,50 @@ 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 +119,18 @@ 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;
+        int result = writeBatches != null ? writeBatches.hashCode() : 0;
+        result = 31 * result + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
+        result = 31 * result + (returnItemCollectionMetrics != null ? returnItemCollectionMetrics.hashCode() : 0);
+        return result;
     }
 
     private static List<WriteBatch> getListIfExist(List<WriteBatch> writeBatches) {
@@ -92,10 +143,54 @@ 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() {
         }
 
+        /**
+         * Sets the ConsumedCapacity via ReturnConsumedCapacity object
+         *
+         * @see Builder#returnConsumedCapacity(ReturnConsumedCapacity)
+         */
+        public Builder returnConsumedCapacity(ReturnConsumedCapacity returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity == null ? null : returnConsumedCapacity.toString();
+            return this;
+        }
+
+        /**
+         * Sets the ConsumedCapacity via String.
+         *
+         * @see Builder#returnConsumedCapacity(String)
+         */
+        public Builder returnConsumedCapacity(String returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity;
+            return this;
+        }
+
+        /**
+         * Sets the item collection metrics.
+         *
+         * @see BatchWriteItemEnhancedRequest.Builder#returnItemCollectionMetrics(ReturnItemCollectionMetrics)
+         */
+        public BatchWriteItemEnhancedRequest.Builder returnItemCollectionMetrics(ReturnItemCollectionMetrics
+                                                                                     returnItemCollectionMetrics) {
+            this.returnItemCollectionMetrics = returnItemCollectionMetrics == null ? null :
+                                               returnItemCollectionMetrics.toString();
+            return this;
+        }
+
+        /**
+         * Sets 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 +214,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