From ffa9a702b7263fdb743f680dc9018c67f2b6482a Mon Sep 17 00:00:00 2001 From: Zelda Hessler Date: Wed, 20 Nov 2024 10:59:28 -0600 Subject: [PATCH 1/2] Only generate the `presign` customize operation method for ops that support it --- aws/rust-runtime/aws-config/Cargo.lock | 2 +- .../smithy/rustsdk/AwsPresigningDecorator.kt | 58 ++++++++++++------- .../client/smithy/ClientCodegenVisitor.kt | 31 ++++++---- .../generators/OperationCustomization.kt | 5 ++ .../CustomizableOperationImplGenerator.kt | 33 +++++++++++ .../core/smithy/customize/Customization.kt | 10 ++++ 6 files changed, 105 insertions(+), 34 deletions(-) create mode 100644 codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt diff --git a/aws/rust-runtime/aws-config/Cargo.lock b/aws/rust-runtime/aws-config/Cargo.lock index d81bd719c0..712530ba64 100644 --- a/aws/rust-runtime/aws-config/Cargo.lock +++ b/aws/rust-runtime/aws-config/Cargo.lock @@ -89,7 +89,7 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.4.3" +version = "1.5.0" dependencies = [ "aws-credential-types", "aws-sigv4", diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt index 0c3cc189b4..53cb2ffcba 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt @@ -20,7 +20,8 @@ import software.amazon.smithy.model.transform.ModelTransformer 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.customize.ClientCodegenDecorator -import software.amazon.smithy.rust.codegen.client.smithy.generators.client.CustomizableOperationSection +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.generators.client.FluentClientCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientSection import software.amazon.smithy.rust.codegen.client.smithy.generators.client.InternalTraitsModule @@ -39,11 +40,8 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope import software.amazon.smithy.rust.codegen.core.smithy.RustCrate import software.amazon.smithy.rust.codegen.core.smithy.contextName -import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization -import software.amazon.smithy.rust.codegen.core.smithy.customize.adhocCustomization import software.amazon.smithy.rust.codegen.core.util.cloneOperation import software.amazon.smithy.rust.codegen.core.util.expectTrait -import software.amazon.smithy.rust.codegen.core.util.thenSingletonListOf import software.amazon.smithy.rustsdk.traits.PresignableTrait private val presigningTypes: Array> = @@ -141,26 +139,42 @@ class AwsPresigningDecorator internal constructor( TopDownIndex.of(ctx.model).getContainedOperations(ctx.serviceShape) .any { presignableOperations.containsKey(it.id) } - override fun extraSections(codegenContext: ClientCodegenContext): List = - anyPresignedShapes(codegenContext).thenSingletonListOf { - adhocCustomization { - rustTemplate( - """ - /// Sends the request and returns the response. - ##[allow(unused_mut)] - pub async fn presigned(mut self, presigning_config: #{PresigningConfig}) -> #{Result}<#{PresignedRequest}, crate::error::SdkError> where - E: std::error::Error + #{Send} + #{Sync} + 'static, - B: #{CustomizablePresigned} - { - self.execute(move |sender, conf|sender.presign(conf, presigning_config)).await + override fun operationCustomizations( + codegenContext: ClientCodegenContext, + operation: OperationShape, + baseCustomizations: List, + ): List { + return baseCustomizations + + object : OperationCustomization() { + override fun section(section: OperationSection): Writable { + return writable { + when (section) { + is OperationSection.CustomizableOperationImpl -> { + if (PRESIGNABLE_OPERATIONS.containsKey(operation.id)) { + rustTemplate( + """ + /// Sends the request and returns the response. + ##[allow(unused_mut)] + pub async fn presigned(mut self, presigning_config: #{PresigningConfig}) -> #{Result}<#{PresignedRequest}, crate::error::SdkError> where + E: std::error::Error + #{Send} + #{Sync} + 'static, + B: #{CustomizablePresigned} + { + self.execute(move |sender, conf|sender.presign(conf, presigning_config)).await + } + """, + *preludeScope, + *presigningTypes, + "CustomizablePresigned" to customizablePresigned, + ) + } + } + + else -> {} + } } - """, - *preludeScope, - *presigningTypes, - "CustomizablePresigned" to customizablePresigned, - ) + } } - } + } override fun extras( codegenContext: ClientCodegenContext, diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt index 7e8737f50d..24f4e6e38e 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt @@ -21,6 +21,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegen import software.amazon.smithy.rust.codegen.client.smithy.generators.ClientEnumGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceGenerator +import software.amazon.smithy.rust.codegen.client.smithy.generators.client.CustomizableOperationImplGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.error.ErrorGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.error.OperationErrorGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.ClientProtocolTestGenerator @@ -311,10 +312,10 @@ class ClientCodegenVisitor( * Generate operations */ override fun operationShape(operationShape: OperationShape) { - rustCrate.useShapeWriter(operationShape) operationWriter@{ + rustCrate.useShapeWriter(operationShape) { // Render the operation shape operationGenerator.renderOperation( - this@operationWriter, + this, operationShape, codegenDecorator, ) @@ -327,16 +328,24 @@ class ClientCodegenVisitor( protocolGeneratorFactory.support(), operationShape, ), - ).render(this@operationWriter) + ).render(this) + } - rustCrate.withModule(symbolProvider.moduleForOperationError(operationShape)) { - OperationErrorGenerator( - model, - symbolProvider, - operationShape, - codegenDecorator.errorCustomizations(codegenContext, emptyList()), - ).render(this) - } + rustCrate.withModule(symbolProvider.moduleForOperationError(operationShape)) { + OperationErrorGenerator( + model, + symbolProvider, + operationShape, + codegenDecorator.errorCustomizations(codegenContext, emptyList()), + ).render(this) + } + + rustCrate.withModule(ClientRustModule.Client.customize) { + CustomizableOperationImplGenerator( + codegenContext, + operationShape, + codegenDecorator.operationCustomizations(codegenContext, operationShape, emptyList()), + ).render(this) } } } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt index 07ed4e51cb..2e31dbac44 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt @@ -125,6 +125,11 @@ sealed class OperationSection(name: String) : Section(name) { writer.rustTemplate(".with_retry_classifier(#{classifier})", "classifier" to classifier) } } + + data class CustomizableOperationImpl( + override val customizations: List, + val operationShape: OperationShape, + ) : OperationSection("CustomizableOperationImpl") } abstract class OperationCustomization : NamedCustomization() diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt new file mode 100644 index 0000000000..7f839be24b --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt @@ -0,0 +1,33 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.rust.codegen.client.smithy.generators.client + +import software.amazon.smithy.model.shapes.OperationShape +import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext +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.core.rustlang.RustWriter +import software.amazon.smithy.rust.codegen.core.rustlang.rust +import software.amazon.smithy.rust.codegen.core.smithy.customize.allCustomizationsAreEmpty +import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations + +class CustomizableOperationImplGenerator( + private val codegenContext: ClientCodegenContext, + private val operation: OperationShape, + private val customizations: List, +) { + fun render(writer: RustWriter) { + val section = OperationSection.CustomizableOperationImpl(customizations, operation) + // When no customizations are set or there is nothing to write, return early. + if (customizations.isEmpty() || allCustomizationsAreEmpty(customizations, section)) { + return + } + + writer.rust("impl CustomizableOperation<#T, E, B> {", codegenContext.symbolProvider.toSymbol(operation)) + writer.writeCustomizations(customizations, section) + writer.rust("}") + } +} diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt index 79ff1bdb6f..3a969bbb08 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt @@ -80,3 +80,13 @@ fun RustWriter.writeCustomizationsOrElse( orElse(this) } } + +fun allCustomizationsAreEmpty( + customizations: List>, + section: T, +): Boolean { + val test = RustWriter.root() + test.writeCustomizations(customizations, section) + // If they're not dirty, then they're empty. + return !test.dirty() +} From 6747d671aeeb4cf9ca9466d88327d1cd0966cee4 Mon Sep 17 00:00:00 2001 From: Zelda Hessler Date: Mon, 2 Dec 2024 10:09:24 -0600 Subject: [PATCH 2/2] undo presignable ops generation changes --- aws/rust-runtime/aws-config/Cargo.lock | 94 +++++++++---------- .../smithy/rustsdk/AwsPresigningDecorator.kt | 58 +++++------- 2 files changed, 69 insertions(+), 83 deletions(-) diff --git a/aws/rust-runtime/aws-config/Cargo.lock b/aws/rust-runtime/aws-config/Cargo.lock index 712530ba64..5f5dea15b2 100644 --- a/aws/rust-runtime/aws-config/Cargo.lock +++ b/aws/rust-runtime/aws-config/Cargo.lock @@ -89,7 +89,7 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.5.0" +version = "1.4.3" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -428,9 +428,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.1.36" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baee610e9452a8f6f0a1b6194ec09ff9e2d85dea54432acdae41aa0761c95d70" +checksum = "fd9de9f2205d5ef3fd67e685b0df337994ddd4495e2a28d185500d0e1edfea47" dependencies = [ "shlex", ] @@ -495,9 +495,9 @@ checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" [[package]] name = "cpufeatures" -version = "0.2.14" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "608697df725056feaccfa42cffdaeeec3fccc4ffc38358ecd19b243e716a78e0" +checksum = "16b80225097f2e5ae4e7179dd2266824648f3e2f49d9134d584b76389d31c4c3" dependencies = [ "libc", ] @@ -576,9 +576,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8c02a5121d4ea3eb16a80748c74f5549a5665e4c21333c6098f283870fbdea6" +checksum = "486f806e73c5707928240ddc295403b1b93c96a02038563881c4a2fd84b81ac4" [[package]] name = "fnv" @@ -692,9 +692,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.15.1" +version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a9bfc1af68b1726ea47d3d5109de126281def866b33970e10fbab11b5dafab3" +checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" [[package]] name = "hermit-abi" @@ -849,9 +849,9 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.11" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" [[package]] name = "lazy_static" @@ -861,9 +861,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.161" +version = "0.2.164" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e9489c2807c139ffd9c1794f4af0ebe86a828db53ecdc7fea2111d0fed085d1" +checksum = "433bfe06b8c75da9b2e3fbea6e5329ff87748f0b144ef75306e674c3f6f7c13f" [[package]] name = "lock_api" @@ -1104,9 +1104,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.89" +version = "1.0.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e" +checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" dependencies = [ "unicode-ident", ] @@ -1137,7 +1137,7 @@ checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" dependencies = [ "aho-corasick", "memchr", - "regex-automata 0.4.8", + "regex-automata 0.4.9", "regex-syntax 0.8.5", ] @@ -1152,9 +1152,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.8" +version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3" +checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" dependencies = [ "aho-corasick", "memchr", @@ -1269,9 +1269,9 @@ checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" [[package]] name = "schannel" -version = "0.1.26" +version = "0.1.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01227be5826fa0690321a2ba6c5cd57a19cf3f6a09e76973b58e61de6ab9d1c1" +checksum = "1f29ebaa345f945cec9fbbc532eb307f0fdad8161f281b6369539c8d84876b3d" dependencies = [ "windows-sys 0.59.0", ] @@ -1307,9 +1307,9 @@ dependencies = [ [[package]] name = "security-framework-sys" -version = "2.12.0" +version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea4a292869320c0272d7bc55a5a6aafaff59b4f63404a003887b679a2e05b4b6" +checksum = "fa39c7303dc58b5543c94d22c1766b0d31f2ee58306363ea622b10bbc075eaa2" dependencies = [ "core-foundation-sys", "libc", @@ -1329,29 +1329,29 @@ checksum = "f97841a747eef040fcd2e7b3b9a220a7205926e60488e673d9e4926d27772ce5" [[package]] name = "serde" -version = "1.0.214" +version = "1.0.215" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f55c3193aca71c12ad7890f1785d2b73e1b9f63a0bbc353c08ef26fe03fc56b5" +checksum = "6513c1ad0b11a9376da888e3e0baa0077f1aed55c17f50e7b2397136129fb88f" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.214" +version = "1.0.215" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de523f781f095e28fa605cdce0f8307e451cc0fd14e2eb4cd2e98a355b147766" +checksum = "ad1e866f866923f252f05c889987993144fb74e722403468a4ebd70c3cd756c0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.89", ] [[package]] name = "serde_json" -version = "1.0.132" +version = "1.0.133" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d726bfaff4b320266d395898905d0eba0345aae23b54aee3a737e260fd46db03" +checksum = "c7fceb2473b9166b2294ef05efcb65a3db80803f0b03ef86a5fc88a2b85ee377" dependencies = [ "indexmap", "itoa", @@ -1445,9 +1445,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.87" +version = "2.0.89" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d" +checksum = "44d46482f1c1c87acd84dea20c1bf5ebff4c757009ed6bf19cfd36fb10e92c4e" dependencies = [ "proc-macro2", "quote", @@ -1456,22 +1456,22 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.68" +version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02dd99dc800bbb97186339685293e1cc5d9df1f8fae2d0aecd9ff1c77efea892" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.68" +version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7c61ec9a6f64d2793d8a45faba21efbe3ced62a886d44c36a009b2b519b4c7e" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.89", ] [[package]] @@ -1531,9 +1531,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.41.0" +version = "1.41.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "145f3413504347a2be84393cc8a7d2fb4d863b375909ea59f2158261aa258bbb" +checksum = "22cfb5bee7a6a52939ca9224d6ac897bb669134078daa8735560897f69de4d33" dependencies = [ "backtrace", "bytes", @@ -1555,7 +1555,7 @@ checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.89", ] [[package]] @@ -1600,20 +1600,20 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.27" +version = "0.1.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.87", + "syn 2.0.89", ] [[package]] name = "tracing-core" -version = "0.1.32" +version = "0.1.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" +checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c" dependencies = [ "once_cell", "valuable", @@ -1679,7 +1679,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04659ddb06c87d233c566112c1c9c5b9e98256d9af50ec3bc9c8327f873a7568" dependencies = [ "quote", - "syn 2.0.87", + "syn 2.0.89", ] [[package]] @@ -1702,9 +1702,9 @@ checksum = "5ab17db44d7388991a428b2ee655ce0c212e862eff1768a455c58f9aad6e7893" [[package]] name = "unicode-ident" -version = "1.0.13" +version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" +checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" [[package]] name = "unicode-normalization" diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt index 53cb2ffcba..0c3cc189b4 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt @@ -20,8 +20,7 @@ import software.amazon.smithy.model.transform.ModelTransformer 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.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.generators.client.CustomizableOperationSection import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientSection import software.amazon.smithy.rust.codegen.client.smithy.generators.client.InternalTraitsModule @@ -40,8 +39,11 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope import software.amazon.smithy.rust.codegen.core.smithy.RustCrate import software.amazon.smithy.rust.codegen.core.smithy.contextName +import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization +import software.amazon.smithy.rust.codegen.core.smithy.customize.adhocCustomization import software.amazon.smithy.rust.codegen.core.util.cloneOperation import software.amazon.smithy.rust.codegen.core.util.expectTrait +import software.amazon.smithy.rust.codegen.core.util.thenSingletonListOf import software.amazon.smithy.rustsdk.traits.PresignableTrait private val presigningTypes: Array> = @@ -139,42 +141,26 @@ class AwsPresigningDecorator internal constructor( TopDownIndex.of(ctx.model).getContainedOperations(ctx.serviceShape) .any { presignableOperations.containsKey(it.id) } - override fun operationCustomizations( - codegenContext: ClientCodegenContext, - operation: OperationShape, - baseCustomizations: List, - ): List { - return baseCustomizations + - object : OperationCustomization() { - override fun section(section: OperationSection): Writable { - return writable { - when (section) { - is OperationSection.CustomizableOperationImpl -> { - if (PRESIGNABLE_OPERATIONS.containsKey(operation.id)) { - rustTemplate( - """ - /// Sends the request and returns the response. - ##[allow(unused_mut)] - pub async fn presigned(mut self, presigning_config: #{PresigningConfig}) -> #{Result}<#{PresignedRequest}, crate::error::SdkError> where - E: std::error::Error + #{Send} + #{Sync} + 'static, - B: #{CustomizablePresigned} - { - self.execute(move |sender, conf|sender.presign(conf, presigning_config)).await - } - """, - *preludeScope, - *presigningTypes, - "CustomizablePresigned" to customizablePresigned, - ) - } - } - - else -> {} - } + override fun extraSections(codegenContext: ClientCodegenContext): List = + anyPresignedShapes(codegenContext).thenSingletonListOf { + adhocCustomization { + rustTemplate( + """ + /// Sends the request and returns the response. + ##[allow(unused_mut)] + pub async fn presigned(mut self, presigning_config: #{PresigningConfig}) -> #{Result}<#{PresignedRequest}, crate::error::SdkError> where + E: std::error::Error + #{Send} + #{Sync} + 'static, + B: #{CustomizablePresigned} + { + self.execute(move |sender, conf|sender.presign(conf, presigning_config)).await } - } + """, + *preludeScope, + *presigningTypes, + "CustomizablePresigned" to customizablePresigned, + ) } - } + } override fun extras( codegenContext: ClientCodegenContext,