From 925e731008b370dfaefb4b006d348629b2b3a16c Mon Sep 17 00:00:00 2001 From: George Fu Date: Mon, 11 Nov 2024 13:17:20 -0500 Subject: [PATCH] feat(codegen): send x-amzn-query-mode header (#6651) --- .../client-sqs/src/protocols/Aws_json1_0.ts | 1 + .../aws/typescript/codegen/AwsJsonRpc1_1.java | 21 ++++++++++++++++ .../codegen/AwsSmithyRpcV2Cbor.java | 19 +++++++++++++++ .../codegen/JsonRpcProtocolGenerator.java | 5 ++++ .../src/middleware-sdk-sqs.integ.spec.ts | 24 +++++++++++++++++++ 5 files changed, 70 insertions(+) diff --git a/clients/client-sqs/src/protocols/Aws_json1_0.ts b/clients/client-sqs/src/protocols/Aws_json1_0.ts index 35d045756d8d..61855c0b3ba5 100644 --- 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/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsJsonRpc1_1.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsJsonRpc1_1.java index 538caab3b549..69eba96658be 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsJsonRpc1_1.java +++ b/codegen/smithy-aws-typescript-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/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsSmithyRpcV2Cbor.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsSmithyRpcV2Cbor.java index 2256aea79c18..d4f6064a58b4 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AwsSmithyRpcV2Cbor.java +++ b/codegen/smithy-aws-typescript-codegen/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/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/JsonRpcProtocolGenerator.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/JsonRpcProtocolGenerator.java index 09504b5de899..81a156cbce1d 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/JsonRpcProtocolGenerator.java +++ b/codegen/smithy-aws-typescript-codegen/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 index 73307c091435..914d962d8ba0 100644 --- 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({