Port base changes from hackathon-2022-extension-methods branch

Author: RanVaknin

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

@@ -0,0 +1,41 @@
+/*
+ * 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;
+
+/**
+ * Whether the service has API extensions defined apart from the service model.
+ */
+public class ClientExtensions {
+
+    private boolean sync = false;
+    private boolean async = false;
+
+    public boolean getSync() {
+        return sync;
+    }
+
+    public void setSync(boolean sync) {
+        this.sync = sync;
+    }
+
+    public boolean getAsync() {
+        return async;
+    }
+
+    public void setAsync(boolean async) {
+        this.async = async;
+    }
+}

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

@@ -249,6 +249,11 @@ public class CustomizationConfig {
      */
     private String asyncClientDecorator;
 
+    /**
+     * If the client has API extensions to include.
+     */
+    private ClientExtensions clientExtensions;
+
     /**
      * Only for s3. A set of customization to related to multipart operations.
      */
@@ -746,6 +751,14 @@ public void setAsyncClientDecorator(String asyncClientDecorator) {
         this.asyncClientDecorator = asyncClientDecorator;
     }
 
+    public ClientExtensions getClientExtensions() {
+        return clientExtensions;
+    }
+
+    public void setClientExtensions(ClientExtensions clientExtensions) {
+        this.clientExtensions = clientExtensions;
+    }
+
     public boolean isDelegateSyncClientClass() {
         return delegateSyncClientClass;
     }

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

@@ -37,6 +37,9 @@ public final class PoetUtils {
                                                                   .addMember("value", "$S", "software.amazon.awssdk:codegen")
                                                                   .build();
 
+    private static final String EXTENSION_PACKAGE = ".extensions";
+    private static final String EXTENSION_INTERFACE_SUFFIX = "SdkExtension";
+
     private PoetUtils() {
     }
 
@@ -108,4 +111,10 @@ public static JavaFile buildJavaFile(ClassSpec spec) {
         spec.staticImports().forEach(i -> i.memberNames().forEach(m -> builder.addStaticImport(i.className(), m)));
         return builder.build();
     }
+
+    public static ClassName extensionInterfaceClassName(ClassName clientInterfaceName) {
+        String packageName = clientInterfaceName.packageName() + EXTENSION_PACKAGE;
+        String className = clientInterfaceName.simpleName() + EXTENSION_INTERFACE_SUFFIX;
+        return ClassName.get(packageName, className);
+    }
 }

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

@@ -21,6 +21,7 @@
 import static javax.lang.model.element.Modifier.PUBLIC;
 import static javax.lang.model.element.Modifier.STATIC;
 import static software.amazon.awssdk.codegen.internal.Constant.ASYNC_STREAMING_INPUT_PARAM;
+import static software.amazon.awssdk.codegen.poet.PoetUtils.extensionInterfaceClassName;
 import static software.amazon.awssdk.codegen.internal.Constant.EVENT_PUBLISHER_PARAM_NAME;
 import static software.amazon.awssdk.codegen.internal.Constant.EVENT_RESPONSE_HANDLER_PARAM_NAME;
 
@@ -35,6 +36,7 @@
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 import org.reactivestreams.Publisher;
@@ -47,6 +49,8 @@
 import software.amazon.awssdk.codegen.docs.SimpleMethodOverload;
 import software.amazon.awssdk.codegen.docs.WaiterDocs;
 import software.amazon.awssdk.codegen.model.config.customization.AdditionalBuilderMethod;
+import software.amazon.awssdk.codegen.model.config.customization.ClientExtensions;
+import software.amazon.awssdk.codegen.model.config.customization.CustomizationConfig;
 import software.amazon.awssdk.codegen.model.config.customization.UtilitiesMethod;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
@@ -111,6 +115,8 @@ protected TypeSpec.Builder createTypeSpec() {
 
     protected void addInterfaceClass(TypeSpec.Builder type) {
         type.addSuperinterface(AwsClient.class);
+        Optional<ClassName> extensionInterface = findExtensionInterface(className, model.getCustomizationConfig());
+        extensionInterface.ifPresent(type::addSuperinterface);
     }
 
     protected void addAnnotations(TypeSpec.Builder type) {
@@ -551,4 +557,13 @@ protected MethodSpec.Builder batchManagerOperationBody(MethodSpec.Builder builde
                       .addStatement("throw new $T()", UnsupportedOperationException.class);
     }
 
+    private static Optional<ClassName> findExtensionInterface(ClassName clientInterfaceName,
+                                                              CustomizationConfig customizationConfig) {
+        ClientExtensions clientExtensions = customizationConfig.getClientExtensions();
+        if (clientExtensions != null && clientExtensions.getAsync()) {
+            return Optional.of(extensionInterfaceClassName(clientInterfaceName));
+        }
+        return Optional.empty();
+    }
+
 }

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

@@ -25,6 +25,7 @@
 import static software.amazon.awssdk.codegen.internal.Constant.SYNC_CLIENT_SOURCE_PATH_PARAM_NAME;
 import static software.amazon.awssdk.codegen.internal.Constant.SYNC_STREAMING_INPUT_PARAM;
 import static software.amazon.awssdk.codegen.internal.Constant.SYNC_STREAMING_OUTPUT_PARAM;
+import static software.amazon.awssdk.codegen.poet.PoetUtils.extensionInterfaceClassName;
 import static software.amazon.awssdk.codegen.poet.client.AsyncClientInterface.STREAMING_TYPE_VARIABLE;
 
 import com.squareup.javapoet.ClassName;
@@ -38,6 +39,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 import software.amazon.awssdk.annotations.SdkPublicApi;
@@ -49,6 +51,8 @@
 import software.amazon.awssdk.codegen.docs.DocConfiguration;
 import software.amazon.awssdk.codegen.docs.SimpleMethodOverload;
 import software.amazon.awssdk.codegen.docs.WaiterDocs;
+import software.amazon.awssdk.codegen.model.config.customization.ClientExtensions;
+import software.amazon.awssdk.codegen.model.config.customization.CustomizationConfig;
 import software.amazon.awssdk.codegen.model.config.customization.UtilitiesMethod;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
@@ -102,6 +106,8 @@ public final TypeSpec poetSpec() {
 
     protected void addInterfaceClass(TypeSpec.Builder type) {
         type.addSuperinterface(AwsClient.class);
+        Optional<ClassName> extensionInterface = findExtensionInterface(className, model.getCustomizationConfig());
+        extensionInterface.ifPresent(type::addSuperinterface);
     }
 
     protected TypeSpec.Builder createTypeSpec() {
@@ -559,4 +565,13 @@ protected MethodSpec.Builder waiterOperationBody(MethodSpec.Builder builder) {
         return builder.addModifiers(DEFAULT, PUBLIC)
                       .addStatement("throw new $T()", UnsupportedOperationException.class);
     }
+
+    private static Optional<ClassName> findExtensionInterface(ClassName clientInterfaceName,
+                                                              CustomizationConfig customizationConfig) {
+        ClientExtensions clientExtensions = customizationConfig.getClientExtensions();
+        if (clientExtensions != null && clientExtensions.getSync()) {
+            return Optional.of(extensionInterfaceClassName(clientInterfaceName));
+        }
+        return Optional.empty();
+    }
 }

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

@@ -606,6 +606,18 @@ public static IntermediateModel batchManagerModels() {
         return new IntermediateModelBuilder(models).build();
     }
 
+    public static IntermediateModel clientExtensionModels() {
+        File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/clientextensions/service-2.json").getFile());
+        File customizationModel = new File(ClientTestModels.class.getResource("client/c2j/clientextensions/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/AsyncClientInterfaceTest.java

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

codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/SyncClientInterfaceTest.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.clientExtensionModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.restJsonServiceModels;
 import static software.amazon.awssdk.codegen.poet.PoetMatchers.generatesTo;
 
@@ -28,4 +29,10 @@ public void syncClientInterface() {
         ClassSpec syncClientInterface = new SyncClientInterface(restJsonServiceModels());
         assertThat(syncClientInterface, generatesTo("test-json-client-interface.java"));
     }
+
+    @Test
+    public void syncClientInterfaceClientExtensions() {
+        ClassSpec syncClientInterface = new SyncClientInterface(clientExtensionModels());
+        assertThat(syncClientInterface, generatesTo("test-clientextensions-sync-interface.java"));
+    }
 }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/utils/PoetUtilsTest.java

@@ -0,0 +1,38 @@
+/*
+ * 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.poet.utils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.squareup.javapoet.ClassName;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.codegen.poet.PoetUtils;
+
+/**
+ * Validate functionality of {@link PoetUtils} methods.
+ */
+public class PoetUtilsTest {
+
+    @Test
+    public void findExtensionInterface_whenNotExists_returnEmpty() {
+        ClassName clientName = ClassName.get("software.amazon.awssdk.sdkservice", "SdkServiceClient");
+        ClassName extensionClass = PoetUtils.extensionInterfaceClassName(clientName);
+        assertThat(extensionClass).isNotNull();
+        assertThat(extensionClass.canonicalName())
+            .isEqualTo("software.amazon.awssdk.sdkservice.extensions.SdkServiceClientSdkExtension");
+    }
+
+}

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

@@ -0,0 +1,6 @@
+{
+    "clientExtensions": {
+      "sync": true,
+      "async": true
+    }
+}