From ffa9a702b7263fdb743f680dc9018c67f2b6482a Mon Sep 17 00:00:00 2001 From: Zelda Hessler Date: Wed, 20 Nov 2024 10:59:28 -0600 Subject: [PATCH] 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() +}