new(StringArrayEndpointParams) Customization of Operation Context par…

…ams and adding customizations for S3

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

@@ -23,6 +23,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.TreeMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -40,6 +41,7 @@
 import software.amazon.awssdk.codegen.model.service.AuthType;
 import software.amazon.awssdk.codegen.model.service.EndpointRuleSetModel;
 import software.amazon.awssdk.codegen.model.service.Operation;
+import software.amazon.awssdk.codegen.model.service.OperationContextParam;
 import software.amazon.awssdk.codegen.model.service.Paginators;
 import software.amazon.awssdk.codegen.model.service.ServiceModel;
 import software.amazon.awssdk.codegen.model.service.Waiters;
@@ -168,9 +170,27 @@ public IntermediateModel build() {
 
         namingStrategy.validateCustomerVisibleNaming(trimmedModel);
         customizeEndpointParameters(fullModel, endpointRuleSet);
+        customizeOperationContextParams(fullModel, trimmedModel);
         return trimmedModel;
     }
 
+    private static void customizeOperationContextParams(IntermediateModel fullModel, IntermediateModel trimmedModel) {
+        if (!CollectionUtils.isNullOrEmpty(fullModel.getCustomizationConfig().getOperationContextParams())) {
+            fullModel.getCustomizationConfig().getOperationContextParams().forEach((operationName, operationContextParamMap) -> {
+                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(operationContextParamMap);
+            });
+        }
+    }
+
     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.OperationContextParam;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.core.traits.PayloadTrait;
 import software.amazon.awssdk.utils.AttributeMap;
@@ -324,6 +325,7 @@ public class CustomizationConfig {
      * This should be removed once the service updates its models
      */
     private Map<String, ParameterModel> endpointParameters;
+    private Map<String, Map<String, OperationContextParam>> operationContextParams;
 
     private CustomizationConfig() {
     }
@@ -858,4 +860,12 @@ public Map<String, ParameterModel> getEndpointParameters() {
     public void setEndpointParameters(Map<String, ParameterModel> endpointParameters) {
         this.endpointParameters = endpointParameters;
     }
+
+    public Map<String, Map<String, OperationContextParam>> getOperationContextParams() {
+        return operationContextParams;
+    }
+
+    public void setOperationContextParams(Map<String, Map<String, OperationContextParam>> operationContextParams) {
+        this.operationContextParams = operationContextParams;
+    }
 }

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

@@ -43,4 +43,15 @@ 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 an operation which already has OperationContextParams.");
+    }
 }
+
+
+

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

@@ -0,0 +1,8 @@
+{
+  "operationContextParams":
+    {
+      "OperationWithOperationContextParam": {
+        "customEndpointArray":{"value":"ListMember.StringList[*].LeafString"}
+      }
+    }
+}

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

@@ -11,5 +11,16 @@
       "documentation": "Parameter from the customization config",
       "type": "StringArray"
     }
+  },
+  "operationContextParams": {
+      "OperationWithCustomizedOperationContextParam": {
+        "customEndpointArray":{"value":"ListMember.StringList[*].LeafString"}
+      }
   }
-}
+  ,
+  "CustomEndpointArray": {
+      "required": false,
+      "documentation": "Parameter from the customization config",
+      "type": "StringArray"
+    }
+}

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