Code generation for Service BatchFunctions class

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

@@ -58,6 +58,7 @@ public static Metadata constructMetadata(ServiceModel serviceModel,
                 .withRequestTransformPackageName(namingStrategy.getRequestTransformPackageName(serviceName))
                 .withPaginatorsPackageName(namingStrategy.getPaginatorsPackageName(serviceName))
                 .withWaitersPackageName(namingStrategy.getWaitersPackageName(serviceName))
+                .withBatchmanagerPackageName(namingStrategy.getBatchManagerPackageName(serviceName))
                 .withServiceAbbreviation(serviceMetadata.getServiceAbbreviation())
                 .withServiceFullName(serviceMetadata.getServiceFullName())
                 .withServiceName(serviceName)

codegen/src/main/java/software/amazon/awssdk/codegen/emitters/GeneratorPathProvider.java

@@ -68,4 +68,12 @@ public String getWaitersDirectory() {
     public String getWaitersInternalDirectory() {
         return sourceDirectory + "/" + Utils.packageToDirectory(model.getMetadata().getFullWaitersInternalPackageName());
     }
+
+    public String getBatchManagerDirectory() {
+        return sourceDirectory + "/" + Utils.packageToDirectory(model.getMetadata().getFullBatchManagerPackageName());
+    }
+
+    public String getBatchManagerInternalDirectory() {
+        return sourceDirectory + "/" + Utils.packageToDirectory(model.getMetadata().getFullBatchManagerInternalPackageName());
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/AwsGeneratorTasks.java

@@ -26,6 +26,7 @@ public AwsGeneratorTasks(GeneratorTaskParams params) {
               new AsyncClientGeneratorTasks(params),
               new PaginatorsGeneratorTasks(params),
               new EventStreamGeneratorTasks(params),
-              new WaitersGeneratorTasks(params));
+              new WaitersGeneratorTasks(params),
+              new BatchManagerGeneratorTasks(params));
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/BatchManagerGeneratorTasks.java

@@ -0,0 +1,55 @@
+/*
+ * 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.emitters.tasks;
+
+import java.util.ArrayList;
+import java.util.List;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.codegen.emitters.GeneratorTask;
+import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
+import software.amazon.awssdk.codegen.emitters.PoetGeneratorTask;
+import software.amazon.awssdk.codegen.poet.batchmanager.BatchFunctionsClassSpec;
+
+@SdkInternalApi
+public class BatchManagerGeneratorTasks extends BaseGeneratorTasks {
+    private final GeneratorTaskParams generatorTaskParams;
+
+    public BatchManagerGeneratorTasks(GeneratorTaskParams dependencies) {
+        super(dependencies);
+        this.generatorTaskParams = dependencies;
+    }
+
+    @Override
+    protected boolean hasTasks() {
+        return model.getCustomizationConfig().getBatchManagerMethods() != null;
+    }
+
+    @Override
+    protected List<GeneratorTask> createTasks() {
+        List<GeneratorTask> generatorTasks = new ArrayList<>();
+        generatorTasks.add(createBatchFunctions());
+        return generatorTasks;
+    }
+
+    private GeneratorTask createBatchFunctions() {
+        return new PoetGeneratorTask(batchManagerInternalClassDir(), model.getFileHeader(),
+                                     new BatchFunctionsClassSpec(model));
+    }
+
+    private String batchManagerInternalClassDir() {
+        return generatorTaskParams.getPathProvider().getBatchManagerInternalDirectory();
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/internal/Constant.java

@@ -72,6 +72,8 @@ public final class Constant {
 
     public static final String PACKAGE_NAME_CUSTOM_AUTH_PATTERN = "%s.auth";
 
+    public static final String PACKAGE_NAME_BATCHMANAGER_PATTERN = "%s.batchmanager";
+
     public static final String AUTH_POLICY_ENUM_CLASS_DIR = "software/amazon/awssdk/auth/policy/actions";
 
     public static final String REQUEST_CLASS_SUFFIX = "Request";

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

@@ -0,0 +1,165 @@
+/*
+ * 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.config.customization;
+
+/**
+ * Config to define a batchable method. The key is a method that has a batch counterpart
+ *
+ * ex. For SQS, we can use a key of: sendMessage, meanwhile the batchFunctionsTypes will store the types
+ * SendMessageRequest, SendMessageResponse, and SendMessageBatchResponse.
+ */
+public class BatchManager {
+
+    /**
+     * The batch equivalent of the request method. This is required.
+     *
+     * Ex. if the request method is sendMessage, batchMethod is sendMessageBatch
+     */
+    private String batchMethod;
+
+    /** Type of a single entry that is contained within a batch request. This is required. */
+    private String batchRequestEntry;
+
+    /**
+     * Type of a successful batch response entry. If a successful and failed batch response entry are the same,
+     * successfulBatchEntry indicates the type of that entry. This is required.
+     *
+     * Ex. SendMessageBatchResultEntry for SQS, PutRecordsResultEntry for kinesis
+     */
+    private String successBatchEntry;
+
+    /** Type of a failed batch response entry. This is optional (depending on service). */
+    private String errorBatchEntry;
+
+    /**
+     * Name of the method used to extract the successful responses from a batch response. If the method to extract successful
+     * and failed entries are the same, successEntriesMethod indicates the name of that method. This is required.
+     *
+     * Ex. successful for SQS, records for kinesis, responses for dynamodb etc.
+     */
+    private String successEntriesMethod;
+
+    /** Name of the method used to extract failed responses from a batch responses. This is optional (depending on service). */
+    private String errorEntriesMethod;
+
+    /**
+     * Name of the method used to get/set the destination for a request. This is required.
+     *
+     * Ex. queueUrl for SQS, streamName for Kinesis
+     */
+    private String batchKey;
+
+    /**
+     * Name of the method used to extract the status code from a failed batch entry.
+     *
+     * Ex. code for SQS, errorCode for Kinesis.
+     */
+    private String errorCodeMethod;
+
+    /**
+     * Name of the method used to extract the status code from a failed batch entry.
+     *
+     * Ex. code for SQS, errorCode for Kinesis.
+     */
+    private String errorMessageMethod;
+
+    /**
+     * Name of the method used to set or extract the request identifier used to identify entries within a batch request.
+     *
+     * Ex. id for SQS
+     */
+    private String batchRequestIdentifier;
+
+    public String getBatchMethod() {
+        return batchMethod;
+    }
+
+    public void setBatchMethod(String batchMethod) {
+        this.batchMethod = batchMethod;
+    }
+
+    public String getBatchRequestEntry() {
+        return batchRequestEntry;
+    }
+
+    public void setBatchRequestEntry(String batchRequestEntry) {
+        this.batchRequestEntry = batchRequestEntry;
+    }
+
+    public String getSuccessBatchEntry() {
+        return successBatchEntry;
+    }
+
+    public void setSuccessBatchEntry(String successBatchEntry) {
+        this.successBatchEntry = successBatchEntry;
+    }
+
+    public String getErrorBatchEntry() {
+        return errorBatchEntry;
+    }
+
+    public void setErrorBatchEntry(String errorBatchEntry) {
+        this.errorBatchEntry = errorBatchEntry;
+    }
+
+    public String getSuccessEntriesMethod() {
+        return successEntriesMethod;
+    }
+
+    public void setSuccessEntriesMethod(String successEntriesMethod) {
+        this.successEntriesMethod = successEntriesMethod;
+    }
+
+    public String getErrorEntriesMethod() {
+        return errorEntriesMethod;
+    }
+
+    public void setErrorEntriesMethod(String errorEntriesMethod) {
+        this.errorEntriesMethod = errorEntriesMethod;
+    }
+
+    public String getBatchKey() {
+        return batchKey;
+    }
+
+    public void setBatchKey(String batchKey) {
+        this.batchKey = batchKey;
+    }
+
+    public String getErrorCodeMethod() {
+        return errorCodeMethod;
+    }
+
+    public void setErrorCodeMethod(String errorCodeMethod) {
+        this.errorCodeMethod = errorCodeMethod;
+    }
+
+    public String getErrorMessageMethod() {
+        return errorMessageMethod;
+    }
+
+    public void setErrorMessageMethod(String errorMessageMethod) {
+        this.errorMessageMethod = errorMessageMethod;
+    }
+
+    public String getBatchRequestIdentifier() {
+        return batchRequestIdentifier;
+    }
+
+    public void setBatchRequestIdentifier(String requestIdentifier) {
+        this.batchRequestIdentifier = requestIdentifier;
+    }
+}

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

@@ -189,7 +189,7 @@ public class CustomizationConfig {
 
     private String userAgent;
 
-    private BatchManagerMethod batchManagerMethod;
+    private Map<String, BatchManager> batchManagerMethods;
 
     private CustomizationConfig() {
     }
@@ -488,11 +488,11 @@ public CustomizationConfig withUserAgent(String userAgent) {
         return this;
     }
 
-    public BatchManagerMethod getBatchManagerMethod() {
-        return batchManagerMethod;
+    public Map<String, BatchManager>  getBatchManagerMethods() {
+        return batchManagerMethods;
     }
 
-    public void setBatchManagerMethod(BatchManagerMethod batchManagerMethod) {
-        this.batchManagerMethod = batchManagerMethod;
+    public void setBatchManager(Map<String, BatchManager> batchManagerMethods) {
+        this.batchManagerMethods = batchManagerMethods;
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/Metadata.java

@@ -69,6 +69,8 @@ public class Metadata {
 
     private String waitersPackageName;
 
+    private String batchManagerPackageName;
+
     private String serviceAbbreviation;
 
     private String serviceFullName;
@@ -683,4 +685,25 @@ public String getFullWaitersPackageName() {
     public String getFullWaitersInternalPackageName() {
         return joinPackageNames(getFullWaitersPackageName(), "internal");
     }
+
+    public String getBatchManagerPackageName() {
+        return batchManagerPackageName;
+    }
+
+    public void setBatchManagerPackageName(String batchManagerPackageName) {
+        this.batchManagerPackageName = batchManagerPackageName;
+    }
+
+    public Metadata withBatchmanagerPackageName(String batchmanagerPackageName) {
+        setBatchManagerPackageName(batchmanagerPackageName);
+        return this;
+    }
+
+    public String getFullBatchManagerPackageName() {
+        return joinPackageNames(rootPackageName, getBatchManagerPackageName());
+    }
+
+    public String getFullBatchManagerInternalPackageName() {
+        return joinPackageNames(getFullBatchManagerPackageName(), "internal");
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/naming/DefaultNamingStrategy.java

@@ -181,6 +181,11 @@ public String getSmokeTestPackageName(String serviceName) {
                                         Constant.PACKAGE_NAME_SMOKE_TEST_PATTERN);
     }
 
+    @Override
+    public String getBatchManagerPackageName(String serviceName) {
+        return getCustomizedPackageName(concatServiceNameIfShareModel(serviceName), Constant.PACKAGE_NAME_BATCHMANAGER_PATTERN);
+    }
+
     /**
      * If the service is sharing models with other services, we need to concatenate its customized package name
      * if provided or service name with the shared service name.

codegen/src/main/java/software/amazon/awssdk/codegen/naming/NamingStrategy.java

@@ -63,6 +63,11 @@ public interface NamingStrategy {
      */
     String getSmokeTestPackageName(String serviceName);
 
+    /**
+     * Retrieve the batchmanager package name that should be used based on the service name.
+     */
+    String getBatchManagerPackageName(String serviceName);
+
     /**
      * @param errorShapeName Name of error shape to derive exception class name from.
      * @return Appropriate name to use for a Java exception class name