Rest Xml marshaller refactor

Author: varunnvs92

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

@@ -28,6 +28,7 @@
 import software.amazon.awssdk.codegen.emitters.GeneratorTask;
 import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
 import software.amazon.awssdk.codegen.model.intermediate.Metadata;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeType;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
@@ -61,15 +62,15 @@ private boolean shouldGenerate(ShapeModel shapeModel) {
             info("Skip generating marshaller class for " + shapeModel.getShapeName());
             return false;
         }
+
         ShapeType shapeType = shapeModel.getShapeType();
         return (ShapeType.Request == shapeType || (ShapeType.Model == shapeType && metadata.isJsonProtocol()))
                // The event stream shape is a container for event subtypes and isn't something that needs to ever be marshalled
                && !shapeModel.isEventStream();
     }
 
     private Stream<GeneratorTask> createTask(String javaShapeName, ShapeModel shapeModel) throws Exception {
-        if (metadata.isJsonProtocol()) {
-
+        if (metadata.isJsonProtocol() || isRestXml()) {
             return ShapeType.Request == shapeModel.getShapeType() ||
                    (ShapeType.Model == shapeModel.getShapeType() && shapeModel.isEvent()
                     && EventStreamUtils.isRequestEvent(model, shapeModel))
@@ -84,6 +85,10 @@ private Stream<GeneratorTask> createTask(String javaShapeName, ShapeModel shapeM
                                  transformClassDir));
     }
 
+    private boolean isRestXml() {
+        return Protocol.REST_XML == metadata.getProtocol();
+    }
+
     private GeneratorTask createMarshallerTask(String javaShapeName, Template template,
                                                String marshallerClassName, String marshallerDirectory) throws IOException {
         Map<String, Object> marshallerDataModel = ImmutableMap.<String, Object>builder()

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

@@ -17,6 +17,8 @@
 
 import static com.squareup.javapoet.TypeSpec.Builder;
 import static java.util.Collections.singletonList;
+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.applySignerOverrideMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.getCustomResponseHandler;
@@ -42,10 +44,12 @@
 import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionJsonMarshaller;
 import software.amazon.awssdk.awscore.internal.client.handler.AwsClientHandlerUtils;
 import software.amazon.awssdk.awscore.protocol.json.AwsJsonProtocolFactory;
+import software.amazon.awssdk.awscore.protocol.xml.AwsXmlProtocolFactory;
 import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
@@ -101,6 +105,10 @@ public TypeSpec poetSpec() {
             classBuilder.addField(AwsJsonProtocolFactory.class, "jsonProtocolFactory", Modifier.PRIVATE, Modifier.FINAL);
         }
 
+        if (model.getMetadata().getProtocol() == Protocol.REST_XML) {
+            classBuilder.addField(AwsXmlProtocolFactory.class, "protocolFactory", PRIVATE, FINAL);
+        }
+
         if (model.hasPaginators()) {
             classBuilder.addMethod(applyPaginatorUserAgentMethod(poetExtensions, model));
         }
@@ -129,6 +137,9 @@ private MethodSpec constructor(Builder classBuilder) {
         if (model.getMetadata().isCborProtocol()) {
             builder.addStatement("this.jsonProtocolFactory = init(false)");
         }
+        if (model.getMetadata().getProtocol() == Protocol.REST_XML) {
+            builder.addStatement("this.protocolFactory = AwsXmlProtocolFactory.builder().build()");
+        }
         if (hasOperationWithEventStreamOutput()) {
             classBuilder.addField(FieldSpec.builder(ClassName.get(Executor.class), "executor",
                                                     Modifier.PRIVATE, Modifier.FINAL)

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

@@ -30,6 +30,7 @@
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.protocol.xml.AwsXmlProtocolFactory;
 import software.amazon.awssdk.codegen.docs.SimpleMethodOverload;
 import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
@@ -84,6 +85,10 @@ public TypeSpec poetSpec() {
 
         classBuilder.addMethod(closeMethod());
 
+        if (model.getMetadata().getProtocol() == Protocol.REST_XML) {
+            classBuilder.addField(AwsXmlProtocolFactory.class, "protocolFactory", PRIVATE, FINAL);
+        }
+
         if (model.hasPaginators()) {
             classBuilder.addMethod(applyPaginatorUserAgentMethod(poetExtensions, model));
         }
@@ -122,6 +127,10 @@ private MethodSpec constructor() {
         } else {
             builder.addStatement("this.$N = init()", protocolSpec.protocolFactory(model).name);
         }
+
+        if (model.getMetadata().getProtocol() == Protocol.REST_XML) {
+            builder.addStatement("this.protocolFactory = AwsXmlProtocolFactory.builder().build()");
+        }
         return builder.build();
     }
 
@@ -186,7 +195,7 @@ static ProtocolSpec getProtocolSpecs(PoetExtensions poetExtensions, Protocol pro
         switch (protocol) {
             case QUERY:
             case REST_XML:
-                return new QueryXmlProtocolSpec(poetExtensions);
+                return new QueryXmlProtocolSpec(poetExtensions, protocol);
             case EC2:
                 return new Ec2ProtocolSpec(poetExtensions);
             case AWS_JSON:

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/Ec2ProtocolSpec.java

@@ -18,12 +18,13 @@
 import com.squareup.javapoet.MethodSpec;
 import java.util.ArrayList;
 import java.util.List;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 
 public class Ec2ProtocolSpec extends QueryXmlProtocolSpec {
 
     public Ec2ProtocolSpec(PoetExtensions poetExtensions) {
-        super(poetExtensions);
+        super(poetExtensions, Protocol.EC2);
     }
 
     @Override

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryXmlProtocolSpec.java

@@ -31,29 +31,32 @@
 import software.amazon.awssdk.awscore.http.response.DefaultErrorResponseHandler;
 import software.amazon.awssdk.awscore.http.response.StaxResponseHandler;
 import software.amazon.awssdk.awscore.protocol.xml.StandardErrorUnmarshaller;
-import software.amazon.awssdk.awscore.protocol.xml.StaxOperationMetadata;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
+import software.amazon.awssdk.core.internal.protocol.restxml.unmarshall.StaxOperationMetadata;
 import software.amazon.awssdk.core.runtime.transform.StreamingRequestMarshaller;
 import software.amazon.awssdk.core.runtime.transform.Unmarshaller;
 import software.amazon.awssdk.utils.StringUtils;
 
 public class QueryXmlProtocolSpec implements ProtocolSpec {
 
     private final PoetExtensions poetExtensions;
+    private final Protocol protocol;
     private final TypeName unmarshallerType = ParameterizedTypeName.get(Unmarshaller.class,
                                                                         AwsServiceException.class,
                                                                         Node.class);
     private final TypeName listOfUnmarshallersType = ParameterizedTypeName.get(ClassName.get("java.util", "List"),
                                                                                unmarshallerType);
 
-    public QueryXmlProtocolSpec(PoetExtensions poetExtensions) {
+    public QueryXmlProtocolSpec(PoetExtensions poetExtensions, Protocol protocol) {
         this.poetExtensions = poetExtensions;
+        this.protocol = protocol;
     }
 
     @Override
@@ -118,7 +121,9 @@ public CodeBlock errorResponseHandler(OperationModel opModel) {
     public CodeBlock executionHandler(OperationModel opModel) {
         TypeName responseType = poetExtensions.getModelClass(opModel.getReturnType().getReturnType());
         ClassName requestType = poetExtensions.getModelClass(opModel.getInput().getVariableType());
+
         ClassName marshaller = poetExtensions.getTransformClass(opModel.getInputShape().getShapeName() + "Marshaller");
+
         CodeBlock.Builder codeBlock = CodeBlock
                 .builder()
                 .add("\n\nreturn clientHandler.execute(new $T<$T, $T>()" +
@@ -131,13 +136,16 @@ public CodeBlock executionHandler(OperationModel opModel) {
                      "responseHandler",
                      "errorResponseHandler",
                      opModel.getInput().getVariableName());
+
         if (opModel.hasStreamingInput()) {
-            return codeBlock.add(".withMarshaller(new $T(new $T(), requestBody)));",
+            return codeBlock.add(".withMarshaller(new $T(new $T($L), requestBody)));",
                                  ParameterizedTypeName.get(ClassName.get(StreamingRequestMarshaller.class), requestType),
-                                 marshaller)
+                                 marshaller,
+                                 protocol == Protocol.REST_XML ? "protocolFactory" : "")
                             .build();
         }
-        return codeBlock.add(".withMarshaller(new $T()) $L);", marshaller,
+        return codeBlock.add(".withMarshaller(new $T($L)) $L);", marshaller,
+                             protocol == Protocol.REST_XML ? "protocolFactory" : "",
                              opModel.hasStreamingOutput() ? ", responseTransformer" : "").build();
     }
 
@@ -150,7 +158,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
         String asyncRequestBody = opModel.hasStreamingInput() ? ".withAsyncRequestBody(requestBody)"
                 : "";
         return CodeBlock.builder().add("\n\nreturn clientHandler.execute(new $T<$T, $T>()\n" +
-                                       ".withMarshaller(new $T())" +
+                                       ".withMarshaller(new $T($L))" +
                                        ".withResponseHandler(responseHandler)" +
                                        ".withErrorResponseHandler($N)\n" +
                                        asyncRequestBody +
@@ -159,6 +167,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                                        requestType,
                                        pojoResponseType,
                                        marshaller,
+                                       protocol == Protocol.REST_XML ? "protocolFactory" : "",
                                        "errorResponseHandler",
                                        opModel.getInput().getVariableName(),
                                        opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "")

codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ShapeModelSpec.java

@@ -239,9 +239,11 @@ private CodeBlock createListTrait(MemberModel m) {
                         .add("$T.builder()\n"
                              + ".memberLocationName($S)\n"
                              + ".memberFieldInfo($L)\n"
+                             + ".flattened($L)\n"
                              + ".build()", ClassName.get(ListTrait.class),
                              m.getListModel().getMemberLocationName(),
-                             containerSdkFieldInitializer(m.getListModel().getListMemberModel()))
+                             containerSdkFieldInitializer(m.getListModel().getListMemberModel()),
+                             m.getHttp().isFlattened())
                         .build();
     }
 

codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/MarshallerSpec.java

@@ -32,6 +32,7 @@
 import software.amazon.awssdk.codegen.poet.transform.protocols.EventStreamJsonMarshallerSpec;
 import software.amazon.awssdk.codegen.poet.transform.protocols.JsonMarshallerSpec;
 import software.amazon.awssdk.codegen.poet.transform.protocols.MarshallerProtocolSpec;
+import software.amazon.awssdk.codegen.poet.transform.protocols.XmlMarshallerSpec;
 import software.amazon.awssdk.core.Request;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.runtime.transform.Marshaller;
@@ -63,6 +64,7 @@ public MarshallerSpec(IntermediateModel intermediateModel, ShapeModel shapeModel
 
     @Override
     public TypeSpec poetSpec() {
+
         return TypeSpec.classBuilder(className)
                        .addJavadoc("{@link $T} Marshaller", requestClassName)
                        .addModifiers(Modifier.PUBLIC)
@@ -121,8 +123,11 @@ private MarshallerProtocolSpec getProtocolSpecs(software.amazon.awssdk.codegen.m
             case ION:
             case AWS_JSON:
                 return getJsonMarshallerSpec();
-            case QUERY:
+
             case REST_XML:
+                return new XmlMarshallerSpec(intermediateModel, shapeModel);
+
+            case QUERY:
             case EC2:
             case API_GATEWAY:
                 throw new UnsupportedOperationException("Not yet supported.");