fix(CopySourceRequestInS3): Updated codegen to customize the request before calling service

codegen/src/main/java/software/amazon/awssdk/codegen/model/config/customization/CustomizationConfig.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.codegen.model.rules.endpoints.ParameterModel;
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
 import software.amazon.awssdk.codegen.model.service.CustomOperationContextParam;
+import software.amazon.awssdk.codegen.model.service.CustomRequestTransformer;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.core.traits.PayloadTrait;
 import software.amazon.awssdk.utils.AttributeMap;
@@ -328,6 +329,13 @@ public class CustomizationConfig {
 
     private List<CustomOperationContextParam> customOperationContextParams;
 
+    /**
+     * A map that associates API names with their respective custom request transformers.
+     * The {@link CustomRequestTransformer} allows for dynamic and specific handling of API requests, ensuring that each
+     * request that requires custom handling can be appropriately transformed based on its corresponding API name.
+     */
+    private Map<String, CustomRequestTransformer> customRequestTransformerMap;
+
     private CustomizationConfig() {
     }
 
@@ -869,4 +877,13 @@ public List<CustomOperationContextParam> getCustomOperationContextParams() {
     public void setCustomOperationContextParams(List<CustomOperationContextParam> customOperationContextParams) {
         this.customOperationContextParams = customOperationContextParams;
     }
+
+    public Map<String, CustomRequestTransformer> getCustomRequestTransformerMap() {
+        return customRequestTransformerMap;
+    }
+
+    public void setCustomRequestTransformerMap(Map<String, CustomRequestTransformer> customRequestTransformerMap) {
+        this.customRequestTransformerMap = customRequestTransformerMap;
+    }
+
 }

codegen/src/main/java/software/amazon/awssdk/codegen/model/service/CustomRequestTransformer.java

@@ -0,0 +1,65 @@
+/*
+ * 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.codegen.model.service;
+
+/**
+ * Represents a custom request transformer for API requests.
+ *
+ * <p>This class allows for dynamic and specific transformation of API requests,
+ * ensuring that each request is appropriately transformed based on the
+ * transformation logic defined in the specified {@link CustomRequestTransformer#getClassName()} and
+ * {@link CustomRequestTransformer#getMethodName()}.
+ *
+ * <p>Example:
+ * <pre>
+ * {
+ *     "methodName": "dummyRequestModifier",
+ *     "className": "software.amazon.awssdk.codegen.internal.UtilsTest"
+ * }
+ * </pre>
+ *
+ * <p>The class should have a public static method   dummyRequestModifier
+ * that takes an input and returns an output of ApiRequest for which Customization is applied.
+ */
+
+public class CustomRequestTransformer {
+
+    /**
+     * The fully qualified name of the class that defines the transformation method. The {@code methodName} is the
+     */
+    private String className;
+
+    /**
+     * The name of the method within that class which will perform the transformation
+     */
+    private String methodName;
+
+    public String getClassName() {
+        return className;
+    }
+
+    public void setClassName(String className) {
+        this.className = className;
+    }
+
+    public String getMethodName() {
+        return methodName;
+    }
+
+    public void setMethodName(String methodName) {
+        this.methodName = methodName;
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java

@@ -25,6 +25,7 @@
 import static software.amazon.awssdk.codegen.internal.Constant.EVENT_PUBLISHER_PARAM_NAME;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.addS3ArnableFieldCode;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
+import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.addRequestModifierCode;
 import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.getProtocolSpecs;
 
 import com.squareup.javapoet.ClassName;
@@ -332,6 +333,7 @@ protected void addCloseMethod(TypeSpec.Builder type) {
     @Override
     protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, OperationModel opModel) {
 
+        addRequestModifierCode(opModel, model).ifPresent(builder::addCode);
         builder.addModifiers(PUBLIC)
                .addAnnotation(Override.class);
         builder.addStatement("$T clientConfiguration = updateSdkClientConfiguration($L, this.clientConfiguration)",

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java

@@ -20,10 +20,12 @@
 import static javax.lang.model.element.Modifier.PROTECTED;
 import static javax.lang.model.element.Modifier.PUBLIC;
 import static javax.lang.model.element.Modifier.STATIC;
+import static software.amazon.awssdk.codegen.poet.PoetUtils.classNameFromFqcn;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.addS3ArnableFieldCode;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
 
 import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.CodeBlock;
 import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
@@ -36,6 +38,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -50,6 +53,7 @@
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
+import software.amazon.awssdk.codegen.model.service.CustomRequestTransformer;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
 import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
@@ -249,6 +253,8 @@ private Stream<MethodSpec> operations(OperationModel opModel) {
     private MethodSpec traditionalMethod(OperationModel opModel) {
         MethodSpec.Builder method = SyncClientInterface.operationMethodSignature(model, opModel)
                                                        .addAnnotation(Override.class);
+
+        addRequestModifierCode(opModel, model).ifPresent(method::addCode);
         if (!useSraAuth) {
             method.addCode(ClientClassUtils.callApplySignerOverrideMethod(opModel));
         }
@@ -341,6 +347,27 @@ private MethodSpec traditionalMethod(OperationModel opModel) {
         return method.build();
     }
 
+    public static Optional<CodeBlock> addRequestModifierCode(OperationModel opModel, IntermediateModel model) {
+
+        Map<String, CustomRequestTransformer> customRequestTransformerMap =
+            model.getCustomizationConfig().getCustomRequestTransformerMap();
+
+        if (!CollectionUtils.isNullOrEmpty(customRequestTransformerMap)) {
+            CustomRequestTransformer requestCustomizer = customRequestTransformerMap.get(opModel.getOperationName());
+            if (requestCustomizer != null) {
+                CodeBlock.Builder builder = CodeBlock.builder();
+                ClassName instanceType = classNameFromFqcn(requestCustomizer.getClassName());
+                builder.addStatement("$L = $T.$N($L)",
+                                     opModel.getInput().getVariableName(),
+                                     instanceType,
+                                     requestCustomizer.getMethodName(),
+                                     opModel.getInput().getVariableName());
+                return Optional.of(builder.build());
+            }
+        }
+        return Optional.empty();
+    }
+
     @Override
     protected void addCloseMethod(TypeSpec.Builder type) {
         MethodSpec method = MethodSpec.methodBuilder("close")

codegen/src/test/java/software/amazon/awssdk/codegen/internal/UtilsTest.java

@@ -43,4 +43,9 @@ public void testUnCapitalize() {
         capitalizedToUncapitalized.forEach((capitalized,unCapitalized) ->
                 assertThat(Utils.unCapitalize(capitalized), is(equalTo(unCapitalized))));
     }
+
+    // Dummy No-op function which just returns the input as the return function.
+    public static <T> T dummyRequestModifier(T input) {
+        return input;
+    }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/query/customization.config

@@ -21,5 +21,11 @@
         }
       }
     }
-  ]
+  ],
+  "customRequestTransformerMap": {
+    "OperationWithCustomMember": {
+      "methodName": "dummyRequestModifier",
+      "className": "software.amazon.awssdk.codegen.internal.UtilsTest"
+    }
+  }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/query/service-2.json

@@ -88,6 +88,14 @@
         "encodings": ["gzip"]
       }
     },
+    "OperationWithCustomMember": {
+      "name": "OperationWithCustomMember",
+      "http": {
+        "method": "POST",
+        "requestUri": "/"
+      },
+      "input":{"shape":"WithCustomMember"}
+    },
     "APostOperation": {
       "name": "APostOperation",
       "http": {
@@ -202,6 +210,17 @@
         }
       }
     },
+    "WithCustomMember": {
+      "type": "structure",
+      "members": {
+        "StringMemberToBeUpdate" : {
+          "shape": "String"
+        },
+        "StringMember": {
+          "shape": "String"
+        }
+      }
+    },
     "WithOperationContextParam": {
       "type": "structure",
       "members": {

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/sra/test-query-async-client-class.java

@@ -14,6 +14,7 @@
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.awscore.internal.AwsProtocolMetadata;
 import software.amazon.awssdk.awscore.internal.AwsServiceProtocol;
+import software.amazon.awssdk.codegen.internal.UtilsTest;
 import software.amazon.awssdk.core.CredentialType;
 import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.SdkPlugin;
@@ -51,6 +52,8 @@
 import software.amazon.awssdk.services.query.model.OperationWithChecksumRequiredResponse;
 import software.amazon.awssdk.services.query.model.OperationWithContextParamRequest;
 import software.amazon.awssdk.services.query.model.OperationWithContextParamResponse;
+import software.amazon.awssdk.services.query.model.OperationWithCustomMemberRequest;
+import software.amazon.awssdk.services.query.model.OperationWithCustomMemberResponse;
 import software.amazon.awssdk.services.query.model.OperationWithCustomizedOperationContextParamRequest;
 import software.amazon.awssdk.services.query.model.OperationWithCustomizedOperationContextParamResponse;
 import software.amazon.awssdk.services.query.model.OperationWithNoneAuthTypeRequest;
@@ -74,6 +77,7 @@
 import software.amazon.awssdk.services.query.transform.GetOperationWithChecksumRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithChecksumRequiredRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithContextParamRequestMarshaller;
+import software.amazon.awssdk.services.query.transform.OperationWithCustomMemberRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithCustomizedOperationContextParamRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithNoneAuthTypeRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithOperationContextParamRequestMarshaller;
@@ -470,6 +474,63 @@ public CompletableFuture<OperationWithContextParamResponse> operationWithContext
         }
     }
 
+    /**
+     * Invokes the OperationWithCustomMember operation asynchronously.
+     *
+     * @param operationWithCustomMemberRequest
+     * @return A Java Future containing the result of the OperationWithCustomMember operation returned by the service.<br/>
+     *         The CompletableFuture returned by this method can be completed exceptionally with the following
+     *         exceptions. The exception returned is wrapped with CompletionException, so you need to invoke
+     *         {@link Throwable#getCause} to retrieve the underlying exception.
+     *         <ul>
+     *         <li>SdkException Base class for all exceptions that can be thrown by the SDK (both service and client).
+     *         Can be used for catch all scenarios.</li>
+     *         <li>SdkClientException If any client side error occurs such as an IO related failure, failure to get
+     *         credentials, etc.</li>
+     *         <li>QueryException Base class for all service exceptions. Unknown exceptions will be thrown as an
+     *         instance of this type.</li>
+     *         </ul>
+     * @sample QueryAsyncClient.OperationWithCustomMember
+     * @see <a href="https://docs.aws.amazon.com/goto/WebAPI/query-service-2010-05-08/OperationWithCustomMember"
+     *      target="_top">AWS API Documentation</a>
+     */
+    @Override
+    public CompletableFuture<OperationWithCustomMemberResponse> operationWithCustomMember(
+        OperationWithCustomMemberRequest operationWithCustomMemberRequest) {
+        operationWithCustomMemberRequest = UtilsTest.dummyRequestModifier(operationWithCustomMemberRequest);
+        SdkClientConfiguration clientConfiguration = updateSdkClientConfiguration(operationWithCustomMemberRequest,
+                                                                                  this.clientConfiguration);
+        List<MetricPublisher> metricPublishers = resolveMetricPublishers(clientConfiguration, operationWithCustomMemberRequest
+            .overrideConfiguration().orElse(null));
+        MetricCollector apiCallMetricCollector = metricPublishers.isEmpty() ? NoOpMetricCollector.create() : MetricCollector
+            .create("ApiCall");
+        try {
+            apiCallMetricCollector.reportMetric(CoreMetric.SERVICE_ID, "Query Service");
+            apiCallMetricCollector.reportMetric(CoreMetric.OPERATION_NAME, "OperationWithCustomMember");
+
+            HttpResponseHandler<OperationWithCustomMemberResponse> responseHandler = protocolFactory
+                .createResponseHandler(OperationWithCustomMemberResponse::builder);
+
+            HttpResponseHandler<AwsServiceException> errorResponseHandler = protocolFactory.createErrorResponseHandler();
+
+            CompletableFuture<OperationWithCustomMemberResponse> executeFuture = clientHandler
+                .execute(new ClientExecutionParams<OperationWithCustomMemberRequest, OperationWithCustomMemberResponse>()
+                             .withOperationName("OperationWithCustomMember").withProtocolMetadata(protocolMetadata)
+                             .withMarshaller(new OperationWithCustomMemberRequestMarshaller(protocolFactory))
+                             .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
+                             .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
+                             .withInput(operationWithCustomMemberRequest));
+            CompletableFuture<OperationWithCustomMemberResponse> whenCompleteFuture = null;
+            whenCompleteFuture = executeFuture.whenComplete((r, e) -> {
+                metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+            });
+            return CompletableFutureUtils.forwardExceptionTo(whenCompleteFuture, executeFuture);
+        } catch (Throwable t) {
+            metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+            return CompletableFutureUtils.failedFuture(t);
+        }
+    }
+
     /**
      * Invokes the OperationWithCustomizedOperationContextParam operation asynchronously.
      *