Wraps s3 client in cross regional client when enabled

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

@@ -19,6 +19,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import software.amazon.awssdk.awscore.internal.client.ClientComposer;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.core.traits.PayloadTrait;
 import software.amazon.awssdk.utils.AttributeMap;
@@ -212,6 +213,20 @@ public class CustomizationConfig {
      */
     private boolean delegateSyncClientClass;
 
+    /**
+     * Fully qualified name of a class that given the default sync client instance can return the final client instance,
+     * for instance by decorating the client with specific-purpose implementations of the client interface.
+     * The class should implement the {@link ClientComposer} interface. See S3 customization.config for an example.
+     */
+    private String syncClientComposer;
+
+    /**
+     * Fully qualified name of a class that given the default async client instance can return the final client instance,
+     * for instance by decorating the client with specific-purpose implementations of the client interface.
+     * The class should implement the {@link ClientComposer} interface. See S3 customization.config for an example.
+     */
+    private String asyncClientComposer;
+
     /**
      * Whether to skip generating endpoint tests from endpoint-tests.json
      */
@@ -561,6 +576,22 @@ public void setDelegateAsyncClientClass(boolean delegateAsyncClientClass) {
         this.delegateAsyncClientClass = delegateAsyncClientClass;
     }
 
+    public String getSyncClientComposer() {
+        return syncClientComposer;
+    }
+
+    public void setSyncClientComposer(String syncClientComposer) {
+        this.syncClientComposer = syncClientComposer;
+    }
+
+    public String getAsyncClientComposer() {
+        return asyncClientComposer;
+    }
+
+    public void setAsyncClientComposer(String asyncClientComposer) {
+        this.asyncClientComposer = asyncClientComposer;
+    }
+
     public boolean isDelegateSyncClientClass() {
         return delegateSyncClientClass;
     }

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

@@ -235,6 +235,21 @@ public String getSdkResponseBaseClassName() {
         }
     }
 
+    public Optional<String> syncClientComposerClassName() {
+        if (customizationConfig.getSyncClientComposer() != null) {
+            return Optional.of(customizationConfig.getSyncClientComposer());
+        }
+        return Optional.empty();
+    }
+
+    public Optional<String> asyncClientComposerClassName() {
+        String asyncClientComposer = customizationConfig.getAsyncClientComposer();
+        if (customizationConfig.getAsyncClientComposer() != null) {
+            return Optional.of(asyncClientComposer);
+        }
+        return Optional.empty();
+    }
+
     public String getFileHeader() {
         return FILE_HEADER;
     }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/AsyncClientBuilderClass.java

@@ -24,6 +24,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.awscore.client.config.AwsClientOption;
+import software.amazon.awssdk.awscore.internal.client.ClientComposer;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
@@ -80,7 +81,10 @@ public TypeSpec poetSpec() {
             builder.addMethod(bearerTokenProviderMethod());
         }
 
-        return builder.addMethod(buildClientMethod()).build();
+        builder.addMethod(buildClientMethod());
+        builder.addMethod(initializeServiceClientConfigMethod());
+
+        return builder.build();
     }
 
     private MethodSpec endpointDiscoveryEnabled() {
@@ -120,30 +124,28 @@ private MethodSpec endpointProviderMethod() {
     }
 
     private MethodSpec buildClientMethod() {
-        return MethodSpec.methodBuilder("buildClient")
-                         .addAnnotation(Override.class)
-                         .addModifiers(Modifier.PROTECTED, Modifier.FINAL)
-                         .returns(clientInterfaceName)
-                         .addStatement("$T clientConfiguration = super.asyncClientConfiguration()", SdkClientConfiguration.class)
-                         .addStatement("this.validateClientOptions(clientConfiguration)")
-                         .addStatement("$T endpointOverride = null", URI.class)
-                         .addStatement("$T endpointProvider = clientConfiguration.option($T.ENDPOINT_PROVIDER)",
-                                       EndpointProvider.class,
-                                       SdkClientOption.class)
-                         .addCode("if (clientConfiguration.option($T.ENDPOINT_OVERRIDDEN) != null"
-                                  + "&& $T.TRUE.equals(clientConfiguration.option($T.ENDPOINT_OVERRIDDEN))) {"
-                                  + "endpointOverride = clientConfiguration.option($T.ENDPOINT);"
-                                  + "}",
-                                  SdkClientOption.class, Boolean.class, SdkClientOption.class, SdkClientOption.class)
-                         .addStatement("$T serviceClientConfiguration = $T.builder()"
-                                       + ".overrideConfiguration(overrideConfiguration())"
-                                       + ".region(clientConfiguration.option($T.AWS_REGION))"
-                                       + ".endpointOverride(endpointOverride)"
-                                       + ".endpointProvider(endpointProvider)"
-                                       + ".build()",
-                                       serviceConfigClassName, serviceConfigClassName, AwsClientOption.class)
-                         .addStatement("return new $T(serviceClientConfiguration, clientConfiguration)", clientClassName)
-                         .build();
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("buildClient")
+                                               .addAnnotation(Override.class)
+                                               .addModifiers(Modifier.PROTECTED, Modifier.FINAL)
+                                               .returns(clientInterfaceName)
+                                               .addStatement("$T clientConfiguration = super.asyncClientConfiguration()",
+                                                             SdkClientConfiguration.class)
+                                               .addStatement("this.validateClientOptions(clientConfiguration)")
+                                               .addStatement("$T serviceClientConfiguration = initializeServiceClientConfig"
+                                                             + "(clientConfiguration)",
+                                                             serviceConfigClassName);
+
+        builder.addStatement("$1T client = new $2T(serviceClientConfiguration, clientConfiguration)",
+                             clientInterfaceName, clientClassName);
+        if (model.asyncClientComposerClassName().isPresent()) {
+            builder.addStatement("$1T composer = new $2T()",
+                                 ClientComposer.class,
+                                 PoetUtils.classNameFromFqcn(model.asyncClientComposerClassName().get()));
+            builder.addStatement("return ($T) composer.compose(client, clientConfiguration)", clientInterfaceName);
+        } else {
+            builder.addStatement("return client");
+        }
+        return builder.build();
     }
 
     private MethodSpec bearerTokenProviderMethod() {
@@ -157,6 +159,29 @@ private MethodSpec bearerTokenProviderMethod() {
                          .build();
     }
 
+    private MethodSpec initializeServiceClientConfigMethod() {
+        return MethodSpec.methodBuilder("initializeServiceClientConfig").addModifiers(Modifier.PRIVATE)
+                         .addParameter(SdkClientConfiguration.class, "clientConfig")
+                         .returns(serviceConfigClassName)
+                         .addStatement("$T endpointOverride = null", URI.class)
+                         .addStatement("$T endpointProvider = clientConfig.option($T.ENDPOINT_PROVIDER)",
+                                       EndpointProvider.class,
+                                       SdkClientOption.class)
+                         .addCode("if (clientConfig.option($T.ENDPOINT_OVERRIDDEN) != null"
+                                  + "&& $T.TRUE.equals(clientConfig.option($T.ENDPOINT_OVERRIDDEN))) {"
+                                  + "endpointOverride = clientConfig.option($T.ENDPOINT);"
+                                  + "}",
+                                  SdkClientOption.class, Boolean.class, SdkClientOption.class, SdkClientOption.class)
+                         .addStatement("return $T.builder()"
+                                       + ".overrideConfiguration(overrideConfiguration())"
+                                       + ".region(clientConfig.option($T.AWS_REGION))"
+                                       + ".endpointOverride(endpointOverride)"
+                                       + ".endpointProvider(endpointProvider)"
+                                       + ".build()",
+                                       serviceConfigClassName, AwsClientOption.class)
+                         .build();
+    }
+
     @Override
     public ClassName className() {
         return builderClassName;

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/SyncClientBuilderClass.java

@@ -24,6 +24,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.awscore.client.config.AwsClientOption;
+import software.amazon.awssdk.awscore.internal.client.ClientComposer;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
@@ -80,7 +81,10 @@ public TypeSpec poetSpec() {
             builder.addMethod(tokenProviderMethodImpl());
         }
 
-        return builder.addMethod(buildClientMethod()).build();
+        builder.addMethod(buildClientMethod());
+        builder.addMethod(initializeServiceClientConfigMethod());
+
+        return builder.build();
     }
 
     private MethodSpec endpointDiscoveryEnabled() {
@@ -118,31 +122,29 @@ private MethodSpec endpointProviderMethod() {
                          .build();
     }
 
-
     private MethodSpec buildClientMethod() {
-        return MethodSpec.methodBuilder("buildClient")
-                         .addAnnotation(Override.class)
-                         .addModifiers(Modifier.PROTECTED, Modifier.FINAL)
-                         .returns(clientInterfaceName)
-                         .addStatement("$T clientConfiguration = super.syncClientConfiguration()", SdkClientConfiguration.class)
-                         .addStatement("this.validateClientOptions(clientConfiguration)")
-                         .addStatement("$T endpointOverride = null", URI.class)
-                         .addStatement("$T endpointProvider = clientConfiguration.option($T.ENDPOINT_PROVIDER)",
-                                       EndpointProvider.class, SdkClientOption.class)
-                         .addCode("if (clientConfiguration.option($T.ENDPOINT_OVERRIDDEN) != null"
-                                  + "&& $T.TRUE.equals(clientConfiguration.option($T.ENDPOINT_OVERRIDDEN))) {"
-                                  + "endpointOverride = clientConfiguration.option($T.ENDPOINT);"
-                                  + "}",
-                                  SdkClientOption.class, Boolean.class, SdkClientOption.class, SdkClientOption.class)
-                         .addStatement("$T serviceClientConfiguration = $T.builder()"
-                                       + ".overrideConfiguration(overrideConfiguration())"
-                                       + ".region(clientConfiguration.option($T.AWS_REGION))"
-                                       + ".endpointOverride(endpointOverride)"
-                                       + ".endpointProvider(endpointProvider)"
-                                       + ".build()",
-                                       serviceConfigClassName, serviceConfigClassName, AwsClientOption.class)
-                         .addStatement("return new $T(serviceClientConfiguration, clientConfiguration)", clientClassName)
-                         .build();
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("buildClient")
+                                               .addAnnotation(Override.class)
+                                               .addModifiers(Modifier.PROTECTED, Modifier.FINAL)
+                                               .returns(clientInterfaceName)
+                                               .addStatement("$T clientConfiguration = super.syncClientConfiguration()",
+                                                             SdkClientConfiguration.class)
+                                               .addStatement("this.validateClientOptions(clientConfiguration)")
+                                               .addStatement("$T serviceClientConfiguration = initializeServiceClientConfig"
+                                                             + "(clientConfiguration)",
+                                                             serviceConfigClassName);
+
+        builder.addStatement("$1T client = new $2T(serviceClientConfiguration, clientConfiguration)",
+                             clientInterfaceName, clientClassName);
+        if (model.syncClientComposerClassName().isPresent()) {
+            builder.addStatement("$1T composer = new $2T()",
+                                 ClientComposer.class,
+                                 PoetUtils.classNameFromFqcn(model.syncClientComposerClassName().get()));
+            builder.addStatement("return ($T) composer.compose(client, clientConfiguration)", clientInterfaceName);
+        } else {
+            builder.addStatement("return client");
+        }
+        return builder.build();
     }
 
     private MethodSpec tokenProviderMethodImpl() {
@@ -156,6 +158,29 @@ private MethodSpec tokenProviderMethodImpl() {
                          .build();
     }
 
+    private MethodSpec initializeServiceClientConfigMethod() {
+        return MethodSpec.methodBuilder("initializeServiceClientConfig").addModifiers(Modifier.PRIVATE)
+                         .addParameter(SdkClientConfiguration.class, "clientConfig")
+                         .returns(serviceConfigClassName)
+                         .addStatement("$T endpointOverride = null", URI.class)
+                         .addStatement("$T endpointProvider = clientConfig.option($T.ENDPOINT_PROVIDER)",
+                                       EndpointProvider.class,
+                                       SdkClientOption.class)
+                         .addCode("if (clientConfig.option($T.ENDPOINT_OVERRIDDEN) != null"
+                                  + "&& $T.TRUE.equals(clientConfig.option($T.ENDPOINT_OVERRIDDEN))) {"
+                                  + "endpointOverride = clientConfig.option($T.ENDPOINT);"
+                                  + "}",
+                                  SdkClientOption.class, Boolean.class, SdkClientOption.class, SdkClientOption.class)
+                         .addStatement("return $T.builder()"
+                                       + ".overrideConfiguration(overrideConfiguration())"
+                                       + ".region(clientConfig.option($T.AWS_REGION))"
+                                       + ".endpointOverride(endpointOverride)"
+                                       + ".endpointProvider(endpointProvider)"
+                                       + ".build()",
+                                       serviceConfigClassName, AwsClientOption.class)
+                         .build();
+    }
+
     @Override
     public ClassName className() {
         return builderClassName;

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

@@ -192,7 +192,7 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
         builder.addModifiers(PUBLIC)
                .addAnnotation(Override.class);
 
-        if (builder.parameters.size() < 1) {
+        if (builder.parameters.isEmpty()) {
             throw new IllegalStateException("All client methods must have an argument");
         }
 

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

@@ -151,7 +151,7 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
         builder.addModifiers(PUBLIC)
                .addAnnotation(Override.class);
 
-        if (builder.parameters.size() < 1) {
+        if (builder.parameters.isEmpty()) {
             throw new IllegalStateException("All client methods must have an argument");
         }
 

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

@@ -146,6 +146,19 @@ public static IntermediateModel customContentTypeModels() {
         return new IntermediateModelBuilder(models).build();
     }
 
+    public static IntermediateModel composedClientJsonServiceModels() {
+        File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/rest-json/service-2.json").getFile());
+        File customizationModel =
+            new File(ClientTestModels.class.getResource("client/c2j/composedclient/customization.config").getFile());
+        CustomizationConfig customizationConfig = getCustomizationConfig(customizationModel);
+        C2jModels models = C2jModels.builder()
+                                    .serviceModel(getServiceModel(serviceModel))
+                                    .customizationConfig(customizationConfig)
+                                    .build();
+
+        return new IntermediateModelBuilder(models).build();
+    }
+
     public static IntermediateModel internalConfigModels() {
         File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/internalconfig/service-2.json").getFile());
         File customizationModel = new File(ClientTestModels.class.getResource("client/c2j/internalconfig/customization.config").getFile());

codegen/src/test/java/software/amazon/awssdk/codegen/poet/builder/BuilderClassTest.java

@@ -68,6 +68,11 @@ public void syncClientBuilderClass() throws Exception {
         validateGeneration(SyncClientBuilderClass::new, "test-sync-client-builder-class.java");
     }
 
+    @Test
+    public void syncComposedClientBuilderClass() throws Exception {
+        validateComposedClientGeneration(SyncClientBuilderClass::new, "test-composed-sync-client-builder-class.java");
+    }
+
     @Test
     public void asyncClientBuilderInterface() throws Exception {
         validateGeneration(AsyncClientBuilderInterface::new, "test-async-client-builder-interface.java");
@@ -78,10 +83,20 @@ public void asyncClientBuilderClass() throws Exception {
         validateGeneration(AsyncClientBuilderClass::new, "test-async-client-builder-class.java");
     }
 
+    @Test
+    public void asyncComposedClientBuilderClass() throws Exception {
+        validateComposedClientGeneration(AsyncClientBuilderClass::new, "test-composed-async-client-builder-class.java");
+    }
+
     private void validateGeneration(Function<IntermediateModel, ClassSpec> generatorConstructor, String expectedClassName) {
         assertThat(generatorConstructor.apply(ClientTestModels.restJsonServiceModels()), generatesTo(expectedClassName));
     }
 
+    private void validateComposedClientGeneration(Function<IntermediateModel, ClassSpec> generatorConstructor,
+                                                  String expectedClassName) {
+        assertThat(generatorConstructor.apply(ClientTestModels.composedClientJsonServiceModels()), generatesTo(expectedClassName));
+    }
+
     private void validateBearerAuthGeneration(Function<IntermediateModel, ClassSpec> generatorConstructor,
                                            String expectedClassName) {
         assertThat(generatorConstructor.apply(ClientTestModels.bearerAuthServiceModels()), generatesTo(expectedClassName));

codegen/src/test/java/software/amazon/awssdk/codegen/poet/common/UserAgentClassSpecTest.java

@@ -26,7 +26,7 @@
 public class UserAgentClassSpecTest {
 
     @Test
-    public void testGeneratedResponseForSyncOperations() {
+    void testUserAgentClass() {
         ClassSpec useragentspec = new UserAgentUtilsSpec(ClientTestModels.restJsonServiceModels());
         assertThat(useragentspec, generatesTo("test-user-agent-class.java"));
     }