feature: Support of ReturnConsumedCapacity in DynamoDBEnhancedClient's TransactWriteItem operation (#5501)

* 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)

* Revert "Added changelog"

This reverts commit 1d72986

* Revert "Add support for ConsumedCapacity in a BatchWrite request"

This reverts commit d185260.

* Refactor the Consumed Capacity changes to make them backward compatible

* Fix doc comments

* Fix formatting and change Objects.hash usage to Objects.hashCode

* dependency fix

* Fix javadoc comments

* pull changes from 4d5c584

---------

Co-authored-by: prateek.vats <[email protected]>

Author: prateek-vats

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedAsyncClient.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
@@ -104,7 +105,7 @@ public CompletableFuture<List<Document>> transactGetItems(
     @Override
     public CompletableFuture<Void> transactWriteItems(TransactWriteItemsEnhancedRequest request) {
         TransactWriteItemsOperation operation = TransactWriteItemsOperation.create(request);
-        return operation.executeAsync(dynamoDbClient, extension);
+        return operation.executeAsync(dynamoDbClient, extension).thenApply(r -> null);
     }
 
     @Override
@@ -116,6 +117,21 @@ public CompletableFuture<Void> transactWriteItems(
         return transactWriteItems(builder.build());
     }
 
+    @Override
+    public CompletableFuture<TransactWriteItemsEnhancedResponse> transactWriteItemsWithResponse(
+        TransactWriteItemsEnhancedRequest request) {
+        TransactWriteItemsOperation operation = TransactWriteItemsOperation.create(request);
+        return operation.executeAsync(dynamoDbClient, extension);
+    }
+
+    @Override
+    public CompletableFuture<TransactWriteItemsEnhancedResponse> transactWriteItemsWithResponse(
+        Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
+        TransactWriteItemsEnhancedRequest.Builder builder = TransactWriteItemsEnhancedRequest.builder();
+        requestConsumer.accept(builder);
+        return transactWriteItemsWithResponse(builder.build());
+    }
+
     public DynamoDbAsyncClient dynamoDbAsyncClient() {
         return dynamoDbClient;
     }
@@ -139,8 +155,8 @@ public boolean equals(Object o) {
 
         DefaultDynamoDbEnhancedAsyncClient that = (DefaultDynamoDbEnhancedAsyncClient) o;
 
-        if (dynamoDbClient != null ? ! dynamoDbClient.equals(that.dynamoDbClient)
-            : that.dynamoDbClient != null) {
+        if (dynamoDbClient != null ? !dynamoDbClient.equals(that.dynamoDbClient)
+                                   : that.dynamoDbClient != null) {
 
             return false;
         }

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedClient.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
@@ -100,7 +101,8 @@ public List<Document> transactGetItems(
     @Override
     public Void transactWriteItems(TransactWriteItemsEnhancedRequest request) {
         TransactWriteItemsOperation operation = TransactWriteItemsOperation.create(request);
-        return operation.execute(dynamoDbClient, extension);
+        operation.execute(dynamoDbClient, extension);
+        return null;
     }
 
     @Override
@@ -110,6 +112,20 @@ public Void transactWriteItems(Consumer<TransactWriteItemsEnhancedRequest.Builde
         return transactWriteItems(builder.build());
     }
 
+    @Override
+    public TransactWriteItemsEnhancedResponse transactWriteItemsWithResponse(TransactWriteItemsEnhancedRequest request) {
+        TransactWriteItemsOperation operation = TransactWriteItemsOperation.create(request);
+        return operation.execute(dynamoDbClient, extension);
+    }
+
+    @Override
+    public TransactWriteItemsEnhancedResponse transactWriteItemsWithResponse(
+        Consumer<TransactWriteItemsEnhancedRequest.Builder> requestConsumer) {
+        TransactWriteItemsEnhancedRequest.Builder builder = TransactWriteItemsEnhancedRequest.builder();
+        requestConsumer.accept(builder);
+        return transactWriteItemsWithResponse(builder.build());
+    }
+
     public DynamoDbClient dynamoDbClient() {
         return dynamoDbClient;
     }
@@ -133,19 +149,19 @@ public boolean equals(Object o) {
 
         DefaultDynamoDbEnhancedClient that = (DefaultDynamoDbEnhancedClient) o;
 
-        if (dynamoDbClient != null ? ! dynamoDbClient.equals(that.dynamoDbClient) : that.dynamoDbClient != null) {
+        if (dynamoDbClient != null ? !dynamoDbClient.equals(that.dynamoDbClient) : that.dynamoDbClient != null) {
             return false;
         }
         return extension != null ?
-            extension.equals(that.extension) :
-            that.extension == null;
+               extension.equals(that.extension) :
+               that.extension == null;
     }
 
     @Override
     public int hashCode() {
         int result = dynamoDbClient != null ? dynamoDbClient.hashCode() : 0;
         result = 31 * result + (extension != null ?
-            extension.hashCode() : 0);
+                                extension.hashCode() : 0);
         return result;
     }
 

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

@@ -20,16 +20,17 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsResponse;
 
 @SdkInternalApi
 public class TransactWriteItemsOperation
-    implements DatabaseOperation<TransactWriteItemsRequest, TransactWriteItemsResponse, Void> {
+    implements DatabaseOperation<TransactWriteItemsRequest, TransactWriteItemsResponse, TransactWriteItemsEnhancedResponse> {
 
-    private TransactWriteItemsEnhancedRequest request;
+    private final TransactWriteItemsEnhancedRequest request;
 
     private TransactWriteItemsOperation(TransactWriteItemsEnhancedRequest request) {
         this.request = request;
@@ -49,18 +50,23 @@ 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
     public Function<TransactWriteItemsRequest, TransactWriteItemsResponse> serviceCall(
         DynamoDbClient dynamoDbClient) {
-
         return dynamoDbClient::transactWriteItems;
     }