Add support for Endpoints 2.0 Parameters

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

@@ -6,8 +6,11 @@
 package software.amazon.smithy.rustsdk
 
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.rulesengine.language.syntax.parameters.Builtins
+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.RustCodegenDecorator
+import software.amazon.smithy.rust.codegen.client.smithy.endpoint.RulesEngineBuiltInResolver
 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.client.smithy.generators.protocol.ClientProtocolGenerator
@@ -97,6 +100,22 @@ class RegionDecorator : RustCodegenDecorator<ClientProtocolGenerator, ClientCode
         return baseCustomizations + PubUseRegion(codegenContext.runtimeConfig)
     }
 
+    override fun builtInResolvers(codegenContext: ClientCodegenContext): List<RulesEngineBuiltInResolver> {
+        return listOf(
+            object : RulesEngineBuiltInResolver {
+                override fun defaultFor(
+                    parameter: Parameter,
+                    configRef: String,
+                ): Writable? {
+                    return when (parameter) {
+                        Builtins.REGION -> writable { rust("$configRef.region.as_ref().map(|r|r.as_ref())") }
+                        else -> null
+                    }
+                }
+            },
+        )
+    }
+
     override fun supportsCodegenContext(clazz: Class<out CodegenContext>): Boolean =
         clazz.isAssignableFrom(ClientCodegenContext::class.java)
 }
@@ -117,8 +136,10 @@ class RegionProviderConfig(codegenContext: CodegenContext) : ConfigCustomization
                 """,
                 *codegenScope,
             )
+
             is ServiceConfig.BuilderStruct ->
                 rustTemplate("region: Option<#{Region}>,", *codegenScope)
+
             ServiceConfig.BuilderImpl ->
                 rustTemplate(
                     """
@@ -162,6 +183,7 @@ class RegionConfigPlugin : OperationCustomization() {
                     """,
                 )
             }
+
             else -> emptySection
         }
     }
@@ -176,6 +198,7 @@ class PubUseRegion(private val runtimeConfig: RuntimeConfig) : LibRsCustomizatio
                     region(runtimeConfig),
                 )
             }
+
             else -> emptySection
         }
     }

codegen-client-test/build.gradle.kts

@@ -88,6 +88,7 @@ val allCodegenTests = "../codegen-core/common-test-models".let { commonModels ->
             """.trimIndent(),
             imports = listOf("$commonModels/naming-obstacle-course-structs.smithy"),
         ),
+        CodegenTest("aws.protocoltests.json#TestService", "endpoint-rules"),
         CodegenTest("com.aws.example.rust#PokemonService", "pokemon-service-client", imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy")),
     )
 }

codegen-client-test/model/endpoint-rules.smithy

@@ -0,0 +1,33 @@
+$version: "1.0"
+
+namespace aws.protocoltests.json
+
+use smithy.rules#endpointRuleSet
+use smithy.rules#endpointTests
+
+use smithy.rules#clientContextParams
+use smithy.rules#staticContextParams
+use smithy.rules#contextParam
+use aws.protocols#awsJson1_1
+
+@awsJson1_1
+@endpointRuleSet({
+    "version": "1.0",
+    "rules": [],
+    "parameters": {
+        "Bucket": { "required": false, "type": "String" },
+        "Region": { "required": false, "type": "String", "builtIn": "AWS::Region" },
+    }
+})
+service TestService {
+    operations: [TestOperation]
+}
+
+operation TestOperation {
+    input: TestOperationInput
+}
+
+structure TestOperationInput {
+    @contextParam(name: "Bucket")
+    bucket: String
+}

codegen-client/build.gradle.kts

@@ -29,6 +29,7 @@ dependencies {
     implementation("software.amazon.smithy:smithy-aws-traits:$smithyVersion")
     implementation("software.amazon.smithy:smithy-protocol-test-traits:$smithyVersion")
     implementation("software.amazon.smithy:smithy-waiters:$smithyVersion")
+    implementation("software.amazon.smithy:smithy-rules-engine:$smithyVersion")
     runtimeOnly(project(":rust-runtime"))
     testImplementation("org.junit.jupiter:junit-jupiter:5.6.1")
     testImplementation("io.kotest:kotest-assertions-core-jvm:$kotestVersion")

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

@@ -8,6 +8,8 @@ package software.amazon.smithy.rust.codegen.client.smithy
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.rust.codegen.client.smithy.customize.RustCodegenDecorator
+import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.ClientProtocolGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget
 import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
@@ -24,6 +26,9 @@ data class ClientCodegenContext(
     override val serviceShape: ServiceShape,
     override val protocol: ShapeId,
     override val settings: ClientRustSettings,
+    // Expose the `rootDecorator`, enabling customizations to compose by referencing information from the root codegen
+    // decorator
+    val rootDecorator: RustCodegenDecorator<ClientProtocolGenerator, ClientCodegenContext>,
 ) : CodegenContext(
     model, symbolProvider, serviceShape, protocol, settings, CodegenTarget.CLIENT,
 )

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

@@ -80,7 +80,7 @@ class CodegenVisitor(
         model = codegenDecorator.transformModel(service, baseModel)
         symbolProvider = RustCodegenPlugin.baseSymbolProvider(model, service, symbolVisitorConfig)
 
-        codegenContext = ClientCodegenContext(model, symbolProvider, service, protocol, settings)
+        codegenContext = ClientCodegenContext(model, symbolProvider, service, protocol, settings, codegenDecorator)
 
         val clientPublicModules = setOf(
             RustModule.Error,

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

@@ -10,6 +10,7 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.rust.codegen.client.smithy.endpoint.RulesEngineBuiltInResolver
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
@@ -69,6 +70,8 @@ interface RustCodegenDecorator<T, C : CodegenContext> {
 
     fun transformModel(service: ServiceShape, model: Model): Model = model
 
+    fun builtInResolvers(codegenContext: C): List<RulesEngineBuiltInResolver> = listOf()
+
     fun supportsCodegenContext(clazz: Class<out CodegenContext>): Boolean
 }
 
@@ -141,6 +144,10 @@ open class CombinedCodegenDecorator<T, C : CodegenContext>(decorators: List<Rust
         }
     }
 
+    override fun builtInResolvers(codegenContext: C): List<RulesEngineBuiltInResolver> {
+        return orderedDecorators.flatMap { it.builtInResolvers(codegenContext) }
+    }
+
     override fun supportsCodegenContext(clazz: Class<out CodegenContext>): Boolean =
         // `CombinedCodegenDecorator` can work with all types of codegen context.
         CodegenContext::class.java.isAssignableFrom(clazz)

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

@@ -0,0 +1,97 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.endpoint
+
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.model.shapes.BooleanShape
+import software.amazon.smithy.model.shapes.ShapeType
+import software.amazon.smithy.model.shapes.StringShape
+import software.amazon.smithy.rulesengine.traits.ClientContextParamDefinition
+import software.amazon.smithy.rulesengine.traits.ClientContextParamsTrait
+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.RustReservedWords
+import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.docs
+import software.amazon.smithy.rust.codegen.core.rustlang.rust
+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.RustSymbolProvider
+import software.amazon.smithy.rust.codegen.core.smithy.makeOptional
+import software.amazon.smithy.rust.codegen.core.util.getTrait
+import software.amazon.smithy.rust.codegen.core.util.orNull
+import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
+
+/**
+ * This decorator adds `ClientContextParams` to the service config.
+ *
+ * This handles injecting parameters like `s3::Accelerate` or `s3::ForcePathStyle`
+ */
+class ClientContextDecorator(ctx: CodegenContext) : ConfigCustomization() {
+    private val contextParams = ctx.serviceShape.getTrait<ClientContextParamsTrait>()?.parameters.orEmpty().toList()
+        .map { (key, value) -> ContextParam.fromClientParam(key, value, ctx.symbolProvider) }
+
+    data class ContextParam(val name: String, val type: Symbol, val docs: String?) {
+        companion object {
+            private fun toSymbol(shapeType: ShapeType, symbolProvider: RustSymbolProvider): Symbol =
+                symbolProvider.toSymbol(
+                    when (shapeType) {
+                        ShapeType.STRING -> StringShape.builder().id("smithy.api#String").build()
+                        ShapeType.BOOLEAN -> BooleanShape.builder().id("smithy.api#Boolean").build()
+                        else -> TODO("unsupported type")
+                    },
+                )
+
+            fun fromClientParam(
+                name: String,
+                definition: ClientContextParamDefinition,
+                symbolProvider: RustSymbolProvider,
+            ): ContextParam {
+                return ContextParam(
+                    RustReservedWords.escapeIfNeeded(name.toSnakeCase()),
+                    toSymbol(definition.type, symbolProvider),
+                    definition.documentation.orNull(),
+                )
+            }
+        }
+    }
+
+    override fun section(section: ServiceConfig): Writable {
+        return when (section) {
+            is ServiceConfig.ConfigStruct -> writable {
+                contextParams.forEach { param ->
+                    rust("pub (crate) ${param.name}: #T,", param.type.makeOptional())
+                }
+            }
+            ServiceConfig.ConfigImpl -> emptySection
+            ServiceConfig.BuilderStruct -> writable {
+                contextParams.forEach { param ->
+                    rust("${param.name}: #T,", param.type.makeOptional())
+                }
+            }
+            ServiceConfig.BuilderImpl -> writable {
+                contextParams.forEach { param ->
+                    param.docs?.also { docs(it) }
+                    rust(
+                        """
+                        pub fn ${param.name}(mut self, ${param.name}: impl Into<#T>) -> Self {
+                            self.${param.name} = Some(${param.name}.into());
+                            self
+                        }
+                        """,
+                        param.type,
+                    )
+                }
+            }
+            ServiceConfig.BuilderBuild -> writable {
+                contextParams.forEach { param ->
+                    rust("${param.name}: self.${param.name},")
+                }
+            }
+            else -> emptySection
+        }
+    }
+}