Send x-amzn-query-mode to inform a service with the `awsQueryCompat…

…ible` trait that SDK is operating in that mode (#3883)

## Motivation and Context
If a client SDK is generated from a service model that has the
[`awsQueryCompatible`](https://smithy.io/2.0/aws/protocols/aws-query-protocol.html#aws-protocols-awsquerycompatible-trait)
trait, the SDK now sends `x-amzn-query-mode` in the request header for
the service.

## Description
The change in the PR itself is pretty simple, as said in the title. It's
more important to understand why we are making these changes. The rest
of the description will focus on the reason driving this change.

The `awsQueryCompatible` trait, added by a service, is specifically for
deserializing errors. It allows for deserializing errors in a backward
compatible manner when the service migrates away from the AWS Query
protocol.

With [the awsQueryError
trait](https://smithy.io/2.0/aws/protocols/aws-query-protocol.html#aws-protocols-awsqueryerror-trait),
the AWS Query supports customizing error codes that is not supported in
any other AWS protocol, e.g.
```
@awsQueryError(
    code: "AWS.SimpleQueueService.NonExistentQueue",
    httpResponseCode: 400
)
@error("client")
structure QueueDoesNotExistException {
    message: String
}
```

In short, the `awsQueryCompatible` trait makes it possible to continue
using the custom error codes even when the service drops support for the
AWS Query protocol and switches to other protocols such as `awsJson1_0`
and `rpcv2Cbor` (see [example
snippet](https://smithy.io/2.0/aws/protocols/aws-query-protocol.html#aws-protocols-awsquerycompatible-trait)
in the Smithy documentation)

The changes in this PR would be unnecessary if a service had originally
supported only `@awsQuery` and had _atomically_ updated its Smithy model
to replace `@awsQuery` with `@awsQueryCompatible` and `@awsJson1_0` in
lockstep. However, that's not always the case in practice.

Consider a service whose Smithy model supports two protocols:
`@awsQuery` and `@awsJson1_0`. While the Rust SDK maintains [an ordered
map of
protocols](https://github.com/smithy-lang/smithy-rs/blob/de4bc4547df15e2472b9bca91c89f963ffad0b03/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/protocols/ClientProtocolLoader.kt#L38-L47)
to determine which one to use, it’s possible for two groups of client
SDKs to be generated over time, depending on which protocol was added
first to the service:
1. Client SDKs that understand the `awsQuery` protocol (this group can
interpret custom error codes sent in responses from the service)
2. Client SDKs that understand the `awsJson1_0` protocol (this group
does not interpret custom error codes)

Now, imagine if the service updated its Smithy model to remove
`@awsQuery` and add `@awsQueryCompatible` (likely it would add the
`awsQueryError` trait to an error structure somewhere as well). The
supported protocols would then be `@awsJson1_0` and
`@awsQueryCompatible`. Group 1 remains unaffected, as they can continue
deserializing responses with custom error codes. However, group 2 would
now be broken, as they would begin receiving custom error codes in
responses due to the `awsQueryCompatible` trait.

To prevent the issue for group 2 above, the `x-amzn-query-mode` header
in this PR informs the service that it should only send back custom
error codes if the client SDK is built with the `awsQueryCompatible`
trait. With this update, client SDKs built only with the `awsJson1_0`
will remain unaffected, as they do not send the `x-amzn-query-mode`
header to the service and, therefore, will not receive custom error
codes in response.

## Testing
- Added a Kotlin test to check for the `x-amzn-query-mode` header
- Added `x-amzn-query-mode` as a required header to a protocol test for
the `awsQueryCompatible`

## Checklist
- [x] For changes to the AWS SDK, generated SDK code, or SDK runtime
crates, I have created a changelog entry Markdown file in the
`.changelog` directory, specifying "aws-sdk-rust" in the `applies_to`
key.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

.changelog/1729271936.md

@@ -0,0 +1,12 @@
+---
+applies_to:
+- aws-sdk-rust
+authors:
+- ysaito1001
+references:
+- smithy-rs#3883
+breaking: false
+new_feature: false
+bug_fix: false
+---
+Client SDKs built with the `awsQueryCompatible` trait now include the `x-amzn-query-mode` header. This header signals the service that the clients are operating in compatible mode.

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/protocols/AwsQueryCompatibleTest.kt

@@ -6,25 +6,29 @@
 package software.amazon.smithy.rust.codegen.client.smithy.protocols
 
 import org.junit.jupiter.api.Test
-import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
-import software.amazon.smithy.rust.codegen.core.util.lookup
+import software.amazon.smithy.rust.codegen.core.testutil.testModule
+import software.amazon.smithy.rust.codegen.core.testutil.tokioTest
+import software.amazon.smithy.rust.codegen.core.util.letIf
 
 class AwsQueryCompatibleTest {
- @Test
- fun `aws-query-compatible json with aws query error should allow for retrieving error code and type from custom header`() {
- val model =
- """
+ companion object {
+ const val prologue = """
  namespace test
  use aws.protocols#awsJson1_0
  use aws.protocols#awsQueryCompatible
  use aws.protocols#awsQueryError
+ """
 
- @awsQueryCompatible
- @awsJson1_0
+ const val awsjson10Trait = "@awsJson1_0"
+ const val awsQueryCompatibleTrait = "@awsQueryCompatible"
+
+ fun testService(withAwsQueryError: Boolean = true) =
+ """
  service TestService {
  version: "2023-02-20",
  operations: [SomeOperation]
@@ -40,36 +44,45 @@ class AwsQueryCompatibleTest {
  a: String,
  b: Integer
  }
-
- @awsQueryError(
- code: "InvalidThing",
- httpResponseCode: 400,
- )
- @error("client")
- structure InvalidThingException {
- message: String
+ """.letIf(withAwsQueryError) {
+ it +
+ """
+ @awsQueryError(
+ code: "InvalidThing",
+ httpResponseCode: 400,
+ )
+ """
+ }.let {
+ it +
+ """
+ @error("client")
+ structure InvalidThingException {
+ message: String
+ }
+ """
  }
-  """.asSmithyModel()
+ }
 
+ @Test
+ fun `aws-query-compatible json with aws query error should allow for retrieving error code and type from custom header`() {
+ val model =
+ (prologue + awsQueryCompatibleTrait + awsjson10Trait + testService()).asSmithyModel(
+ smithyVersion = "2",
+ )
  clientIntegrationTest(model) { context, rustCrate ->
- val operation: OperationShape = context.model.lookup("test#SomeOperation")
- rustCrate.withModule(context.symbolProvider.moduleForShape(operation)) {
- rustTemplate(
- """
- ##[cfg(test)]
- ##[#{tokio}::test]
- async fn should_parse_code_and_type_fields() {
- use aws_smithy_types::body::SdkBody;
-
- let response = |_: http::Request<SdkBody>| {
+ rustCrate.testModule {
+ tokioTest("should_parse_code_and_type_fields") {
+ rustTemplate(
+ """
+ let response = |_: http::Request<#{SdkBody}>| {
  http::Response::builder()
  .header(
  "x-amzn-query-error",
  http::HeaderValue::from_static("AWS.SimpleQueueService.NonExistentQueue;Sender"),
  )
  .status(400)
  .body(
- SdkBody::from(
+ #{SdkBody}::from(
  r##"{
  "__type": "com.amazonaws.sqs##QueueDoesNotExist",
  "message": "Some user-visible message"
@@ -86,68 +99,38 @@ class AwsQueryCompatibleTest {
  );
  let error = dbg!(client.some_operation().send().await).err().unwrap().into_service_error();
  assert_eq!(
- Some("AWS.SimpleQueueService.NonExistentQueue"),
+ #{Some}("AWS.SimpleQueueService.NonExistentQueue"),
  error.meta().code(),
  );
- assert_eq!(Some("Sender"), error.meta().extra("type"));
- }
- """,
- "infallible_client_fn" to
- CargoDependency.smithyRuntimeTestUtil(context.runtimeConfig)
- .toType().resolve("client::http::test_util::infallible_client_fn"),
- "tokio" to CargoDependency.Tokio.toType(),
- )
+ assert_eq!(#{Some}("Sender"), error.meta().extra("type"));
+ """,
+ *RuntimeType.preludeScope,
+ "SdkBody" to RuntimeType.sdkBody(context.runtimeConfig),
+ "infallible_client_fn" to
+ CargoDependency.smithyRuntimeTestUtil(context.runtimeConfig)
+ .toType().resolve("client::http::test_util::infallible_client_fn"),
+ )
+ }
  }
  }
  }
 
  @Test
  fun `aws-query-compatible json without aws query error should allow for retrieving error code from payload`() {
  val model =
- """
- namespace test
- use aws.protocols#awsJson1_0
- use aws.protocols#awsQueryCompatible
-
- @awsQueryCompatible
- @awsJson1_0
- service TestService {
- version: "2023-02-20",
- operations: [SomeOperation]
- }
-
- operation SomeOperation {
- input: SomeOperationInputOutput,
- output: SomeOperationInputOutput,
- errors: [InvalidThingException],
- }
-
- structure SomeOperationInputOutput {
- a: String,
- b: Integer
- }
-
- @error("client")
- structure InvalidThingException {
- message: String
- }
- """.asSmithyModel()
-
+ (prologue + awsQueryCompatibleTrait + awsjson10Trait + testService(withAwsQueryError = false)).asSmithyModel(
+ smithyVersion = "2",
+ )
  clientIntegrationTest(model) { context, rustCrate ->
- val operation: OperationShape = context.model.lookup("test#SomeOperation")
- rustCrate.withModule(context.symbolProvider.moduleForShape(operation)) {
- rustTemplate(
- """
- ##[cfg(test)]
- ##[#{tokio}::test]
- async fn should_parse_code_from_payload() {
- use aws_smithy_types::body::SdkBody;
-
- let response = |_: http::Request<SdkBody>| {
+ rustCrate.testModule {
+ tokioTest("should_parse_code_from_payload") {
+ rustTemplate(
+ """
+ let response = |_: http::Request<#{SdkBody}>| {
  http::Response::builder()
  .status(400)
  .body(
- SdkBody::from(
+ #{SdkBody}::from(
  r##"{
  "__type": "com.amazonaws.sqs##QueueDoesNotExist",
  "message": "Some user-visible message"
@@ -163,15 +146,45 @@ class AwsQueryCompatibleTest {
  .build()
  );
  let error = dbg!(client.some_operation().send().await).err().unwrap().into_service_error();
- assert_eq!(Some("QueueDoesNotExist"), error.meta().code());
- assert_eq!(None, error.meta().extra("type"));
- }
- """,
- "infallible_client_fn" to
- CargoDependency.smithyRuntimeTestUtil(context.runtimeConfig)
- .toType().resolve("client::http::test_util::infallible_client_fn"),
- "tokio" to CargoDependency.Tokio.toType(),
- )
+ assert_eq!(#{Some}("QueueDoesNotExist"), error.meta().code());
+ assert_eq!(#{None}, error.meta().extra("type"));
+ """,
+ *RuntimeType.preludeScope,
+ "SdkBody" to RuntimeType.sdkBody(context.runtimeConfig),
+ "infallible_client_fn" to
+ CargoDependency.smithyRuntimeTestUtil(context.runtimeConfig)
+ .toType().resolve("client::http::test_util::infallible_client_fn"),
+ )
+ }
+ }
+ }
+ }
+
+ @Test
+ fun `request header should include x-amzn-query-mode when the service has the awsQueryCompatible trait`() {
+ val model =
+ (prologue + awsQueryCompatibleTrait + awsjson10Trait + testService()).asSmithyModel(
+ smithyVersion = "2",
+ )
+ clientIntegrationTest(model) { context, rustCrate ->
+ rustCrate.testModule {
+ tokioTest("test_request_header_should_include_x_amzn_query_mode") {
+ rustTemplate(
+ """
+ let (http_client, rx) = #{capture_request}(#{None});
+ let config = crate::Config::builder()
+ .http_client(http_client)
+ .endpoint_url("http://localhost:1234/SomeOperation")
+ .build();
+ let client = crate::Client::from_conf(config);
+ let _ = dbg!(client.some_operation().send().await);
+ let request = rx.expect_request();
+ assert_eq!("true", request.headers().get("x-amzn-query-mode").unwrap());
+ """,
+ *RuntimeType.preludeScope,
+ "capture_request" to RuntimeType.captureRequest(context.runtimeConfig),
+ )
+ }
  }
  }
  }

codegen-core/common-test-models/aws-json-query-compat.smithy

@@ -24,8 +24,10 @@ service QueryCompatService {
  params: {
  message: "hello!"
  },
- headers: { "x-amz-target": "QueryCompatService.Operation"}
-
+ headers: {
+ "x-amz-target": "QueryCompatService.Operation",
+ "x-amzn-query-mode": "true",
+ }
  }
 ])
 operation Operation {

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/AwsQueryCompatible.kt

@@ -97,5 +97,8 @@ class AwsQueryCompatible(
  awsJson.parseEventStreamErrorMetadata(operationShape)
 
  override fun additionalRequestHeaders(operationShape: OperationShape): List<Pair<String, String>> =
- listOf("x-amz-target" to "${codegenContext.serviceShape.id.name}.${operationShape.id.name}")
+ listOf(
+ "x-amz-target" to "${codegenContext.serviceShape.id.name}.${operationShape.id.name}",
+ "x-amzn-query-mode" to "true",
+ )
 }