Refact xml marshallers/unmarshallers

Author: varunnvs92

codegen/pom.xml

@@ -130,6 +130,11 @@
             <artifactId>aws-query-protocol</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>aws-xml-protocol</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>protocol-core</artifactId>

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

@@ -15,37 +15,26 @@
 
 package software.amazon.awssdk.codegen.emitters.tasks;
 
-import static software.amazon.awssdk.codegen.model.intermediate.Protocol.AWS_JSON;
 import static software.amazon.awssdk.utils.FunctionalUtils.safeFunction;
 
-import freemarker.template.Template;
-import java.io.IOException;
 import java.util.List;
-import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import software.amazon.awssdk.codegen.emitters.FreemarkerGeneratorTask;
 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;
 import software.amazon.awssdk.codegen.poet.transform.MarshallerSpec;
-import software.amazon.awssdk.utils.ImmutableMap;
 
 public class MarshallerGeneratorTasks extends BaseGeneratorTasks {
 
-    private final String transformClassDir;
     private final Metadata metadata;
-    private final Map<String, ShapeModel> shapes;
 
     public MarshallerGeneratorTasks(GeneratorTaskParams dependencies) {
         super(dependencies);
-        this.transformClassDir = dependencies.getPathProvider().getTransformDirectory();
         this.metadata = model.getMetadata();
-        this.shapes = model.getShapes();
     }
 
     @Override
@@ -62,56 +51,18 @@ 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() || metadata.getProtocol() == Protocol.QUERY || metadata.getProtocol() == Protocol.EC2) {
-            return ShapeType.Request == shapeModel.getShapeType() ||
-                   (ShapeType.Model == shapeModel.getShapeType() && shapeModel.isEvent()
-                    && EventStreamUtils.isRequestEvent(model, shapeModel))
+        return ShapeType.Request == shapeModel.getShapeType() || (ShapeType.Model == shapeModel.getShapeType()
+                                                                  && shapeModel.isEvent()
+                                                                  && EventStreamUtils.isRequestEvent(model, shapeModel))
                    ? Stream.of(createPoetGeneratorTask(new MarshallerSpec(model, shapeModel)))
                    : Stream.empty();
-        }
-
-        return Stream.of(
-            createMarshallerTask(javaShapeName,
-                                 freemarker.getModelMarshallerTemplate(),
-                                 javaShapeName + "Marshaller",
-                                 transformClassDir));
-    }
-
-    private GeneratorTask createMarshallerTask(String javaShapeName, Template template,
-                                               String marshallerClassName, String marshallerDirectory) throws IOException {
-        Map<String, Object> marshallerDataModel = ImmutableMap.<String, Object>builder()
-            .put("fileHeader", model.getFileHeader())
-            .put("shapeName", javaShapeName)
-            .put("shapes", shapes)
-            .put("metadata", metadata)
-            .put("transformPackage", model.getMetadata().getFullTransformPackageName())
-            .put("requestTransformPackage", model.getMetadata().getFullRequestTransformPackageName())
-            .put("customConfig", model.getCustomizationConfig())
-            .put("className", marshallerClassName)
-            .put("protocolEnum", getProtocolEnumName())
-            .build();
-
-        return new FreemarkerGeneratorTask(marshallerDirectory,
-                                           marshallerClassName,
-                                           template,
-                                           marshallerDataModel);
-    }
-
-    private String getProtocolEnumName() {
-        switch (metadata.getProtocol()) {
-            case CBOR:
-            case ION:
-            case AWS_JSON:
-                return AWS_JSON.name();
-            default:
-                return metadata.getProtocol().name();
-        }
     }
 }

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

@@ -17,7 +17,6 @@
 
 import static software.amazon.awssdk.utils.FunctionalUtils.safeFunction;
 
-import freemarker.template.Template;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -54,7 +53,6 @@ protected List<GeneratorTask> createTasks() throws Exception {
     }
 
     private GeneratorTask createTask(String javaShapeName, ShapeModel shapeModel) throws Exception {
-        Template template = freemarker.getModelUnmarshallerTemplate();
         ShapeType shapeType = shapeModel.getShapeType();
         Map<String, Object> dataModel = ImmutableMap.<String, Object>builder()
                 .put("fileHeader", model.getFileHeader())
@@ -65,16 +63,9 @@ private GeneratorTask createTask(String javaShapeName, ShapeModel shapeModel) th
                 .build();
 
         switch (shapeType) {
-            case Response:
-            case Model: {
-                return new FreemarkerGeneratorTask(transformClassDir,
-                                         javaShapeName + "Unmarshaller",
-                                                   template,
-                                                   dataModel);
-            }
             case Exception: {
                 return new FreemarkerGeneratorTask(transformClassDir,
-                                         javaShapeName + "Unmarshaller",
+                                                   javaShapeName + "Unmarshaller",
                                                    freemarker.getExceptionUnmarshallerTemplate(),
                                                    dataModel);
             }
@@ -90,10 +81,6 @@ private boolean shouldGenerate(ShapeModel shapeModel) {
             return false;
         }
         switch (shapeModel.getShapeType()) {
-            case Response:
-            case Model:
-                // The event stream shape is a container for event subtypes and isn't something that needs to ever be unmarshalled
-                return !shapeModel.isEventStream();
             case Exception:
                 // Generating Exception Unmarshallers is not required for the JSON protocol
                 return !metadata.isJsonProtocol();

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

@@ -140,6 +140,11 @@ public class CustomizationConfig {
 
     private boolean skipSyncClientGeneration;
 
+    /**
+     * List of output shapes for which the root xml element should be used while unmarshalling the response
+     */
+    private List<String> useRootXmlElementForResult = new ArrayList<>();
+
     /**
      * Custom Response metadata
      */
@@ -389,6 +394,14 @@ public void setSkipSyncClientGeneration(boolean skipSyncClientGeneration) {
         this.skipSyncClientGeneration = skipSyncClientGeneration;
     }
 
+    public List<String> getUseRootXmlElementForResult() {
+        return useRootXmlElementForResult;
+    }
+
+    public void setUseRootXmlElementForResult(List<String> useRootXmlElementForResult) {
+        this.useRootXmlElementForResult = useRootXmlElementForResult;
+    }
+
     public Map<String, String> getCustomResponseMetadata() {
         return customResponseMetadata;
     }

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

@@ -42,7 +42,7 @@
 import software.amazon.awssdk.codegen.poet.client.specs.Ec2ProtocolSpec;
 import software.amazon.awssdk.codegen.poet.client.specs.JsonProtocolSpec;
 import software.amazon.awssdk.codegen.poet.client.specs.ProtocolSpec;
-import software.amazon.awssdk.codegen.poet.client.specs.QueryXmlProtocolSpec;
+import software.amazon.awssdk.codegen.poet.client.specs.QueryProtocolSpec;
 import software.amazon.awssdk.codegen.poet.client.specs.XmlProtocolSpec;
 import software.amazon.awssdk.codegen.utils.PaginatorUtils;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
@@ -191,7 +191,7 @@ static ProtocolSpec getProtocolSpecs(PoetExtensions poetExtensions, Intermediate
         Protocol protocol = model.getMetadata().getProtocol();
         switch (protocol) {
             case QUERY:
-                return new QueryXmlProtocolSpec(poetExtensions);
+                return new QueryProtocolSpec(poetExtensions);
             case REST_XML:
                 return new XmlProtocolSpec(poetExtensions);
             case EC2:

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

@@ -21,7 +21,7 @@
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.protocols.query.AwsEc2ProtocolFactory;
 
-public class Ec2ProtocolSpec extends QueryXmlProtocolSpec {
+public class Ec2ProtocolSpec extends QueryProtocolSpec {
 
     public Ec2ProtocolSpec(PoetExtensions poetExtensions) {
         super(poetExtensions);
@@ -90,4 +90,4 @@ private MethodSpec dryRunMethod() {
 
     }
     */
-}
+}

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

@@ -44,16 +44,16 @@
 import software.amazon.awssdk.protocols.query.AwsQueryProtocolFactory;
 import software.amazon.awssdk.utils.StringUtils;
 
-public class QueryXmlProtocolSpec implements ProtocolSpec {
+public class QueryProtocolSpec implements ProtocolSpec {
 
-    private final PoetExtensions poetExtensions;
+    protected final PoetExtensions poetExtensions;
     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 QueryProtocolSpec(PoetExtensions poetExtensions) {
         this.poetExtensions = poetExtensions;
     }
 
@@ -63,6 +63,10 @@ public FieldSpec protocolFactory(IntermediateModel model) {
                         .addModifiers(Modifier.PRIVATE, Modifier.FINAL).build();
     }
 
+    protected Class<?> protocolFactoryClass() {
+        return AwsQueryProtocolFactory.class;
+    }
+
     @Override
     public List<FieldSpec> additionalFields() {
         return singletonList(FieldSpec.builder(listOfUnmarshallersType, "exceptionUnmarshallers")
@@ -83,15 +87,11 @@ public MethodSpec initProtocolFactory(IntermediateModel model) {
                                                    poetExtensions.getModelClass(model.getSdkModeledExceptionBaseClassName()))
                                     .build());
         methodSpec.addStatement("this.exceptionUnmarshallers = unmarshallers");
-        methodSpec.addStatement("return new $T()", protocolFactoryClass());
+        methodSpec.addStatement("return $T.builder().build()", protocolFactoryClass());
 
         return methodSpec.build();
     }
 
-    protected Class<?> protocolFactoryClass() {
-        return AwsQueryProtocolFactory.class;
-    }
-
     private ClassName getErrorUnmarshallerClass(IntermediateModel model) {
         return StringUtils.isNotBlank(model.getExceptionUnmarshallerImpl()) ?
                PoetUtils.classNameFromFqcn(model.getExceptionUnmarshallerImpl()) :
@@ -126,17 +126,17 @@ public CodeBlock executionHandler(OperationModel opModel) {
         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>()" +
-                     ".withResponseHandler($N)" +
-                     ".withErrorResponseHandler($N)" +
-                     ".withInput($L)",
-                     ClientExecutionParams.class,
-                     requestType,
-                     responseType,
-                     "responseHandler",
-                     "errorResponseHandler",
-                     opModel.getInput().getVariableName());
+            .builder()
+            .add("\n\nreturn clientHandler.execute(new $T<$T, $T>()" +
+                 ".withResponseHandler($N)" +
+                 ".withErrorResponseHandler($N)" +
+                 ".withInput($L)",
+                 ClientExecutionParams.class,
+                 requestType,
+                 responseType,
+                 "responseHandler",
+                 "errorResponseHandler",
+                 opModel.getInput().getVariableName());
         if (opModel.hasStreamingInput()) {
             return codeBlock.add(".withMarshaller(new $T(new $T(protocolFactory), requestBody)));",
                                  ParameterizedTypeName.get(ClassName.get(StreamingRequestMarshaller.class), requestType),
@@ -154,7 +154,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
         ClassName marshaller = poetExtensions.getRequestTransformClass(opModel.getInputShape().getShapeName() + "Marshaller");
 
         String asyncRequestBody = opModel.hasStreamingInput() ? ".withAsyncRequestBody(requestBody)"
-                : "";
+                                                              : "";
         return CodeBlock.builder().add("\n\nreturn clientHandler.execute(new $T<$T, $T>()\n" +
                                        ".withMarshaller(new $T(protocolFactory))" +
                                        ".withResponseHandler(responseHandler)" +
@@ -190,4 +190,4 @@ public List<CodeBlock> errorUnmarshallers(IntermediateModel model) {
                             .add("unmarshallers.add(new $T());", exceptionClass).build();
         }).collect(Collectors.toList());
     }
-}
+}