Support modeled exceptions in event streams

Author: varunnvs92

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

@@ -16,15 +16,13 @@
 package software.amazon.awssdk.codegen;
 
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import software.amazon.awssdk.codegen.internal.Utils;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeType;
-import software.amazon.awssdk.codegen.model.service.ErrorMap;
 import software.amazon.awssdk.codegen.model.service.ErrorTrait;
-import software.amazon.awssdk.codegen.model.service.Operation;
+import software.amazon.awssdk.codegen.model.service.Shape;
 
 /**
  * Constructs the exception shapes for the intermediate model. Analyzes the operations in the
@@ -46,28 +44,21 @@ private Map<String, ShapeModel> constructExceptionShapes() {
         // Java shape models, to be constructed
         final Map<String, ShapeModel> javaShapes = new HashMap<String, ShapeModel>();
 
-        for (Map.Entry<String, Operation> entry : getServiceModel().getOperations().entrySet()) {
+        for (Map.Entry<String, Shape> shape : getServiceModel().getShapes().entrySet()) {
+            if (shape.getValue().isException()) {
+                String errorShapeName = shape.getKey();
+                String javaClassName = getNamingStrategy().getExceptionName(errorShapeName);
 
-            Operation operation = entry.getValue();
-            List<ErrorMap> operationErrors = operation.getErrors();
+                ShapeModel exceptionShapeModel = generateShapeModel(javaClassName,
+                                                                    errorShapeName);
 
-            if (operationErrors != null) {
-                for (ErrorMap error : operationErrors) {
-
-                    String errorShapeName = error.getShape();
-                    String javaClassName = getNamingStrategy().getExceptionName(errorShapeName);
-
-                    ShapeModel exceptionShapeModel = generateShapeModel(javaClassName,
-                                                                        errorShapeName);
-
-                    exceptionShapeModel.setType(ShapeType.Exception.getValue());
-                    exceptionShapeModel.setErrorCode(getErrorCode(errorShapeName));
-                    if (exceptionShapeModel.getDocumentation() == null) {
-                        exceptionShapeModel.setDocumentation(error.getDocumentation());
-                    }
-
-                    javaShapes.put(javaClassName, exceptionShapeModel);
+                exceptionShapeModel.setType(ShapeType.Exception.getValue());
+                exceptionShapeModel.setErrorCode(getErrorCode(errorShapeName));
+                if (exceptionShapeModel.getDocumentation() == null) {
+                    exceptionShapeModel.setDocumentation(shape.getValue().getDocumentation());
                 }
+
+                javaShapes.put(javaClassName, exceptionShapeModel);
             }
         }
 

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

@@ -28,9 +28,7 @@
 import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
-
 import software.amazon.awssdk.awscore.eventstream.EventStreamAsyncResponseTransformer;
-import software.amazon.awssdk.awscore.eventstream.EventStreamExceptionJsonUnmarshaller;
 import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionJsonUnmarshaller;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.awscore.internal.protocol.json.AwsJsonProtocol;
@@ -159,19 +157,6 @@ public CodeBlock responseHandler(IntermediateModel model, OperationModel opModel
                             });
             builder.add(".defaultUnmarshaller((in) -> $T.UNKNOWN)\n"
                         + ".build());\n", eventStreamUtils.eventStreamBaseClass());
-
-            builder.add("\n\n$T<$T> exceptionHandler = $L.createResponseHandler(\n" +
-                        "    new JsonOperationMetadata().withPayloadJson(true).withHasStreamingSuccessResponse(false),\n" +
-                        "    $T.builder()\n" +
-                        "        .defaultUnmarshaller(x -> $T.populateDefaultException($T::builder, x))\n" +
-                        "        .build());",
-                        HttpResponseHandler.class,
-                        WildcardTypeName.subtypeOf(Throwable.class),
-                        protocolFactory,
-                        EventStreamExceptionJsonUnmarshaller.class,
-                        EventStreamExceptionJsonUnmarshaller.class,
-                        baseExceptionClassName(model));
-
         }
         return builder.build();
     }
@@ -229,7 +214,7 @@ public CodeBlock asyncExecutionHandler(OperationModel opModel) {
         CodeBlock.Builder builder = CodeBlock.builder();
         if (opModel.hasEventStreamOutput()) {
             builder.add("$T<$T, $T> asyncResponseTransformer = new $T<>(\n"
-                        + "asyncResponseHandler, responseHandler, eventResponseHandler, exceptionHandler);\n",
+                        + "asyncResponseHandler, responseHandler, eventResponseHandler, errorResponseHandler, serviceName());\n",
                         ClassName.get(AsyncResponseTransformer.class),
                         ClassName.get(SdkResponse.class),
                         ClassName.get(Void.class),

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-async-client-class.java

@@ -11,7 +11,6 @@
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.client.handler.AwsAsyncClientHandler;
 import software.amazon.awssdk.awscore.eventstream.EventStreamAsyncResponseTransformer;
-import software.amazon.awssdk.awscore.eventstream.EventStreamExceptionJsonUnmarshaller;
 import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionJsonUnmarshaller;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.awscore.internal.protocol.json.AwsJsonProtocol;
@@ -42,7 +41,6 @@
 import software.amazon.awssdk.services.json.model.GetWithoutRequiredMembersRequest;
 import software.amazon.awssdk.services.json.model.GetWithoutRequiredMembersResponse;
 import software.amazon.awssdk.services.json.model.InvalidInputException;
-import software.amazon.awssdk.services.json.model.JsonException;
 import software.amazon.awssdk.services.json.model.JsonRequest;
 import software.amazon.awssdk.services.json.model.PaginatedOperationWithResultKeyRequest;
 import software.amazon.awssdk.services.json.model.PaginatedOperationWithResultKeyResponse;
@@ -222,18 +220,9 @@ public CompletableFuture<Void> eventStreamOperation(EventStreamOperationRequest
                             .addUnmarshaller("EventTwo", EventTwoUnmarshaller.getInstance())
                             .defaultUnmarshaller((in) -> EventStream.UNKNOWN).build());
 
-            HttpResponseHandler<? extends Throwable> exceptionHandler = jsonProtocolFactory
-                    .createResponseHandler(
-                            new JsonOperationMetadata().withPayloadJson(true).withHasStreamingSuccessResponse(false),
-                            EventStreamExceptionJsonUnmarshaller
-                                    .builder()
-                                    .defaultUnmarshaller(
-                                            x -> EventStreamExceptionJsonUnmarshaller.populateDefaultException(
-                                                    JsonException::builder, x)).build());
-
             HttpResponseHandler<AwsServiceException> errorResponseHandler = createErrorResponseHandler(jsonProtocolFactory);
             AsyncResponseTransformer<SdkResponse, Void> asyncResponseTransformer = new EventStreamAsyncResponseTransformer<>(
-                    asyncResponseHandler, responseHandler, eventResponseHandler, exceptionHandler);
+                    asyncResponseHandler, responseHandler, eventResponseHandler, errorResponseHandler, serviceName());
 
             return clientHandler.execute(
                     new ClientExecutionParams<EventStreamOperationRequest, SdkResponse>()

core/aws-core/src/main/java/software/amazon/awssdk/awscore/eventstream/EventStreamAsyncResponseTransformer.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.awscore.eventstream;
 
 import static java.util.Collections.singletonList;
+import static software.amazon.awssdk.core.http.HttpResponseHandler.X_AMZN_REQUEST_ID_HEADER;
 import static software.amazon.awssdk.utils.FunctionalUtils.runAndLogError;
 
 import java.io.ByteArrayInputStream;
@@ -37,6 +38,7 @@
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
 import software.amazon.awssdk.core.internal.util.ThrowableUtils;
 import software.amazon.awssdk.http.AbortableInputStream;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
@@ -106,6 +108,18 @@ public class EventStreamAsyncResponseTransformer<ResponseT, EventT>
      */
     private final AtomicReference<Throwable> error = new AtomicReference<>();
 
+    /**
+     * The name of the aws service
+     */
+    private final String serviceName;
+
+    /**
+     * Request Id for the streaming request. The value is populated when the initial response is received from the service.
+     * As request id is not sent in event messages (including exceptions), this can be returned by the SDK along with
+     * received exception details.
+     */
+    private String requestId = null;
+
     /**
      * @param eventStreamResponseTransformer Response transformer provided by customer.
      * @param initialResponseUnmarshaller Unmarshaller for the initial-response event stream message.
@@ -115,19 +129,26 @@ public EventStreamAsyncResponseTransformer(
         EventStreamResponseHandler<ResponseT, EventT> eventStreamResponseTransformer,
         HttpResponseHandler<? extends ResponseT> initialResponseUnmarshaller,
         HttpResponseHandler<? extends EventT> eventUnmarshaller,
-        HttpResponseHandler<? extends Throwable> exceptionUnmarshaller) {
+        HttpResponseHandler<? extends Throwable> exceptionUnmarshaller,
+        String serviceName) {
 
         this.eventStreamResponseTransformer = eventStreamResponseTransformer;
         this.initialResponseUnmarshaller = initialResponseUnmarshaller;
         this.eventUnmarshaller = eventUnmarshaller;
         this.exceptionUnmarshaller = exceptionUnmarshaller;
+        this.serviceName = serviceName;
     }
 
     @Override
     public void responseReceived(SdkResponse response) {
         // We use a void unmarshaller and unmarshall the actual response in the message
         // decoder when we receive the initial-response frame. TODO not clear
         // how we would handle REST protocol which would unmarshall the response from the HTTP headers
+        if (response != null && response.sdkHttpResponse() != null) {
+            this.requestId = response.sdkHttpResponse()
+                                     .firstMatchingHeader(X_AMZN_REQUEST_ID_HEADER)
+                                     .orElse(null);
+        }
     }
 
     @Override
@@ -205,8 +226,18 @@ private MessageDecoder createDecoder() {
                                                                             EMPTY_EXECUTION_ATTRIBUTES));
                     }
                 } else if (isError(m) || isException(m)) {
-                    Throwable exception = exceptionUnmarshaller.handle(adaptMessageToResponse(m, true),
-                                                                       EMPTY_EXECUTION_ATTRIBUTES);
+                    SdkHttpFullResponse errorResponse = adaptMessageToResponse(m, true);
+                    if (!errorResponse.firstMatchingHeader(X_AMZN_REQUEST_ID_HEADER).isPresent() &&
+                        requestId != null) {
+                        errorResponse = errorResponse.toBuilder()
+                                                     .putHeader(X_AMZN_REQUEST_ID_HEADER, requestId)
+                                                     .build();
+                    }
+
+                    Throwable exception = exceptionUnmarshaller.handle(errorResponse,
+                                                                       new ExecutionAttributes()
+                                                                           .putAttribute(SdkExecutionAttribute.SERVICE_NAME,
+                                                                                         serviceName));
                     runAndLogError(log, "Error thrown from exceptionOccurred, ignoring.", () -> exceptionOccurred(exception));
                 }
             } catch (Exception e) {