Add support for request event streams

Author: varunnvs92

build-tools/src/main/resources/software/amazon/awssdk/checkstyle-suppressions.xml

@@ -42,4 +42,4 @@
     <!-- Assumes getter/setter match JSON property -->
     <suppress checks="AbbreviationAsWordInName"
               files=".*[\\/]software[\\/]amazon[\\/]awssdk[\\/]services[\\/]s3[\\/]event[\\/]S3EventNotification.java$"/>
-</suppressions>
+</suppressions>

build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml

@@ -104,4 +104,4 @@
         <Bug pattern="BC_UNCONFIRMED_CAST,SIC_INNER_SHOULD_BE_STATIC_ANON,DLS_DEAD_LOCAL_STORE,DM_CONVERT_CASE,NM_CLASS_NOT_EXCEPTION,NP_NULL_ON_SOME_PATH"/>
     </Match>
 
-</FindBugsFilter>
+</FindBugsFilter>

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

@@ -173,6 +173,8 @@ private MemberModel generateMemberModel(String c2jMemberName, Member c2jMemberDe
                 .withBeanStyleGetterMethodName(namingStrategy.getBeanStyleGetterMethodName(c2jMemberName))
                 .withBeanStyleSetterMethodName(namingStrategy.getBeanStyleSetterMethodName(c2jMemberName));
         memberModel.setIdempotencyToken(c2jMemberDefinition.isIdempotencyToken());
+        memberModel.setEventPayload(c2jMemberDefinition.isEventPayload());
+        memberModel.setEventHeader(c2jMemberDefinition.isEventHeader());
 
         // Pass the xmlNameSpace from the member reference
         if (c2jMemberDefinition.getXmlNamespace() != null) {

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

@@ -25,7 +25,6 @@
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamResponseHandlerBuilderImplSpec;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamResponseHandlerSpec;
-import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamVisitorBuilderImplSpec;
 
 /**
@@ -54,10 +53,9 @@ protected List<GeneratorTask> createTasks() throws Exception {
     }
 
     private Stream<ClassSpec> eventStreamClassSpecs(OperationModel opModel) {
-        EventStreamUtils eventStreamUtils = EventStreamUtils.create(params.getPoetExtensions(), opModel);
         return Stream.of(
-            new EventStreamResponseHandlerSpec(params, eventStreamUtils),
-            new EventStreamResponseHandlerBuilderImplSpec(params, eventStreamUtils),
-            new EventStreamVisitorBuilderImplSpec(params, eventStreamUtils));
+            new EventStreamResponseHandlerSpec(params, opModel),
+            new EventStreamResponseHandlerBuilderImplSpec(params, opModel),
+            new EventStreamVisitorBuilderImplSpec(params, opModel));
     }
 }

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

@@ -74,6 +74,10 @@ public class MemberModel extends DocumentationModel {
 
     private String beanStyleSetterName;
 
+    private boolean eventPayload;
+
+    private boolean eventHeader;
+
     public String getName() {
         return name;
     }
@@ -274,6 +278,22 @@ public void setDeprecated(boolean deprecated) {
         this.deprecated = deprecated;
     }
 
+    public boolean isEventPayload() {
+        return eventPayload;
+    }
+
+    public void setEventPayload(boolean eventPayload) {
+        this.eventPayload = eventPayload;
+    }
+
+    public boolean isEventHeader() {
+        return eventHeader;
+    }
+
+    public void setEventHeader(boolean eventHeader) {
+        this.eventHeader = eventHeader;
+    }
+
     public ListModel getListModel() {
         return listModel;
     }
@@ -448,7 +468,8 @@ public void setIdempotencyToken(boolean idempotencyToken) {
     }
 
     public boolean getIsBinary() {
-        return http.getIsStreaming() || (http.getIsPayload() && isSdkBytesType());
+        return http.getIsStreaming() ||
+               (isSdkBytesType() && (http.getIsPayload() || isEventPayload()));
     }
 
     /**

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

@@ -193,11 +193,21 @@ public void setPaginated(boolean paginated) {
      * @return True if the operation has an event stream member in the output shape. False otherwise.
      */
     public boolean hasEventStreamOutput() {
-        return outputShape != null
-               && outputShape.getMembers() != null
-               && outputShape.getMembers().stream()
-                             .filter(m -> m.getShape() != null)
-                             .anyMatch(m -> m.getShape().isEventStream());
+        return containsEventStream(outputShape);
     }
 
+    /**
+     * @return True if the operation has an event stream member in the input shape. False otherwise.
+     */
+    public boolean hasEventStreamInput() {
+        return containsEventStream(inputShape);
+    }
+
+    private boolean containsEventStream(ShapeModel shapeModel) {
+        return shapeModel != null
+               && shapeModel.getMembers() != null
+               && shapeModel.getMembers().stream()
+                            .filter(m -> m.getShape() != null)
+                            .anyMatch(m -> m.getShape().isEventStream());
+    }
 }

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

@@ -175,11 +175,53 @@ public List<MemberModel> getUnboundMembers() {
         return unboundMembers;
     }
 
+    /**
+     * @return The list of members whose are not marked with either eventheader or eventpayload trait.
+     */
+    @JsonIgnore
+    public List<MemberModel> getUnboundEventMembers() {
+        if (members == null) {
+            return new ArrayList<>();
+        }
+
+        return members.stream()
+                      .filter(m -> !m.isEventHeader())
+                      .filter(m -> !m.isEventPayload())
+                      .collect(Collectors.toList());
+    }
+
     /**
      * @return True if the shape has an explicit payload member or implicit payload member(s).
      */
     public boolean hasPayloadMembers() {
-        return hasPayloadMember || getUnboundMembers().size() > 0;
+        return hasPayloadMember ||
+               getExplicitEventPayloadMember() != null ||
+               getUnboundMembers().size() > 0 ||
+               getUnboundEventMembers().size() > 0;
+    }
+
+    /**
+     * Explicit event payload member will have "eventpayload" trait set to true.
+     * There can be at most only one member that can be declared as explicit payload.
+     *
+     * @return the member that has the 'eventpayload' trait set to true. If none found, return null.
+     */
+    public MemberModel getExplicitEventPayloadMember() {
+        if (members == null) {
+            return null;
+        }
+
+        return members.stream()
+                      .filter(m -> m.isEventPayload())
+                      .findFirst()
+                      .orElse(null);
+    }
+
+    /**
+     * If all members in shape have eventheader trait, then there is no payload
+     */
+    public boolean hasNoEventPayload() {
+        return members == null || members.stream().allMatch(m -> m.isEventHeader());
     }
 
     public boolean isHasStreamingMember() {

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

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.codegen.model.service;
 
+import com.fasterxml.jackson.annotation.JsonProperty;
+
 public class Member {
 
     private String shape;
@@ -39,6 +41,12 @@ public class Member {
 
     private boolean deprecated;
 
+    @JsonProperty(value = "eventpayload")
+    private boolean eventPayload;
+
+    @JsonProperty(value = "eventheader")
+    private boolean eventHeader;
+
     public String getShape() {
         return shape;
     }
@@ -126,4 +134,20 @@ public boolean isDeprecated() {
     public void setDeprecated(boolean deprecated) {
         this.deprecated = deprecated;
     }
+
+    public boolean isEventPayload() {
+        return eventPayload;
+    }
+
+    public void setEventPayload(boolean eventPayload) {
+        this.eventPayload = eventPayload;
+    }
+
+    public boolean isEventHeader() {
+        return eventHeader;
+    }
+
+    public void setEventHeader(boolean eventHeader) {
+        this.eventHeader = eventHeader;
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/PoetExtensions.java

@@ -16,7 +16,10 @@
 package software.amazon.awssdk.codegen.poet;
 
 import com.squareup.javapoet.ClassName;
+import software.amazon.awssdk.codegen.internal.Utils;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 
 /**
  * Extension and convenience methods to Poet that use the intermediate model.
@@ -90,4 +93,55 @@ public ClassName getResponseMetadataClass() {
         return ClassName.get(model.getMetadata().getFullModelPackageName(),
                              model.getSdkResponseBaseClassName() + "Metadata");
     }
+
+    /**
+     * @return The correctly cased name of the API.
+     */
+    public String getApiName(OperationModel operation) {
+        return Utils.capitalize(operation.getOperationName());
+    }
+
+    /**
+     * @return The {@link ClassName} for the response pojo.
+     */
+    public ClassName responsePojoType(OperationModel operation) {
+        return getModelClass(operation.getOutputShape().getShapeName());
+    }
+
+    // TODO Should we move the event stream specific methods to a new class
+    /**
+     * @return {@link ClassName} for generated event stream response handler interface.
+     */
+    public ClassName eventStreamResponseHandlerType(OperationModel operation) {
+        return getModelClass(getApiName(operation) + "ResponseHandler");
+    }
+
+    /**
+     * @return {@link ClassName} for the builder interface for the response handler interface
+     */
+    public ClassName eventStreamResponseHandlerBuilderType(OperationModel operation) {
+        return eventStreamResponseHandlerType(operation).nestedClass("Builder");
+    }
+
+    /**
+     * @return {@link ClassName} for the event stream visitor interface.
+     */
+    public ClassName eventStreamResponseHandlerVisitorType(OperationModel operation) {
+        return eventStreamResponseHandlerType(operation).nestedClass("Visitor");
+    }
+
+    /**
+     * @return {@link ClassName} for the builder interface for the event stream visitor interface.
+     */
+    public ClassName eventStreamResponseHandlerVisitorBuilderType(OperationModel operation) {
+        return eventStreamResponseHandlerVisitorType(operation).nestedClass("Builder");
+    }
+
+    /**
+     * @param shapeModel shape model for the class in model package
+     * @return {@link ClassName} for the shape represented by the given {@link ShapeModel}.
+     */
+    public ClassName getModelClassFromShape(ShapeModel shapeModel) {
+        return getModelClass(shapeModel.getShapeName());
+    }
 }

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

@@ -23,23 +23,35 @@
 import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.getProtocolSpecs;
 
 import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.CodeBlock;
 import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
+import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
+import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.concurrent.Executor;
+import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
+import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.awscore.client.handler.AwsAsyncClientHandler;
+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.codegen.emitters.GeneratorTaskParams;
 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.codegen.poet.PoetUtils;
 import software.amazon.awssdk.codegen.poet.StaticImport;
 import software.amazon.awssdk.codegen.poet.client.specs.ProtocolSpec;
+import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
+import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
 import software.amazon.awssdk.core.client.handler.AsyncClientHandler;
 import software.amazon.awssdk.core.internal.client.config.SdkClientConfiguration;
@@ -158,7 +170,8 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
                           .addCode(getCustomResponseHandler(opModel, returnType)
                                        .orElseGet(() -> protocolSpec.responseHandler(model, opModel)))
                           .addCode(protocolSpec.errorResponseHandler(opModel))
-                          .addCode(protocolSpec.asyncExecutionHandler(opModel))
+                          .addCode(eventToByteBufferPublisher(opModel))
+                          .addCode(protocolSpec.asyncExecutionHandler(model, opModel))
                .endControlFlow()
                .beginControlFlow("catch ($T t)", Throwable.class);
 
@@ -190,4 +203,46 @@ public ClassName className() {
     public Iterable<StaticImport> staticImports() {
         return singletonList(StaticImport.staticMethodImport(FunctionalUtils.class, "runAndLogError"));
     }
+
+    private CodeBlock eventToByteBufferPublisher(OperationModel opModel) {
+        if (!opModel.hasEventStreamInput()) {
+            return CodeBlock.builder().build();
+        }
+
+        ShapeModel eventStreamShape = EventStreamUtils.getEventStreamInRequest(opModel.getInputShape());
+        CodeBlock code = CodeBlock.builder()
+                                  .add(createEventStreamTaggedUnionJsonMarshaller(eventStreamShape))
+                                  .addStatement("$1T eventPublisher = $2T.adapt($3L)",
+                                                ParameterizedTypeName.get(
+                                                    ClassName.get(SdkPublisher.class),
+                                                    eventStreamType(eventStreamShape)),
+                                                SdkPublisher.class,
+                                                EVENT_PUBLISHER_PARAM_NAME)
+                                  .add("$T adapted = eventPublisher.map(event -> eventMarshaller.marshall(event))",
+                                       ParameterizedTypeName.get(Publisher.class, ByteBuffer.class))
+                                  .add(".map($T::encodeEventStreamRequestToByteBuffer);", AwsClientHandlerUtils.class)
+                                  .build();
+
+        return code;
+    }
+
+    private CodeBlock createEventStreamTaggedUnionJsonMarshaller(ShapeModel eventStreamShape) {
+        CodeBlock.Builder builder = CodeBlock.builder().add("$1T eventMarshaller = $1T.builder()",
+                                                            EventStreamTaggedUnionJsonMarshaller.class);
+
+        List<String> eventNames = EventStreamUtils.getEvents(eventStreamShape)
+                                                  .map(shape -> shape.getShapeName())
+                                                  .collect(Collectors.toList());
+
+        eventNames.forEach(event -> builder.add(".putMarshaller($T.class, new $T(protocolFactory))",
+                                                      poetExtensions.getModelClass(event),
+                                                      poetExtensions.getTransformClass(event + "Marshaller")));
+
+        builder.add(".build();");
+        return builder.build();
+    }
+
+    private TypeName eventStreamType(ShapeModel shapeModel) {
+        return poetExtensions.getModelClass(shapeModel.getShapeName());
+    }
 }