Add AwsEndpointResolver when constructing operations (#198)

* Add AwsEndpointResolver when constructing operations

This commit adds a customization for AWS-services to allow specifying an EndpointResolver, with a default fallback provided.

* Enforce no doc warnings, fix bug, add cargoCheck to CI

.github/workflows/ci.yaml

@@ -147,6 +147,8 @@ jobs:
       # docs are not included in the artifact; this step validates that they can be generated
     - name: Generate docs
       run: ./gradlew :aws:sdk:cargoDocs
+    - name: Run tests
+      run: ./gradlew :aws:sdk:cargoTest
     - name: Get current date
       id: date
       run: echo "name=${GITHUB_REF##*/}-$(date +'%Y-%m-%d')" >> $GITHUB_ENV

aws/rust-runtime/aws-endpoint/src/lib.rs

@@ -26,6 +26,12 @@ pub struct AwsEndpoint {
     signing_region: Option<SigningRegion>,
 }
 
+impl AwsEndpoint {
+    pub fn set_endpoint(&self, mut uri: &mut http::Uri, endpoint_prefix: Option<&EndpointPrefix>) {
+        self.endpoint.set_endpoint(&mut uri, endpoint_prefix);
+    }
+}
+
 pub type BoxError = Box<dyn Error + Send + Sync + 'static>;
 
 /// Resolve the AWS Endpoint for a given region
@@ -109,7 +115,7 @@ fn get_endpoint_resolver(config: &PropertyBag) -> Option<&AwsEndpointResolver> {
     config.get()
 }
 
-pub fn set_endpoint_resolver(provider: AwsEndpointResolver, config: &mut PropertyBag) {
+pub fn set_endpoint_resolver(config: &mut PropertyBag, provider: AwsEndpointResolver) {
     config.insert(provider);
 }
 
@@ -117,7 +123,7 @@ pub fn set_endpoint_resolver(provider: AwsEndpointResolver, config: &mut Propert
 ///
 /// AwsEndpointStage implements [`MapRequest`](smithy_http::middleware::MapRequest). It will:
 /// 1. Load an endpoint provider from the property bag.
-/// 2. Load an endpoint given the [`Region`](aws_types::Region) in the property bag.
+/// 2. Load an endpoint given the [`Region`](aws_types::region::Region) in the property bag.
 /// 3. Apply the endpoint to the URI in the request
 /// 4. Set the `SigningRegion` and `SigningService` in the property bag to drive downstream
 /// signing middleware.
@@ -188,7 +194,7 @@ mod test {
         {
             let mut conf = req.config_mut();
             conf.insert(region.clone());
-            set_endpoint_resolver(provider, &mut conf);
+            set_endpoint_resolver(&mut conf, provider);
         };
         let req = AwsEndpointStage.apply(req).expect("should succeed");
         assert_eq!(

aws/rust-runtime/aws-sig-auth/src/middleware.rs

@@ -140,7 +140,7 @@ mod test {
             .augment(|req, conf| {
                 conf.insert(region.clone());
                 conf.insert(UNIX_EPOCH + Duration::new(1611160427, 0));
-                set_endpoint_resolver(provider, conf);
+                set_endpoint_resolver(conf, provider);
                 Result::<_, Infallible>::Ok(req)
             })
             .expect("succeeds");

aws/rust-runtime/aws-types/src/region.rs

@@ -28,6 +28,12 @@ impl Region {
     }
 }
 
+impl From<&str> for Region {
+    fn from(region: &str) -> Self {
+        Region(Arc::new(region.to_string()))
+    }
+}
+
 /// Provide a [`Region`](Region) to use with AWS requests
 ///
 /// For most cases [`default_provider`](default_provider) will be the best option, implementing

aws/sdk-codegen/build.gradle.kts

@@ -17,6 +17,7 @@ group = "software.amazon.software.amazon.smithy.rust.codegen.smithy"
 version = "0.1.0"
 
 val smithyVersion: String by project
+val kotestVersion: String by project
 
 dependencies {
     implementation(project(":codegen"))
@@ -24,6 +25,7 @@ dependencies {
     implementation("software.amazon.smithy:smithy-protocol-test-traits:$smithyVersion")
     implementation("software.amazon.smithy:smithy-aws-traits:$smithyVersion")
     testImplementation("org.junit.jupiter:junit-jupiter:5.6.1")
+    testImplementation("io.kotest:kotest-assertions-core-jvm:$kotestVersion")
 }
 
 tasks.compileKotlin {

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

@@ -9,7 +9,8 @@ import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecor
 
 val DECORATORS = listOf(
     CredentialsProviderDecorator(),
-    RegionDecorator()
+    RegionDecorator(),
+    AwsEndpointDecorator()
 )
 
 class AwsCodegenDecorator : CombinedCodegenDecorator(DECORATORS) {

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

@@ -0,0 +1,95 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package software.amazon.smithy.rustsdk
+
+import software.amazon.smithy.aws.traits.ServiceTrait
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.rustlang.Local
+import software.amazon.smithy.rust.codegen.rustlang.Writable
+import software.amazon.smithy.rust.codegen.rustlang.asType
+import software.amazon.smithy.rust.codegen.rustlang.rust
+import software.amazon.smithy.rust.codegen.rustlang.writable
+import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
+import software.amazon.smithy.rust.codegen.smithy.generators.OperationCustomization
+import software.amazon.smithy.rust.codegen.smithy.generators.OperationSection
+import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig
+import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustomization
+import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig
+import software.amazon.smithy.rust.codegen.util.dq
+
+class AwsEndpointDecorator : RustCodegenDecorator {
+    override val name: String = "AwsEndpoint"
+    override val order: Byte = 0
+
+    override fun configCustomizations(
+        protocolConfig: ProtocolConfig,
+        baseCustomizations: List<ConfigCustomization>
+    ): List<ConfigCustomization> {
+        return baseCustomizations + EndpointConfigCustomization(protocolConfig.runtimeConfig, protocolConfig.serviceShape)
+    }
+
+    override fun operationCustomizations(
+        protocolConfig: ProtocolConfig,
+        operation: OperationShape,
+        baseCustomizations: List<OperationCustomization>
+    ): List<OperationCustomization> {
+        return baseCustomizations + EndpointResolverFeature(protocolConfig.runtimeConfig, operation)
+    }
+}
+
+class EndpointConfigCustomization(private val runtimeConfig: RuntimeConfig, serviceShape: ServiceShape) : ConfigCustomization() {
+    private val endpointPrefix = serviceShape.expectTrait(ServiceTrait::class.java).endpointPrefix
+    private val resolveAwsEndpoint = runtimeConfig.awsEndpointDependency().asType().copy(name = "ResolveAwsEndpoint")
+    override fun section(section: ServiceConfig): Writable = writable {
+        when (section) {
+            is ServiceConfig.ConfigStruct -> rust("pub endpoint_resolver: ::std::sync::Arc<dyn #T>,", resolveAwsEndpoint)
+            is ServiceConfig.ConfigImpl -> emptySection
+            is ServiceConfig.BuilderStruct ->
+                rust("endpoint_resolver: Option<::std::sync::Arc<dyn #T>>,", resolveAwsEndpoint)
+            ServiceConfig.BuilderImpl ->
+                rust(
+                    """
+            pub fn endpoint_resolver(mut self, endpoint_resolver: impl #T + 'static) -> Self {
+                self.endpoint_resolver = Some(::std::sync::Arc::new(endpoint_resolver));
+                self
+            }
+            """,
+                    resolveAwsEndpoint
+                )
+            ServiceConfig.BuilderBuild -> rust(
+                """endpoint_resolver: self.endpoint_resolver.unwrap_or_else(||
+                                ::std::sync::Arc::new(
+                                    #T::DefaultAwsEndpointResolver::for_service(${endpointPrefix.dq()})
+                                )
+                         ),""",
+                runtimeConfig.awsEndpointDependency().asType()
+            )
+        }
+    }
+}
+
+// This is an experiment in a slightly different way to create runtime types. All code MAY be refactored to use this pattern
+fun RuntimeConfig.awsEndpointDependency() = CargoDependency("aws-endpoint", Local(this.relativePath))
+
+class EndpointResolverFeature(private val runtimeConfig: RuntimeConfig, private val operationShape: OperationShape) :
+    OperationCustomization() {
+    override fun section(section: OperationSection): Writable {
+        return when (section) {
+            OperationSection.ImplBlock -> emptySection
+            is OperationSection.Feature -> writable {
+                rust(
+                    """
+                #T::set_endpoint_resolver(&mut ${section.request}.config_mut(), ${section.config}.endpoint_resolver.clone());
+                """,
+                    runtimeConfig.awsEndpointDependency().asType()
+                )
+            }
+        }
+    }
+}

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

@@ -1,34 +0,0 @@
-/*
- * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
- * SPDX-License-Identifier: Apache-2.0.
- */
-
-package software.amazon.smithy.rustsdk
-
-import software.amazon.smithy.rust.codegen.rustlang.Attribute
-import software.amazon.smithy.rust.codegen.rustlang.rust
-import software.amazon.smithy.rust.codegen.rustlang.writable
-import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustomization
-import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig
-
-/**
- * Just a Stub
- *
- * Augment the config object with the AWS-specific fields like service and region
- */
-class BaseAwsConfig : ConfigCustomization() {
-    override fun section(section: ServiceConfig) = writable {
-        when (section) {
-            ServiceConfig.ConfigStruct -> {
-                Attribute.AllowUnused.render(this)
-                rust("pub(crate) region: String,")
-            }
-            ServiceConfig.BuilderBuild -> rust("region: \"todo\".to_owned(),")
-            else -> {}
-            /*ServiceConfig.ConfigImpl -> TODO()
-            ServiceConfig.BuilderStruct -> TODO()
-            ServiceConfig.BuilderImpl -> TODO()
-            ServiceConfig.BuilderBuild -> TODO()*/
-        }
-    }
-}

aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/EndpointConfigCustomizationTest.kt

@@ -0,0 +1,69 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package software.amazon.smithy.rustsdk
+
+import io.kotest.matchers.shouldBe
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.aws.traits.ServiceTrait
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.testutil.TestRuntimeConfig
+import software.amazon.smithy.rust.codegen.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.testutil.compileAndTest
+import software.amazon.smithy.rust.codegen.testutil.stubConfigProject
+import software.amazon.smithy.rust.codegen.testutil.unitTest
+import software.amazon.smithy.rust.codegen.testutil.validateConfigCustomizations
+import software.amazon.smithy.rust.codegen.util.lookup
+
+internal class EndpointConfigCustomizationTest {
+
+    private val model = """
+    namespace test
+    @aws.api#service(sdkId: "Test", endpointPrefix: "differentprefix")
+    service TestService {
+        version: "123"
+    }
+
+    @aws.api#service(sdkId: "Test")
+    service NoEndpointPrefix {
+        version: "123"
+    }
+    """.asSmithyModel()
+
+    @Test
+    fun `generates valid code`() {
+        validateConfigCustomizations(EndpointConfigCustomization(TestRuntimeConfig, model.lookup("test#TestService")))
+    }
+
+    @Test
+    fun `generates valid code when no endpoint prefix is provided`() {
+        val serviceShape = model.lookup<ServiceShape>("test#NoEndpointPrefix")
+        validateConfigCustomizations(EndpointConfigCustomization(TestRuntimeConfig, serviceShape))
+        serviceShape.expectTrait(ServiceTrait::class.java).endpointPrefix shouldBe "noendpointprefix"
+    }
+
+    @Test
+    fun `write an endpoint into the config`() {
+        val project = stubConfigProject(EndpointConfigCustomization(TestRuntimeConfig, model.lookup("test#TestService")))
+        project.useFileWriter("src/lib.rs", "crate") {
+            it.addDependency(awsTypes(TestRuntimeConfig))
+            it.addDependency(CargoDependency.Http)
+            it.unitTest(
+                """
+                use aws_types::region::Region;
+                use http::Uri;
+                let conf = crate::config::Config::builder().build();
+                let endpoint = conf.endpoint_resolver
+                    .endpoint(&Region::from("us-east-1")).expect("default resolver produces a valid endpoint");
+                let mut uri = Uri::from_static("/?k=v");
+                endpoint.set_endpoint(&mut uri, None);
+                assert_eq!(uri, Uri::from_static("https://us-east-1.differentprefix.amazonaws.com/?k=v"));
+            """
+            )
+        }
+        project.compileAndTest()
+    }
+}

aws/sdk/build.gradle.kts

@@ -179,7 +179,7 @@ tasks.register<Exec>("cargoTest") {
 tasks.register<Exec>("cargoDocs") {
     workingDir(sdkOutputDir)
     // disallow warnings
-    environment("RUSTFLAGS", "-D warnings")
+    environment("RUSTDOCFLAGS", "-D warnings")
     commandLine("cargo", "doc", "--no-deps")
     dependsOn("assemble")
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/CargoDependency.kt

@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.rustlang
 import software.amazon.smithy.codegen.core.SymbolDependency
 import software.amazon.smithy.codegen.core.SymbolDependencyContainer
 import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.util.dq
 
 sealed class DependencyScope
@@ -101,6 +102,9 @@ class InlineDependency(
     }
 }
 
+fun CargoDependency.asType(): RuntimeType =
+    RuntimeType(null, dependency = this, namespace = this.name.replace("-", "_"))
+
 /**
  * A dependency on an internal or external Cargo Crate
  */

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/HttpProtocolGenerator.kt

@@ -10,6 +10,7 @@ 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.model.shapes.StructureShape
+import software.amazon.smithy.model.traits.EndpointTrait
 import software.amazon.smithy.rust.codegen.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.rustlang.documentShape
 import software.amazon.smithy.rust.codegen.rustlang.rustBlock
@@ -48,6 +49,9 @@ abstract class HttpProtocolGenerator(
     private val symbolProvider = protocolConfig.symbolProvider
     private val model = protocolConfig.model
     fun renderOperation(operationWriter: RustWriter, inputWriter: RustWriter, operationShape: OperationShape, customizations: List<OperationCustomization>) {
+        if (operationShape.hasTrait(EndpointTrait::class.java)) {
+            TODO("https://github.com/awslabs/smithy-rs/issues/197")
+        }
         val inputShape = operationShape.inputShape(model)
         val inputSymbol = symbolProvider.toSymbol(inputShape)
         val builderGenerator = OperationInputBuilderGenerator(model, symbolProvider, operationShape, customizations)

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/testutil/Rust.kt

@@ -165,7 +165,7 @@ fun TestWriterDelegator.compileAndTest() {
             model = stubModel
         )
     )
-    "cargo test".runCommand(baseDir)
+    "cargo test".runCommand(baseDir, mapOf("RUSTFLAGS" to "-A dead_code"))
 }
 
 // TODO: unify these test helpers a bit

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/testutil/TestConfigCustomization.kt

@@ -42,7 +42,13 @@ fun stubCustomization(name: String): ConfigCustomization {
  * This test is not comprehensive, but it ensures that your customization generates Rust code that compiles and correctly
  * composes with other customizations.
  * */
-fun validateConfigCustomizations(vararg customization: ConfigCustomization) {
+fun validateConfigCustomizations(vararg customization: ConfigCustomization): TestWriterDelegator {
+    val project = stubConfigProject(*customization)
+    project.compileAndTest()
+    return project
+}
+
+fun stubConfigProject(vararg customization: ConfigCustomization): TestWriterDelegator {
     val customizations = listOf(stubCustomization("a")) + customization.toList() + stubCustomization("b")
     val generator = ServiceConfigGenerator(customizations = customizations.toList())
     val project = TestWorkspace.testProject()
@@ -56,5 +62,5 @@ fun validateConfigCustomizations(vararg customization: ConfigCustomization) {
         """
         )
     }
-    project.compileAndTest()
+    return project
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/util/Exec.kt

@@ -11,15 +11,18 @@ import java.util.concurrent.TimeUnit
 
 class CommandFailed(output: String) : Exception("Command Failed\n$output")
 
-fun String.runCommand(workdir: Path? = null): String {
+fun String.runCommand(workdir: Path? = null, environment: Map<String, String> = mapOf()): String {
     val parts = this.split("\\s".toRegex())
-    val proc = ProcessBuilder(*parts.toTypedArray())
+    val builder = ProcessBuilder(*parts.toTypedArray())
         .redirectOutput(ProcessBuilder.Redirect.PIPE)
         .redirectError(ProcessBuilder.Redirect.PIPE)
         .letIf(workdir != null) {
             it.directory(workdir?.toFile())
         }
-        .start()
+
+    val env = builder.environment()
+    environment.forEach { (k, v) -> env[k] = v }
+    val proc = builder.start()
 
     proc.waitFor(60, TimeUnit.MINUTES)
     val stdErr = proc.errorStream.bufferedReader().readText()

design/src/operation.md

@@ -48,7 +48,7 @@ pub fn build(self, config: &dynamodb::config::Config) -> Operation<BatchExecuteS
     let mut req = operation::Request::new(req);
     let mut conf = req.config_mut();
     conf.insert_signing_config(config.signing_service());
-    conf.insert_endpoint_provider(config.endpoint_provider.clone());
+    conf.insert_endpoint_resolver(config.endpoint_resolver.clone());
     Operation::new(req)
 }
 ```