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

aws · Nov 11, 2024 · 925e731 · 925e731
1 parent 4e2f525
commit 925e731
Show file tree

Hide file tree

Showing 5 changed files with 70 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-mode": "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,10 @@
 package software.amazon.smithy.aws.typescript.codegen;
 
 import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait;
+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.utils.SmithyInternalApi;
 
 /**
@@ -43,4 +46,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-mode": "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-mode": "true",
+                        """);
+                    }
                 }
         );
     }

diff --git a/packages/middleware-sdk-sqs/src/middleware-sdk-sqs.integ.spec.ts b/packages/middleware-sdk-sqs/src/middleware-sdk-sqs.integ.spec.ts
@@ -9,6 +9,8 @@ import { requireRequestsFrom } from "../../../private/aws-util-test/src";
 
 const sqsModel: any = require("../../../codegen/sdk-codegen/aws-models/sqs.json");
 const useAwsQuery = !!sqsModel.shapes["com.amazonaws.sqs#AmazonSQS"].traits["aws.protocols#awsQuery"];
+const isAwsQueryCompatible =
+  !!sqsModel.shapes["com.amazonaws.sqs#AmazonSQS"].traits["aws.protocols#awsQueryCompatible"];
 
 let hashError = "";
 const md5 = (str: string) =>
@@ -319,6 +321,28 @@ describe("middleware-sdk-sqs", () => {
     });
   });
 
+  it("should send the x-amzn-query-mode header when in awsQueryCompatible mode", async () => {
+    const client = new SQS({
+      region: "us-west-2",
+      credentials: mockCredentials,
+      logger,
+    });
+
+    requireRequestsFrom(client).toMatch({
+      hostname: "abc.com",
+      protocol: "https:",
+      path: "/",
+      headers: {
+        "x-amzn-query-mode": "true",
+      },
+    });
+
+    await client.sendMessage({
+      QueueUrl: "https://abc.com/123/MyQueue",
+      MessageBody: "hello",
+    });
+  });
+
   describe("queue-url", () => {
     it("should override resolved endpoint by default", async () => {
       const client = new SQS({