Feature/ddb enhanced add override configuration to enhanced req

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

@@ -63,6 +63,7 @@ public BatchGetItemRequest generateRequest(DynamoDbEnhancedClientExtension exten
         return BatchGetItemRequest.builder()
                                   .requestItems(Collections.unmodifiableMap(requestItems))
                                   .returnConsumedCapacity(request.returnConsumedCapacityAsString())
+                                  .overrideConfiguration(request.overrideConfiguration())
                                   .build();
     }
 

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

@@ -69,6 +69,7 @@ public BatchWriteItemRequest generateRequest(DynamoDbEnhancedClientExtension ext
                                         Collections.unmodifiableMap(CollectionUtils.deepCopyMap(allRequestItems)))
                                     .returnConsumedCapacity(request.returnConsumedCapacity())
                                     .returnItemCollectionMetrics(request.returnItemCollectionMetrics())
+                                    .overrideConfiguration(request.overrideConfiguration())
                                     .build();
     }
 

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

@@ -141,6 +141,7 @@ public CreateTableRequest generateRequest(TableSchema<T> tableSchema,
                                  .billingMode(billingMode)
                                  .provisionedThroughput(this.request.provisionedThroughput())
                                  .streamSpecification(this.request.streamSpecification())
+                                 .overrideConfiguration(this.request.overrideConfiguration())
                                  .build();
     }
 

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

@@ -20,6 +20,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
 import software.amazon.awssdk.enhanced.dynamodb.Expression;
 import software.amazon.awssdk.enhanced.dynamodb.Key;
@@ -94,6 +95,8 @@ public DeleteItemRequest generateRequest(TableSchema<T> tableSchema,
 
         requestBuilder = addExpressionsIfExist(requestBuilder);
 
+        applyOverrideConfiguration(requestBuilder, request);
+
         return requestBuilder.build();
     }
 
@@ -186,4 +189,19 @@ private DeleteItemRequest.Builder addPlainDeleteItemParameters(DeleteItemRequest
             requestBuilder.returnValuesOnConditionCheckFailure(enhancedRequest.returnValuesOnConditionCheckFailureAsString());
         return requestBuilder;
     }
+
+    private DeleteItemRequest.Builder applyOverrideConfiguration(DeleteItemRequest.Builder requestBuilder,
+                                                              Either<DeleteItemEnhancedRequest,
+                                                                  TransactDeleteItemEnhancedRequest> req) {
+        AwsRequestOverrideConfiguration overrideConfiguration = null;
+        if (req.left().isPresent()) {
+            overrideConfiguration = req.left().get().overrideConfiguration();
+        } else if (req.right().isPresent()) {
+            overrideConfiguration = req.right().get().overrideConfiguration();
+        }
+        if (overrideConfiguration != null) {
+            requestBuilder.overrideConfiguration(overrideConfiguration);
+        }
+        return requestBuilder;
+    }
 }

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

@@ -76,6 +76,7 @@ public GetItemRequest generateRequest(TableSchema<T> tableSchema,
                              .key(this.request.key().keyMap(tableSchema, context.indexName()))
                              .consistentRead(this.request.consistentRead())
                              .returnConsumedCapacity(this.request.returnConsumedCapacityAsString())
+                             .overrideConfiguration(this.request.overrideConfiguration())
                              .build();
     }
 

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

@@ -20,6 +20,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
 import software.amazon.awssdk.enhanced.dynamodb.Expression;
 import software.amazon.awssdk.enhanced.dynamodb.OperationContext;
@@ -108,6 +109,8 @@ public PutItemRequest generateRequest(TableSchema<T> tableSchema,
                                                               .tableName(operationContext.tableName())
                                                               .item(itemMap);
 
+        applyOverrideConfiguration(requestBuilder, request);
+
         if (request.left().isPresent()) {
             requestBuilder = addPlainPutItemParameters(requestBuilder, request.left().get());
         }
@@ -228,4 +231,19 @@ private PutItemRequest.Builder addPlainPutItemParameters(PutItemRequest.Builder
             requestBuilder.returnValuesOnConditionCheckFailure(enhancedRequest.returnValuesOnConditionCheckFailureAsString());
         return requestBuilder;
     }
+
+    private PutItemRequest.Builder applyOverrideConfiguration(PutItemRequest.Builder requestBuilder,
+                                                              Either<PutItemEnhancedRequest<T>,
+                                                                  TransactPutItemEnhancedRequest<T>> req) {
+        AwsRequestOverrideConfiguration overrideConfiguration = null;
+        if (req.left().isPresent()) {
+            overrideConfiguration = req.left().get().overrideConfiguration();
+        } else if (req.right().isPresent()) {
+            overrideConfiguration = req.right().get().overrideConfiguration();
+        }
+        if (overrideConfiguration != null) {
+            requestBuilder.overrideConfiguration(overrideConfiguration);
+        }
+        return requestBuilder;
+    }
 }

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

@@ -85,7 +85,8 @@ public QueryRequest generateRequest(TableSchema<T> tableSchema,
                                                         .exclusiveStartKey(this.request.exclusiveStartKey())
                                                         .consistentRead(this.request.consistentRead())
                                                         .returnConsumedCapacity(this.request.returnConsumedCapacity())
-                                                        .projectionExpression(projectionExpressionAsString);
+                                                        .projectionExpression(projectionExpressionAsString)
+                                                        .overrideConfiguration(this.request.overrideConfiguration());
 
         if (!TableMetadata.primaryIndexName().equals(operationContext.indexName())) {
             queryRequest = queryRequest.indexName(operationContext.indexName());

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

@@ -85,7 +85,8 @@ public ScanRequest generateRequest(TableSchema<T> tableSchema,
             .select(this.request.select())
             .expressionAttributeValues(expressionValues)
             .expressionAttributeNames(expressionNames)
-            .projectionExpression(projectionExpressionAsString);
+            .projectionExpression(projectionExpressionAsString)
+            .overrideConfiguration(this.request.overrideConfiguration());
 
         if (!TableMetadata.primaryIndexName().equals(operationContext.indexName())) {
             scanRequest = scanRequest.indexName(operationContext.indexName());

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

@@ -52,6 +52,7 @@ public OperationName operationName() {
     public TransactGetItemsRequest generateRequest(DynamoDbEnhancedClientExtension extension) {
         return TransactGetItemsRequest.builder()
                                       .transactItems(this.request.transactGetItems())
+                                      .overrideConfiguration(this.request.overrideConfiguration())
                                       .build();
     }
 

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

@@ -52,6 +52,7 @@ public TransactWriteItemsRequest generateRequest(DynamoDbEnhancedClientExtension
                                         .clientRequestToken(this.request.clientRequestToken())
                                         .returnConsumedCapacity(this.request.returnConsumedCapacity())
                                         .returnItemCollectionMetrics(this.request.returnItemCollectionMetrics())
+                                        .overrideConfiguration(this.request.overrideConfiguration())
                                         .build();
     }
 

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

@@ -28,6 +28,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
 import software.amazon.awssdk.enhanced.dynamodb.Expression;
 import software.amazon.awssdk.enhanced.dynamodb.OperationContext;
@@ -158,6 +159,8 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
             requestBuilder = requestBuilder.expressionAttributeValues(expressionValues);
         }
 
+        applyOverrideConfiguration(requestBuilder, request);
+
         return requestBuilder.build();
     }
 
@@ -355,4 +358,19 @@ private static Map<String, AttributeValue> coalesceExpressionValues(Expression f
         }
         return expressionValues;
     }
+
+    private UpdateItemRequest.Builder applyOverrideConfiguration(UpdateItemRequest.Builder requestBuilder,
+                                                                 Either<UpdateItemEnhancedRequest<T>,
+                                                                     TransactUpdateItemEnhancedRequest<T>> req) {
+        AwsRequestOverrideConfiguration overrideConfiguration = null;
+        if (req.left().isPresent()) {
+            overrideConfiguration = req.left().get().overrideConfiguration();
+        } else if (req.right().isPresent()) {
+            overrideConfiguration = req.right().get().overrideConfiguration();
+        }
+        if (overrideConfiguration != null) {
+            requestBuilder.overrideConfiguration(overrideConfiguration);
+        }
+        return requestBuilder;
+    }
 }

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

@@ -21,9 +21,11 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
 import software.amazon.awssdk.services.dynamodb.model.BatchGetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
@@ -41,10 +43,12 @@ public final class BatchGetItemEnhancedRequest {
 
     private final List<ReadBatch> readBatches;
     private final String returnConsumedCapacity;
+    private final AwsRequestOverrideConfiguration overrideConfiguration;
 
     private BatchGetItemEnhancedRequest(Builder builder) {
         this.readBatches = getListIfExist(builder.readBatches);
         this.returnConsumedCapacity = builder.returnConsumedCapacity;
+        this.overrideConfiguration = builder.overrideConfiguration;
     }
 
     /**
@@ -58,7 +62,9 @@ public static Builder builder() {
      * Returns a builder initialized with all existing values on the request object.
      */
     public Builder toBuilder() {
-        return new Builder().readBatches(readBatches).returnConsumedCapacity(this.returnConsumedCapacity);
+        return new Builder().readBatches(readBatches)
+                            .returnConsumedCapacity(this.returnConsumedCapacity)
+                            .overrideConfiguration(this.overrideConfiguration);
     }
 
     /**
@@ -87,6 +93,18 @@ public String returnConsumedCapacityAsString() {
         return returnConsumedCapacity;
     }
 
+    /**
+     * Returns the override configuration to apply to the low-level {@link BatchGetItemRequest}.
+     * <p>
+     * This can be used to customize the request, such as adding custom headers, MetricPublisher or AwsCredentialsProvider.
+     * </p>
+     *
+     * @return the {@link AwsRequestOverrideConfiguration} to apply to the underlying service call.
+     */
+    public AwsRequestOverrideConfiguration overrideConfiguration() {
+        return overrideConfiguration;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -99,13 +117,15 @@ public boolean equals(Object o) {
         BatchGetItemEnhancedRequest that = (BatchGetItemEnhancedRequest) o;
 
         return Objects.equals(this.readBatches, that.readBatches) &&
-               Objects.equals(this.returnConsumedCapacity, that.returnConsumedCapacity);
+               Objects.equals(this.returnConsumedCapacity, that.returnConsumedCapacity) &&
+               Objects.equals(this.overrideConfiguration, that.overrideConfiguration);
     }
 
     @Override
     public int hashCode() {
         int hc = readBatches != null ? readBatches.hashCode() : 0;
         hc = 31 * hc + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
+        hc = 31 * hc + (overrideConfiguration != null ? overrideConfiguration.hashCode() : 0);
         return hc;
     }
 
@@ -120,6 +140,7 @@ private static List<ReadBatch> getListIfExist(List<ReadBatch> readBatches) {
     public static final class Builder {
         private List<ReadBatch> readBatches;
         private String returnConsumedCapacity;
+        private AwsRequestOverrideConfiguration overrideConfiguration;
 
         private Builder() {
         }
@@ -183,6 +204,30 @@ public Builder returnConsumedCapacity(String returnConsumedCapacity) {
             return this;
         }
 
+        /**
+         * Sets the override configuration to apply to the low-level {@link BatchGetItemRequest}.
+         *
+         * @see BatchGetItemRequest.Builder#overrideConfiguration(AwsRequestOverrideConfiguration) 
+         * @return a builder of this type
+         */
+        public Builder overrideConfiguration(AwsRequestOverrideConfiguration overrideConfiguration) {
+            this.overrideConfiguration = overrideConfiguration;
+            return this;
+        }
+
+        /**
+         * Sets the override configuration to apply to the low-level {@link BatchGetItemRequest}.
+         *
+         * @see BatchGetItemRequest.Builder#overrideConfiguration(Consumer)
+         * @return a builder of this type
+         */
+        public Builder overrideConfiguration(Consumer<AwsRequestOverrideConfiguration.Builder> overrideConfigurationBuilder) {
+            AwsRequestOverrideConfiguration.Builder builder = AwsRequestOverrideConfiguration.builder();
+            overrideConfigurationBuilder.accept(builder);
+            this.overrideConfiguration = builder.build();
+            return this;
+        }
+
         public BatchGetItemEnhancedRequest build() {
             return new BatchGetItemEnhancedRequest(this);
         }