Generate marshallers for h2 request events

varunnvs92 · varunnvs92 · commit 28a5f480a6b2 · 2018-10-12T11:39:34.000-07:00
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/MarshallerGeneratorTasks.java b/codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/MarshallerGeneratorTasks.java
@@ -30,6 +30,7 @@
 import software.amazon.awssdk.codegen.model.intermediate.Metadata;
 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;
 
@@ -68,9 +69,12 @@ private boolean shouldGenerate(ShapeModel shapeModel) {
 
     private Stream<GeneratorTask> createTask(String javaShapeName, ShapeModel shapeModel) throws Exception {
         if (metadata.isJsonProtocol()) {
-            return ShapeType.Request == shapeModel.getShapeType() ?
-                   Stream.of(createPoetGeneratorTask(new MarshallerSpec(model, shapeModel))) :
-                   Stream.empty();
+
+            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(
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ShapeModelSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ShapeModelSpec.java
@@ -150,7 +150,7 @@ private CodeBlock traits(MemberModel m) {
         } else if (m.isMap()) {
             traits.add(createMapTrait(m));
         }
-        if (m.getHttp().getIsPayload()) {
+        if (m.getHttp().getIsPayload() || m.isEventPayload()) {
             traits.add(createPayloadTrait());
         }
         if (m.isJsonValue()) {
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/MarshallerSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/MarshallerSpec.java
@@ -29,6 +29,7 @@
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
+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.core.Request;
@@ -119,7 +120,7 @@ private MarshallerProtocolSpec getProtocolSpecs(software.amazon.awssdk.codegen.m
             case CBOR:
             case ION:
             case AWS_JSON:
-                return new JsonMarshallerSpec(intermediateModel, shapeModel);
+                return getJsonMarshallerSpec();
             case QUERY:
             case REST_XML:
             case EC2:
@@ -129,4 +130,11 @@ private MarshallerProtocolSpec getProtocolSpecs(software.amazon.awssdk.codegen.m
                 throw new RuntimeException("Unknown protocol: " + protocol.name());
         }
     }
+
+    private MarshallerProtocolSpec getJsonMarshallerSpec() {
+        if (shapeModel.isEvent()) {
+            return new EventStreamJsonMarshallerSpec(intermediateModel, shapeModel);
+        }
+        return new JsonMarshallerSpec(intermediateModel, shapeModel);
+    }
 }
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/protocols/EventStreamJsonMarshallerSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/protocols/EventStreamJsonMarshallerSpec.java
@@ -0,0 +1,117 @@
+/*
+ * Copyright 2010-2018 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.transform.protocols;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.CodeBlock;
+import com.squareup.javapoet.FieldSpec;
+import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
+import software.amazon.awssdk.core.Request;
+import software.amazon.awssdk.core.protocol.OperationInfo;
+import software.amazon.awssdk.core.protocol.ProtocolMarshaller;
+
+/**
+ * MarshallerSpec for event shapes in Json protocol
+ */
+public final class EventStreamJsonMarshallerSpec extends JsonMarshallerSpec {
+
+    private static final String JSON_CONTENT_TYPE = "application/json";
+
+    private final IntermediateModel intermediateModel;
+
+    public EventStreamJsonMarshallerSpec(IntermediateModel model, ShapeModel shapeModel) {
+        super(model, shapeModel);
+        this.intermediateModel = model;
+    }
+
+    @Override
+    public CodeBlock marshalCodeBlock(ClassName requestClassName) {
+        String variableName = shapeModel.getVariable().getVariableName();
+        CodeBlock.Builder builder =
+            CodeBlock.builder()
+                     .addStatement("$T<$T<$T>> protocolMarshaller = protocolFactory.createProtocolMarshaller"
+                                   + "(SDK_OPERATION_BINDING, $L)",
+                                   ProtocolMarshaller.class, Request.class,
+                                   requestClassName, variableName)
+                     .addStatement("$T<$T> request = protocolMarshaller.marshall($L)",
+                                   Request.class, requestClassName, variableName)
+                     .addStatement("request.addHeader(\":message-type\", \"event\")")
+                     .addStatement("request.addHeader(\":event-type\", \"$L\")", getMemberNameFromEventStream());
+
+        // Add :content-type header only if payload is present
+        if (!shapeModel.hasNoEventPayload()) {
+            builder.addStatement("request.addHeader(\":content-type\", \"$L\")", determinePayloadContentType());
+        }
+
+        builder.addStatement("return request");
+
+        return builder.build();
+    }
+
+    @Override
+    protected FieldSpec operationInfoField() {
+        CodeBlock.Builder builder =
+            CodeBlock.builder()
+                     .add("$T.builder()", OperationInfo.class)
+                     .add(".hasExplicitPayloadMember($L)", shapeModel.isHasPayloadMember() ||
+                                                           shapeModel.getExplicitEventPayloadMember() != null)
+                     .add(".hasPayloadMembers($L)", shapeModel.hasPayloadMembers())
+                     .add(".build()");
+
+
+        return FieldSpec.builder(ClassName.get(OperationInfo.class), "SDK_OPERATION_BINDING")
+                        .addModifiers(Modifier.PRIVATE, Modifier.FINAL, Modifier.STATIC)
+                        .initializer(builder.build())
+                        .build();
+    }
+
+    private String getMemberNameFromEventStream() {
+        ShapeModel eventStream = EventStreamUtils.getBaseEventStreamShape(intermediateModel, shapeModel);
+        return eventStream.getMembers().stream()
+                          .filter(memberModel -> memberModel.getShape().equals(shapeModel))
+                          .findAny()
+                          .map(MemberModel::getC2jName)
+                          .orElseThrow(() -> new IllegalStateException(
+                              String.format("Unable to find %s from its parent event stream", shapeModel.getC2jName())));
+    }
+
+    private String determinePayloadContentType() {
+        MemberModel explicitEventPayload = shapeModel.getExplicitEventPayloadMember();
+        if (explicitEventPayload != null) {
+            return getPayloadContentType(explicitEventPayload);
+        }
+
+        return JSON_CONTENT_TYPE;
+    }
+
+    private String getPayloadContentType(MemberModel memberModel) {
+        final String blobContentType = "application/octet-stream";
+        final String stringContentType = "text/plain";
+        final String variableType = memberModel.getVariable().getVariableType();
+
+        if ("software.amazon.awssdk.core.SdkBytes".equals(variableType)) {
+            return blobContentType;
+        } else if ("String".equals(variableType)) {
+            return stringContentType;
+        }
+
+        return JSON_CONTENT_TYPE;
+    }
+}
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/protocols/JsonMarshallerSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/transform/protocols/JsonMarshallerSpec.java
@@ -39,8 +39,8 @@
  */
 public class JsonMarshallerSpec implements MarshallerProtocolSpec {
 
-    private final Metadata metadata;
-    private final ShapeModel shapeModel;
+    protected final Metadata metadata;
+    protected final ShapeModel shapeModel;
 
     public JsonMarshallerSpec(IntermediateModel model, ShapeModel shapeModel) {
         this.metadata = model.getMetadata();
@@ -82,7 +82,12 @@ public FieldSpec protocolFactory() {
     @Override
     public List<FieldSpec> memberVariables() {
         List<FieldSpec> fields = new ArrayList<>();
+        fields.add(operationInfoField());
+        fields.add(protocolFactory());
+        return fields;
+    }
 
+    protected FieldSpec operationInfoField() {
         CodeBlock.Builder initializationCodeBlockBuilder = CodeBlock.builder()
                                                                     .add("$T.builder()", OperationInfo.class);
         initializationCodeBlockBuilder.add(".requestUri($S)", shapeModel.getMarshaller().getRequestUri())
@@ -102,12 +107,9 @@ public List<FieldSpec> memberVariables() {
 
         CodeBlock codeBlock = initializationCodeBlockBuilder.add(".build()").build();
 
-        FieldSpec.Builder instance = FieldSpec.builder(ClassName.get(OperationInfo.class), "SDK_OPERATION_BINDING")
-                                              .addModifiers(Modifier.PRIVATE, Modifier.FINAL, Modifier.STATIC)
-                                              .initializer(codeBlock);
-
-        fields.add(instance.build());
-        fields.add(protocolFactory());
-        return fields;
+        return FieldSpec.builder(ClassName.get(OperationInfo.class), "SDK_OPERATION_BINDING")
+                        .addModifiers(Modifier.PRIVATE, Modifier.FINAL, Modifier.STATIC)
+                        .initializer(codeBlock)
+                        .build();
     }
 }
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/inputevent.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/inputevent.java
@@ -16,6 +16,7 @@
 import software.amazon.awssdk.core.protocol.SdkField;
 import software.amazon.awssdk.core.protocol.SdkPojo;
 import software.amazon.awssdk.core.protocol.traits.LocationTrait;
+import software.amazon.awssdk.core.protocol.traits.PayloadTrait;
 import software.amazon.awssdk.utils.ToString;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
 import software.amazon.awssdk.utils.builder.ToCopyableBuilder;
@@ -25,10 +26,11 @@
 @Generated("software.amazon.awssdk:codegen")
 public final class InputEvent implements SdkPojo, ToCopyableBuilder<InputEvent.Builder, InputEvent>, InputEventStream {
     private static final SdkField<SdkBytes> EXPLICIT_PAYLOAD_MEMBER_FIELD = SdkField
-        .<SdkBytes> builder(MarshallingType.SDK_BYTES).getter(getter(InputEvent::explicitPayloadMember))
-                                                      .setter(setter(Builder::explicitPayloadMember))
-                                                      .traits(LocationTrait.builder().location(MarshallLocation.PAYLOAD).locationName("ExplicitPayloadMember").build())
-                                                      .build();
+        .<SdkBytes> builder(MarshallingType.SDK_BYTES)
+        .getter(getter(InputEvent::explicitPayloadMember))
+        .setter(setter(Builder::explicitPayloadMember))
+        .traits(LocationTrait.builder().location(MarshallLocation.PAYLOAD).locationName("ExplicitPayloadMember").build(),
+                PayloadTrait.create()).build();
 
     private static final List<SdkField<?>> SDK_FIELDS = Collections
         .unmodifiableList(Arrays.asList(EXPLICIT_PAYLOAD_MEMBER_FIELD));
@@ -156,4 +158,3 @@ public List<SdkField<?>> sdkFields() {
         }
     }
 }
-
diff --git a/core/aws-core/src/main/java/software/amazon/awssdk/awscore/protocol/json/AwsJsonProtocolFactory.java b/core/aws-core/src/main/java/software/amazon/awssdk/awscore/protocol/json/AwsJsonProtocolFactory.java
@@ -171,7 +171,7 @@ private AwsStructuredJsonFactory getSdkFactory() {
         }
     }
 
-    public <T extends software.amazon.awssdk.awscore.AwsRequest> ProtocolMarshaller<Request<T>> createProtocolMarshaller(
+    public <T> ProtocolMarshaller<Request<T>> createProtocolMarshaller(
         OperationInfo operationInfo, T origRequest) {
         return JsonProtocolMarshallerBuilder.<T>standard()
             .jsonGenerator(createGenerator(operationInfo))
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/protocol/json/JsonProtocolMarshaller.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/protocol/json/JsonProtocolMarshaller.java
@@ -23,7 +23,6 @@
 import software.amazon.awssdk.core.DefaultRequest;
 import software.amazon.awssdk.core.Request;
 import software.amazon.awssdk.core.SdkBytes;
-import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.protocol.MarshallingType;
 import software.amazon.awssdk.core.protocol.OperationInfo;
 import software.amazon.awssdk.core.protocol.ProtocolRequestMarshaller;
@@ -40,7 +39,7 @@
  * @param <OrigRequestT> Type of the original request object.
  */
 @SdkInternalApi
-public class JsonProtocolMarshaller<OrigRequestT extends SdkRequest> implements ProtocolRequestMarshaller<OrigRequestT> {
+public class JsonProtocolMarshaller<OrigRequestT> implements ProtocolRequestMarshaller<OrigRequestT> {
 
     private static final MarshallerRegistry MARSHALLER_REGISTRY = createMarshallerRegistry();
 
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/protocol/json/JsonProtocolMarshallerBuilder.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/protocol/json/JsonProtocolMarshallerBuilder.java
@@ -16,7 +16,6 @@
 package software.amazon.awssdk.core.protocol.json;
 
 import software.amazon.awssdk.annotations.SdkProtectedApi;
-import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.internal.protocol.json.JsonProtocolMarshaller;
 import software.amazon.awssdk.core.internal.protocol.json.NullAsEmptyBodyProtocolRequestMarshaller;
 import software.amazon.awssdk.core.protocol.OperationInfo;
@@ -28,16 +27,16 @@
  * @param <T> Type of the original request object.
  */
 @SdkProtectedApi
-public class JsonProtocolMarshallerBuilder<T extends SdkRequest> {
+public class JsonProtocolMarshallerBuilder<T> {
 
     private StructuredJsonGenerator jsonGenerator;
     private String contentType;
     private OperationInfo operationInfo;
     private boolean sendExplicitNullForPayload;
     private T originalRequest;
 
-    public static <T extends SdkRequest> JsonProtocolMarshallerBuilder<T> standard() {
-        return new JsonProtocolMarshallerBuilder<T>();
+    public static <T> JsonProtocolMarshallerBuilder<T> standard() {
+        return new JsonProtocolMarshallerBuilder<>();
     }
 
     public JsonProtocolMarshallerBuilder<T> jsonGenerator(StructuredJsonGenerator jsonGenerator) {

Original file line number	Diff line number	Diff line change
`@@ -150,7 +150,7 @@ private CodeBlock traits(MemberModel m) {`
`150`	`150`	`} else if (m.isMap()) {`
`151`	`151`	`traits.add(createMapTrait(m));`
`152`	`152`	`}`
`153`		`- if (m.getHttp().getIsPayload()) {`
	`153`	`+ if (m.getHttp().getIsPayload() \|\| m.isEventPayload()) {`
`154`	`154`	`traits.add(createPayloadTrait());`
`155`	`155`	`}`
`156`	`156`	`if (m.isJsonValue()) {`
Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ private AwsStructuredJsonFactory getSdkFactory() {`
`171`	`171`	`}`
`172`	`172`	`}`
`173`	`173`
`174`		`- public <T extends software.amazon.awssdk.awscore.AwsRequest> ProtocolMarshaller<Request<T>> createProtocolMarshaller(`
	`174`	`+ public <T> ProtocolMarshaller<Request<T>> createProtocolMarshaller(`
`175`	`175`	`OperationInfo operationInfo, T origRequest) {`
`176`	`176`	`return JsonProtocolMarshallerBuilder.<T>standard()`
`177`	`177`	`.jsonGenerator(createGenerator(operationInfo))`