Various refactoring for JSON (un)marshallers and fixes. Tests and support for Timestamp format trait

Author: shorea

codegen/src/main/java/software/amazon/awssdk/codegen/AddShapes.java

@@ -163,6 +163,7 @@ private MemberModel generateMemberModel(String c2jMemberName, Member c2jMemberDe
                                          .withDocumentation(c2jMemberDefinition.getDocumentation()))
                    .withSetterModel(new VariableModel(variableName, variableType, variableDeclarationType))
                    .withGetterModel(new ReturnTypeModel(variableType))
+                   .withTimestampFormat(resolveTimestampFormat(c2jMemberDefinition, c2jShape))
                    .withJsonValue(c2jMemberDefinition.getJsonValue());
         memberModel.setDocumentation(c2jMemberDefinition.getDocumentation());
         memberModel.setDeprecated(c2jMemberDefinition.isDeprecated());
@@ -204,6 +205,11 @@ private MemberModel generateMemberModel(String c2jMemberName, Member c2jMemberDe
         return memberModel;
     }
 
+    private String resolveTimestampFormat(Member c2jMemberDefinition, Shape c2jShape) {
+        return c2jMemberDefinition.getTimestampFormat() != null ?
+               c2jMemberDefinition.getTimestampFormat() : c2jShape.getTimestampFormat();
+    }
+
     private ParameterHttpMapping generateParameterHttpMapping(Shape parentShape,
                                                               String memberName,
                                                               Member member,

codegen/src/main/java/software/amazon/awssdk/codegen/IntermediateModelBuilder.java

@@ -18,7 +18,6 @@
 import static software.amazon.awssdk.codegen.AddMetadata.constructMetadata;
 import static software.amazon.awssdk.codegen.RemoveUnusedShapes.removeUnusedShapes;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/IntermediateModel.java

@@ -31,7 +31,6 @@
 import software.amazon.awssdk.codegen.internal.Utils;
 import software.amazon.awssdk.codegen.model.config.customization.CustomizationConfig;
 import software.amazon.awssdk.codegen.model.service.PaginatorDefinition;
-import software.amazon.awssdk.codegen.naming.DefaultNamingStrategy;
 import software.amazon.awssdk.codegen.naming.NamingStrategy;
 import software.amazon.awssdk.core.SdkResponseMetadata;
 import software.amazon.awssdk.utils.IoUtils;

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/MemberModel.java

@@ -76,6 +76,8 @@ public class MemberModel extends DocumentationModel {
 
     private boolean isJsonValue;
 
+    private String timestampFormat;
+
     public String getName() {
         return name;
     }
@@ -466,6 +468,19 @@ public MemberModel withJsonValue(boolean jsonValue) {
         return this;
     }
 
+    public String getTimestampFormat() {
+        return timestampFormat;
+    }
+
+    public void setTimestampFormat(String timestampFormat) {
+        this.timestampFormat = timestampFormat;
+    }
+
+    public MemberModel withTimestampFormat(String timestampFormat) {
+        setTimestampFormat(timestampFormat);
+        return this;
+    }
+
     /**
      * @return Implementation of {@link PathMarshaller} to use if this member is bound the the URI.
      * @throws IllegalStateException If this member is not bound to the URI. Templates should first check

codegen/src/main/java/software/amazon/awssdk/codegen/model/service/Member.java

@@ -44,6 +44,8 @@ public class Member {
     @JsonProperty("jsonvalue")
     private boolean jsonValue;
 
+    private String timestampFormat;
+
     public String getShape() {
         return shape;
     }
@@ -139,4 +141,12 @@ public boolean getJsonValue() {
     public void setJsonValue(boolean jsonValue) {
         this.jsonValue = jsonValue;
     }
+
+    public String getTimestampFormat() {
+        return timestampFormat;
+    }
+
+    public void setTimestampFormat(String timestampFormat) {
+        this.timestampFormat = timestampFormat;
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/model/service/Shape.java

@@ -67,6 +67,8 @@ public class Shape {
     @JsonProperty(value = "event")
     private boolean isEvent;
 
+    private String timestampFormat;
+
     public boolean isFault() {
         return fault;
     }
@@ -238,4 +240,12 @@ public boolean isEvent() {
     public void setIsEvent(boolean event) {
         isEvent = event;
     }
+
+    public String getTimestampFormat() {
+        return timestampFormat;
+    }
+
+    public void setTimestampFormat(String timestampFormat) {
+        this.timestampFormat = timestampFormat;
+    }
 }

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

@@ -30,7 +30,7 @@
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.awscore.eventstream.EventStreamAsyncResponseTransformer;
-import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionJsonUnmarshaller;
+import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionPojoSupplier;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.awscore.internal.protocol.json.AwsJsonProtocol;
 import software.amazon.awssdk.awscore.protocol.json.AwsJsonProtocolFactory;
@@ -145,16 +145,14 @@ public CodeBlock responseHandler(IntermediateModel model, OperationModel opModel
                      JsonOperationMetadata.class,
                      true,
                      false,
-                     ClassName.get(EventStreamTaggedUnionJsonUnmarshaller.class));
+                     ClassName.get(EventStreamTaggedUnionPojoSupplier.class));
 
             eventStreamUtils.getEventStreamMembers()
                             .forEach(m -> {
-                                String unmarshallerClassName = m.getShape().getVariable().getVariableType() + "Unmarshaller";
-                                builder.add(".addUnmarshaller(\"$L\", $T.getInstance())\n",
-                                            m.getC2jName(),
-                                            poetExtensions.getTransformClass(unmarshallerClassName));
+                                builder.add(".addSdkPojoSupplier(\"$L\", $T::builder)\n",
+                                            m.getC2jName(), poetExtensions.getModelClass(m.getName()));
                             });
-            builder.add(".defaultUnmarshaller((in) -> $T.UNKNOWN)\n"
+            builder.add(".defaultSdkPojoSupplier(() -> $T.UNKNOWN)\n"
                         + ".build());\n", eventStreamUtils.eventStreamBaseClass());
         }
         return builder.build();

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

@@ -25,6 +25,7 @@
 import com.squareup.javapoet.TypeVariableName;
 import com.squareup.javapoet.WildcardTypeName;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -88,16 +89,23 @@ public TypeSpec poetSpec() {
             ClassName responseHandlerClass = eventStreamUtils.responseHandlerType();
             return PoetUtils.createInterfaceBuilder(modelClass)
                             .addAnnotation(SdkPublicApi.class)
+                            .addSuperinterface(ClassName.get(SdkPojo.class))
                             .addJavadoc("Base interface for all event types of the $L API.", eventStreamUtils.getApiName())
                             .addField(FieldSpec.builder(modelClass, "UNKNOWN")
                                                .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                                                .initializer(CodeBlock.builder()
                                                                      .add("new $T() {\n"
+                                                                          + "        @Override\n"
+                                                                          + "        public $T<$T<?>> sdkFields() {\n"
+                                                                          + "            return $T.emptyList();\n"
+                                                                          + "        }\n"
                                                                           + "        @Override\n"
                                                                           + "        public void accept($T.Visitor visitor) {\n"
                                                                           + "            visitor.visitDefault(this);\n"
                                                                           + "        }\n"
-                                                                          + "    };\n", modelClass, responseHandlerClass
+                                                                          + "    };\n",
+                                                                          modelClass, List.class, SdkField.class,
+                                                                          Collections.class, responseHandlerClass
                                                                      )
                                                                      .build())
                                                .addJavadoc("Special type of {@link $T} for unknown types of events that this "
@@ -206,7 +214,7 @@ public ClassName className() {
         return shapeModelSpec.className();
     }
 
-    private ClassName builderClassName(){
+    private ClassName builderClassName() {
         return className().nestedClass("Builder");
     }
 

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

@@ -31,15 +31,16 @@
 import software.amazon.awssdk.codegen.model.intermediate.ShapeType;
 import software.amazon.awssdk.codegen.naming.NamingStrategy;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
+import software.amazon.awssdk.core.protocol.MarshallLocation;
+import software.amazon.awssdk.core.protocol.MarshallingType;
+import software.amazon.awssdk.core.protocol.SdkField;
 import software.amazon.awssdk.core.protocol.traits.IdempotencyTrait;
 import software.amazon.awssdk.core.protocol.traits.JsonValueTrait;
 import software.amazon.awssdk.core.protocol.traits.ListTrait;
 import software.amazon.awssdk.core.protocol.traits.LocationTrait;
 import software.amazon.awssdk.core.protocol.traits.MapTrait;
-import software.amazon.awssdk.core.protocol.MarshallLocation;
-import software.amazon.awssdk.core.protocol.MarshallingType;
-import software.amazon.awssdk.core.protocol.SdkField;
 import software.amazon.awssdk.core.protocol.traits.PayloadTrait;
+import software.amazon.awssdk.core.protocol.traits.TimestampFormatTrait;
 
 /**
  * Provides Poet specs related to shape models.
@@ -148,12 +149,15 @@ private CodeBlock traits(MemberModel m) {
         if (m.getHttp().getIsPayload()) {
             traits.add(createPayloadTrait());
         }
-        if(m.isJsonValue()) {
+        if (m.isJsonValue()) {
             traits.add(createJsonValueTrait());
         }
-        if(m.isIdempotencyToken()) {
+        if (m.isIdempotencyToken()) {
             traits.add(createIdempotencyTrait());
         }
+        if (m.getTimestampFormat() != null) {
+            traits.add(createTimestampFormatTrait(m));
+        }
         if (!traits.isEmpty()) {
             return CodeBlock.builder()
                             .add(".traits(" + traits.stream().map(t -> "$L").collect(Collectors.joining(", ")) + ")",
@@ -164,6 +168,15 @@ private CodeBlock traits(MemberModel m) {
         }
     }
 
+    private CodeBlock createTimestampFormatTrait(MemberModel m) {
+        TimestampFormatTrait.Format format = TimestampFormatTrait.Format.fromString(m.getTimestampFormat());
+        ClassName traitClass = ClassName.get(TimestampFormatTrait.class);
+        ClassName formatClass = ClassName.get(TimestampFormatTrait.Format.class);
+        return CodeBlock.builder()
+                        .add("$T.create($T.$L)", traitClass, formatClass, format.name())
+                        .build();
+    }
+
     private CodeBlock createLocationTrait(MemberModel m) {
         return CodeBlock.builder()
                         // TODO will marshall and unmarshall location name ever differ?

codegen/src/test/java/software/amazon/awssdk/codegen/naming/DefaultNamingStrategyTest.java

@@ -276,19 +276,19 @@ public void getServiceName_ThrowsException_WhenServiceIdHasWhiteSpace() {
     @Test
     public void getSdkFieldFieldName_SingleWord() {
         assertThat(strat.getSdkFieldFieldName(new MemberModel().withName("foo")))
-            .isEqualTo("FOO");
+            .isEqualTo("FOO_FIELD");
     }
 
     @Test
     public void getSdkFieldFieldName_CamalCaseConvertedToScreamCase() {
         assertThat(strat.getSdkFieldFieldName(new MemberModel().withName("fooBar")))
-            .isEqualTo("FOO_BAR");
+            .isEqualTo("FOO_BAR_FIELD");
     }
 
     @Test
     public void getSdkFieldFieldName_PascalCaseConvertedToScreamCase() {
         assertThat(strat.getSdkFieldFieldName(new MemberModel().withName("FooBar")))
-            .isEqualTo("FOO_BAR");
+            .isEqualTo("FOO_BAR_FIELD");
     }
 
     private void validateConversion(String input, String expectedOutput) {