Event transformation based on type

dagnir · dagnir · commit 65534411b910 · 2020-12-07T11:24:19.000-08:00
Fixes event marshalling and unmarshalling:

 - Fix unmarshalling by mapping event type to the correct builder static
   constructor on the eventstream interface.

 - Fix marshalling by retrieving the value for the :event-type header from the
   event using getSdkEventType()
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java
@@ -41,6 +41,7 @@
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumRequiredTrait;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
+import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
 import software.amazon.awssdk.core.SdkPojoBuilder;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
@@ -419,9 +420,14 @@ private void responseHandlersForEventStreaming(OperationModel opModel, TypeName
                  protocolFactory,
                  JsonOperationMetadata.class,
                  ClassName.get(EventStreamTaggedUnionPojoSupplier.class));
+
+        EventStreamSpecHelper eventStreamSpecHelper = new EventStreamSpecHelper(eventStream, model);
         EventStreamUtils.getEventMembers(eventStream)
-                        .forEach(m -> builder.add(".putSdkPojoSupplier(\"$L\", $T::builder)\n",
-                                                  m.getC2jName(), poetExtensions.getModelClass(m.getShape().getC2jName())));
+                        .forEach(m -> {
+                            String builderMethod = eventStreamSpecHelper.eventBuilderMethodName(m);
+                            builder.add(".putSdkPojoSupplier($S, $T::$N)",
+                                        m.getC2jName(), eventStreamBaseClass, builderMethod);
+                        });
         builder.add(".defaultSdkPojoSupplier(() -> new $T($T.UNKNOWN))\n"
                     + ".build());\n", SdkPojoBuilder.class, eventStreamBaseClass);
     }
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/AwsServiceModel.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/AwsServiceModel.java
@@ -699,9 +699,7 @@ private MethodSpec eventBuilderMethod(MemberModel event) {
             returnType = baseClass.nestedClass("Builder");
         }
 
-        String eventClass = intermediateModel.getNamingStrategy().getShapeClassName(event.getName());
-        String methodName = String.format("%sBuilder", StringUtils.uncapitalize(eventClass));
-
+        String methodName = specHelper.eventBuilderMethodName(event);
         return MethodSpec.methodBuilder(methodName)
                 .addModifiers(PUBLIC, Modifier.STATIC)
                 .returns(returnType)
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/EventStreamSpecHelper.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/EventStreamSpecHelper.java
@@ -29,6 +29,7 @@
 import software.amazon.awssdk.codegen.naming.NamingStrategy;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.eventstream.EventTypeEnumSpec;
+import software.amazon.awssdk.utils.StringUtils;
 import software.amazon.awssdk.utils.internal.CodegenNamingUtils;
 
 public final class EventStreamSpecHelper {
@@ -91,4 +92,8 @@ public String eventTypeEnumValue(MemberModel eventModel) {
         NamingStrategy namingStrategy = intermediateModel.getNamingStrategy();
         return namingStrategy.getEnumValueName(eventModel.getName());
     }
+
+    public String eventBuilderMethodName(MemberModel eventModel) {
+        return String.format("%sBuilder", StringUtils.uncapitalize(eventModel.getName()));
+    }
 }
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
@@ -22,7 +22,6 @@
 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.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpMethod;
 import software.amazon.awssdk.protocols.core.OperationInfo;
@@ -35,11 +34,8 @@ 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(shapeModel);
-        this.intermediateModel = model;
     }
 
     @Override
@@ -50,8 +46,8 @@ public CodeBlock marshalCodeBlock(ClassName requestClassName) {
                      .addStatement("$T<$T> protocolMarshaller = protocolFactory.createProtocolMarshaller(SDK_OPERATION_BINDING)",
                                    ProtocolMarshaller.class, SdkHttpFullRequest.class)
                      .add("return protocolMarshaller.marshall($L).toBuilder()", variableName)
-                     .add(".putHeader(\":message-type\", \"event\")")
-                     .add(".putHeader(\":event-type\", \"$L\")", getMemberNameFromEventStream());
+                     .add(".putHeader($S, $S)", ":message-type", "event")
+                     .add(".putHeader($S, $N.sdkEventType().toString())", ":event-type", variableName);
 
         // Add :content-type header only if payload is present
         if (!shapeModel.hasNoEventPayload()) {
@@ -83,18 +79,6 @@ protected FieldSpec operationInfoField() {
                         .build();
     }
 
-    private String getMemberNameFromEventStream() {
-        ShapeModel eventStream = EventStreamUtils.getBaseEventStreamShape(intermediateModel, shapeModel)
-            .orElseThrow(() -> new IllegalStateException("Could not find associated event stream spec for "
-                                                         + shapeModel.getC2jName()));
-        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) {
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-async-client-class.java
@@ -57,7 +57,6 @@
 import software.amazon.awssdk.services.json.model.APostOperationResponse;
 import software.amazon.awssdk.services.json.model.APostOperationWithOutputRequest;
 import software.amazon.awssdk.services.json.model.APostOperationWithOutputResponse;
-import software.amazon.awssdk.services.json.model.EventOne;
 import software.amazon.awssdk.services.json.model.EventStream;
 import software.amazon.awssdk.services.json.model.EventStreamOperationRequest;
 import software.amazon.awssdk.services.json.model.EventStreamOperationResponse;
@@ -67,7 +66,6 @@
 import software.amazon.awssdk.services.json.model.EventStreamOperationWithOnlyOutputRequest;
 import software.amazon.awssdk.services.json.model.EventStreamOperationWithOnlyOutputResponse;
 import software.amazon.awssdk.services.json.model.EventStreamOperationWithOnlyOutputResponseHandler;
-import software.amazon.awssdk.services.json.model.EventTwo;
 import software.amazon.awssdk.services.json.model.GetWithoutRequiredMembersRequest;
 import software.amazon.awssdk.services.json.model.GetWithoutRequiredMembersResponse;
 import software.amazon.awssdk.services.json.model.InputEvent;
@@ -304,9 +302,9 @@ public CompletableFuture<Void> eventStreamOperation(EventStreamOperationRequest
 
             HttpResponseHandler<? extends EventStream> eventResponseHandler = protocolFactory.createResponseHandler(
                     JsonOperationMetadata.builder().isPayloadJson(true).hasStreamingSuccessResponse(false).build(),
-                    EventStreamTaggedUnionPojoSupplier.builder().putSdkPojoSupplier("EventOne", EventOne::builder)
-                            .putSdkPojoSupplier("EventTheSecond", EventTwo::builder)
-                            .putSdkPojoSupplier("secondEventOne", EventOne::builder)
+                    EventStreamTaggedUnionPojoSupplier.builder().putSdkPojoSupplier("EventOne", EventStream::eventOneBuilder)
+                            .putSdkPojoSupplier("EventTheSecond", EventStream::eventTheSecondBuilder)
+                            .putSdkPojoSupplier("secondEventOne", EventStream::secondEventOneBuilder)
                             .defaultSdkPojoSupplier(() -> new SdkPojoBuilder(EventStream.UNKNOWN)).build());
 
             HttpResponseHandler<AwsServiceException> errorResponseHandler = createErrorResponseHandler(protocolFactory,
@@ -467,9 +465,9 @@ public CompletableFuture<Void> eventStreamOperationWithOnlyOutput(
 
             HttpResponseHandler<? extends EventStream> eventResponseHandler = protocolFactory.createResponseHandler(
                     JsonOperationMetadata.builder().isPayloadJson(true).hasStreamingSuccessResponse(false).build(),
-                    EventStreamTaggedUnionPojoSupplier.builder().putSdkPojoSupplier("EventOne", EventOne::builder)
-                            .putSdkPojoSupplier("EventTheSecond", EventTwo::builder)
-                            .putSdkPojoSupplier("secondEventOne", EventOne::builder)
+                    EventStreamTaggedUnionPojoSupplier.builder().putSdkPojoSupplier("EventOne", EventStream::eventOneBuilder)
+                            .putSdkPojoSupplier("EventTheSecond", EventStream::eventTheSecondBuilder)
+                            .putSdkPojoSupplier("secondEventOne", EventStream::secondEventOneBuilder)
                             .defaultSdkPojoSupplier(() -> new SdkPojoBuilder(EventStream.UNKNOWN)).build());
 
             HttpResponseHandler<AwsServiceException> errorResponseHandler = createErrorResponseHandler(protocolFactory,
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/eventonemarshaller.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/eventonemarshaller.java
@@ -19,7 +19,7 @@
 @SdkInternalApi
 public class EventOneMarshaller implements Marshaller<EventOne> {
     private static final OperationInfo SDK_OPERATION_BINDING = OperationInfo.builder().hasExplicitPayloadMember(false)
-                                                                            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
+            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
 
     private final BaseAwsJsonProtocolFactory protocolFactory;
 
@@ -32,11 +32,13 @@ public SdkHttpFullRequest marshall(EventOne eventOne) {
         Validate.paramNotNull(eventOne, "eventOne");
         try {
             ProtocolMarshaller<SdkHttpFullRequest> protocolMarshaller = protocolFactory
-                .createProtocolMarshaller(SDK_OPERATION_BINDING);
+                    .createProtocolMarshaller(SDK_OPERATION_BINDING);
             return protocolMarshaller.marshall(eventOne).toBuilder().putHeader(":message-type", "event")
-                                     .putHeader(":event-type", "EventOne").putHeader(":content-type", "application/json").build();
+                    .putHeader(":event-type", eventOne.sdkEventType().toString()).putHeader(":content-type", "application/json")
+                    .build();
         } catch (Exception e) {
             throw SdkClientException.builder().message("Unable to marshall request to JSON: " + e.getMessage()).cause(e).build();
         }
     }
 }
+
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/eventthreemarshaller.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/eventthreemarshaller.java
@@ -19,7 +19,8 @@
 @SdkInternalApi
 public class EventThreeMarshaller implements Marshaller<EventThree> {
     private static final OperationInfo SDK_OPERATION_BINDING = OperationInfo.builder().hasExplicitPayloadMember(false)
-                                                                            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
+            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
+
     private final BaseAwsJsonProtocolFactory protocolFactory;
 
     public EventThreeMarshaller(BaseAwsJsonProtocolFactory protocolFactory) {
@@ -33,7 +34,8 @@ public SdkHttpFullRequest marshall(EventThree eventThree) {
             ProtocolMarshaller<SdkHttpFullRequest> protocolMarshaller = protocolFactory
                     .createProtocolMarshaller(SDK_OPERATION_BINDING);
             return protocolMarshaller.marshall(eventThree).toBuilder().putHeader(":message-type", "event")
-                    .putHeader(":event-type", "EventThree").putHeader(":content-type", "application/json").build();
+                    .putHeader(":event-type", eventThree.sdkEventType().toString())
+                    .putHeader(":content-type", "application/json").build();
         } catch (Exception e) {
             throw SdkClientException.builder().message("Unable to marshall request to JSON: " + e.getMessage()).cause(e).build();
         }
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/eventtwomarshaller.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/eventtwomarshaller.java
@@ -19,7 +19,7 @@
 @SdkInternalApi
 public class EventTwoMarshaller implements Marshaller<EventTwo> {
     private static final OperationInfo SDK_OPERATION_BINDING = OperationInfo.builder().hasExplicitPayloadMember(false)
-                                                                            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
+            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
 
     private final BaseAwsJsonProtocolFactory protocolFactory;
 
@@ -32,11 +32,13 @@ public SdkHttpFullRequest marshall(EventTwo eventTwo) {
         Validate.paramNotNull(eventTwo, "eventTwo");
         try {
             ProtocolMarshaller<SdkHttpFullRequest> protocolMarshaller = protocolFactory
-                .createProtocolMarshaller(SDK_OPERATION_BINDING);
+                    .createProtocolMarshaller(SDK_OPERATION_BINDING);
             return protocolMarshaller.marshall(eventTwo).toBuilder().putHeader(":message-type", "event")
-                                     .putHeader(":event-type", "EventTwo").putHeader(":content-type", "application/json").build();
+                    .putHeader(":event-type", eventTwo.sdkEventType().toString()).putHeader(":content-type", "application/json")
+                    .build();
         } catch (Exception e) {
             throw SdkClientException.builder().message("Unable to marshall request to JSON: " + e.getMessage()).cause(e).build();
         }
     }
 }
+
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/inputeventmarshaller.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/inputeventmarshaller.java
@@ -19,7 +19,7 @@
 @SdkInternalApi
 public class InputEventMarshaller implements Marshaller<InputEvent> {
     private static final OperationInfo SDK_OPERATION_BINDING = OperationInfo.builder().hasExplicitPayloadMember(true)
-                                                                            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
+            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
 
     private final BaseAwsJsonProtocolFactory protocolFactory;
 
@@ -32,11 +32,13 @@ public SdkHttpFullRequest marshall(InputEvent inputEvent) {
         Validate.paramNotNull(inputEvent, "inputEvent");
         try {
             ProtocolMarshaller<SdkHttpFullRequest> protocolMarshaller = protocolFactory
-                .createProtocolMarshaller(SDK_OPERATION_BINDING);
+                    .createProtocolMarshaller(SDK_OPERATION_BINDING);
             return protocolMarshaller.marshall(inputEvent).toBuilder().putHeader(":message-type", "event")
-                                     .putHeader(":event-type", "InputEvent").putHeader(":content-type", "application/octet-stream").build();
+                    .putHeader(":event-type", inputEvent.sdkEventType().toString())
+                    .putHeader(":content-type", "application/octet-stream").build();
         } catch (Exception e) {
             throw SdkClientException.builder().message("Unable to marshall request to JSON: " + e.getMessage()).cause(e).build();
         }
     }
 }
+
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/inputeventtwomarshaller.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/transform/inputeventtwomarshaller.java
@@ -19,7 +19,7 @@
 @SdkInternalApi
 public class InputEventTwoMarshaller implements Marshaller<InputEventTwo> {
     private static final OperationInfo SDK_OPERATION_BINDING = OperationInfo.builder().hasExplicitPayloadMember(false)
-                                                                            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
+            .hasPayloadMembers(true).httpMethod(SdkHttpMethod.GET).hasEvent(true).build();
 
     private final BaseAwsJsonProtocolFactory protocolFactory;
 
@@ -32,11 +32,13 @@ public SdkHttpFullRequest marshall(InputEventTwo inputEventTwo) {
         Validate.paramNotNull(inputEventTwo, "inputEventTwo");
         try {
             ProtocolMarshaller<SdkHttpFullRequest> protocolMarshaller = protocolFactory
-                .createProtocolMarshaller(SDK_OPERATION_BINDING);
+                    .createProtocolMarshaller(SDK_OPERATION_BINDING);
             return protocolMarshaller.marshall(inputEventTwo).toBuilder().putHeader(":message-type", "event")
-                                     .putHeader(":event-type", "InputEventTwo").putHeader(":content-type", "application/json").build();
+                    .putHeader(":event-type", inputEventTwo.sdkEventType().toString())
+                    .putHeader(":content-type", "application/json").build();
         } catch (Exception e) {
             throw SdkClientException.builder().message("Unable to marshall request to JSON: " + e.getMessage()).cause(e).build();
         }
     }
 }
+
