Reorganize credentials, endpoints, region, and op customization

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

@@ -9,8 +9,9 @@ import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.traits.TitleTrait
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
-import software.amazon.smithy.rust.codegen.client.smithy.generators.client.CustomizableOperationGenerator
+import software.amazon.smithy.rust.codegen.client.smithy.featureGatedCustomizeModule
 import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientGenerator
 import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientGenerics
@@ -102,10 +103,10 @@ class AwsFluentClientDecorator : ClientCodegenDecorator {
             ),
             retryClassifier = AwsRuntimeType.awsHttp(runtimeConfig).resolve("retry::AwsResponseRetryClassifier"),
         ).render(rustCrate)
-        rustCrate.withModule(CustomizableOperationGenerator.CustomizeModule) {
+        rustCrate.withModule(codegenContext.featureGatedCustomizeModule()) {
             renderCustomizableOperationSendMethod(runtimeConfig, generics, this)
         }
-        rustCrate.withModule(FluentClientGenerator.clientModule) {
+        rustCrate.withModule(ClientRustModule.client) {
             AwsFluentClientExtensions(types).render(this)
         }
         val awsSmithyClient = "aws-smithy-client"

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

@@ -7,10 +7,10 @@ package software.amazon.smithy.rustsdk
 
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.generators.error.ErrorCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.error.ErrorSection
-import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
@@ -66,7 +66,7 @@ abstract class BaseRequestIdDecorator : ClientCodegenDecorator {
     ): List<BuilderCustomization> = baseCustomizations + listOf(RequestIdBuilderCustomization())
 
     override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
-        rustCrate.withModule(RustModule.Types) {
+        rustCrate.withModule(ClientRustModule.Types) {
             // Re-export RequestId in generated crate
             rust("pub use #T;", accessorTrait(codegenContext))
         }

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

@@ -8,9 +8,9 @@ package software.amazon.smithy.rustsdk
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customize.TestUtilFeature
+import software.amazon.smithy.rust.codegen.client.smithy.featureGatedConfigModule
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
-import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
@@ -19,8 +19,6 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 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.smithy.generators.LibRsCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsSection
 
 class CredentialsProviderDecorator : ClientCodegenDecorator {
     override val name: String = "CredentialsProvider"
@@ -33,13 +31,6 @@ class CredentialsProviderDecorator : ClientCodegenDecorator {
         return baseCustomizations + CredentialProviderConfig(codegenContext.runtimeConfig)
     }
 
-    override fun libRsCustomizations(
-        codegenContext: ClientCodegenContext,
-        baseCustomizations: List<LibRsCustomization>,
-    ): List<LibRsCustomization> {
-        return baseCustomizations + PubUseCredentials(codegenContext.runtimeConfig)
-    }
-
     override fun extraSections(codegenContext: ClientCodegenContext): List<AdHocCustomization> =
         listOf(
             adhocCustomization<SdkConfigSection.CopySdkConfigToClientConfig> { section ->
@@ -49,6 +40,13 @@ class CredentialsProviderDecorator : ClientCodegenDecorator {
 
     override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
         rustCrate.mergeFeature(TestUtilFeature.copy(deps = listOf("aws-credential-types/test-util")))
+
+        rustCrate.withModule(codegenContext.featureGatedConfigModule()) {
+            rust(
+                "pub use #T::Credentials;",
+                AwsRuntimeType.awsCredentialTypes(codegenContext.runtimeConfig),
+            )
+        }
     }
 }
 
@@ -95,20 +93,5 @@ class CredentialProviderConfig(runtimeConfig: RuntimeConfig) : ConfigCustomizati
     }
 }
 
-class PubUseCredentials(private val runtimeConfig: RuntimeConfig) : LibRsCustomization() {
-    override fun section(section: LibRsSection): Writable {
-        return when (section) {
-            is LibRsSection.Body -> writable {
-                rust(
-                    "pub use #T::Credentials;",
-                    AwsRuntimeType.awsCredentialTypes(runtimeConfig),
-                )
-            }
-
-            else -> emptySection
-        }
-    }
-}
-
 fun defaultProvider() =
     RuntimeType.forInlineDependency(InlineAwsDependency.forRustFile("no_credentials")).resolve("NoCredentials")

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

@@ -12,6 +12,7 @@ import software.amazon.smithy.rulesengine.language.syntax.parameters.Parameter
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.featureGatedConfigModule
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
@@ -20,12 +21,11 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection
 import software.amazon.smithy.rust.codegen.core.smithy.customize.adhocCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsSection
 import software.amazon.smithy.rust.codegen.core.util.dq
 import software.amazon.smithy.rust.codegen.core.util.extendIf
 import software.amazon.smithy.rust.codegen.core.util.thenSingletonListOf
@@ -101,13 +101,6 @@ class RegionDecorator : ClientCodegenDecorator {
         return baseCustomizations.extendIf(usesRegion(codegenContext)) { RegionConfigPlugin() }
     }
 
-    override fun libRsCustomizations(
-        codegenContext: ClientCodegenContext,
-        baseCustomizations: List<LibRsCustomization>,
-    ): List<LibRsCustomization> {
-        return baseCustomizations.extendIf(usesRegion(codegenContext)) { PubUseRegion(codegenContext.runtimeConfig) }
-    }
-
     override fun extraSections(codegenContext: ClientCodegenContext): List<AdHocCustomization> {
         return usesRegion(codegenContext).thenSingletonListOf {
             adhocCustomization<SdkConfigSection.CopySdkConfigToClientConfig> { section ->
@@ -121,6 +114,14 @@ class RegionDecorator : ClientCodegenDecorator {
         }
     }
 
+    override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
+        if (usesRegion(codegenContext)) {
+            rustCrate.withModule(codegenContext.featureGatedConfigModule()) {
+                rust("pub use #T::Region;", region(codegenContext.runtimeConfig))
+            }
+        }
+    }
+
     override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> {
         if (!usesRegion(codegenContext)) {
             return listOf()
@@ -221,19 +222,4 @@ class RegionConfigPlugin : OperationCustomization() {
     }
 }
 
-class PubUseRegion(private val runtimeConfig: RuntimeConfig) : LibRsCustomization() {
-    override fun section(section: LibRsSection): Writable {
-        return when (section) {
-            is LibRsSection.Body -> writable {
-                rust(
-                    "pub use #T::Region;",
-                    region(runtimeConfig),
-                )
-            }
-
-            else -> emptySection
-        }
-    }
-}
-
 fun region(runtimeConfig: RuntimeConfig) = AwsRuntimeType.awsTypes(runtimeConfig).resolve("region")

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

@@ -18,6 +18,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointTypesGenerator
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.EndpointsModule
+import software.amazon.smithy.rust.codegen.client.smithy.featureGatedConfigModule
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
@@ -27,12 +28,9 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
-import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 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.smithy.generators.LibRsCustomization
-import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsSection
 import software.amazon.smithy.rust.codegen.core.util.extendIf
 import software.amazon.smithy.rust.codegen.core.util.letIf
 import software.amazon.smithy.rust.codegen.core.util.thenSingletonListOf
@@ -91,14 +89,14 @@ class AwsEndpointDecorator : ClientCodegenDecorator {
         }
     }
 
-    override fun libRsCustomizations(
-        codegenContext: ClientCodegenContext,
-        baseCustomizations: List<LibRsCustomization>,
-    ): List<LibRsCustomization> {
-        return baseCustomizations + PubUseEndpoint(codegenContext.runtimeConfig)
-    }
-
     override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
+        rustCrate.withModule(codegenContext.featureGatedConfigModule()) {
+            rust(
+                "pub use #T::endpoint::Endpoint;",
+                CargoDependency.smithyHttp(codegenContext.runtimeConfig).toType(),
+            )
+        }
+
         val epTypes = EndpointTypesGenerator.fromContext(codegenContext)
         if (epTypes.defaultResolver() == null) {
             throw CodegenException(
@@ -257,21 +255,6 @@ class AwsEndpointDecorator : ClientCodegenDecorator {
             }
         }
     }
-
-    class PubUseEndpoint(private val runtimeConfig: RuntimeConfig) : LibRsCustomization() {
-        override fun section(section: LibRsSection): Writable {
-            return when (section) {
-                is LibRsSection.Body -> writable {
-                    rust(
-                        "pub use #T::endpoint::Endpoint;",
-                        CargoDependency.smithyHttp(runtimeConfig).toType(),
-                    )
-                }
-
-                else -> emptySection
-            }
-        }
-    }
 }
 
 fun ClientCodegenContext.isRegionalized() = getBuiltIn(Builtins.REGION) != null

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt

@@ -20,6 +20,9 @@ import software.amazon.smithy.rust.codegen.core.util.hasTrait
  * Modules for code generated client crates.
  */
 object ClientRustModule {
+    /** crate */
+    val root = RustModule.LibRs
+
     /** crate::client */
     val client = Client.self
     object Client {
@@ -57,3 +60,19 @@ object ClientModuleProvider : ModuleProvider {
     override fun moduleForEventStreamError(eventStream: UnionShape): RustModule.LeafModule =
         ClientRustModule.Error
 }
+
+// TODO(CrateReorganization): Remove when cleaning up `enableNewCrateOrganizationScheme`
+fun ClientCodegenContext.featureGatedConfigModule() = when (settings.codegenConfig.enableNewCrateOrganizationScheme) {
+    true -> ClientRustModule.Config
+    else -> ClientRustModule.root
+}
+
+// TODO(CrateReorganization): Remove when cleaning up `enableNewCrateOrganizationScheme`
+fun ClientCodegenContext.featureGatedCustomizeModule() = when (settings.codegenConfig.enableNewCrateOrganizationScheme) {
+    true -> ClientRustModule.Client.customize
+    else -> RustModule.public(
+        "customize",
+        "Operation customization and supporting types",
+        parent = ClientRustModule.Operation,
+    )
+}

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt

@@ -6,6 +6,7 @@
 package software.amazon.smithy.rust.codegen.client.smithy.generators
 
 import software.amazon.smithy.model.shapes.StringShape
+import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
@@ -86,7 +87,7 @@ data class InfallibleEnumType(
     }
 
     private fun unknownVariantValue(context: EnumGeneratorContext): RuntimeType {
-        return RuntimeType.forInlineFun(UnknownVariantValue, RustModule.Types) {
+        return RuntimeType.forInlineFun(UnknownVariantValue, ClientRustModule.Types) {
             docs(
                 """
                 Opaque struct used as inner data for the `Unknown` variant defined in enums in
@@ -167,4 +168,4 @@ data class InfallibleEnumType(
 }
 
 class ClientEnumGenerator(codegenContext: CodegenContext, shape: StringShape) :
-    EnumGenerator(codegenContext.model, codegenContext.symbolProvider, shape, InfallibleEnumType(RustModule.Types))
+    EnumGenerator(codegenContext.model, codegenContext.symbolProvider, shape, InfallibleEnumType(ClientRustModule.Types))

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt

@@ -5,14 +5,13 @@
 
 package software.amazon.smithy.rust.codegen.client.smithy.generators.client
 
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.featureGatedCustomizeModule
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.GenericTypeArg
 import software.amazon.smithy.rust.codegen.core.rustlang.RustGenerics
-import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
-import software.amazon.smithy.rust.codegen.core.rustlang.Visibility
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
-import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 
@@ -21,20 +20,16 @@ import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
  * fluent client builders.
  */
 class CustomizableOperationGenerator(
-    private val runtimeConfig: RuntimeConfig,
+    private val codegenContext: ClientCodegenContext,
     private val generics: FluentClientGenerics,
-    private val includeFluentClient: Boolean,
 ) {
-
-    companion object {
-        val CustomizeModule = RustModule.public("customize", "Operation customization and supporting types", parent = RustModule.operation(Visibility.PUBLIC))
-    }
-
+    private val includeFluentClient = codegenContext.settings.codegenConfig.includeFluentClient
+    private val runtimeConfig = codegenContext.runtimeConfig
     private val smithyHttp = CargoDependency.smithyHttp(runtimeConfig).toType()
     private val smithyTypes = CargoDependency.smithyTypes(runtimeConfig).toType()
 
     fun render(crate: RustCrate) {
-        crate.withModule(CustomizeModule) {
+        crate.withModule(codegenContext.featureGatedCustomizeModule()) {
             rustTemplate(
                 """
                 pub use #{Operation};
@@ -67,6 +62,7 @@ class CustomizableOperationGenerator(
             "handle_generics_bounds" to handleGenerics.bounds(),
             "operation_generics_decl" to operationGenerics.declaration(),
             "combined_generics_decl" to combinedGenerics.declaration(),
+            "customize_module" to codegenContext.featureGatedCustomizeModule(),
         )
 
         writer.rustTemplate(
@@ -81,7 +77,7 @@ class CustomizableOperationGenerator(
 
             /// A wrapper type for [`Operation`](aws_smithy_http::operation::Operation)s that allows for
             /// customization of the operation before it is sent. A `CustomizableOperation` may be sent
-            /// by calling its [`.send()`][crate::operation::customize::CustomizableOperation::send] method.
+            /// by calling its [`.send()`][#{customize_module}::CustomizableOperation::send] method.
             ##[derive(Debug)]
             pub struct CustomizableOperation#{combined_generics_decl:W} {
                 pub(crate) handle: Arc<Handle#{handle_generics_decl:W}>,