Add support for consumed capacity in transactWriteItems

Author: ragnar9780

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

@@ -1,6 +1,6 @@
 {
     "type": "feature",
     "category": "DynamoDBEnhancedClient",
-    "description": "This commit introduces ConsumedCapacity in the response of a BatchWrite response",
+    "description": "Support of ReturnConsumedCapacity in DynamoDBEnhancedClient's BatchWrite and TransctWriteItems operations",
     "contributor": "prateek-vats"
 }

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

@@ -35,6 +35,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactGetItemsEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 
@@ -380,7 +381,7 @@ default CompletableFuture<List<Document>> transactGetItems(
      * @param request A {@link BatchWriteItemEnhancedRequest} containing keys grouped by tables.
      * @return a {@link CompletableFuture} of {@link Void}.
      */
-    default CompletableFuture<Void> transactWriteItems(TransactWriteItemsEnhancedRequest request) {
+    default CompletableFuture<TransactWriteItemsEnhancedResponse> transactWriteItems(TransactWriteItemsEnhancedRequest request) {
         throw new UnsupportedOperationException();
     }
 
@@ -425,7 +426,7 @@ default CompletableFuture<Void> transactWriteItems(TransactWriteItemsEnhancedReq
      * tables.
      * @return a {@link CompletableFuture} of {@link Void}.
      */
-    default CompletableFuture<Void> transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
+    default CompletableFuture<TransactWriteItemsEnhancedResponse> transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
         throw new UnsupportedOperationException();
     }
 

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

@@ -34,6 +34,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactGetItemsEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.BatchGetItemRequest;
@@ -394,7 +395,7 @@ default List<Document> transactGetItems(Consumer<TransactGetItemsEnhancedRequest
      *
      * @param request A {@link BatchWriteItemEnhancedRequest} containing keys grouped by tables.
      */
-    default Void transactWriteItems(TransactWriteItemsEnhancedRequest request) {
+    default TransactWriteItemsEnhancedResponse transactWriteItems(TransactWriteItemsEnhancedRequest request) {
         throw new UnsupportedOperationException();
     }
 
@@ -439,7 +440,7 @@ default Void transactWriteItems(TransactWriteItemsEnhancedRequest request) {
      * @param requestConsumer a {@link Consumer} of {@link TransactWriteItemsEnhancedRequest} containing keys and items grouped
      * by tables.
      */
-    default Void transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
+    default TransactWriteItemsEnhancedResponse transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
         throw new UnsupportedOperationException();
     }
 

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/client/DefaultDynamoDbEnhancedAsyncClient.java

@@ -36,6 +36,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.BatchWriteResult;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactGetItemsEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 
 @SdkInternalApi
@@ -102,13 +103,13 @@ public CompletableFuture<List<Document>> transactGetItems(
     }
 
     @Override
-    public CompletableFuture<Void> transactWriteItems(TransactWriteItemsEnhancedRequest request) {
+    public CompletableFuture<TransactWriteItemsEnhancedResponse> transactWriteItems(TransactWriteItemsEnhancedRequest request) {
         TransactWriteItemsOperation operation = TransactWriteItemsOperation.create(request);
         return operation.executeAsync(dynamoDbClient, extension);
     }
 
     @Override
-    public CompletableFuture<Void> transactWriteItems(
+    public CompletableFuture<TransactWriteItemsEnhancedResponse> transactWriteItems(
         Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
 
         TransactWriteItemsEnhancedRequest.Builder builder = TransactWriteItemsEnhancedRequest.builder();

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/client/DefaultDynamoDbEnhancedClient.java

@@ -35,6 +35,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.BatchWriteResult;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactGetItemsEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 
 @SdkInternalApi
@@ -98,13 +99,13 @@ public List<Document> transactGetItems(
     }
 
     @Override
-    public Void transactWriteItems(TransactWriteItemsEnhancedRequest request) {
+    public TransactWriteItemsEnhancedResponse transactWriteItems(TransactWriteItemsEnhancedRequest request) {
         TransactWriteItemsOperation operation = TransactWriteItemsOperation.create(request);
         return operation.execute(dynamoDbClient, extension);
     }
 
     @Override
-    public Void transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
+    public TransactWriteItemsEnhancedResponse transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
         TransactWriteItemsEnhancedRequest.Builder builder = TransactWriteItemsEnhancedRequest.builder();
         requestConsumer.accept(builder);
         return transactWriteItems(builder.build());

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

@@ -19,6 +19,7 @@
 import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -27,7 +28,7 @@
 
 @SdkInternalApi
 public class TransactWriteItemsOperation
-    implements DatabaseOperation<TransactWriteItemsRequest, TransactWriteItemsResponse, Void> {
+    implements DatabaseOperation<TransactWriteItemsRequest, TransactWriteItemsResponse, TransactWriteItemsEnhancedResponse> {
 
     private TransactWriteItemsEnhancedRequest request;
 
@@ -49,12 +50,17 @@ public TransactWriteItemsRequest generateRequest(DynamoDbEnhancedClientExtension
         return TransactWriteItemsRequest.builder()
                                         .transactItems(this.request.transactWriteItems())
                                         .clientRequestToken(this.request.clientRequestToken())
+            .returnConsumedCapacity(this.request.returnConsumedCapacity())
+            .returnItemCollectionMetrics(this.request.returnItemCollectionMetrics())
                                         .build();
     }
 
     @Override
-    public Void transformResponse(TransactWriteItemsResponse response, DynamoDbEnhancedClientExtension extension) {
-        return null;        // this operation does not return results
+    public TransactWriteItemsEnhancedResponse transformResponse(TransactWriteItemsResponse response, DynamoDbEnhancedClientExtension extension) {
+        return TransactWriteItemsEnhancedResponse.builder()
+            .consumedCapacity(response.consumedCapacity())
+            .itemCollectionMetrics(response.itemCollectionMetrics())
+                                                 .build();
     }
 
     @Override

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

@@ -19,6 +19,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.NotThreadSafe;
@@ -33,7 +34,12 @@
 import software.amazon.awssdk.enhanced.dynamodb.internal.operations.PutItemOperation;
 import software.amazon.awssdk.enhanced.dynamodb.internal.operations.TransactableWriteOperation;
 import software.amazon.awssdk.enhanced.dynamodb.internal.operations.UpdateItemOperation;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
+import software.amazon.awssdk.services.dynamodb.model.ReturnItemCollectionMetrics;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
+import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 
 /**
  * Defines parameters used for the transaction operation transactWriteItems() (such as
@@ -56,10 +62,14 @@ public final class TransactWriteItemsEnhancedRequest {
     private final List<TransactWriteItem> transactWriteItems;
 
     private final String clientRequestToken;
+    private final String returnConsumedCapacity;
+    private final String returnItemCollectionMetrics;
 
     private TransactWriteItemsEnhancedRequest(Builder builder) {
         this.transactWriteItems = getItemsFromSupplier(builder.itemSupplierList);
         this.clientRequestToken = builder.clientRequestToken;
+        this.returnConsumedCapacity = builder.returnConsumedCapacity;
+        this.returnItemCollectionMetrics = builder.returnItemCollectionMetrics;
     }
 
     /**
@@ -89,6 +99,34 @@ public String clientRequestToken() {
         return clientRequestToken;
     }
 
+    /**
+     * Whether to return the capacity consumed by this operation.
+     *
+     * @see PutItemRequest#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 DeleteItemRequest#returnItemCollectionMetrics()
+     */
+    public ReturnItemCollectionMetrics returnItemCollectionMetrics() {
+        return ReturnItemCollectionMetrics.fromValue(returnItemCollectionMetrics);
+    }
+
     /**
      * Returns the list of {@link TransactWriteItem} that represents all actions in the request.
      */
@@ -104,15 +142,13 @@ public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) {
             return false;
         }
-
         TransactWriteItemsEnhancedRequest that = (TransactWriteItemsEnhancedRequest) o;
-
-        return transactWriteItems != null ? transactWriteItems.equals(that.transactWriteItems) : that.transactWriteItems == null;
+        return Objects.equals(transactWriteItems, that.transactWriteItems) && Objects.equals(clientRequestToken, that.clientRequestToken) && Objects.equals(returnConsumedCapacity, that.returnConsumedCapacity) && Objects.equals(returnItemCollectionMetrics, that.returnItemCollectionMetrics);
     }
 
     @Override
     public int hashCode() {
-        return transactWriteItems != null ? transactWriteItems.hashCode() : 0;
+        return Objects.hash(transactWriteItems, clientRequestToken, returnConsumedCapacity, returnItemCollectionMetrics);
     }
 
     /**
@@ -125,6 +161,8 @@ public static final class Builder {
         private List<Supplier<TransactWriteItem>> itemSupplierList = new ArrayList<>();
 
         private String clientRequestToken;
+        private String returnConsumedCapacity;
+        private String returnItemCollectionMetrics;
 
         private Builder() {
         }
@@ -337,6 +375,37 @@ public <T> Builder addUpdateItem(MappedTableResource<T> mappedTableResource, T i
                                                  .build());
         }
 
+        /**
+         * Whether to return the capacity consumed by this operation.
+         *
+         * @see TransactWriteItemsEnhancedRequest.Builder#returnConsumedCapacity(ReturnConsumedCapacity)
+         */
+        public TransactWriteItemsEnhancedRequest.Builder returnConsumedCapacity(ReturnConsumedCapacity returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity == null ? null : returnConsumedCapacity.toString();
+            return this;
+        }
+
+        /**
+         * Whether to return the capacity consumed by this operation.
+         *
+         * @see TransactWriteItemsEnhancedRequest.Builder#returnConsumedCapacity(String)
+         */
+        public TransactWriteItemsEnhancedRequest.Builder returnConsumedCapacity(String returnConsumedCapacity) {
+            this.returnConsumedCapacity = returnConsumedCapacity;
+            return this;
+        }
+
+        /**
+         * Whether to return the item collection metrics.
+         *
+         * @see UpdateItemRequest.Builder#returnItemCollectionMetrics(ReturnItemCollectionMetrics)
+         */
+        public TransactWriteItemsEnhancedRequest.Builder returnItemCollectionMetrics(ReturnItemCollectionMetrics returnItemCollectionMetrics) {
+            this.returnItemCollectionMetrics = returnItemCollectionMetrics == null ? null :
+                                               returnItemCollectionMetrics.toString();
+            return this;
+        }
+
         /**
          * Sets the clientRequestToken in this builder.
          *

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

@@ -0,0 +1,105 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.model;
+
+import java.util.List;
+import java.util.Map;
+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.ConsumedCapacity;
+import software.amazon.awssdk.services.dynamodb.model.ItemCollectionMetrics;
+
+/**
+ * Defines the elements returned by DynamoDB from a {@code TransactWriteItemsOperation} operation, such as
+ * {@link DynamoDbEnhancedClient#transactWriteItems(TransactWriteItemsEnhancedRequest)}
+ *
+ */
+@SdkPublicApi
+@ThreadSafe
+public final class TransactWriteItemsEnhancedResponse {
+    private final List<ConsumedCapacity> consumedCapacity;
+    private final Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics;
+
+    private TransactWriteItemsEnhancedResponse(Builder builder) {
+        this.consumedCapacity = builder.consumedCapacity;
+        this.itemCollectionMetrics = builder.itemCollectionMetrics;
+    }
+
+    /**
+     * The capacity units consumed by the {@code UpdateItem} operation.
+     *
+     * @see TransactWriteItemsEnhancedResponse#consumedCapacity() for more information.
+     */
+    public List<ConsumedCapacity> consumedCapacity() {
+        return consumedCapacity;
+    }
+
+    /**
+     * Information about item collections, if any, that were affected by the {@code UpdateItem} operation.
+     *
+     * @see TransactWriteItemsEnhancedResponse#itemCollectionMetrics() for more information.
+     */
+    public Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics() {
+        return itemCollectionMetrics;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        TransactWriteItemsEnhancedResponse that = (TransactWriteItemsEnhancedResponse) o;
+        return Objects.equals(consumedCapacity, that.consumedCapacity) && Objects.equals(itemCollectionMetrics, that.itemCollectionMetrics);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(consumedCapacity, itemCollectionMetrics);
+    }
+
+    /**
+     * Creates a newly initialized builder for a request object.
+     */
+    public static TransactWriteItemsEnhancedResponse.Builder builder() {
+        return new TransactWriteItemsEnhancedResponse.Builder();
+    }
+
+    @NotThreadSafe
+    public static final class Builder{
+        private List<ConsumedCapacity> consumedCapacity;
+        private Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics;
+
+        public Builder consumedCapacity(List<ConsumedCapacity> consumedCapacity) {
+            this.consumedCapacity = consumedCapacity;
+            return this;
+        }
+
+        public Builder itemCollectionMetrics(Map<String, List<ItemCollectionMetrics>> itemCollectionMetrics) {
+            this.itemCollectionMetrics = itemCollectionMetrics;
+            return this;
+        }
+
+        public TransactWriteItemsEnhancedResponse build() {
+            return new TransactWriteItemsEnhancedResponse(this);
+        }
+    }
+}