new(StringArrayEndpointParams) Customization of Operation Context par…

…ams and adding customizations for S3 (#5159)

new(StringArrayEndpointParams) Customization of Operation Context params and adding customizations for S3

codegen/src/main/java/software/amazon/awssdk/codegen/IntermediateModelBuilder.java

@@ -38,6 +38,7 @@
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.model.rules.endpoints.EndpointTestSuiteModel;
 import software.amazon.awssdk.codegen.model.service.AuthType;
+import software.amazon.awssdk.codegen.model.service.CustomOperationContextParam;
 import software.amazon.awssdk.codegen.model.service.EndpointRuleSetModel;
 import software.amazon.awssdk.codegen.model.service.Operation;
 import software.amazon.awssdk.codegen.model.service.Paginators;
@@ -168,9 +169,31 @@ public IntermediateModel build() {
 
         namingStrategy.validateCustomerVisibleNaming(trimmedModel);
         customizeEndpointParameters(fullModel, endpointRuleSet);
+        customizeOperationContextParams(trimmedModel, fullModel.getCustomizationConfig().getCustomOperationContextParams());
         return trimmedModel;
     }
 
+    private static void customizeOperationContextParams(IntermediateModel trimmedModel,
+                                                        List<CustomOperationContextParam> customOperationContextParams) {
+
+        if (CollectionUtils.isNullOrEmpty(customOperationContextParams)) {
+            return;
+        }
+        customOperationContextParams.forEach(customOperationContextParam -> {
+            String operationName = customOperationContextParam.getOperationName();
+            OperationModel operation = trimmedModel.getOperation(operationName);
+            if (operation == null) {
+                throw new IllegalStateException(
+                    "Could not find operation " + operationName + " to customize Operation Context Params.");
+            }
+            if (operation.getOperationContextParams() != null) {
+                throw new IllegalStateException(
+                    "Cannot customize operation " + operationName + " which already has OperationContextParams.");
+            }
+            operation.setOperationContextParams(customOperationContextParam.getOperationContextParamsMap());
+        });
+    }
+
     private void customizeEndpointParameters(IntermediateModel fullModel, EndpointRuleSetModel endpointRuleSet) {
         if (fullModel.getCustomizationConfig().getEndpointParameters() != null) {
             fullModel.getCustomizationConfig().getEndpointParameters().keySet().forEach(key -> {

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

@@ -22,6 +22,7 @@
 import java.util.Map;
 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.core.retry.RetryMode;
 import software.amazon.awssdk.core.traits.PayloadTrait;
 import software.amazon.awssdk.utils.AttributeMap;
@@ -325,6 +326,8 @@ public class CustomizationConfig {
      */
     private Map<String, ParameterModel> endpointParameters;
 
+    private List<CustomOperationContextParam> customOperationContextParams;
+
     private CustomizationConfig() {
     }
 
@@ -858,4 +861,12 @@ public Map<String, ParameterModel> getEndpointParameters() {
     public void setEndpointParameters(Map<String, ParameterModel> endpointParameters) {
         this.endpointParameters = endpointParameters;
     }
+
+    public List<CustomOperationContextParam> getCustomOperationContextParams() {
+        return customOperationContextParams;
+    }
+
+    public void setCustomOperationContextParams(List<CustomOperationContextParam> customOperationContextParams) {
+        this.customOperationContextParams = customOperationContextParams;
+    }
 }

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

@@ -0,0 +1,43 @@
+/*
+ * 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;
+
+
+import java.util.Map;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+@SdkInternalApi
+public class CustomOperationContextParam {
+    Map<String, OperationContextParam> operationContextParamsMap;
+
+    String operationName;
+
+    public Map<String, OperationContextParam> getOperationContextParamsMap() {
+        return operationContextParamsMap;
+    }
+
+    public void setOperationContextParamsMap(Map<String, OperationContextParam> operationContextParamsMap) {
+        this.operationContextParamsMap = operationContextParamsMap;
+    }
+
+    public String getOperationName() {
+        return operationName;
+    }
+
+    public void setOperationName(String operationName) {
+        this.operationName = operationName;
+    }
+}

codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules/EndpointProviderInterfaceSpecTest.java

@@ -43,4 +43,23 @@ public void endpointParametersWithDuplicatesInCustomizationConfig() {
                 ClientTestModels.queryServiceModelWithSpecialCustomization("customization-with-duplicate-endpointparameter.config")))
             .withMessageContaining("Duplicate parameters found in customizationConfig");
     }
+
+    @Test
+    public void endpointParametersWithDuplicatesOperationContextInCustomizationConfig() {
+        assertThatIllegalStateException()
+            .isThrownBy(() -> new EndpointParametersClassSpec(
+                ClientTestModels.queryServiceModelWithSpecialCustomization("customization-with-duplicate-operationcontextparams.config")))
+            .withMessageContaining("Cannot customize operation OperationWithOperationContextParam which already has OperationContextParams.");
+    }
+
+    @Test
+    public void endpointParametersWithIncorrectNameOperationContextInCustomizationConfig() {
+        assertThatIllegalStateException()
+            .isThrownBy(() -> new EndpointParametersClassSpec(
+                ClientTestModels.queryServiceModelWithSpecialCustomization("customization-with-incorrectName-operationcontextparams.config")))
+            .withMessageContaining("Could not find operation RandomOperationName to customize Operation Context Params.");
+    }
 }
+
+
+

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/query/customization-with-duplicate-operationcontextparams.config

@@ -0,0 +1,12 @@
+{
+  "customOperationContextParams": [
+    {
+      "operationName": "OperationWithOperationContextParam",
+      "operationContextParamsMap": {
+        "customEndpointArray": {
+          "value": "ListMember.StringList[*].LeafString"
+        }
+      }
+    }
+  ]
+}

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/query/customization-with-incorrectName-operationcontextparams.config

@@ -0,0 +1,12 @@
+{
+  "customOperationContextParams": [
+    {
+      "operationName": "RandomOperationName",
+      "operationContextParamsMap": {
+        "customEndpointArray": {
+          "value": "ListMember.StringList[*].LeafString"
+        }
+      }
+    }
+  ]
+}

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

@@ -11,5 +11,15 @@
       "documentation": "Parameter from the customization config",
       "type": "StringArray"
     }
-  }
-}
+  },
+  "customOperationContextParams": [
+    {
+      "operationName": "OperationWithCustomizedOperationContextParam",
+      "operationContextParamsMap": {
+        "customEndpointArray": {
+          "value": "ListMember.StringList[*].LeafString"
+        }
+      }
+    }
+  ]
+}

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

@@ -38,7 +38,15 @@
         "requestUri": "/"
       },
       "operationContextParams":{
-        "deleteKeys":{"value":"ListMember.StringList[*].LeafString"}
+        "customEndpointArray":{"value":"ListMember.StringList[*].LeafString"}
+      },
+      "input":{"shape":"WithOperationContextParam"}
+    },
+    "OperationWithCustomizedOperationContextParam": {
+      "name": "OperationWithCustomizedOperationContextParam",
+      "http": {
+        "method": "POST",
+        "requestUri": "/"
       },
       "input":{"shape":"WithOperationContextParam"}
     },

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

@@ -51,6 +51,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.OperationWithCustomizedOperationContextParamRequest;
+import software.amazon.awssdk.services.query.model.OperationWithCustomizedOperationContextParamResponse;
 import software.amazon.awssdk.services.query.model.OperationWithNoneAuthTypeRequest;
 import software.amazon.awssdk.services.query.model.OperationWithNoneAuthTypeResponse;
 import software.amazon.awssdk.services.query.model.OperationWithOperationContextParamRequest;
@@ -72,6 +74,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.OperationWithCustomizedOperationContextParamRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithNoneAuthTypeRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithOperationContextParamRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithRequestCompressionRequestMarshaller;
@@ -467,6 +470,65 @@ public CompletableFuture<OperationWithContextParamResponse> operationWithContext
         }
     }
 
+    /**
+     * Invokes the OperationWithCustomizedOperationContextParam operation asynchronously.
+     *
+     * @param operationWithCustomizedOperationContextParamRequest
+     * @return A Java Future containing the result of the OperationWithCustomizedOperationContextParam 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.OperationWithCustomizedOperationContextParam
+     * @see <a
+     *      href="https://docs.aws.amazon.com/goto/WebAPI/query-service-2010-05-08/OperationWithCustomizedOperationContextParam"
+     *      target="_top">AWS API Documentation</a>
+     */
+    @Override
+    public CompletableFuture<OperationWithCustomizedOperationContextParamResponse> operationWithCustomizedOperationContextParam(
+        OperationWithCustomizedOperationContextParamRequest operationWithCustomizedOperationContextParamRequest) {
+        SdkClientConfiguration clientConfiguration = updateSdkClientConfiguration(
+            operationWithCustomizedOperationContextParamRequest, this.clientConfiguration);
+        List<MetricPublisher> metricPublishers = resolveMetricPublishers(clientConfiguration,
+                                                                         operationWithCustomizedOperationContextParamRequest.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, "OperationWithCustomizedOperationContextParam");
+
+            HttpResponseHandler<OperationWithCustomizedOperationContextParamResponse> responseHandler = protocolFactory
+                .createResponseHandler(OperationWithCustomizedOperationContextParamResponse::builder);
+
+            HttpResponseHandler<AwsServiceException> errorResponseHandler = protocolFactory.createErrorResponseHandler();
+
+            CompletableFuture<OperationWithCustomizedOperationContextParamResponse> executeFuture = clientHandler
+                .execute(new ClientExecutionParams<OperationWithCustomizedOperationContextParamRequest, OperationWithCustomizedOperationContextParamResponse>()
+                             .withOperationName("OperationWithCustomizedOperationContextParam")
+                             .withProtocolMetadata(protocolMetadata)
+                             .withMarshaller(new OperationWithCustomizedOperationContextParamRequestMarshaller(protocolFactory))
+                             .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
+                             .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
+                             .withInput(operationWithCustomizedOperationContextParamRequest));
+            CompletableFuture<OperationWithCustomizedOperationContextParamResponse> 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 OperationWithNoneAuthType operation asynchronously.
      *

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

@@ -45,6 +45,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.OperationWithCustomizedOperationContextParamRequest;
+import software.amazon.awssdk.services.query.model.OperationWithCustomizedOperationContextParamResponse;
 import software.amazon.awssdk.services.query.model.OperationWithNoneAuthTypeRequest;
 import software.amazon.awssdk.services.query.model.OperationWithNoneAuthTypeResponse;
 import software.amazon.awssdk.services.query.model.OperationWithOperationContextParamRequest;
@@ -66,6 +68,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.OperationWithCustomizedOperationContextParamRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithNoneAuthTypeRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithOperationContextParamRequestMarshaller;
 import software.amazon.awssdk.services.query.transform.OperationWithRequestCompressionRequestMarshaller;
@@ -401,6 +404,55 @@ public OperationWithContextParamResponse operationWithContextParam(
         }
     }
 
+    /**
+     * Invokes the OperationWithCustomizedOperationContextParam operation.
+     *
+     * @param operationWithCustomizedOperationContextParamRequest
+     * @return Result of the OperationWithCustomizedOperationContextParam operation returned by the service.
+     * @throws SdkException
+     *         Base class for all exceptions that can be thrown by the SDK (both service and client). Can be used for
+     *         catch all scenarios.
+     * @throws SdkClientException
+     *         If any client side error occurs such as an IO related failure, failure to get credentials, etc.
+     * @throws QueryException
+     *         Base class for all service exceptions. Unknown exceptions will be thrown as an instance of this type.
+     * @sample QueryClient.OperationWithCustomizedOperationContextParam
+     * @see <a
+     *      href="https://docs.aws.amazon.com/goto/WebAPI/query-service-2010-05-08/OperationWithCustomizedOperationContextParam"
+     *      target="_top">AWS API Documentation</a>
+     */
+    @Override
+    public OperationWithCustomizedOperationContextParamResponse operationWithCustomizedOperationContextParam(
+        OperationWithCustomizedOperationContextParamRequest operationWithCustomizedOperationContextParamRequest)
+        throws AwsServiceException, SdkClientException, QueryException {
+
+        HttpResponseHandler<OperationWithCustomizedOperationContextParamResponse> responseHandler = protocolFactory
+            .createResponseHandler(OperationWithCustomizedOperationContextParamResponse::builder);
+
+        HttpResponseHandler<AwsServiceException> errorResponseHandler = protocolFactory.createErrorResponseHandler();
+        SdkClientConfiguration clientConfiguration = updateSdkClientConfiguration(
+            operationWithCustomizedOperationContextParamRequest, this.clientConfiguration);
+        List<MetricPublisher> metricPublishers = resolveMetricPublishers(clientConfiguration,
+                                                                         operationWithCustomizedOperationContextParamRequest.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, "OperationWithCustomizedOperationContextParam");
+
+            return clientHandler
+                .execute(new ClientExecutionParams<OperationWithCustomizedOperationContextParamRequest, OperationWithCustomizedOperationContextParamResponse>()
+                             .withOperationName("OperationWithCustomizedOperationContextParam")
+                             .withProtocolMetadata(protocolMetadata).withResponseHandler(responseHandler)
+                             .withErrorResponseHandler(errorResponseHandler).withRequestConfiguration(clientConfiguration)
+                             .withInput(operationWithCustomizedOperationContextParamRequest)
+                             .withMetricCollector(apiCallMetricCollector)
+                             .withMarshaller(new OperationWithCustomizedOperationContextParamRequestMarshaller(protocolFactory)));
+        } finally {
+            metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+        }
+    }
+
     /**
      * Invokes the OperationWithNoneAuthType operation.
      *