Codegenerate BatchManager API under AsyncClient and Initail Interfaces for BatchManager

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

@@ -60,6 +60,7 @@ public static Metadata constructMetadata(ServiceModel serviceModel,
                 .withBaseBuilder(String.format(Constant.BASE_BUILDER_CLASS_NAME_PATTERN, serviceName))
                 .withDocumentation(serviceModel.getDocumentation())
                 .withServiceAbbreviation(serviceMetadata.getServiceAbbreviation())
+                .withBatchmanagerPackageName(namingStrategy.getBatchManagerPackageName(serviceName))
                 .withServiceFullName(serviceMetadata.getServiceFullName())
                 .withServiceName(serviceName)
                 .withSyncClient(String.format(Constant.SYNC_CLIENT_CLASS_NAME_PATTERN, serviceName))

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

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

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

@@ -342,6 +342,11 @@ public class CustomizationConfig {
      */
     private Map<String, PreClientExecutionRequestCustomizer> preClientExecutionRequestCustomizer;
 
+    /**
+     * A boolean flag to indicate if Automatic Batch Request is supported.
+     */
+    private boolean batchManagerSupported;
+
     private CustomizationConfig() {
     }
 
@@ -901,4 +906,12 @@ public void setPreClientExecutionRequestCustomizer(Map<String, PreClientExecutio
         this.preClientExecutionRequestCustomizer = preClientExecutionRequestCustomizer;
     }
 
+    public boolean getBatchManagerSupported() {
+        return batchManagerSupported;
+    }
+
+    public void setBatchManagerSupported(boolean batchManagerSupported) {
+        this.batchManagerSupported = batchManagerSupported;
+    }
+
 }

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

@@ -71,6 +71,8 @@ public class Metadata {
 
     private String waitersPackageName;
 
+    private String batchManagerPackageName;
+
     private String endpointRulesPackageName;
 
     private String authSchemePackageName;
@@ -789,4 +791,22 @@ public String getFullInternalJmesPathPackageName() {
         return joinPackageNames(getFullJmesPathPackageName(), "internal");
     }
 
+    public Metadata withBatchmanagerPackageName(String batchmanagerPackageName) {
+        setBatchManagerPackageName(batchmanagerPackageName);
+        return this;
+    }
+
+
+    public String getBatchManagerPackageName() {
+        return batchManagerPackageName;
+    }
+
+    public void setBatchManagerPackageName(String batchManagerPackageName) {
+        this.batchManagerPackageName = batchManagerPackageName;
+    }
+
+    public String getFullBatchManagerPackageName() {
+        return joinPackageNames(rootPackageName, getBatchManagerPackageName());
+    }
+
 }

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

@@ -189,6 +189,11 @@ public String getJmesPathPackageName(String serviceName) {
         return getCustomizedPackageName(concatServiceNameIfShareModel(serviceName), Constant.PACKAGE_NAME_JMESPATH_PATTERN);
     }
 
+    @Override
+    public String getBatchManagerPackageName(String serviceName) {
+        return getCustomizedPackageName(concatServiceNameIfShareModel(serviceName), Constant.PACKAGE_NAME_BATCHMANAGER_PATTERN);
+    }
+
     @Override
     public String getSmokeTestPackageName(String serviceName) {
 

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

@@ -74,6 +74,11 @@ public interface NamingStrategy {
      */
     String getJmesPathPackageName(String serviceName);
 
+    /**
+     * Retrieve the batchManager package name that should be used based on the service name.
+     */
+    String getBatchManagerPackageName(String serviceName);
+
     /**
      * Retrieve the smote test package name that should be used based on the service name.
      */

codegen/src/main/java/software/amazon/awssdk/codegen/poet/PoetExtension.java

@@ -196,4 +196,9 @@ public boolean isResponse(ShapeModel shapeModel) {
     public boolean isRequest(ShapeModel shapeModel) {
         return shapeModel.getShapeType() == ShapeType.Request;
     }
+
+    public ClassName getBatchManagerAsyncInterface() {
+        return ClassName.get(model.getMetadata().getFullBatchManagerPackageName(),
+                             model.getMetadata().getServiceName() + "AsyncBatchManager");
+    }
 }

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

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.codegen.poet.client;
 
+import static com.squareup.javapoet.TypeSpec.Builder;
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 import static javax.lang.model.element.Modifier.FINAL;
@@ -104,6 +105,7 @@ public final class AsyncClientClass extends AsyncClientInterface {
     private final ClassName serviceClientConfigurationClassName;
     private final ServiceClientConfigurationUtils configurationUtils;
     private final boolean useSraAuth;
+    private boolean hasScheduledExecutor;
 
     public AsyncClientClass(GeneratorTaskParams dependencies) {
         super(dependencies.getModel());
@@ -139,7 +141,7 @@ protected void addModifiers(TypeSpec.Builder type) {
     }
 
     @Override
-    protected void addFields(TypeSpec.Builder type) {
+    protected void addFields(Builder type) {
         type.addField(FieldSpec.builder(ClassName.get(Logger.class), "log")
                                .addModifiers(PRIVATE, STATIC, FINAL)
                                .initializer("$T.getLogger($T.class)", LoggerFactory.class,
@@ -155,6 +157,10 @@ protected void addFields(TypeSpec.Builder type) {
             type.addField(AwsJsonProtocolFactory.class, "jsonProtocolFactory", PRIVATE, FINAL);
         }
 
+        if (shouldAddScheduledExecutor()) {
+            addScheduledExecutorIfNeeded(type);
+        }
+
         model.getEndpointOperation().ifPresent(
             o -> type.addField(EndpointDiscoveryRefreshCache.class, "endpointDiscoveryCache", PRIVATE));
     }
@@ -180,9 +186,6 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
 
     @Override
     protected void addWaiterMethod(TypeSpec.Builder type) {
-        type.addField(FieldSpec.builder(ClassName.get(ScheduledExecutorService.class), "executorService")
-                                       .addModifiers(PRIVATE, FINAL)
-                                       .build());
 
         MethodSpec waiter = MethodSpec.methodBuilder("waiter")
                                       .addModifiers(PUBLIC)
@@ -263,14 +266,18 @@ private MethodSpec constructor(TypeSpec.Builder classBuilder) {
             builder.endControlFlow();
         }
 
-        if (model.hasWaiters()) {
+        if (shouldAddScheduledExecutor()) {
             builder.addStatement("this.executorService = clientConfiguration.option($T.SCHEDULED_EXECUTOR_SERVICE)",
                                  SdkClientOption.class);
         }
 
         return builder.build();
     }
 
+    private boolean shouldAddScheduledExecutor() {
+        return model.hasWaiters() || model.getCustomizationConfig().getBatchManagerSupported();
+    }
+
     private boolean hasOperationWithEventStreamOutput() {
         return model.getOperations().values().stream().anyMatch(OperationModel::hasEventStreamOutput);
     }
@@ -547,6 +554,26 @@ protected MethodSpec utilitiesMethod() {
                          .build();
     }
 
+    @Override
+    protected void addBatchManagerMethod(Builder type) {
+
+        String scheduledExecutor = "executorService";
+        ClassName returnType;
+
+        returnType = poetExtensions.getBatchManagerAsyncInterface();
+
+        MethodSpec batchManager = MethodSpec.methodBuilder("batchManager")
+                                            .addModifiers(PUBLIC)
+                                            .addAnnotation(Override.class)
+                                            .returns(returnType)
+                                            .addStatement("return $T.builder().client(this).scheduledExecutor($N).build()",
+                                                          returnType, scheduledExecutor)
+                                            .build();
+
+
+        type.addMethod(batchManager);
+    }
+
     private MethodSpec resolveMetricPublishersMethod() {
         String clientConfigName = "clientConfiguration";
         String requestOverrideConfigName = "requestOverrideConfiguration";
@@ -621,4 +648,13 @@ private boolean hasStreamingV4AuthOperations() {
         return model.getOperations().values().stream()
                 .anyMatch(this::shouldUseAsyncWithBodySigner);
     }
+
+    private void addScheduledExecutorIfNeeded(Builder classBuilder) {
+        if (!hasScheduledExecutor) {
+            classBuilder.addField(FieldSpec.builder(ClassName.get(ScheduledExecutorService.class), "executorService")
+                                           .addModifiers(PRIVATE, FINAL)
+                                           .build());
+            hasScheduledExecutor = true;
+        }
+    }
 }

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

@@ -96,6 +96,9 @@ public TypeSpec poetSpec() {
         if (model.hasWaiters()) {
             addWaiterMethod(result);
         }
+        if (model.getCustomizationConfig().getBatchManagerSupported()) {
+            addBatchManagerMethod(result);
+        }
         result.addMethod(serviceClientConfigMethod());
         addAdditionalMethods(result);
         addCloseMethod(result);
@@ -162,6 +165,16 @@ protected void addWaiterMethod(TypeSpec.Builder type) {
         type.addMethod(waiterOperationBody(builder).build());
     }
 
+    protected void addBatchManagerMethod(TypeSpec.Builder type) {
+        ClassName returnType = poetExtensions.getBatchManagerAsyncInterface();
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("batchManager")
+                                               .addModifiers(PUBLIC)
+                                               .returns(returnType)
+                                               .addJavadoc("Creates an instance of {@link $T} object with the "
+                                                           + "configuration set on this client.", returnType);
+        type.addMethod(batchManagerOperationBody(builder).build());
+    }
+
     @Override
     public ClassName className() {
         return className;
@@ -532,4 +545,10 @@ protected MethodSpec.Builder waiterOperationBody(MethodSpec.Builder builder) {
         return builder.addModifiers(DEFAULT, PUBLIC)
                       .addStatement("throw new $T()", UnsupportedOperationException.class);
     }
+
+    protected MethodSpec.Builder batchManagerOperationBody(MethodSpec.Builder builder) {
+        return builder.addModifiers(DEFAULT, PUBLIC)
+                      .addStatement("throw new $T()", UnsupportedOperationException.class);
+    }
+
 }

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

@@ -217,4 +217,9 @@ protected MethodSpec.Builder utilitiesOperationBody(MethodSpec.Builder builder)
     protected MethodSpec.Builder waiterOperationBody(MethodSpec.Builder builder) {
         return builder.addAnnotation(Override.class).addStatement("return delegate.waiter()");
     }
+
+    @Override
+    protected MethodSpec.Builder batchManagerOperationBody(MethodSpec.Builder builder) {
+        return builder.addAnnotation(Override.class).addStatement("return delegate.batchManager()");
+    }
 }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/ClientTestModels.java

@@ -380,6 +380,18 @@ public static IntermediateModel internalConfigModels() {
         return new IntermediateModelBuilder(models).build();
     }
 
+    public static IntermediateModel batchManagerModels() {
+        File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/batchmanager/service-2.json").getFile());
+        File customizationModel = new File(ClientTestModels.class.getResource("client/c2j/batchmanager/customization.config").getFile());
+
+        C2jModels models = C2jModels.builder()
+                                    .serviceModel(getServiceModel(serviceModel))
+                                    .customizationConfig(getCustomizationConfig(customizationModel))
+                                    .build();
+
+        return new IntermediateModelBuilder(models).build();
+    }
+
     private static ServiceModel getServiceModel(File file) {
         return ModelLoaderUtils.loadModel(ServiceModel.class, file);
     }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClassTest.java

@@ -18,6 +18,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.awsJsonServiceModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.awsQueryCompatibleJsonServiceModels;
+import static software.amazon.awssdk.codegen.poet.ClientTestModels.batchManagerModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.customContentTypeModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.customPackageModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.endpointDiscoveryModels;
@@ -92,6 +93,12 @@ public void asyncClientCustomPackageName() {
         assertThat(syncClientCustomServiceMetaData, generatesTo("test-custompackage-async.java"));
     }
 
+    @Test
+    public void asyncClientBatchManager() {
+        ClassSpec aSyncClientBatchManager = createAsyncClientClass(batchManagerModels());
+        assertThat(aSyncClientBatchManager, generatesTo("test-batchmanager-async.java"));
+    }
+
     private AsyncClientClass createAsyncClientClass(IntermediateModel model) {
         return new AsyncClientClass(GeneratorTaskParams.create(model, "sources/", "tests/", "resources/"));
     }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/AsyncClientInterfaceTest.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.codegen.poet.client;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static software.amazon.awssdk.codegen.poet.ClientTestModels.batchManagerModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.restJsonServiceModels;
 import static software.amazon.awssdk.codegen.poet.PoetMatchers.generatesTo;
 
@@ -28,4 +29,10 @@ public void asyncClientInterface() {
         ClassSpec asyncClientInterface = new AsyncClientInterface(restJsonServiceModels());
         assertThat(asyncClientInterface, generatesTo("test-json-async-client-interface.java"));
     }
+
+    @Test
+    public void asyncClientInterfaceWithBatchManager() {
+        ClassSpec asyncClientInterface = new AsyncClientInterface(batchManagerModels());
+        assertThat(asyncClientInterface, generatesTo("test-json-async-client-interface-batchmanager.java"));
+    }
 }

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

@@ -0,0 +1,3 @@
+{
+    "batchManagerSupported": true
+}