From fb882233f163a0648a412a27bdc7f3254765605c Mon Sep 17 00:00:00 2001
From: Matas Lauzadis <lauzmata@amazon.com>
Date: Wed, 30 Oct 2024 10:28:46 -0400
Subject: [PATCH] Add x-amzn-query-mode customization

---
 .../20e287a5-cee0-4a5f-9d08-4022dcdee843.json |  5 ++++
 .../AwsQueryModeCustomization.kt              | 28 +++++++++++++++++++
 ...tlin.codegen.integration.KotlinIntegration |  1 +
 3 files changed, 34 insertions(+)
 create mode 100644 .changes/20e287a5-cee0-4a5f-9d08-4022dcdee843.json
 create mode 100644 codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/AwsQueryModeCustomization.kt

diff --git a/.changes/20e287a5-cee0-4a5f-9d08-4022dcdee843.json b/.changes/20e287a5-cee0-4a5f-9d08-4022dcdee843.json
new file mode 100644
index 00000000000..37442f4cdd2
--- /dev/null
+++ b/.changes/20e287a5-cee0-4a5f-9d08-4022dcdee843.json
@@ -0,0 +1,5 @@
+{
+    "id": "20e287a5-cee0-4a5f-9d08-4022dcdee843",
+    "type": "misc",
+    "description": "Send x-amzn-query-mode=true for services with query-compatible trait"
+}
\ No newline at end of file
diff --git a/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/AwsQueryModeCustomization.kt b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/AwsQueryModeCustomization.kt
new file mode 100644
index 00000000000..4e8de297cd5
--- /dev/null
+++ b/codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/AwsQueryModeCustomization.kt
@@ -0,0 +1,28 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization
+
+import software.amazon.smithy.aws.traits.protocols.AwsQueryCompatibleTrait
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.MutateHeadersMiddleware
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
+import software.amazon.smithy.model.Model
+
+/**
+ * Send an extra `x-amzn-query-mode` header with a value of `true` for services which have the [AwsQueryCompatibleTrait] applied.
+ */
+class AwsQueryModeCustomization : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        model
+            .getShape(settings.service)
+            .get()
+            .hasTrait<AwsQueryCompatibleTrait>()
+
+    override fun customizeMiddleware(ctx: ProtocolGenerator.GenerationContext, resolved: List<ProtocolMiddleware>): List<ProtocolMiddleware> =
+        resolved + MutateHeadersMiddleware(extraHeaders = mapOf("x-amzn-query-mode" to "true"))
+}
\ No newline at end of file
diff --git a/codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration b/codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
index 3ef264034df..6d8199d3072 100644
--- a/codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+++ b/codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
@@ -43,3 +43,4 @@ aws.sdk.kotlin.codegen.customization.s3.express.S3ExpressIntegration
 aws.sdk.kotlin.codegen.customization.s3.S3ExpiresIntegration
 aws.sdk.kotlin.codegen.BusinessMetricsIntegration
 aws.sdk.kotlin.codegen.smoketests.SmokeTestsDenyListIntegration
+aws.sdk.kotlin.codegen.customization.AwsQueryModeCustomization