Only generate the presign customize operation method for ops that s…

…upport it
smithy-lang · Nov 20, 2024 · ffa9a70 · ffa9a70
1 parent df77d5f
commit ffa9a70
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 34 deletions.
diff --git a/aws/rust-runtime/aws-config/Cargo.lock b/aws/rust-runtime/aws-config/Cargo.lock
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
@@ -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<Pair<String, Any>> =
@@ -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<AdHocCustomization> =
-        anyPresignedShapes(codegenContext).thenSingletonListOf {
-            adhocCustomization<CustomizableOperationSection.CustomizableOperationImpl> {
-                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<E>> where
-                        E: std::error::Error + #{Send} + #{Sync} + 'static,
-                        B: #{CustomizablePresigned}<E>
-                    {
-                        self.execute(move |sender, conf|sender.presign(conf, presigning_config)).await
+    override fun operationCustomizations(
+        codegenContext: ClientCodegenContext,
+        operation: OperationShape,
+        baseCustomizations: List<OperationCustomization>,
+    ): List<OperationCustomization> {
+        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<E>> where
+                                            E: std::error::Error + #{Send} + #{Sync} + 'static,
+                                            B: #{CustomizablePresigned}<E>
+                                        {
+                                            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/...src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt b/...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/...in/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt b/...in/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<OperationCustomization>,
+        val operationShape: OperationShape,
+    ) : OperationSection("CustomizableOperationImpl")
 }
 
 abstract class OperationCustomization : NamedCustomization<OperationSection>()
diff --git a/...smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt b/...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<OperationCustomization>,
+) {
+    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<E, B> CustomizableOperation<#T, E, B> {", codegenContext.symbolProvider.toSymbol(operation))
+        writer.writeCustomizations(customizations, section)
+        writer.rust("}")
+    }
+}
diff --git a/...rc/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt b/...rc/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt
@@ -80,3 +80,13 @@ fun <T : Section> RustWriter.writeCustomizationsOrElse(
         orElse(this)
     }
 }
+
+fun <T : Section> allCustomizationsAreEmpty(
+    customizations: List<NamedCustomization<T>>,
+    section: T,
+): Boolean {
+    val test = RustWriter.root()
+    test.writeCustomizations(customizations, section)
+    // If they're not dirty, then they're empty.
+    return !test.dirty()
+}