complete event deserializer and event union deserializer

Author: AllanZhengYP

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/HttpBindingProtocolGenerator.java

@@ -54,6 +54,8 @@
 import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.EndpointTrait;
 import software.amazon.smithy.model.traits.ErrorTrait;
+import software.amazon.smithy.model.traits.EventHeaderTrait;
+import software.amazon.smithy.model.traits.EventPayloadTrait;
 import software.amazon.smithy.model.traits.EventStreamTrait;
 import software.amazon.smithy.model.traits.HttpTrait;
 import software.amazon.smithy.model.traits.StreamingTrait;
@@ -74,8 +76,8 @@ public abstract class HttpBindingProtocolGenerator implements ProtocolGenerator
     private final Set<Shape> serializingDocumentShapes = new TreeSet<>();
     private final Set<Shape> deserializingDocumentShapes = new TreeSet<>();
     private final Set<StructureShape> deserializingErrorShapes = new TreeSet<>();
-    private final Set<Shape> serializeEventShapes = new TreeSet<>();
-    private final Set<Shape> deserializingEventShapes = new TreeSet<>();
+    private final Set<StructureShape> serializeEventShapes = new TreeSet<>();
+    private final Set<StructureShape> deserializingEventShapes = new TreeSet<>();
     private final Set<UnionShape> serializeEventUnions = new TreeSet<>();
     private final Set<UnionShape> deserializeEventUnions = new TreeSet<>();
     private final boolean isErrorCodeInBody;
@@ -134,12 +136,14 @@ public ApplicationProtocol getApplicationProtocol() {
     @Override
     public void generateSharedComponents(GenerationContext context) {
         deserializeEventUnions.forEach(eventUnion -> {
-            Set<StructureShape> events = HttpProtocolGeneratorUtils.generateSerializingEventUnion(context, eventUnion);
+            Set<StructureShape> events =
+                    HttpProtocolGeneratorUtils.generateDeserializingEventUnion(context, eventUnion);
             deserializingEventShapes.addAll(events);
         });
         serializeEventUnions.forEach(eventUnion -> {
 
         });
+        deserializingEventShapes.forEach(event -> generateEventDeserializer(context, event));
 
         deserializingErrorShapes.forEach(error -> generateErrorDeserializer(context, error));
         generateDocumentBodyShapeSerializers(context, serializingDocumentShapes);
@@ -618,7 +622,7 @@ protected void serializeInputPayload(
     private void collectSerializingEventShapes(GenerationContext context, Shape target) {
         if (target instanceof StructureShape) {
             // Single-event event stream
-            this.serializeEventShapes.add(target);
+            this.serializeEventShapes.add(target.asStructureShape().get());
         } else if (target instanceof UnionShape) {
             // Multi-event event stream. Save the dispatcher
             this.serializeEventUnions.add(target.asUnionShape().get());
@@ -844,7 +848,8 @@ private List<HttpBinding> readResponsePayload(
             eventStreamDeserializerWrapper(context, binding.getMember(), target);
             writer.write("").dedent();
             writer.write("contents.$L = data;", binding.getMemberName());
-            return payloadBindings;
+            //Not to generate non-eventstream payload shape again
+            return ListUtils.of();
         } else if (hasStreamingComponent) {
             // If payload is streaming, return raw low-level stream directly.
             writer.write("const data: any = output.body;");
@@ -890,7 +895,7 @@ private void eventStreamDeserializerWrapper(GenerationContext context, MemberSha
                 });
                 if (target instanceof  StructureShape) {
                     // Single-event event stream
-                    this.deserializingEventShapes.add(target);
+                    this.deserializingEventShapes.add(target.asStructureShape().get());
                     writer.write("return $L;",
                             getOutputValue(context, Location.PAYLOAD, "eventMessage", member, target));
                 } else if (target instanceof UnionShape) {
@@ -899,12 +904,90 @@ private void eventStreamDeserializerWrapper(GenerationContext context, MemberSha
                     Symbol symbol = context.getSymbolProvider().toSymbol(target);
                     String deserFunction = ProtocolGenerator.getDeserFunctionName(symbol, context.getProtocolName())
                             + "_event(parsedEvent, context)";
-                    writer.write("return $L;", deserFunction);
+                    writer.write("return await $L;", deserFunction);
                 }
             });
         });
     }
 
+    private void generateEventDeserializer(
+            GenerationContext context,
+            StructureShape event
+    ) {
+        TypeScriptWriter writer = context.getWriter();
+        SymbolProvider symbolProvider = context.getSymbolProvider();
+        Symbol symbol = symbolProvider.toSymbol(event);
+        String methodName = ProtocolGenerator.getDeserFunctionName(symbol, context.getProtocolName()) + "_event";
+        Model model = context.getModel();
+        // Handle the general response.
+        writer.openBlock("const $L = async (\n"
+                        + "  output: any,\n"
+                        + "  context: __SerdeContext\n"
+                        + "): Promise<$T> => {", "}",
+            methodName,
+            symbol, () -> {
+                writer.openBlock("let contents: $L = {", "} as any;", symbol.getName(), () -> {
+                    if (!event.getAllMembers().values().isEmpty()) {
+                        writer.write("__type: $S,", event.getId().getName());
+                    }
+                });
+                List<MemberShape> payloadMembers = event.getAllMembers().values().stream()
+                        .filter(member -> member.hasTrait(EventPayloadTrait.class)).collect(Collectors.toList());
+                List<MemberShape> headerMembers = event.getAllMembers().values().stream()
+                        .filter(member -> member.hasTrait(EventHeaderTrait.class)).collect(Collectors.toList());
+                List<MemberShape> documentMembers = event.getAllMembers().values().stream()
+                        .filter(member -> !member.hasTrait(EventHeaderTrait.class)
+                                && !member.hasTrait(EventPayloadTrait.class))
+                        .collect(Collectors.toList());
+                if (!payloadMembers.isEmpty()) {
+                    //if EventPayload trait exists, all other traits are EventHeader
+                    headerMembers.addAll(documentMembers);
+                    documentMembers = ListUtils.of();
+                }
+                for (MemberShape headerMember : headerMembers) {
+                    String memberName = headerMember.getMemberName();
+                    writer.openBlock("if (output.headers[$S] !== undefined) {", "}", memberName, () -> {
+                        Shape target = model.expectShape(headerMember.getTarget());
+                        String headerValue = getOutputValue(context, Location.HEADER,
+                                "output.headers['" + memberName + "']", headerMember, target);
+                        writer.write("contents.$L = $L;", memberName, headerValue);
+                    });
+                }
+                if (!documentMembers.isEmpty()) {
+                    // If response has document binding, the body can be parsed to JavaScript object.
+                    writer.write("const data: any = await parseBody(output.body, context);");
+                    // Deser the event document with the original shape deser function
+                    writer.openBlock("contents = {", "} as any", () -> {
+                       writer.write("...contents,");
+                       writer.write("...$L(data, context)",
+                               ProtocolGenerator.getDeserFunctionName(symbol, context.getProtocolName()));
+                    });
+                    //need original structure shape deserializer to deserialize event body.
+                    deserializingDocumentShapes.add(event);
+                }
+                if (!payloadMembers.isEmpty()) {
+                    //There's only one event payload member
+                    MemberShape payloadMember = payloadMembers.get(0);
+                    Shape payloadTarget = model.expectShape(payloadMember.getTarget());
+                    if (payloadTarget instanceof BlobShape) {
+                        // If event payload is a blob, only need to collect stream to binary data(Uint8Array).
+                        writer.write("const data: any = output.body");
+                    } else if (payloadTarget instanceof StructureShape || payloadTarget instanceof UnionShape) {
+                        // If body is Structure or Union, they we need to parse the string into JavaScript object.
+                        writer.write("const data: any = await parseBody(output.body, context);");
+                    } else if (payloadTarget instanceof StringShape) {
+                        // If payload is string, we need to collect body and encode binary to string.
+                        writer.write("const data: any = await collectBodyString(output.body, context);");
+                    } else {
+                        throw new CodegenException(String.format("Unexpected shape type bound to payload: `%s`",
+                                payloadTarget.getType()));
+                    }
+                    writer.write("contents.$L = data", payloadMember.getMemberName());
+                }
+                writer.write("return contents;");
+            });
+    }
+
     /**
      * Given context and a source of data, generate an output value provider for the
      * shape. This may use native types (like generating a Date for timestamps,)

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/HttpProtocolGeneratorUtils.java

@@ -20,6 +20,7 @@
 import java.util.TreeSet;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
@@ -128,7 +129,10 @@ static void generateCollectBody(GenerationContext context) {
         writer.write("// Collect low-level response body stream to Uint8Array.");
         writer.openBlock("const collectBody = (streamBody: any, context: __SerdeContext): Promise<Uint8Array> => {",
                 "};", () -> {
-            writer.write("return context.streamCollector(streamBody) || Promise.resolve(new Uint8Array());");
+            writer.openBlock("if (streamBody instanceof Uint8Array) {", "}", () -> {
+                writer.write("return Promise.resolve(streamBody);");
+            });
+            writer.write("return context.streamCollector(streamBody) || new Uint8Array();");
         });
 
         writer.write("");
@@ -289,7 +293,7 @@ static void writeHostPrefix(GenerationContext context, OperationShape operation)
         });
     }
 
-    static Set<StructureShape> generateSerializingEventUnion(
+    static Set<StructureShape> generateDeserializingEventUnion(
             GenerationContext context,
             UnionShape events
     ) {
@@ -300,10 +304,10 @@ static Set<StructureShape> generateSerializingEventUnion(
         String methodName = ProtocolGenerator.getDeserFunctionName(symbol, protocolName) + "_event";
         Model model = context.getModel();
         Set<StructureShape> targets = new TreeSet<>();
-        writer.openBlock("const $L = (\n"
+        writer.openBlock("const $L = async (\n"
                 + "  output: any,\n"
                 + "  context: __SerdeContext\n"
-                + "): $T => {", "}", methodName, symbol, () -> {
+                + "): Promise<$T> => {", "}", methodName, symbol, () -> {
             events.getAllMembers().forEach((name, member) -> {
                 Shape target = model.expectShape(member.getTarget());
                 targets.add(target.asStructureShape().orElseThrow(
@@ -318,10 +322,11 @@ static Set<StructureShape> generateSerializingEventUnion(
                         String eventDeserMethodName =
                                 ProtocolGenerator.getDeserFunctionName(eventSymbol, protocolName) + "_event";
                         String statement = eventDeserMethodName + "(output['" + name + "'], context)";
-                        writer.write("$L: $L", name, statement);
+                        writer.write("$L: await $L", name, statement);
                     });
                 });
             });
+            writer.write("return {$$unknown: output}");
         });
         return targets;
     }