Apply AwsChunkedContentEncodingDecorator only to S3

This commit addresses #4504 (comment)

Author: ysaito1001

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AwsChunkedContentEncodingDecorator.kt

@@ -31,8 +31,19 @@ class AwsChunkedContentEncodingDecorator : ClientCodegenDecorator {
     // This decorator must decorate after any of the following:
     // - HttpRequestChecksumDecorator
     // - HttpRequestCompressionDecorator
+    //
+    // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Change ORDER to -1 once
+    //  a dedicated Smithy trait is available.
+    //  Why ORDER -2 is needed temporarily:
+    //  - AwsChunkedContentEncodingDecorator is part of AwsSdkCodegenDecorator (ORDER -1),
+    //    which only applies to S3
+    //  - HttpChecksumTest needs aws-chunked for flexible checksums but uses a non-S3 model
+    //  - To test this, AwsChunkedContentEncodingDecorator must be passed separately to
+    //    awsSdkIntegrationTest alongside AwsSdkCodegenDecorator
+    //  - Since HttpRequestChecksumDecorator is in AwsSdkCodegenDecorator (ORDER -1),
+    //    we need ORDER -2 to run before it
     override val order: Byte =
-        (minOf(HttpRequestChecksumDecorator.ORDER, HttpRequestCompressionDecorator.ORDER) - 1).toByte()
+        (minOf(HttpRequestChecksumDecorator.ORDER, HttpRequestCompressionDecorator.ORDER) - 2).toByte()
 
     override fun configCustomizations(
         codegenContext: ClientCodegenContext,
@@ -46,38 +57,16 @@ class AwsChunkedContentEncodingDecorator : ClientCodegenDecorator {
     ) = baseCustomizations + AwsChunkedOperationCustomization(codegenContext, operation)
 }
 
-// TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Replace this heuristic with a dedicated
-//  Smithy trait once available to determine whether operations require aws-chunked encoding.
-private fun operationRequiresAwsChunked(
-    codegenContext: ClientCodegenContext,
-    operation: OperationShape,
-): Boolean {
-    val checksumTrait = operation.getTrait<HttpChecksumTrait>() ?: return false
-    val requestAlgorithmMember =
-        checksumTrait.requestAlgorithmMemberShape(codegenContext, operation) ?: return false
-    requestAlgorithmMember.getTrait<HttpHeaderTrait>()?.value ?: return false
-    val input = codegenContext.model.expectShape(operation.inputShape, StructureShape::class.java)
-    return input.hasStreamingMember(codegenContext.model)
-}
-
-private fun serviceRequiresAwsChunked(codegenContext: ClientCodegenContext): Boolean =
-    codegenContext.serviceShape.allOperations.any { operationId ->
-        val operation = codegenContext.model.expectShape(operationId, OperationShape::class.java)
-        operationRequiresAwsChunked(codegenContext, operation)
-    }
-
 private class AwsChunkedConfigCustomization(
     codegenContext: ClientCodegenContext,
 ) : ConfigCustomization() {
     private val runtimeConfig = codegenContext.runtimeConfig
     private val moduleUseName = codegenContext.moduleUseName()
-    private val serviceRequiresChunking = serviceRequiresAwsChunked(codegenContext)
 
     override fun section(section: ServiceConfig) =
         writable {
             when (section) {
                 ServiceConfig.BuilderImpl -> {
-                    if (!serviceRequiresChunking) return@writable
                     rustTemplate(
                         """
                         /// Sets the chunk size for [`aws-chunked encoding`].
@@ -144,6 +133,20 @@ private class AwsChunkedConfigCustomization(
         }
 }
 
+// TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Replace this condition with a dedicated
+//  Smithy trait once available to determine whether operations require aws-chunked encoding.
+private fun operationRequiresAwsChunked(
+    codegenContext: ClientCodegenContext,
+    operation: OperationShape,
+): Boolean {
+    val checksumTrait = operation.getTrait<HttpChecksumTrait>() ?: return false
+    val requestAlgorithmMember =
+        checksumTrait.requestAlgorithmMemberShape(codegenContext, operation) ?: return false
+    requestAlgorithmMember.getTrait<HttpHeaderTrait>()?.value ?: return false
+    val input = codegenContext.model.expectShape(operation.inputShape, StructureShape::class.java)
+    return input.hasStreamingMember(codegenContext.model)
+}
+
 private class AwsChunkedOperationCustomization(
     private val codegenContext: ClientCodegenContext,
     private val operation: OperationShape,

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt

@@ -58,7 +58,6 @@ val DECORATORS: List<ClientCodegenDecorator> =
             SdkConfigDecorator(),
             ServiceConfigDecorator(),
             AwsPresigningDecorator(),
-            AwsChunkedContentEncodingDecorator(),
             AwsCrateDocsDecorator(),
             AwsEndpointsStdLib(),
             *PromotedBuiltInsDecorators,
@@ -99,6 +98,9 @@ val DECORATORS: List<ClientCodegenDecorator> =
             S3ExtendedRequestIdDecorator(),
             IsTruncatedPaginatorDecorator(),
             S3ExpiresDecorator(),
+            // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Apply this decorator based on
+            //  a dedicated Smithy trait once available.
+            AwsChunkedContentEncodingDecorator(),
         ),
         S3ControlDecorator().onlyApplyTo("com.amazonaws.s3control#AWSS3ControlServiceV20180820"),
         STSDecorator().onlyApplyTo("com.amazonaws.sts#AWSSecurityTokenServiceV20110615"),

aws/codegen-aws-sdk/src/test/kotlin/software/amazon/smithy/rustsdk/HttpChecksumTest.kt

@@ -144,8 +144,13 @@ internal class HttpChecksumTest {
     }
 
     @Test
-    fun requestChecksumWorks() {
-        awsSdkIntegrationTest(model) { context, rustCrate ->
+    fun requestResponseChecksumWorks() {
+        awsSdkIntegrationTest(
+            model,
+            // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Remove this additional decorator
+            //  and update the test model above with a dedicated Smithy trait once available.
+            additionalDecorators = listOf(AwsChunkedContentEncodingDecorator()),
+        ) { context, rustCrate ->
             // Allows us to use the user-agent test-utils in aws-runtime
             rustCrate.mergeFeature(Feature("test-util", true, listOf("aws-runtime/test-util")))
             val rc = context.runtimeConfig

aws/codegen-aws-sdk/src/test/kotlin/software/amazon/smithy/rustsdk/TestUtil.kt

@@ -10,6 +10,7 @@ import software.amazon.smithy.model.node.ObjectNode
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
 import software.amazon.smithy.rust.codegen.client.smithy.RustClientCodegenPlugin
+import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.testutil.ClientDecoratableBuildPlugin
 import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
 import software.amazon.smithy.rust.codegen.client.testutil.testClientCodegenContext
@@ -44,12 +45,14 @@ fun awsTestCodegenContext(
 fun awsSdkIntegrationTest(
     model: Model,
     params: IntegrationTestParams = awsIntegrationTestParams(),
+    additionalDecorators: List<ClientCodegenDecorator> = listOf(),
     buildPlugin: ClientDecoratableBuildPlugin = RustClientCodegenPlugin(),
     environment: Map<String, String> = mapOf(),
     test: (ClientCodegenContext, RustCrate) -> Unit = { _, _ -> },
 ) = clientIntegrationTest(
     model,
     params,
+    additionalDecorators = additionalDecorators,
     buildPlugin = buildPlugin,
     environment = environment,
     test = test,