Append paginator user agent for pagination methods

Author: zoewangg

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

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.codegen.poet.client;
 
 import static com.squareup.javapoet.TypeSpec.Builder;
+import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applyPaginatorUserAgentMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.getCustomResponseHandler;
 import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.getProtocolSpecs;
 
@@ -67,6 +68,9 @@ public TypeSpec poetSpec() {
                                         .addMethods(protocolSpec.additionalMethods())
                                         .addMethod(protocolSpec.initProtocolFactory(model));
 
+        if (model.hasPaginators()) {
+            classBuilder.addMethod(applyPaginatorUserAgentMethod(poetExtensions, model));
+        }
         protocolSpec.createErrorResponseHandler().ifPresent(classBuilder::addMethod);
 
         return classBuilder.build();
@@ -119,7 +123,7 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
     @Override
     protected MethodSpec.Builder paginatedMethodBody(MethodSpec.Builder builder, OperationModel opModel) {
         return builder.addModifiers(Modifier.PUBLIC)
-                      .addStatement("return new $T(this, $L)",
+                      .addStatement("return new $T(this, applyPaginatorUserAgent($L))",
                                     poetExtensions.getResponseClassForPaginatedAsyncOperation(opModel.getOperationName()),
                                     opModel.getInput().getVariableName());
     }

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

@@ -21,15 +21,23 @@
 import com.squareup.javapoet.ParameterSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
+import com.squareup.javapoet.TypeVariableName;
 import java.util.Optional;
 import java.util.function.Consumer;
 import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.poet.PoetExtensions;
+import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
+import software.amazon.awssdk.core.util.VersionInfo;
 import software.amazon.awssdk.utils.Validate;
 
 final class ClientClassUtils {
+    private static final String PAGINATOR_USER_AGENT = "PAGINATED";
+
     private ClientClassUtils() {
     }
 
@@ -79,4 +87,39 @@ static MethodSpec consumerBuilderVariant(MethodSpec spec, String javadoc) {
 
         return result.build();
     }
+
+    static MethodSpec applyPaginatorUserAgentMethod(PoetExtensions poetExtensions, IntermediateModel model) {
+
+        TypeVariableName typeVariableName =
+            TypeVariableName.get("T", poetExtensions.getModelClass(model.getSdkRequestBaseClassName()));
+
+        ParameterizedTypeName parameterizedTypeName = ParameterizedTypeName
+            .get(ClassName.get(Consumer.class), ClassName.get(AwsRequestOverrideConfiguration.Builder.class));
+
+        CodeBlock codeBlock = CodeBlock.builder()
+                                       .addStatement("$T userAgentApplier = b -> b.addApiName($T.builder().version"
+                                                     + "($T.SDK_VERSION).name($S).build())",
+                                                     parameterizedTypeName, ApiName.class,
+                                                     VersionInfo.class,
+                                                     PAGINATOR_USER_AGENT)
+                                       .addStatement("$T overrideConfiguration =\n"
+                                                     + "            request.overrideConfiguration().map(c -> c.toBuilder()"
+                                                     + ".applyMutation"
+                                                     + "(userAgentApplier).build())\n"
+                                                     + "            .orElse((AwsRequestOverrideConfiguration.builder()"
+                                                     + ".applyMutation"
+                                                     + "(userAgentApplier).build()))", AwsRequestOverrideConfiguration.class)
+                                       .addStatement("return (T) request.toBuilder().overrideConfiguration"
+                                                     + "(overrideConfiguration).build()")
+                                       .build();
+
+        return MethodSpec.methodBuilder("applyPaginatorUserAgent")
+                         .addModifiers(Modifier.PRIVATE)
+                         .addParameter(typeVariableName, "request")
+                         .addTypeVariable(typeVariableName)
+                         .addCode(codeBlock)
+                         .returns(typeVariableName)
+                         .build();
+    }
+
 }

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

@@ -17,6 +17,7 @@
 
 import static javax.lang.model.element.Modifier.FINAL;
 import static javax.lang.model.element.Modifier.PRIVATE;
+import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applyPaginatorUserAgentMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.getCustomResponseHandler;
 
 import com.squareup.javapoet.ClassName;
@@ -81,7 +82,9 @@ public TypeSpec poetSpec() {
 
         classBuilder.addMethod(closeMethod());
 
-        classBuilder.addMethods(protocolSpec.additionalMethods());
+        if (model.hasPaginators()) {
+            classBuilder.addMethod(applyPaginatorUserAgentMethod(poetExtensions, model));
+        }
 
         return classBuilder.build();
     }
@@ -147,7 +150,7 @@ private List<MethodSpec> paginatedMethods(OperationModel opModel) {
                                                         .addAnnotation(Override.class)
                                                         .returns(poetExtensions.getResponseClassForPaginatedSyncOperation(
                                                             opModel.getOperationName()))
-                                                        .addStatement("return new $T(this, $L)",
+                                                        .addStatement("return new $T(this, applyPaginatorUserAgent($L))",
                                                                       poetExtensions.getResponseClassForPaginatedSyncOperation(
                                                                           opModel.getOperationName()),
                                                                       opModel.getInput().getVariableName())