Support event streams that are shared between two operations.

Author: adamthom-amzn

.changes/next-release/bugfix-AWSSDKforJavav2-66d1f62.json

@@ -0,0 +1,5 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "type": "bugfix", 
+    "description": "Support event streams that are shared between two operations."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/eventstream/EventStreamUtils.java

@@ -15,7 +15,9 @@
 
 package software.amazon.awssdk.codegen.poet.eventstream;
 
+import java.util.Collection;
 import java.util.Objects;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
@@ -116,16 +118,21 @@ public static Stream<MemberModel> getEventMembers(ShapeModel eventStreamShape) {
     }
 
     /**
-     * Returns the first operation that contains the given event stream shape. The event stream can be in operation
+     * Returns the all operations that contain the given event stream shape. The event stream can be in operation
      * request or response shape.
      */
-    public static OperationModel findOperationWithEventStream(IntermediateModel model, ShapeModel eventStreamShape) {
-        return model.getOperations().values()
+    public static Collection<OperationModel> findOperationsWithEventStream(IntermediateModel model, ShapeModel eventStreamShape) {
+        Collection<OperationModel> operations = model.getOperations().values()
                     .stream()
                     .filter(op -> operationContainsEventStream(op, eventStreamShape))
-                    .findFirst()
-                    .orElseThrow(() -> new IllegalStateException(String.format(
-                        "%s is an event shape but has no corresponding operation in the model", eventStreamShape.getC2jName())));
+                    .collect(Collectors.toList());
+
+        if (operations.isEmpty()) {
+            throw new IllegalStateException(String.format(
+                "%s is an event shape but has no corresponding operation in the model", eventStreamShape.getC2jName()));
+        }
+
+        return operations;
     }
 
     /**

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

@@ -28,6 +28,7 @@
 import com.squareup.javapoet.WildcardTypeName;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
@@ -84,46 +85,40 @@ public AwsServiceModel(IntermediateModel intermediateModel, ShapeModel shapeMode
     @Override
     public TypeSpec poetSpec() {
         if (shapeModel.isEventStream()) {
-            OperationModel opModel = EventStreamUtils.findOperationWithEventStream(intermediateModel,
-                                                                                   shapeModel);
-            String apiName = poetExtensions.getApiName(opModel);
-            ClassName modelClass = poetExtensions.getModelClassFromShape(shapeModel);
+            Collection<OperationModel> opModels = EventStreamUtils.findOperationsWithEventStream(intermediateModel,
+                                                                                                 shapeModel);
 
-            if (EventStreamUtils.doesShapeContainsEventStream(opModel.getOutputShape(), shapeModel)) {
-                ClassName responseHandlerClass = poetExtensions.eventStreamResponseHandlerType(opModel);
-                return PoetUtils.createInterfaceBuilder(modelClass)
-                                .addAnnotation(SdkPublicApi.class)
-                                .addSuperinterface(ClassName.get(SdkPojo.class))
-                                .addJavadoc("Base interface for all event types of the $L API.", apiName)
-                                .addField(FieldSpec.builder(modelClass, "UNKNOWN")
-                                                   .addModifiers(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) {"
-                                                                              + "            \nvisitor.visitDefault(this);\n"
-                                                                              + "        }\n"
-                                                                              + "    };\n",
-                                                                              modelClass, List.class, SdkField.class,
-                                                                              Collections.class, responseHandlerClass
-                                                                         )
-                                                                         .build())
-                                                   .addJavadoc("Special type of {@link $T} for unknown types of events that this "
-                                                               + "version of the SDK does not know about", modelClass)
-                                                   .build())
-                                .addMethod(acceptMethodSpec(modelClass, responseHandlerClass)
-                                               .addModifiers(Modifier.ABSTRACT)
-                                               .build())
-                                .build();
+            Collection<OperationModel> outputOperations = findOutputEventStreamOperations(opModels, shapeModel);
 
-            } else if (EventStreamUtils.doesShapeContainsEventStream(opModel.getInputShape(), shapeModel)) {
+            ClassName modelClass = poetExtensions.getModelClassFromShape(shapeModel);
+
+            if (!outputOperations.isEmpty()) {
+                CodeBlock unknownInitializer = buildUnknownEventStreamInitializer(outputOperations,
+                                                                                  modelClass);
+
+                TypeSpec.Builder builder =
+                    PoetUtils.createInterfaceBuilder(modelClass)
+                              .addAnnotation(SdkPublicApi.class)
+                              .addSuperinterface(ClassName.get(SdkPojo.class))
+                              .addJavadoc("Base interface for all event types in $L.", shapeModel.getShapeName())
+                              .addField(FieldSpec.builder(modelClass, "UNKNOWN")
+                                       .addModifiers(PUBLIC, Modifier.STATIC, Modifier.FINAL)
+                                       .initializer(unknownInitializer)
+                                       .addJavadoc("Special type of {@link $T} for unknown types of events that this "
+                                                   + "version of the SDK does not know about", modelClass)
+                                       .build());
+
+                for (OperationModel opModel : outputOperations) {
+                    ClassName responseHandlerClass = poetExtensions.eventStreamResponseHandlerType(opModel);
+                    builder.addMethod(acceptMethodSpec(modelClass, responseHandlerClass)
+                                          .addModifiers(Modifier.ABSTRACT)
+                                          .build());
+                }
+                return builder.build();
+            } else if (hasInputStreamOperations(opModels, shapeModel)) {
                 return PoetUtils.createInterfaceBuilder(modelClass)
                                 .addAnnotation(SdkPublicApi.class)
-                                .addJavadoc("Base interface for all event types of the $L API.", apiName)
+                                .addJavadoc("Base interface for all event types in $L.", shapeModel.getShapeName())
                                 .build();
             }
 
@@ -160,21 +155,25 @@ public TypeSpec poetSpec() {
             if (this.shapeModel.isEvent()) {
                 ShapeModel eventStream = EventStreamUtils.getBaseEventStreamShape(intermediateModel, shapeModel);
                 ClassName eventStreamClassName = poetExtensions.getModelClassFromShape(eventStream);
-                OperationModel opModel = EventStreamUtils.findOperationWithEventStream(intermediateModel,
+                Collection<OperationModel> opModels = EventStreamUtils.findOperationsWithEventStream(intermediateModel,
                                                                                        eventStream);
 
-                if (EventStreamUtils.doesShapeContainsEventStream(opModel.getOutputShape(), eventStream)) {
+                Collection<OperationModel> outputOperations = findOutputEventStreamOperations(opModels, eventStream);
+
+                if (!outputOperations.isEmpty()) {
                     ClassName modelClass = poetExtensions.getModelClass(shapeModel.getShapeName());
-                    ClassName responseHandlerClass = poetExtensions.eventStreamResponseHandlerType(opModel);
                     specBuilder.addSuperinterface(eventStreamClassName);
-                    specBuilder.addMethod(acceptMethodSpec(modelClass, responseHandlerClass)
-                                              .addAnnotation(Override.class)
-                                              .addCode(CodeBlock.builder()
-                                                                .addStatement("visitor.visit(this)")
-                                                                .build())
-                                              .build());
-
-                } else if (EventStreamUtils.doesShapeContainsEventStream(opModel.getInputShape(), eventStream)) {
+                    for (OperationModel opModel : outputOperations) {
+                        ClassName responseHandlerClass = poetExtensions.eventStreamResponseHandlerType(opModel);
+                        specBuilder.addMethod(acceptMethodSpec(modelClass, responseHandlerClass)
+                                                  .addAnnotation(Override.class)
+                                                  .addCode(CodeBlock.builder()
+                                                                    .addStatement("visitor.visit(this)")
+                                                                    .build())
+                                                  .build());
+                    }
+
+                } else if (hasInputStreamOperations(opModels, eventStream)) {
                     specBuilder.addSuperinterface(eventStreamClassName);
                 } else {
                     throw new IllegalArgumentException(shapeModel.getC2jName() + " event shape is not a member in any "
@@ -190,6 +189,44 @@ public TypeSpec poetSpec() {
         }
     }
 
+    private boolean hasInputStreamOperations(Collection<OperationModel> opModels, ShapeModel eventStream) {
+        return opModels.stream()
+                       .anyMatch(op -> EventStreamUtils.doesShapeContainsEventStream(op.getInputShape(), eventStream));
+    }
+
+    private List<OperationModel> findOutputEventStreamOperations(Collection<OperationModel> opModels,
+                                                                 ShapeModel eventStream) {
+        return opModels
+            .stream()
+            .filter(opModel -> EventStreamUtils.doesShapeContainsEventStream(opModel.getOutputShape(), eventStream))
+            .collect(Collectors.toList());
+    }
+
+    private CodeBlock buildUnknownEventStreamInitializer(Collection<OperationModel> outputOperations,
+                                                         ClassName eventStreamModelClass) {
+        CodeBlock.Builder builder = CodeBlock.builder()
+                                             .add("new $T() {\n"
+                                                  + "        @Override\n"
+                                                  + "        public $T<$T<?>> sdkFields() {\n"
+                                                  + "            return $T.emptyList();\n"
+                                                  + "        }\n",
+                                                  eventStreamModelClass, List.class, SdkField.class,
+                                                  Collections.class
+                                             );
+
+        for (OperationModel opModel : outputOperations) {
+            ClassName responseHandlerClass = poetExtensions.eventStreamResponseHandlerType(opModel);
+            builder.add("        @Override\n"
+                      + "        public void accept($T.Visitor visitor) {"
+                      + "            \nvisitor.visitDefault(this);\n"
+                      + "        }\n", responseHandlerClass);
+        }
+
+        builder.add("    }\n");
+
+        return builder.build();
+    }
+
     private MethodSpec sdkFieldsMethod() {
         ParameterizedTypeName sdkFieldType = ParameterizedTypeName.get(ClassName.get(SdkField.class),
                                                                        WildcardTypeName.subtypeOf(ClassName.get(Object.class)));

codegen/src/test/java/software/amazon/awssdk/codegen/poet/model/SharedStreamAwsModelSpecTest.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 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.model;
+
+import static java.util.stream.Collectors.toList;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static software.amazon.awssdk.codegen.poet.PoetMatchers.generatesTo;
+import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Locale;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import software.amazon.awssdk.codegen.C2jModels;
+import software.amazon.awssdk.codegen.IntermediateModelBuilder;
+import software.amazon.awssdk.codegen.model.config.customization.CustomizationConfig;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.model.service.ServiceModel;
+import software.amazon.awssdk.codegen.utils.ModelLoaderUtils;
+
+@RunWith(Parameterized.class)
+public class SharedStreamAwsModelSpecTest {
+    private static IntermediateModel intermediateModel;
+
+    private final ShapeModel shapeModel;
+
+    @Parameterized.Parameters(name = "{0}")
+    public static Collection<Object[]> data() {
+        invokeSafely(SharedStreamAwsModelSpecTest::setUp);
+        return intermediateModel.getShapes().values().stream().map(shape -> new Object[] { shape }).collect(toList());
+    }
+
+    public SharedStreamAwsModelSpecTest(ShapeModel shapeModel) {
+        this.shapeModel = shapeModel;
+    }
+
+    @Test
+    public void basicGeneration() throws Exception {
+        assertThat(new AwsServiceModel(intermediateModel, shapeModel), generatesTo(referenceFileForShape()));
+    }
+
+    private String referenceFileForShape() {
+        return "sharedstream/" + shapeModel.getShapeName().toLowerCase(Locale.ENGLISH) + ".java";
+    }
+
+    private static void setUp() throws IOException {
+        File serviceModelFile = new File(SharedStreamAwsModelSpecTest.class.getResource("sharedstream/service-2.json").getFile());
+        ServiceModel serviceModel = ModelLoaderUtils.loadModel(ServiceModel.class, serviceModelFile);
+
+        intermediateModel = new IntermediateModelBuilder(
+            C2jModels.builder()
+                     .serviceModel(serviceModel)
+                     .customizationConfig(CustomizationConfig.create())
+                     .build())
+            .build();
+    }
+}

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/eventstream.java

@@ -8,7 +8,7 @@
 import software.amazon.awssdk.core.SdkPojo;
 
 /**
- * Base interface for all event types of the EventStreamOperation API.
+ * Base interface for all event types in EventStream.
  */
 @Generated("software.amazon.awssdk:codegen")
 @SdkPublicApi
@@ -26,7 +26,7 @@ public List<SdkField<?>> sdkFields() {
         public void accept(EventStreamOperationResponseHandler.Visitor visitor) {
             visitor.visitDefault(this);
         }
-    };;
+    };
 
     /**
      * Calls the appropriate visit method depending on the subtype of {@link EventStream}.

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/inputeventstream.java

@@ -4,7 +4,7 @@
 import software.amazon.awssdk.annotations.SdkPublicApi;
 
 /**
- * Base interface for all event types of the EventStreamOperation API.
+ * Base interface for all event types in InputEventStream.
  */
 @Generated("software.amazon.awssdk:codegen")
 @SdkPublicApi

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/inputeventstreamtwo.java

@@ -4,7 +4,7 @@
 import software.amazon.awssdk.annotations.SdkPublicApi;
 
 /**
- * Base interface for all event types of the EventStreamOperationWithOnlyInput API.
+ * Base interface for all event types in InputEventStreamTwo.
  */
 @Generated("software.amazon.awssdk:codegen")
 @SdkPublicApi

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/sharedstream/eventstream.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 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.services.sharedeventstream.model;
+
+import java.util.Collections;
+import java.util.List;
+import software.amazon.awssdk.annotations.Generated;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.core.SdkField;
+import software.amazon.awssdk.core.SdkPojo;
+
+/**
+ * Base interface for all event types in EventStream.
+ */
+@Generated("software.amazon.awssdk:codegen")
+@SdkPublicApi
+public interface EventStream extends SdkPojo {
+    /**
+     * Special type of {@link EventStream} for unknown types of events that this version of the SDK does not know about
+     */
+    EventStream UNKNOWN = new EventStream() {
+        @Override
+        public List<SdkField<?>> sdkFields() {
+            return Collections.emptyList();
+        }
+
+        @Override
+        public void accept(StreamBirthsResponseHandler.Visitor visitor) {
+            visitor.visitDefault(this);
+        }
+
+        @Override
+        public void accept(StreamDeathsResponseHandler.Visitor visitor) {
+            visitor.visitDefault(this);
+        }
+    };
+
+    /**
+     * Calls the appropriate visit method depending on the subtype of {@link EventStream}.
+     *
+     * @param visitor Visitor to invoke.
+     */
+    void accept(StreamBirthsResponseHandler.Visitor visitor);
+
+    /**
+     * Calls the appropriate visit method depending on the subtype of {@link EventStream}.
+     *
+     * @param visitor Visitor to invoke.
+     */
+    void accept(StreamDeathsResponseHandler.Visitor visitor);
+}