Add x-amzn-query-mode customization

Author: lauzadis

.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"
+}

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"))
+}

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