Move endpoint module and add endpoint re-exports (#2822)

jdisanti · web-flow · commit 71b401f5bf22 · 2023-06-30T17:48:41.000Z
All the types in the `crate::endpoint` module are related to endpoint
configuration, so they should be moved into `crate::config::endpoint` to
be consistent with the rest of the generated crate organization. This PR
moves them for the orchestrator implementation, but retains them in
`crate::endpoint` when generating for middleware.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt
@@ -60,6 +60,9 @@ object ClientRustModule {
         /** crate::client */
         val self = RustModule.public("config")
 
+        /** crate::config::endpoint */
+        val endpoint = RustModule.public("endpoint", parent = self)
+
         /** crate::config::retry */
         val retry = RustModule.public("retry", parent = self)
 
@@ -70,8 +73,18 @@ object ClientRustModule {
         val interceptors = RustModule.public("interceptors", parent = self)
     }
 
-    val Error = RustModule.public("error")
+    // TODO(enableNewSmithyRuntimeCleanup): Delete this root endpoint module
+    @Deprecated(message = "use the endpoint() method to get the endpoint module for now")
     val Endpoint = RustModule.public("endpoint")
+
+    // TODO(enableNewSmithyRuntimeCleanup): Just use Config.endpoint directly and delete this function
+    fun endpoint(codegenContext: ClientCodegenContext): RustModule.LeafModule = if (codegenContext.smithyRuntimeMode.defaultToMiddleware) {
+        Endpoint
+    } else {
+        Config.endpoint
+    }
+
+    val Error = RustModule.public("error")
     val Operation = RustModule.public("operation")
     val Meta = RustModule.public("meta")
     val Input = RustModule.public("input")
@@ -99,6 +112,7 @@ class ClientModuleDocProvider(
             ClientRustModule.client -> clientModuleDoc()
             ClientRustModule.Client.customize -> customizeModuleDoc()
             ClientRustModule.config -> strDoc("Configuration for $serviceName.")
+            ClientRustModule.Config.endpoint -> strDoc("Types needed to configure endpoint resolution.")
             ClientRustModule.Config.retry -> strDoc("Retry configuration.")
             ClientRustModule.Config.timeout -> strDoc("Timeout configuration.")
             ClientRustModule.Config.interceptors -> strDoc("Types needed to implement [`Interceptor`](crate::config::Interceptor).")
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointConfigCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointConfigCustomization.kt
@@ -19,7 +19,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.pre
  * Customization which injects an Endpoints 2.0 Endpoint Resolver into the service config struct
  */
 internal class EndpointConfigCustomization(
-    codegenContext: ClientCodegenContext,
+    private val codegenContext: ClientCodegenContext,
     private val typesGenerator: EndpointTypesGenerator,
 ) :
     ConfigCustomization() {
@@ -86,6 +86,8 @@ internal class EndpointConfigCustomization(
                 ServiceConfig.BuilderImpl -> {
                     // if there are no rules, we don't generate a default resolver—we need to also suppress those docs.
                     val defaultResolverDocs = if (typesGenerator.defaultResolver() != null) {
+                        val endpointModule = ClientRustModule.endpoint(codegenContext).fullyQualifiedPath()
+                            .replace("crate::", "$moduleUseName::")
                         """
                         ///
                         /// When unset, the client will used a generated endpoint resolver based on the endpoint resolution
@@ -94,7 +96,7 @@ internal class EndpointConfigCustomization(
                         /// ## Examples
                         /// ```no_run
                         /// use aws_smithy_http::endpoint;
-                        /// use $moduleUseName::endpoint::{Params as EndpointParams, DefaultResolver};
+                        /// use $endpointModule::{Params as EndpointParams, DefaultResolver};
                         /// /// Endpoint resolver which adds a prefix to the generated endpoint
                         /// ##[derive(Debug)]
                         /// struct PrefixResolver {
@@ -193,7 +195,7 @@ internal class EndpointConfigCustomization(
                         }
                     } else {
                         val alwaysFailsResolver =
-                            RuntimeType.forInlineFun("MissingResolver", ClientRustModule.Endpoint) {
+                            RuntimeType.forInlineFun("MissingResolver", ClientRustModule.endpoint(codegenContext)) {
                                 rustTemplate(
                                     """
                                     ##[derive(Debug)]
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointTypesGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointTypesGenerator.kt
@@ -40,10 +40,10 @@ class EndpointTypesGenerator(
         }
     }
 
-    fun paramsStruct(): RuntimeType = EndpointParamsGenerator(params).paramsStruct()
-    fun paramsBuilder(): RuntimeType = EndpointParamsGenerator(params).paramsBuilder()
+    fun paramsStruct(): RuntimeType = EndpointParamsGenerator(codegenContext, params).paramsStruct()
+    fun paramsBuilder(): RuntimeType = EndpointParamsGenerator(codegenContext, params).paramsBuilder()
     fun defaultResolver(): RuntimeType? =
-        rules?.let { EndpointResolverGenerator(stdlib, runtimeConfig).defaultEndpointResolver(it) }
+        rules?.let { EndpointResolverGenerator(codegenContext, stdlib).defaultEndpointResolver(it) }
 
     fun testGenerator(): Writable =
         defaultResolver()?.let {
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt
@@ -18,7 +18,7 @@ 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.endpoint.generators.CustomRuntimeFunction
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.EndpointParamsGenerator
-import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.EndpointTests
+import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.endpointTestsModule
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.rulesgen.SmithyEndpointsStdLib
 import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection
@@ -134,8 +134,8 @@ class EndpointsDecorator : ClientCodegenDecorator {
 
     override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
         val generator = EndpointTypesGenerator.fromContext(codegenContext)
-        rustCrate.withModule(ClientRustModule.Endpoint) {
-            withInlineModule(EndpointTests, rustCrate.moduleDocProvider) {
+        rustCrate.withModule(ClientRustModule.endpoint(codegenContext)) {
+            withInlineModule(endpointTestsModule(codegenContext), rustCrate.moduleDocProvider) {
                 generator.testGenerator()(this)
             }
         }
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointParamsGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointParamsGenerator.kt
@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators
 import software.amazon.smithy.rulesengine.language.eval.Value
 import software.amazon.smithy.rulesengine.language.syntax.Identifier
 import software.amazon.smithy.rulesengine.language.syntax.parameters.Parameters
+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.endpoint.memberName
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.rustName
@@ -37,12 +38,12 @@ import software.amazon.smithy.rust.codegen.core.util.dq
 import software.amazon.smithy.rust.codegen.core.util.orNull
 
 // internals contains the actual resolver function
-val EndpointImpl = RustModule.private("internals", parent = ClientRustModule.Endpoint)
+fun endpointImplModule(codegenContext: ClientCodegenContext) = RustModule.private("internals", parent = ClientRustModule.endpoint(codegenContext))
 
-val EndpointTests = RustModule.new(
+fun endpointTestsModule(codegenContext: ClientCodegenContext) = RustModule.new(
     "test",
     visibility = Visibility.PRIVATE,
-    parent = ClientRustModule.Endpoint,
+    parent = ClientRustModule.endpoint(codegenContext),
     inline = true,
     documentationOverride = "",
 ).cfgTest()
@@ -108,22 +109,24 @@ val EndpointStdLib = RustModule.private("endpoint_lib")
  *  ```
  */
 
-internal class EndpointParamsGenerator(private val parameters: Parameters) {
-
+internal class EndpointParamsGenerator(
+    private val codegenContext: ClientCodegenContext,
+    private val parameters: Parameters,
+) {
     companion object {
         fun memberName(parameterName: String) = Identifier.of(parameterName).rustName()
         fun setterName(parameterName: String) = "set_${memberName(parameterName)}"
     }
 
-    fun paramsStruct(): RuntimeType = RuntimeType.forInlineFun("Params", ClientRustModule.Endpoint) {
+    fun paramsStruct(): RuntimeType = RuntimeType.forInlineFun("Params", ClientRustModule.endpoint(codegenContext)) {
         generateEndpointsStruct(this)
     }
 
-    internal fun paramsBuilder(): RuntimeType = RuntimeType.forInlineFun("ParamsBuilder", ClientRustModule.Endpoint) {
+    internal fun paramsBuilder(): RuntimeType = RuntimeType.forInlineFun("ParamsBuilder", ClientRustModule.endpoint(codegenContext)) {
         generateEndpointParamsBuilder(this)
     }
 
-    private fun paramsError(): RuntimeType = RuntimeType.forInlineFun("InvalidParams", ClientRustModule.Endpoint) {
+    private fun paramsError(): RuntimeType = RuntimeType.forInlineFun("InvalidParams", ClientRustModule.endpoint(codegenContext)) {
         rust(
             """
             /// An error that occurred during endpoint resolution
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointResolverGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointResolverGenerator.kt
@@ -14,6 +14,7 @@ import software.amazon.smithy.rulesengine.language.syntax.fn.IsSet
 import software.amazon.smithy.rulesengine.language.syntax.rule.Condition
 import software.amazon.smithy.rulesengine.language.syntax.rule.Rule
 import software.amazon.smithy.rulesengine.language.visit.RuleValueVisitor
+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.endpoint.Context
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.Types
@@ -33,7 +34,6 @@ 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.toType
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
-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.util.dq
 import software.amazon.smithy.rust.codegen.core.util.orNull
@@ -119,7 +119,11 @@ class FunctionRegistry(private val functions: List<CustomRuntimeFunction>) {
  *
  */
 
-internal class EndpointResolverGenerator(stdlib: List<CustomRuntimeFunction>, runtimeConfig: RuntimeConfig) {
+internal class EndpointResolverGenerator(
+    private val codegenContext: ClientCodegenContext,
+    stdlib: List<CustomRuntimeFunction>,
+) {
+    private val runtimeConfig = codegenContext.runtimeConfig
     private val registry: FunctionRegistry = FunctionRegistry(stdlib)
     private val types = Types(runtimeConfig)
     private val codegenScope = arrayOf(
@@ -164,7 +168,7 @@ internal class EndpointResolverGenerator(stdlib: List<CustomRuntimeFunction>, ru
 
         // Now that we rendered the rules once (and then threw it away) we can see what functions we actually used!
         val fnsUsed = registry.fnsUsed()
-        return RuntimeType.forInlineFun("DefaultResolver", ClientRustModule.Endpoint) {
+        return RuntimeType.forInlineFun("DefaultResolver", ClientRustModule.endpoint(codegenContext)) {
             rustTemplate(
                 """
                 /// The default endpoint resolver
@@ -190,7 +194,7 @@ internal class EndpointResolverGenerator(stdlib: List<CustomRuntimeFunction>, ru
                 """,
                 "custom_fields" to fnsUsed.mapNotNull { it.structField() }.join(","),
                 "custom_fields_init" to fnsUsed.mapNotNull { it.structFieldInit() }.join(","),
-                "Params" to EndpointParamsGenerator(endpointRuleSet.parameters).paramsStruct(),
+                "Params" to EndpointParamsGenerator(codegenContext, endpointRuleSet.parameters).paramsStruct(),
                 "additional_args" to fnsUsed.mapNotNull { it.additionalArgsInvocation("self") }.join(","),
                 "resolver_fn" to resolverFn(endpointRuleSet, fnsUsed),
                 *codegenScope,
@@ -202,7 +206,7 @@ internal class EndpointResolverGenerator(stdlib: List<CustomRuntimeFunction>, ru
         endpointRuleSet: EndpointRuleSet,
         fnsUsed: List<CustomRuntimeFunction>,
     ): RuntimeType {
-        return RuntimeType.forInlineFun("resolve_endpoint", EndpointImpl) {
+        return RuntimeType.forInlineFun("resolve_endpoint", endpointImplModule(codegenContext)) {
             Attribute(allow(allowLintsForResolver)).render(this)
             rustTemplate(
                 """
@@ -212,7 +216,7 @@ internal class EndpointResolverGenerator(stdlib: List<CustomRuntimeFunction>, ru
 
                 """,
                 *codegenScope,
-                "Params" to EndpointParamsGenerator(endpointRuleSet.parameters).paramsStruct(),
+                "Params" to EndpointParamsGenerator(codegenContext, endpointRuleSet.parameters).paramsStruct(),
                 "additional_args" to fnsUsed.mapNotNull { it.additionalArgsSignature() }.join(","),
                 "body" to resolverFnBody(endpointRuleSet),
             )
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientRuntimeTypesReExportGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientRuntimeTypesReExportGenerator.kt
@@ -33,6 +33,16 @@ class ClientRuntimeTypesReExportGenerator(
                 "Interceptor" to RuntimeType.interceptor(rc),
             )
         }
+        rustCrate.withModule(ClientRustModule.endpoint(codegenContext)) {
+            rustTemplate(
+                """
+                pub use #{ResolveEndpoint};
+                pub use #{SharedEndpointResolver};
+                """,
+                "ResolveEndpoint" to RuntimeType.smithyHttp(rc).resolve("endpoint::ResolveEndpoint"),
+                "SharedEndpointResolver" to RuntimeType.smithyHttp(rc).resolve("endpoint::SharedEndpointResolver"),
+            )
+        }
         rustCrate.withModule(ClientRustModule.Config.retry) {
             rustTemplate(
                 """
diff --git a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointParamsGeneratorTest.kt b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointParamsGeneratorTest.kt
@@ -9,6 +9,7 @@ import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.MethodSource
 import software.amazon.smithy.rulesengine.testutil.TestDiscovery
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.EndpointParamsGenerator
+import software.amazon.smithy.rust.codegen.client.testutil.testClientCodegenContext
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.testutil.TestWorkspace
 import software.amazon.smithy.rust.codegen.core.testutil.compileAndTest
@@ -25,14 +26,15 @@ internal class EndpointParamsGeneratorTest {
     @MethodSource("testSuites")
     fun `generate endpoint params for provided test suites`(testSuite: TestDiscovery.RulesTestSuite) {
         val project = TestWorkspace.testProject()
+        val context = testClientCodegenContext()
         project.lib {
             unitTest("params_work") {
                 rustTemplate(
                     """
                     // this might fail if there are required fields
                     let _ = #{Params}::builder().build();
                     """,
-                    "Params" to EndpointParamsGenerator(testSuite.ruleSet().parameters).paramsStruct(),
+                    "Params" to EndpointParamsGenerator(context, testSuite.ruleSet().parameters).paramsStruct(),
                 )
             }
         }
diff --git a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointResolverGeneratorTest.kt b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointResolverGeneratorTest.kt
@@ -53,15 +53,16 @@ class EndpointResolverGeneratorTest {
             // return
         }
         val project = TestWorkspace.testProject()
+        val context = testClientCodegenContext()
         suite.ruleSet().typecheck()
         project.lib {
             val ruleset = EndpointResolverGenerator(
+                context,
                 SmithyEndpointsStdLib + awsStandardLib(TestRuntimeConfig, partitionsJson),
-                TestRuntimeConfig,
             ).defaultEndpointResolver(suite.ruleSet())
             val testGenerator = EndpointTestGenerator(
                 suite.testSuite().testCases,
-                paramsType = EndpointParamsGenerator(suite.ruleSet().parameters).paramsStruct(),
+                paramsType = EndpointParamsGenerator(context, suite.ruleSet().parameters).paramsStruct(),
                 resolverType = ruleset,
                 suite.ruleSet().parameters,
                 codegenContext = testClientCodegenContext(model = Model.builder().build()),
@@ -79,15 +80,16 @@ class EndpointResolverGeneratorTest {
             testSuites().filter { it.ruleSet().sourceLocation.filename.endsWith("/uri-encode.json") }.findFirst()
                 .orElseThrow()
         val project = TestWorkspace.testProject()
+        val context = testClientCodegenContext()
         suite.ruleSet().typecheck()
         project.lib {
             val ruleset = EndpointResolverGenerator(
+                context,
                 SmithyEndpointsStdLib + awsStandardLib(TestRuntimeConfig, partitionsJson),
-                TestRuntimeConfig,
             ).defaultEndpointResolver(suite.ruleSet())
             val testGenerator = EndpointTestGenerator(
                 suite.testSuite().testCases,
-                paramsType = EndpointParamsGenerator(suite.ruleSet().parameters).paramsStruct(),
+                paramsType = EndpointParamsGenerator(context, suite.ruleSet().parameters).paramsStruct(),
                 resolverType = ruleset,
                 suite.ruleSet().parameters,
                 codegenContext = testClientCodegenContext(Model.builder().build()),
@@ -115,7 +117,8 @@ class EndpointResolverGeneratorTest {
         val scope = Scope<Type>()
         scope.insert("Region", Type.string())
         endpoint.typeCheck(scope)
-        val generator = EndpointResolverGenerator(listOf(), TestRuntimeConfig)
+        val context = testClientCodegenContext()
+        val generator = EndpointResolverGenerator(context, listOf())
         TestWorkspace.testProject().unitTest {
             rustTemplate(
                 """

Original file line number	Diff line number	Diff line change
`@@ -40,10 +40,10 @@ class EndpointTypesGenerator(`
`40`	`40`	`}`
`41`	`41`	`}`
`42`	`42`
`43`		`- fun paramsStruct(): RuntimeType = EndpointParamsGenerator(params).paramsStruct()`
`44`		`- fun paramsBuilder(): RuntimeType = EndpointParamsGenerator(params).paramsBuilder()`
	`43`	`+ fun paramsStruct(): RuntimeType = EndpointParamsGenerator(codegenContext, params).paramsStruct()`
	`44`	`+ fun paramsBuilder(): RuntimeType = EndpointParamsGenerator(codegenContext, params).paramsBuilder()`
`45`	`45`	`fun defaultResolver(): RuntimeType? =`
`46`		`- rules?.let { EndpointResolverGenerator(stdlib, runtimeConfig).defaultEndpointResolver(it) }`
	`46`	`+ rules?.let { EndpointResolverGenerator(codegenContext, stdlib).defaultEndpointResolver(it) }`
`47`	`47`
`48`	`48`	`fun testGenerator(): Writable =`
`49`	`49`	`defaultResolver()?.let {`