feat(codegen): send x-amzn-query-mode header

aws · Nov 11, 2024 · b3c4421 · b3c4421
1 parent 17b37b7
commit b3c4421
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 0 deletions.
diff --git a/clients/client-sqs/src/protocols/Aws_json1_0.ts b/clients/client-sqs/src/protocols/Aws_json1_0.ts
@@ -2087,6 +2087,7 @@ function sharedHeaders(operation: string): __HeaderBag {
   return {
     "content-type": "application/x-amz-json-1.0",
     "x-amz-target": `AmazonSQS.${operation}`,
+    "x-amzn-query-error": "true",
   };
 }
 

diff --git a/...pt-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsJsonRpc1_1.java b/...pt-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsJsonRpc1_1.java
@@ -16,7 +16,12 @@
 package software.amazon.smithy.aws.typescript.codegen;
 
 import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait;
+import software.amazon.smithy.aws.traits.protocols.AwsQueryCompatibleTrait;
+import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
+import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
+import software.amazon.smithy.typescript.codegen.integration.HttpRpcProtocolGenerator;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
 /**
@@ -43,4 +48,22 @@ public ShapeId getProtocol() {
     public String getName() {
         return "aws.json-1.1";
     }
+
+    /**
+     * This override exists because the "x-amzn-query-error" header is only
+     * sent in AwsJsonRpc1_0,
+     */
+    @Override
+    protected void writeSharedRequestHeaders(GenerationContext context) {
+        ServiceShape serviceShape = context.getService();
+        TypeScriptWriter writer = context.getWriter();
+        writer.addImport("HeaderBag", "__HeaderBag", TypeScriptDependency.SMITHY_TYPES);
+        String targetHeader = serviceShape.getId().getName(serviceShape) + ".${operation}";
+        writer.openBlock("function sharedHeaders(operation: string): __HeaderBag { return {", "}};",
+            () -> {
+                writer.write("'content-type': $S,", getDocumentContentType());
+                writer.write("'x-amz-target': `$L`,", targetHeader);
+            }
+        );
+    }
 }
diff --git a/...degen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsSmithyRpcV2Cbor.java b/...degen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsSmithyRpcV2Cbor.java
@@ -8,6 +8,7 @@
 import software.amazon.smithy.aws.traits.protocols.AwsQueryCompatibleTrait;
 import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
+import software.amazon.smithy.typescript.codegen.integration.ProtocolGenerator;
 import software.amazon.smithy.typescript.codegen.protocols.cbor.SmithyRpcV2Cbor;
 
 /**
@@ -34,6 +35,24 @@ public void generateSharedComponents(GenerationContext context) {
         }
     }
 
+    @Override
+    protected void writeSharedRequestHeaders(ProtocolGenerator.GenerationContext context) {
+        TypeScriptWriter writer = context.getWriter();
+        writer.addImport("HeaderBag", "__HeaderBag", TypeScriptDependency.SMITHY_TYPES);
+        writer.openBlock("const SHARED_HEADERS: __HeaderBag = {", "};", () -> {
+            writer.write("'content-type': $S,", getDocumentContentType());
+            writer.write("""
+                "smithy-protocol": "rpc-v2-cbor",
+                "accept": "application/cbor",
+                """);
+            if (context.getService().hasTrait(AwsQueryCompatibleTrait.class)) {
+                writer.write("""
+                    "x-amzn-query-error": "true",
+                    """);
+            }
+        });
+    }
+
     @Override
     protected void writeErrorCodeParser(GenerationContext generationContext) {
         super.writeErrorCodeParser(generationContext);

diff --git a/...src/main/java/software/amazon/smithy/aws/typescript/codegen/JsonRpcProtocolGenerator.java b/...src/main/java/software/amazon/smithy/aws/typescript/codegen/JsonRpcProtocolGenerator.java
@@ -144,6 +144,11 @@ protected void writeSharedRequestHeaders(GenerationContext context) {
                     // AWS JSON RPC protocols use a combination of the service and operation shape names,
                     // separated by a '.' character, for the target header.
                     writer.write("'x-amz-target': `$L`,", targetHeader);
+                    if (serviceShape.hasTrait(AwsQueryCompatibleTrait.class)) {
+                        writer.write("""
+                            "x-amzn-query-error": "true",
+                        """);
+                    }
                 }
         );
     }