diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt index 8af54cfe6e86ad0a61c0638bca9b54be010cac61..6f3577fb98225a3dff154eb105c3976239ac376c 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt @@ -27,6 +27,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCus import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection import software.amazon.smithy.rust.codegen.client.smithy.protocols.ClientRestXmlFactory +import software.amazon.smithy.rust.codegen.client.smithy.traits.IncompatibleWithStalledStreamProtectionTrait import software.amazon.smithy.rust.codegen.core.rustlang.Writable import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate @@ -59,6 +60,10 @@ class S3Decorator : ClientCodegenDecorator { // API returns ListAllMyDirectoryBucketsResult instead of ListDirectoryBucketsOutput ShapeId.from("com.amazonaws.s3#ListDirectoryBucketsOutput"), ) + private val operationsIncompatibleWithStalledStreamProtection = + setOf( + ShapeId.from("com.amazonaws.s3#CopyObject"), + ) override fun protocols( serviceId: ShapeId, @@ -81,6 +86,9 @@ class S3Decorator : ClientCodegenDecorator { shape.letIf(isInInvalidXmlRootAllowList(shape)) { logger.info("Adding AllowInvalidXmlRoot trait to $it") (it as StructureShape).toBuilder().addTrait(AllowInvalidXmlRoot()).build() + }.letIf(operationsIncompatibleWithStalledStreamProtection.contains(shape.id)) { + logger.info("Adding IncompatibleWithStalledStreamProtection trait to $it") + (it as OperationShape).toBuilder().addTrait(IncompatibleWithStalledStreamProtectionTrait()).build() } } // the model has the bucket in the path diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/StalledStreamProtectionConfigCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/StalledStreamProtectionConfigCustomization.kt index c702c6cb0206a8622c8c8ef62dfc39b2b1c8e20f..58eb2c8a64ca5cdeeb160b03a3ee3884b3870797 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/StalledStreamProtectionConfigCustomization.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/StalledStreamProtectionConfigCustomization.kt @@ -11,6 +11,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.configReexport import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection +import software.amazon.smithy.rust.codegen.client.smithy.traits.IncompatibleWithStalledStreamProtectionTrait import software.amazon.smithy.rust.codegen.core.rustlang.Writable import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate import software.amazon.smithy.rust.codegen.core.rustlang.writable @@ -34,7 +35,7 @@ class StalledStreamProtectionDecorator : ClientCodegenDecorator { operation: OperationShape, baseCustomizations: List, ): List { - return baseCustomizations + StalledStreamProtectionOperationCustomization(codegenContext) + return baseCustomizations + StalledStreamProtectionOperationCustomization(codegenContext, operation) } } @@ -111,11 +112,17 @@ class StalledStreamProtectionConfigCustomization(codegenContext: ClientCodegenCo class StalledStreamProtectionOperationCustomization( codegenContext: ClientCodegenContext, + private val operationShape: OperationShape, ) : OperationCustomization() { private val rc = codegenContext.runtimeConfig override fun section(section: OperationSection): Writable = writable { + // Don't add the stalled stream protection interceptor if the operation is incompatible with it + if (operationShape.hasTrait(IncompatibleWithStalledStreamProtectionTrait.ID)) { + return@writable + } + when (section) { is OperationSection.AdditionalInterceptors -> { val stalledStreamProtectionModule = RuntimeType.smithyRuntime(rc).resolve("client::stalled_stream_protection") diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IncompatibleWithStalledStreamProtectionTrait.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IncompatibleWithStalledStreamProtectionTrait.kt new file mode 100644 index 0000000000000000000000000000000000000000..658309653b8f1e45f987b0f64d257ae28ab7a0f9 --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IncompatibleWithStalledStreamProtectionTrait.kt @@ -0,0 +1,19 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.rust.codegen.client.smithy.traits + +import software.amazon.smithy.model.node.Node +import software.amazon.smithy.model.shapes.ShapeId +import software.amazon.smithy.model.traits.AnnotationTrait + +/** + * Indicates that an operation shape is incompatible with stalled stream protection. + */ +class IncompatibleWithStalledStreamProtectionTrait : AnnotationTrait(ID, Node.objectNode()) { + companion object { + val ID: ShapeId = ShapeId.from("software.amazon.smithy.rust.codegen.client.smithy.traits#incompatibleWithStalledStreamProtectionTrait") + } +}