Support awsQueryError trait

Author: skmcgrail

codegen/protocol-test-codegen/smithy-build.json

@@ -13,8 +13,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.ec2#AwsEc2",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/ec2query",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/ec2query"
         }
       }
     },
@@ -30,8 +29,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.json#JsonProtocol",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/jsonrpc",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/jsonrpc"
         }
       }
     },
@@ -47,8 +45,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.json10#JsonRpc10",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/jsonrpc10",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/jsonrpc10"
         }
       }
     },
@@ -64,8 +61,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.query#AwsQuery",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/query",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/query"
         }
       }
     },
@@ -81,8 +77,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.restjson#RestJson",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/awsrestjson",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/awsrestjson"
         }
       }
     },
@@ -98,8 +93,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.restxml#RestXml",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/restxml",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/restxml"
         }
       }
     },
@@ -115,8 +109,7 @@
       "plugins": {
         "go-codegen": {
           "service": "aws.protocoltests.restxml.xmlns#RestXmlWithNamespace",
-          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/restxmlwithnamespace",
-          "moduleVersion": "1.0"
+          "module": "github.com/aws/aws-sdk-go-v2/internal/protocoltest/restxmlwithnamespace"
         }
       }
     }

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsQuery.java

@@ -4,7 +4,11 @@
 import static software.amazon.smithy.aws.go.codegen.AwsProtocolUtils.handleDecodeError;
 import static software.amazon.smithy.aws.go.codegen.XmlProtocolUtils.initializeXmlDecoder;
 
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.TreeMap;
+import software.amazon.smithy.aws.traits.protocols.AwsQueryErrorTrait;
 import software.amazon.smithy.aws.traits.protocols.AwsQueryTrait;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.go.codegen.GoWriter;
@@ -25,7 +29,6 @@
  * Handles generating the aws query protocol for services.
  *
  * @inheritDoc
- *
  * @see HttpRpcProtocolGenerator
  */
 class AwsQuery extends HttpRpcProtocolGenerator {
@@ -81,8 +84,8 @@ protected void serializeInputDocument(GenerationContext context, OperationShape
 
         writer.openBlock("if request, err = request.SetStream(bytes.NewReader(bodyWriter.Bytes())); err != nil {",
                 "}", () -> {
-            writer.write("return out, metadata, &smithy.SerializationError{Err: err}");
-        });
+                    writer.write("return out, metadata, &smithy.SerializationError{Err: err}");
+                });
     }
 
     @Override
@@ -97,7 +100,7 @@ protected void deserializeOutputDocument(GenerationContext context, OperationSha
         StructureShape output = ProtocolUtils.expectOutput(context.getModel(), operation);
         String functionName = ProtocolGenerator.getDocumentDeserializerFunctionName(
                 output, context.getService(), getProtocolName());
-        initializeXmlDecoder(writer, "response.Body", "out, metadata, ","nil");
+        initializeXmlDecoder(writer, "response.Body", "out, metadata, ", "nil");
         unwrapOutputDocument(context, operation);
         writer.write("err = $L(&output, decoder)", functionName);
         handleDecodeError(writer, "out, metadata, ");
@@ -152,4 +155,21 @@ protected void writeErrorMessageCodeDeserializer(GenerationContext context) {
     public void generateProtocolTests(GenerationContext context) {
         AwsProtocolUtils.generateHttpProtocolTests(context);
     }
+
+    @Override
+    public Map<String, ShapeId> getOperationErrors(GenerationContext context, OperationShape operation) {
+        Map<String, ShapeId> errors = new TreeMap<>();
+
+        operation.getErrors().forEach(shapeId -> {
+            Shape errorShape = context.getModel().expectShape(shapeId);
+            String errorName = shapeId.getName(context.getService());
+
+            errors.put(errorName, shapeId);
+
+            errorShape.getTrait(AwsQueryErrorTrait.class).ifPresent(awsQueryErrorTrait ->
+                    errors.put(awsQueryErrorTrait.getCode(), shapeId));
+        });
+
+        return errors;
+    }
 }