Adding functionality to config preferred authschemeProvider

.changes/next-release/feature-AWSSDKforJavav2-8e1c19d.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Added ability to configure preferred authentication schemes when multiple auth options are available."
+}

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

@@ -28,6 +28,7 @@
 import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointAwareAuthSchemeParamsSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointBasedAuthSchemeProviderSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.ModelBasedAuthSchemeProviderSpec;
+import software.amazon.awssdk.codegen.poet.auth.scheme.PreferredAuthSchemeProviderSpec;
 
 public final class AuthSchemeGeneratorTasks extends BaseGeneratorTasks {
     private final GeneratorTaskParams generatorTaskParams;
@@ -45,6 +46,7 @@ protected List<GeneratorTask> createTasks() {
         tasks.add(generateProviderInterface());
         tasks.add(generateDefaultParamsImpl());
         tasks.add(generateModelBasedProvider());
+        tasks.add(generatePreferenceProvider());
         tasks.add(generateAuthSchemeInterceptor());
         if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
             tasks.add(generateEndpointBasedProvider());
@@ -69,6 +71,10 @@ private GeneratorTask generateModelBasedProvider() {
         return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new ModelBasedAuthSchemeProviderSpec(model));
     }
 
+    private GeneratorTask generatePreferenceProvider() {
+        return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new PreferredAuthSchemeProviderSpec(model));
+    }
+
     private GeneratorTask generateEndpointBasedProvider() {
         return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(),
                                      new EndpointBasedAuthSchemeProviderSpec(model));

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeProviderSpec.java

@@ -17,12 +17,16 @@
 
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.CodeBlock;
+import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.Consumer;
 import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
@@ -54,6 +58,9 @@ public TypeSpec poetSpec() {
                         .addMethod(resolveAuthSchemeMethod())
                         .addMethod(resolveAuthSchemeConsumerBuilderMethod())
                         .addMethod(defaultProviderMethod())
+                        .addMethod(staticBuilderMethodSpec())
+                        .addType(builderInterfaceSpec())
+                        .addType(builderClassSpec())
                         .build();
     }
 
@@ -104,4 +111,74 @@ private CodeBlock interfaceJavadoc() {
 
         return b.build();
     }
+
+    private MethodSpec staticBuilderMethodSpec() {
+        return MethodSpec.methodBuilder("builder")
+                         .addJavadoc("Create a builder for the auth scheme provider.")
+                         .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                         .returns(className().nestedClass("Builder"))
+                         .addStatement("return new $T()", ClassName.get(className().packageName(),
+                                                                      className().simpleName(),
+                                                                      className().simpleName() + "Builder")
+                         )
+                         .build();
+    }
+
+
+    private TypeSpec builderInterfaceSpec() {
+        return TypeSpec.interfaceBuilder("Builder")
+                       .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                       .addMethod(MethodSpec.methodBuilder("build")
+                                            .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                                            .addJavadoc("Returns a {@link $T} object that is created from the "
+                                                        + "properties that have been set on the builder.",
+                                                        className())
+                                            .returns(className())
+                                            .build())
+
+                       .addMethod(MethodSpec.methodBuilder("withPreferredAuthSchemes")
+                                            .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                                            .addJavadoc("Set the preferred auth schemes in order of preference.")
+                                            .returns(className().nestedClass("Builder"))
+                                            .addParameter(
+                                                ParameterizedTypeName.get(List.class, String.class),
+                                                "authSchemePreference"
+                                            )
+                                            .build())
+                       .build();
+    }
+
+    private TypeSpec builderClassSpec() {
+        return TypeSpec.classBuilder(authSchemeSpecUtils.authSchemeProviderBuilderName())
+                       .addAnnotation(SdkInternalApi.class)
+                       .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
+                       .addSuperinterface(className().nestedClass("Builder"))
+                       .addField(
+                           FieldSpec
+                               .builder(ParameterizedTypeName.get(List.class, String.class), "authSchemePreference")
+                               .addModifiers(Modifier.PRIVATE)
+                               .build())
+                       .addMethod(
+                           MethodSpec
+                               .methodBuilder("withPreferredAuthSchemes").addAnnotation(Override.class)
+                               .addModifiers(Modifier.PUBLIC)
+                               .addParameter(
+                                   ParameterizedTypeName.get(List.class, String.class),
+                                   "authSchemePreference"
+                               )
+                               .returns(className().nestedClass("Builder"))
+                               .addStatement("this.authSchemePreference = new $T<>(authSchemePreference)", ArrayList.class)
+                               .addStatement("return this")
+                               .build())
+                       .addMethod(
+                           MethodSpec
+                               .methodBuilder("build").addAnnotation(Override.class)
+                               .addModifiers(Modifier.PUBLIC)
+                               .returns(className())
+                               .addStatement("return new $T(defaultProvider(), authSchemePreference)",
+                                             authSchemeSpecUtils.preferredAuthSchemeProviderName())
+                               .build())
+                       .build();
+    }
 }
+

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java

@@ -97,6 +97,15 @@ public ClassName modeledAuthSchemeProviderName() {
         return ClassName.get(internalPackage(), "Modeled" + providerInterfaceName().simpleName());
     }
 
+    public ClassName preferredAuthSchemeProviderName() {
+        return ClassName.get(internalPackage(), "Preferred" + providerInterfaceName().simpleName());
+    }
+
+    public ClassName authSchemeProviderBuilderName() {
+        return ClassName.get(basePackage(),
+                             intermediateModel.getMetadata().getServiceName() + "AuthSchemeProviderBuilder");
+    }
+
     public ClassName authSchemeInterceptor() {
         return ClassName.get(internalPackage(), intermediateModel.getMetadata().getServiceName() + "AuthSchemeInterceptor");
     }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/PreferredAuthSchemeProviderSpec.java

@@ -0,0 +1,101 @@
+/*
+ * 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.auth.scheme;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
+import com.squareup.javapoet.TypeSpec;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.poet.ClassSpec;
+import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.utils.CollectionUtils;
+
+public class PreferredAuthSchemeProviderSpec implements ClassSpec {
+    private final AuthSchemeSpecUtils authSchemeSpecUtils;
+
+    public PreferredAuthSchemeProviderSpec(IntermediateModel intermediateModel) {
+        this.authSchemeSpecUtils = new AuthSchemeSpecUtils(intermediateModel);
+    }
+
+    @Override
+    public ClassName className() {
+        return authSchemeSpecUtils.preferredAuthSchemeProviderName();
+    }
+
+    @Override
+    public TypeSpec poetSpec() {
+        return PoetUtils.createClassBuilder(className())
+                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+                        .addAnnotation(SdkInternalApi.class)
+                        .addField(
+                            authSchemeSpecUtils.providerInterfaceName(), "delegate",
+                            Modifier.PRIVATE, Modifier.FINAL)
+                        .addField(
+                            ParameterizedTypeName.get(List.class, String.class), "authSchemePreference",
+                            Modifier.PRIVATE, Modifier.FINAL)
+                        .addSuperinterface(authSchemeSpecUtils.providerInterfaceName())
+                        .addMethod(constructor())
+                        .addMethod(resolveAuthSchemeMethod())
+                        .build();
+    }
+
+    private MethodSpec constructor() {
+        return MethodSpec
+            .constructorBuilder()
+            .addModifiers(Modifier.PUBLIC)
+            .addParameter(authSchemeSpecUtils.providerInterfaceName(), "delegate")
+            .addParameter(ParameterizedTypeName.get(List.class, String.class), "authSchemePreference")
+            .addStatement("this.delegate = delegate")
+            .addStatement("this.authSchemePreference = authSchemePreference != null ? authSchemePreference : $T.emptyList()",
+                          Collections.class)
+            .build();
+    }
+
+    private MethodSpec resolveAuthSchemeMethod() {
+        MethodSpec.Builder b = MethodSpec.methodBuilder("resolveAuthScheme")
+                                         .addModifiers(Modifier.PUBLIC)
+                                         .addAnnotation(Override.class)
+                                         .returns(authSchemeSpecUtils.resolverReturnType())
+                                         .addParameter(authSchemeSpecUtils.parametersInterfaceName(), "params");
+        b.addJavadoc("Resolve the auth schemes based on the given set of parameters.");
+        b.addStatement("$T candidateAuthSchemes = delegate.resolveAuthScheme(params)",
+                       authSchemeSpecUtils.resolverReturnType());
+        b.beginControlFlow("if ($T.isNullOrEmpty(authSchemePreference))", CollectionUtils.class)
+            .addStatement("return candidateAuthSchemes")
+            .endControlFlow();
+
+        b.addStatement("$T authSchemes = new $T<>()", authSchemeSpecUtils.resolverReturnType(), ArrayList.class);
+        b.beginControlFlow("authSchemePreference.forEach( preferredSchemeId -> ")
+            .addStatement("candidateAuthSchemes.stream().filter(a -> a.schemeId().equals(preferredSchemeId)).findFirst()"
+                       + ".ifPresent(a -> authSchemes.add(a))")
+            .endControlFlow(")");
+
+        b.beginControlFlow("candidateAuthSchemes.forEach(candidate -> ")
+         .beginControlFlow("if (!authSchemes.contains(candidate))")
+         .addStatement("authSchemes.add(candidate)")
+         .endControlFlow()
+            .endControlFlow(")");
+
+        b.addStatement("return authSchemes");
+        return b.build();
+    }
+}

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-auth-scheme-provider.java

@@ -1,27 +1,15 @@
-/*
- * 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.services.query.auth.scheme;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.Generated;
+import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
 import software.amazon.awssdk.services.query.auth.scheme.internal.DefaultQueryAuthSchemeProvider;
+import software.amazon.awssdk.services.query.auth.scheme.internal.PreferredQueryAuthSchemeProvider;
 
 /**
  * An auth scheme provider for Query service. The auth scheme provider takes a set of parameters using
@@ -50,4 +38,40 @@ default List<AuthSchemeOption> resolveAuthScheme(Consumer<QueryAuthSchemeParams.
     static QueryAuthSchemeProvider defaultProvider() {
         return DefaultQueryAuthSchemeProvider.create();
     }
+
+    /**
+     * Create a builder for the auth scheme provider.
+     */
+    static Builder builder() {
+        return new QueryAuthSchemeProviderBuilder();
+    }
+
+    interface Builder {
+        /**
+         * Returns a {@link QueryAuthSchemeProvider} object that is created from the properties that have been set on
+         * the builder.
+         */
+        QueryAuthSchemeProvider build();
+
+        /**
+         * Set the preferred auth schemes in order of preference.
+         */
+        Builder withPreferredAuthSchemes(List<String> authSchemePreference);
+    }
+
+    @SdkInternalApi
+    final class QueryAuthSchemeProviderBuilder implements Builder {
+        private List<String> authSchemePreference;
+
+        @Override
+        public Builder withPreferredAuthSchemes(List<String> authSchemePreference) {
+            this.authSchemePreference = new ArrayList<>(authSchemePreference);
+            return this;
+        }
+
+        @Override
+        public QueryAuthSchemeProvider build() {
+            return new PreferredQueryAuthSchemeProvider(defaultProvider(), authSchemePreference);
+        }
+    }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-provider.java

@@ -1,27 +1,15 @@
-/*
- * 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.services.query.auth.scheme;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.Generated;
+import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
 import software.amazon.awssdk.services.query.auth.scheme.internal.DefaultQueryAuthSchemeProvider;
+import software.amazon.awssdk.services.query.auth.scheme.internal.PreferredQueryAuthSchemeProvider;
 
 /**
  * An auth scheme provider for Query service. The auth scheme provider takes a set of parameters using
@@ -50,4 +38,40 @@ default List<AuthSchemeOption> resolveAuthScheme(Consumer<QueryAuthSchemeParams.
     static QueryAuthSchemeProvider defaultProvider() {
         return DefaultQueryAuthSchemeProvider.create();
     }
+
+    /**
+     * Create a builder for the auth scheme provider.
+     */
+    static Builder builder() {
+        return new QueryAuthSchemeProviderBuilder();
+    }
+
+    interface Builder {
+        /**
+         * Returns a {@link QueryAuthSchemeProvider} object that is created from the properties that have been set on
+         * the builder.
+         */
+        QueryAuthSchemeProvider build();
+
+        /**
+         * Set the preferred auth schemes in order of preference.
+         */
+        Builder withPreferredAuthSchemes(List<String> authSchemePreference);
+    }
+
+    @SdkInternalApi
+    final class QueryAuthSchemeProviderBuilder implements Builder {
+        private List<String> authSchemePreference;
+
+        @Override
+        public Builder withPreferredAuthSchemes(List<String> authSchemePreference) {
+            this.authSchemePreference = new ArrayList<>(authSchemePreference);
+            return this;
+        }
+
+        @Override
+        public QueryAuthSchemeProvider build() {
+            return new PreferredQueryAuthSchemeProvider(defaultProvider(), authSchemePreference);
+        }
+    }
 }