Additional changes for updating non-streaming error unmarshaller

S-Saranya1 · S-Saranya1 · commit d065b95d592d · 2025-04-09T04:39:58.000+05:30
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
@@ -189,7 +189,7 @@ public Optional<CodeBlock> errorResponseHandler(OperationModel opModel) {
             ShapeModel exceptionShape = model.getShapes().get(exceptionName);
             String errorCode = exceptionShape.getErrorCode();
 
-            builder.add("case $S:\n", exceptionShape.getC2jName());
+            builder.add("case $S:\n", errorCode);
             builder.add("return $T.of($T.builder()\n", Optional.class, ExceptionMetadata.class)
                    .add(".errorCode($S)\n", errorCode);
             builder.add(populateHttpStatusCode(exceptionShape, model));
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/ProtocolSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/ProtocolSpec.java
@@ -24,18 +24,21 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import software.amazon.awssdk.awscore.client.handler.AwsSyncClientHandler;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeType;
 import software.amazon.awssdk.codegen.model.service.AuthType;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.CredentialType;
 import software.amazon.awssdk.core.client.handler.SyncClientHandler;
 import software.amazon.awssdk.core.runtime.transform.AsyncStreamingRequestMarshaller;
 import software.amazon.awssdk.core.runtime.transform.StreamingRequestMarshaller;
+import software.amazon.awssdk.protocols.core.ExceptionMetadata;
 import software.amazon.awssdk.utils.StringUtils;
 
 public interface ProtocolSpec {
@@ -71,6 +74,23 @@ default List<MethodSpec> additionalMethods() {
         return new ArrayList<>();
     }
 
+    default List<CodeBlock> registerModeledExceptions(IntermediateModel model, PoetExtension poetExtensions) {
+        return model.getShapes().values().stream()
+                    .filter(s -> s.getShapeType() == ShapeType.Exception)
+                    .map(e -> CodeBlock.builder()
+                                       .add(".registerModeledException($T.builder()"
+                                            + ".errorCode($S)"
+                                            + ".exceptionBuilderSupplier($T::builder)"
+                                            + "$L" // populateHttpStatusCode
+                                            + ".build())",
+                                            ExceptionMetadata.class,
+                                            e.getErrorCode(),
+                                            poetExtensions.getModelClass(e.getShapeName()),
+                                            populateHttpStatusCode(e, model))
+                                       .build())
+                    .collect(Collectors.toList());
+    }
+
     default String populateHttpStatusCode(ShapeModel shapeModel, IntermediateModel model) {
         Integer statusCode = shapeModel.getHttpStatusCode();
 
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryProtocolSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryProtocolSpec.java
@@ -114,7 +114,7 @@ public Optional<CodeBlock> errorResponseHandler(OperationModel opModel) {
             }
             ShapeModel exceptionShape = intermediateModel.getShapes().get(exceptionName);
             String errorCode = exceptionShape.getErrorCode();
-            builder.add("case $S:\n", exceptionShape.getC2jName());
+            builder.add("case $S:\n", errorCode);
             builder.add("return $T.of($T.builder()\n", Optional.class, ExceptionMetadata.class)
                    .add(".errorCode($S)\n", errorCode);
             builder.add(populateHttpStatusCode(exceptionShape, intermediateModel));
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/XmlProtocolSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/XmlProtocolSpec.java
@@ -23,8 +23,11 @@
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.WildcardTypeName;
 import java.util.Map;
+import java.util.HashSet;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
 import software.amazon.awssdk.awscore.eventstream.EventStreamAsyncResponseTransformer;
 import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionPojoSupplier;
 import software.amazon.awssdk.awscore.eventstream.RestEventStreamAsyncResponseTransformer;
@@ -43,6 +46,7 @@
 import software.amazon.awssdk.core.SdkPojoBuilder;
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
+import software.amazon.awssdk.protocols.core.ExceptionMetadata;
 import software.amazon.awssdk.protocols.xml.AwsXmlProtocolFactory;
 import software.amazon.awssdk.protocols.xml.XmlOperationMetadata;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
@@ -108,7 +112,46 @@ private CodeBlock streamingResponseHandler(OperationModel opModel) {
 
     @Override
     public Optional<CodeBlock> errorResponseHandler(OperationModel opModel) {
-        return opModel.hasStreamingOutput() ? streamingErrorResponseHandler(opModel) : Optional.empty();
+        if (opModel.hasStreamingOutput()) {
+            return streamingErrorResponseHandler(opModel);
+        }
+
+        CodeBlock.Builder builder = CodeBlock.builder();
+        ParameterizedTypeName metadataMapperType = ParameterizedTypeName.get(
+            ClassName.get(Function.class),
+            ClassName.get(String.class),
+            ParameterizedTypeName.get(Optional.class, ExceptionMetadata.class));
+
+        builder.add("\n$T exceptionMetadataMapper = errorCode -> {\n", metadataMapperType);
+        builder.add("switch (errorCode) {\n");
+        Set<String> processedExceptions = new HashSet<>();
+
+        opModel.getExceptions().forEach(exception -> {
+            String exceptionName = exception.getExceptionName();
+            if (!processedExceptions.add(exceptionName)) {
+                return;
+            }
+            ShapeModel exceptionShape = model.getShapes().get(exceptionName);
+            String errorCode = exceptionShape.getErrorCode();
+
+            builder.add("case $S:\n", errorCode);
+            builder.add("return $T.of($T.builder()\n", Optional.class, ExceptionMetadata.class)
+                   .add(".errorCode($S)\n", errorCode);
+            builder.add(populateHttpStatusCode(exceptionShape, model));
+            builder.add(".exceptionBuilderSupplier($T::builder)\n",
+                        poetExtensions.getModelClassFromShape(exceptionShape))
+                   .add(".build());\n");
+        });
+
+        builder.add("default: return $T.empty();\n", Optional.class);
+        builder.add("}\n");
+        builder.add("};\n\n");
+
+        builder.add("$T<$T> errorResponseHandler = protocolFactory.createErrorResponseHandler(exceptionMetadataMapper);",
+                    HttpResponseHandler.class,
+                    AwsServiceException.class);
+
+        return Optional.of(builder.build());
     }
 
     private Optional<CodeBlock> streamingErrorResponseHandler(OperationModel opModel) {
@@ -288,9 +331,9 @@ private CodeBlock eventStreamResponseHandlers(OperationModel opModel, TypeName p
                              streamResponseOpMd);
 
         // Response handler responsible for errors for the API call itself, as well as errors sent over the event stream
-        builder.addStatement("$T errorResponseHandler = protocolFactory"
-                             + ".createErrorResponseHandler()", ParameterizedTypeName.get(HttpResponseHandler.class,
-                                                                                          AwsServiceException.class));
+        // builder.addStatement("$T errorResponseHandler = protocolFactory"
+        //                      + ".createErrorResponseHandler()", ParameterizedTypeName.get(HttpResponseHandler.class,
+        //                                                                                   AwsServiceException.class));
 
 
         ShapeModel eventStreamShape = EventStreamUtils.getEventStreamInResponse(opModel.getOutputShape());
diff --git a/core/protocols/aws-query-protocol/src/main/java/software/amazon/awssdk/protocols/query/internal/unmarshall/AwsXmlErrorUnmarshaller.java b/core/protocols/aws-query-protocol/src/main/java/software/amazon/awssdk/protocols/query/internal/unmarshall/AwsXmlErrorUnmarshaller.java
@@ -58,10 +58,8 @@ private AwsXmlErrorUnmarshaller(Builder builder) {
     }
 
     private Optional<ExceptionMetadata> defaultMapToExceptionMetadata(String errorCode) {
-        return Optional.ofNullable(exceptions)
-                       .flatMap(exs -> exs.stream()
-                                          .filter(e -> e.errorCode().equals(errorCode))
-                                          .findFirst());
+        return exceptions.stream().filter(e -> e.errorCode().equals(errorCode))
+                                          .findAny();
     }
 
     /**
diff --git a/core/protocols/aws-query-protocol/src/main/java/software/amazon/awssdk/protocols/query/unmarshall/AwsXmlErrorProtocolUnmarshaller.java b/core/protocols/aws-query-protocol/src/main/java/software/amazon/awssdk/protocols/query/unmarshall/AwsXmlErrorProtocolUnmarshaller.java
@@ -91,6 +91,7 @@ private AwsXmlErrorProtocolUnmarshaller(Builder builder) {
         this.exceptionMetadataMapper = builder.exceptionMetadataMapper;
         this.awsXmlErrorUnmarshaller = AwsXmlErrorUnmarshaller.builder()
                                                               .defaultExceptionSupplier(builder.defaultExceptionSupplier)
+                                                              .exceptions(builder.exceptions)
                                                               .exceptionMetadataMapper(this.exceptionMetadataMapper)
                                                               .errorUnmarshaller(builder.errorUnmarshaller)
                                                               .build();
diff --git a/core/protocols/aws-xml-protocol/src/main/java/software/amazon/awssdk/protocols/xml/AwsXmlProtocolFactory.java b/core/protocols/aws-xml-protocol/src/main/java/software/amazon/awssdk/protocols/xml/AwsXmlProtocolFactory.java
@@ -174,7 +174,9 @@ public <T extends AwsResponse> HttpResponseHandler<Response<T>> createCombinedRe
 
         return new CombinedResponseHandler<>
             (createResponseHandler(pojoSupplier, staxOperationMetadata),
-             createErrorResponseHandler(this.exceptionMetadataMapper));
+             this.exceptionMetadataMapper != null
+             ? createErrorResponseHandler(this.exceptionMetadataMapper)
+             : createErrorResponseHandler());
     }
 
     /**

Original file line number	Diff line number	Diff line change
`@@ -114,7 +114,7 @@ public Optional<CodeBlock> errorResponseHandler(OperationModel opModel) {`
`114`	`114`	`}`
`115`	`115`	`ShapeModel exceptionShape = intermediateModel.getShapes().get(exceptionName);`
`116`	`116`	`String errorCode = exceptionShape.getErrorCode();`
`117`		`- builder.add("case $S:\n", exceptionShape.getC2jName());`
	`117`	`+ builder.add("case $S:\n", errorCode);`
`118`	`118`	`builder.add("return $T.of($T.builder()\n", Optional.class, ExceptionMetadata.class)`
`119`	`119`	`.add(".errorCode($S)\n", errorCode);`
`120`	`120`	`builder.add(populateHttpStatusCode(exceptionShape, intermediateModel));`
Original file line number	Diff line number	Diff line change
`@@ -58,10 +58,8 @@ private AwsXmlErrorUnmarshaller(Builder builder) {`
`58`	`58`	`}`
`59`	`59`
`60`	`60`	`private Optional<ExceptionMetadata> defaultMapToExceptionMetadata(String errorCode) {`
`61`		`- return Optional.ofNullable(exceptions)`
`62`		`- .flatMap(exs -> exs.stream()`
`63`		`- .filter(e -> e.errorCode().equals(errorCode))`
`64`		`- .findFirst());`
	`61`	`+ return exceptions.stream().filter(e -> e.errorCode().equals(errorCode))`
	`62`	`+ .findAny();`
`65`	`63`	`}`
`66`	`64`
`67`	`65`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -174,7 +174,9 @@ public <T extends AwsResponse> HttpResponseHandler<Response<T>> createCombinedRe`
`174`	`174`
`175`	`175`	`return new CombinedResponseHandler<>`
`176`	`176`	`(createResponseHandler(pojoSupplier, staxOperationMetadata),`
`177`		`- createErrorResponseHandler(this.exceptionMetadataMapper));`
	`177`	`+ this.exceptionMetadataMapper != null`
	`178`	`+ ? createErrorResponseHandler(this.exceptionMetadataMapper)`
	`179`	`+ : createErrorResponseHandler());`
`178`	`180`	`}`
`179`	`181`
`180`	`182`	`/**`