Update response and error response parsing (smithy-lang#66)

Currently when response body is a payload member, the SDK either
keeps reponse intact if it has a streaming trait, or parse the body
to JS object.

Actually when response payload is not streaming, it is
NOT always parsable(not a valid JSON or XML string). Specificly,
when payload is String, we shouldn't
parse it with XML or JSON parser; We should also treat Blob shape
differently as we should encode the binaries into string; Only when
shape is Union or Structure can we assume payload is parsable
by JSON or XML parser.

For some protocols, error type flag exists in error response body,
then we need to collect response stream to JS object and parse the
error type; For other protocols, error type flag doesn't exist in
error response body, then we don't need to collect the response
stream in error dispatcher. Instead, we can treat the error like
normal response. So that error shape supports the same traits as
normal responses like streaming, payload etc.

This is done by add a new flag in Protocol generator--
isErrorCodeInBody. When it return true, it means error type flag
exists in error response body, then body is parsed in errors
dispatcher, and each error deser only need to deal with parsed
response body in JS object format. When it returns false, it means
error type can be inferred without touching response body, then
error deser can access the error response intact.

Author: AllanZhengYP

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

@@ -48,7 +48,6 @@
 import software.amazon.smithy.model.shapes.NumberShape;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.model.shapes.ShapeType;
 import software.amazon.smithy.model.shapes.StringShape;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.shapes.TimestampShape;
@@ -72,6 +71,17 @@ 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 boolean isErrorCodeInBody;
+
+    /**
+     * Creates a Http binding protocol generator.
+     *
+     * @param isErrorCodeInBody A boolean that indicates if the error code for the implementing protocol is located in
+     *   the error response body, meaning this generator will parse the body before attempting to load an error code.
+     */
+    public HttpBindingProtocolGenerator(boolean isErrorCodeInBody) {
+        this.isErrorCodeInBody = isErrorCodeInBody;
+    }
 
     @Override
     public ApplicationProtocol getApplicationProtocol() {
@@ -120,6 +130,8 @@ public void generateSharedComponents(GenerationContext context) {
         generateDocumentBodyShapeSerializers(context, serializingDocumentShapes);
         generateDocumentBodyShapeDeserializers(context, deserializingDocumentShapes);
         HttpProtocolGeneratorUtils.generateMetadataDeserializer(context, getApplicationProtocol().getResponseType());
+        HttpProtocolGeneratorUtils.generateCollectBody(context);
+        HttpProtocolGeneratorUtils.generateCollectBodyString(context);
     }
 
     /**
@@ -594,7 +606,7 @@ private void generateOperationDeserializer(
 
         // Write out the error deserialization dispatcher.
         Set<StructureShape> errorShapes = HttpProtocolGeneratorUtils.generateErrorDispatcher(
-                context, operation, responseType, this::writeErrorCodeParser);
+                context, operation, responseType, this::writeErrorCodeParser, isErrorCodeInBody);
         deserializingErrorShapes.addAll(errorShapes);
     }
 
@@ -608,9 +620,10 @@ private void generateErrorDeserializer(GenerationContext context, StructureShape
                 context.getProtocolName()) + "Response";
 
         writer.openBlock("const $L = async (\n"
-                       + "  output: any,\n"
+                       + "  $L: any,\n"
                        + "  context: __SerdeContext\n"
-                       + "): Promise<$T> => {", "};", errorDeserMethodName, errorSymbol, () -> {
+                       + "): Promise<$T> => {", "};",
+                errorDeserMethodName, isErrorCodeInBody ? "parsedOutput" : "output", errorSymbol, () -> {
             writer.openBlock("const contents: $T = {", "};", errorSymbol, () -> {
                 writer.write("__type: $S,", error.getId().getName());
                 writer.write("$$fault: $S,", error.getTrait(ErrorTrait.class).get().getValue());
@@ -621,7 +634,7 @@ private void generateErrorDeserializer(GenerationContext context, StructureShape
             });
 
             readHeaders(context, error, bindingIndex);
-            List<HttpBinding> documentBindings = readResponseBody(context, error, bindingIndex);
+            List<HttpBinding> documentBindings = readErrorResponseBody(context, error, bindingIndex);
             // Track all shapes bound to the document so their deserializers may be generated.
             documentBindings.forEach(binding -> {
                 Shape target = model.expectShape(binding.getMember().getTarget());
@@ -633,6 +646,24 @@ private void generateErrorDeserializer(GenerationContext context, StructureShape
         writer.write("");
     }
 
+    private List<HttpBinding> readErrorResponseBody(
+            GenerationContext context,
+            Shape error,
+            HttpBindingIndex bindingIndex
+    ) {
+        TypeScriptWriter writer = context.getWriter();
+        if (isErrorCodeInBody) {
+            // Body is already parsed in error dispatcher, simply assign body to data.
+            writer.write("const data: any = output.body;");
+            List<HttpBinding> responseBindings = bindingIndex.getResponseBindings(error, Location.DOCUMENT);
+            responseBindings.sort(Comparator.comparing(HttpBinding::getMemberName));
+            return responseBindings;
+        } else {
+            // Deserialize response body just like in normal response.
+            return readResponseBody(context, error, bindingIndex);
+        }
+    }
+
     private void readHeaders(
             GenerationContext context,
             Shape operationOrError,
@@ -691,42 +722,48 @@ private List<HttpBinding> readResponseBody(
         documentBindings.sort(Comparator.comparing(HttpBinding::getMemberName));
         List<HttpBinding> payloadBindings = bindingIndex.getResponseBindings(operationOrError, Location.PAYLOAD);
 
+        // Detect if operation output or error shape contains a streaming member.
+        OperationIndex operationIndex = context.getModel().getKnowledge(OperationIndex.class);
+        StructureShape operationOutputOrError = operationOrError.asStructureShape()
+                .orElseGet(() -> operationIndex.getOutput(operationOrError).orElse(null));
+        boolean hasStreamingComponent = Optional.ofNullable(operationOutputOrError)
+                .map(structure -> structure.getAllMembers().values().stream()
+                        .anyMatch(memberShape -> memberShape.hasTrait(StreamingTrait.class)))
+                .orElse(false);
+
         if (!documentBindings.isEmpty()) {
-            readReponseBodyData(context, operationOrError);
+            // If response has document binding, the body can be parsed to JavaScript object.
+            writer.write("const data: any = await parseBody(output.body, context);");
             deserializeOutputDocument(context, operationOrError, documentBindings);
             return documentBindings;
         }
         if (!payloadBindings.isEmpty()) {
-            readReponseBodyData(context, operationOrError);
             // There can only be one payload binding.
             HttpBinding binding = payloadBindings.get(0);
             Shape target = context.getModel().expectShape(binding.getMember().getTarget());
+            if (hasStreamingComponent) {
+                // If payload is streaming, return raw low-level stream directly.
+                writer.write("const data: any = output.body;");
+            } else if (target instanceof BlobShape) {
+                // If payload is non-streaming blob, only need to collect stream to binary data(Uint8Array).
+                writer.write("const data: any = await collectBody(output.body, context);");
+            } else if (target instanceof StructureShape || target 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 (target 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`",
+                        target.getType()));
+            }
             writer.write("contents.$L = $L;", binding.getMemberName(), getOutputValue(context,
                     Location.PAYLOAD, "data", binding.getMember(), target));
             return payloadBindings;
         }
         return ListUtils.of();
     }
 
-    private void readReponseBodyData(GenerationContext context, Shape operationOrError) {
-        TypeScriptWriter writer = context.getWriter();
-        // Prepare response body for deserializing.
-        OperationIndex operationIndex = context.getModel().getKnowledge(OperationIndex.class);
-        StructureShape operationOutputOrError = operationOrError.asStructureShape()
-                .orElseGet(() -> operationIndex.getOutput(operationOrError).orElse(null));
-        boolean hasStreamingComponent = Optional.ofNullable(operationOutputOrError)
-                .map(structure -> structure.getAllMembers().values().stream()
-                        .anyMatch(memberShape -> memberShape.hasTrait(StreamingTrait.class)))
-                .orElse(false);
-        if (hasStreamingComponent || operationOrError.getType().equals(ShapeType.STRUCTURE)) {
-            // For operations with streaming output or errors with streaming body we keep the body intact.
-            writer.write("const data: any = output.body;");
-        } else {
-            // Otherwise, we collect the response body to structured object with parseBody().
-            writer.write("const data: any = await parseBody(output.body, context);");
-        }
-    }
-
     /**
      * 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,)
@@ -890,10 +927,16 @@ private String getNumberOutputParam(Location bindingType, String dataSource, Sha
      * Writes the code that loads an {@code errorCode} String with the content used
      * to dispatch errors to specific serializers.
      *
-     * <p>Three variables will be in scope:
+     * <p>Two variables will be in scope:
      *   <ul>
-     *       <li>{@code output}: a value of the HttpResponse type.</li>
-     *       <li>{@code data}: the contents of the response body.</li>
+     *       <li>{@code output} or {@code parsedOutput}: a value of the HttpResponse type.
+     *          <ul>
+     *              <li>{@code output} is a raw HttpResponse, available when {@code isErrorCodeInBody} is set to
+     *              {@code false}</li>
+     *              <li>{@code parsedOutput} is a HttpResponse type with body parsed to JavaScript object, available
+     *              when {@code isErrorCodeInBody} is set to {@code true}</li>
+     *          </ul>
+     *       </li>
      *       <li>{@code context}: the SerdeContext.</li>
      *   </ul>
      *

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

@@ -108,6 +108,44 @@ static void generateMetadataDeserializer(GenerationContext context, SymbolRefere
         writer.write("");
     }
 
+    /**
+     * Writes a response body stream collector. This function converts the low-level response body stream to
+     * Uint8Array binary data.
+     *
+     * @param context The generation context.
+     */
+    static void generateCollectBody(GenerationContext context) {
+        TypeScriptWriter writer = context.getWriter();
+
+        writer.addImport("SerdeContext", "__SerdeContext", "@aws-sdk/types");
+        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) || new Uint8Array();");
+        });
+
+        writer.write("");
+    }
+
+    /**
+     * Writes a function converting the low-level response body stream to utf-8 encoded string. It depends on
+     * response body stream collector {@link #generateCollectBody(GenerationContext)}.
+     *
+     * @param context The generation context
+     */
+    static void generateCollectBodyString(GenerationContext context) {
+        TypeScriptWriter writer = context.getWriter();
+
+        writer.addImport("SerdeContext", "__SerdeContext", "@aws-sdk/types");
+        writer.write("// Encode Uint8Array data into string with utf-8.");
+        writer.openBlock("const collectBodyString = (streamBody: any, context: __SerdeContext): Promise<string> => {",
+                "};", () -> {
+            writer.write("return collectBody(streamBody, context).then(body => context.utf8Encoder(body));");
+        });
+
+        writer.write("");
+    }
+
     /**
      * Writes a function used to dispatch to the proper error deserializer
      * for each error that the operation can return. The generated function
@@ -118,13 +156,15 @@ static void generateMetadataDeserializer(GenerationContext context, SymbolRefere
      * @param operation The operation to generate for.
      * @param responseType The response type for the HTTP protocol.
      * @param errorCodeGenerator A consumer
+     * @param shouldParseErrorBody Flag indicating whether need to parse response body in this dispatcher function
      * @return A set of all error structure shapes for the operation that were dispatched to.
      */
     static Set<StructureShape> generateErrorDispatcher(
             GenerationContext context,
             OperationShape operation,
             SymbolReference responseType,
-            Consumer<GenerationContext> errorCodeGenerator
+            Consumer<GenerationContext> errorCodeGenerator,
+            boolean shouldParseErrorBody
     ) {
         TypeScriptWriter writer = context.getWriter();
         SymbolProvider symbolProvider = context.getSymbolProvider();
@@ -138,14 +178,14 @@ static Set<StructureShape> generateErrorDispatcher(
                        + "  output: $T,\n"
                        + "  context: __SerdeContext,\n"
                        + "): Promise<$T> {", "}", errorMethodName, responseType, outputType, () -> {
-            writer.write("const data: any = await parseBody(output.body, context);");
-            // We only consume the parsedOutput if we're dispatching, so only generate if we will.
-            if (!operation.getErrors().isEmpty()) {
-                // Create a holding object since we have already parsed the body, but retain the rest of the output.
-                writer.openBlock("const parsedOutput: any = {", "};", () -> {
-                    writer.write("...output,");
-                    writer.write("body: data,");
-                });
+            // Prepare error response for parsing error code. If error code needs to be parsed from response body
+            // then we collect body and parse it to JS object, otherwise leave the response body as is.
+            if (shouldParseErrorBody) {
+                writer.openBlock("const parsedOutput: any = {", "};",
+                        () -> {
+                            writer.write("...output,");
+                            writer.write("body: await parseBody(output.body, context)");
+                        });
             }
 
             // Error responses must be at least SmithyException and MetadataBearer implementations.
@@ -167,7 +207,8 @@ static Set<StructureShape> generateErrorDispatcher(
                             context.getProtocolName()) + "Response";
                     writer.openBlock("case $S:\ncase $S:", "  break;", errorId.getName(), errorId.toString(), () -> {
                         // Dispatch to the error deserialization function.
-                        writer.write("response = await $L(parsedOutput, context);", errorDeserMethodName);
+                        writer.write("response = await $L($L, context);",
+                                errorDeserMethodName, shouldParseErrorBody ? "parsedOutput" : "output");
                     });
                 });