Feature: Customizable Operations

CHANGELOG.next.toml

@@ -87,7 +87,7 @@ author = "ogudavid"
 [[smithy-rs]]
 message = "Add codegen version to generated package metadata"
 references = ["smithy-rs#1612"]
-meta = { "breaking" = false, "tada" = false, "bug" = false }
+meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "all" }
 author = "unexge"
 
 [[smithy-rs]]
@@ -131,7 +131,7 @@ Then all runtime crates are set with version 0.47.0 by default unless overridden
 implies that we're using `aws-smithy-http` 0.47.1 specifically. For the rest of the crates, it will default to 0.47.0.
 """
 references = ["smithy-rs#1635", "smithy-rs#1416"]
-meta = { "breaking" = true, "tada" = true, "bug" = false }
+meta = { "breaking" = true, "tada" = true, "bug" = false, "target" = "all" }
 author = "weihanglo"
 
 [[smithy-rs]]
@@ -145,6 +145,102 @@ references = ["smithy-rs#1544"]
 meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "server" }
 author = "82marbag"
 
+[[aws-sdk-rust]]
+message = """
+Implemented customizable operations per [RFC-0017](https://awslabs.github.io/smithy-rs/design/rfcs/rfc0017_customizable_client_operations.html).
+
+Before this change, modifying operations before sending them required using lower-level APIs:
+
+```rust
+let input = SomeOperationInput::builder().some_value(5).build()?;
+
+let operation = {
+    let op = input.make_operation(&service_config).await?;
+    let (request, response) = op.into_request_response();
+
+    let request = request.augment(|req, _props| {
+        req.headers_mut().insert(
+            HeaderName::from_static("x-some-header"),
+            HeaderValue::from_static("some-value")
+        );
+        Result::<_, Infallible>::Ok(req)
+    })?;
+
+    Operation::from_parts(request, response)
+};
+
+let response = smithy_client.call(operation).await?;
+```
+
+Now, users may easily modify operations before sending with the `customize` method:
+
+```rust
+let response = client.some_operation()
+    .some_value(5)
+    .customize()
+    .await?
+    .mutate_request(|mut req| {
+        req.headers_mut().insert(
+            HeaderName::from_static("x-some-header"),
+            HeaderValue::from_static("some-value")
+        );
+    })
+    .send()
+    .await?;
+```
+"""
+references = ["smithy-rs#1647", "smithy-rs#1112"]
+meta = { "breaking" = false, "tada" = true, "bug" = false }
+author = "Velfi"
+
+[[smithy-rs]]
+message = """
+Implemented customizable operations per [RFC-0017](https://awslabs.github.io/smithy-rs/design/rfcs/rfc0017_customizable_client_operations.html).
+
+Before this change, modifying operations before sending them required using lower-level APIs:
+
+```rust
+let input = SomeOperationInput::builder().some_value(5).build()?;
+
+let operation = {
+    let op = input.make_operation(&service_config).await?;
+    let (request, response) = op.into_request_response();
+
+    let request = request.augment(|req, _props| {
+        req.headers_mut().insert(
+            HeaderName::from_static("x-some-header"),
+            HeaderValue::from_static("some-value")
+        );
+        Result::<_, Infallible>::Ok(req)
+    })?;
+
+    Operation::from_parts(request, response)
+};
+
+let response = smithy_client.call(operation).await?;
+```
+
+Now, users may easily modify operations before sending with the `customize` method:
+
+```rust
+let response = client.some_operation()
+    .some_value(5)
+    .customize()
+    .await?
+    .mutate_request(|mut req| {
+        req.headers_mut().insert(
+            HeaderName::from_static("x-some-header"),
+            HeaderValue::from_static("some-value")
+        );
+    })
+    .send()
+    .await?;
+```
+"""
+references = ["smithy-rs#1647", "smithy-rs#1112"]
+meta = { "breaking" = false, "tada" = true, "bug" = false, "target" = "client"}
+author = "Velfi"
+
 [[smithy-rs]]
 message = """
 There is a canonical and easier way to run smithy-rs on Lambda [see example].

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

@@ -25,6 +25,7 @@ import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
+import software.amazon.smithy.rust.codegen.smithy.generators.GenericsGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsSection
 import software.amazon.smithy.rust.codegen.smithy.generators.client.FluentClientCustomization
@@ -33,6 +34,7 @@ import software.amazon.smithy.rust.codegen.smithy.generators.client.FluentClient
 import software.amazon.smithy.rust.codegen.smithy.generators.client.FluentClientSection
 import software.amazon.smithy.rust.codegen.util.expectTrait
 import software.amazon.smithy.rustsdk.AwsRuntimeType.defaultMiddleware
+import java.util.Optional
 
 private class Types(runtimeConfig: RuntimeConfig) {
     private val smithyClientDep = CargoDependency.SmithyClient(runtimeConfig)
@@ -71,6 +73,10 @@ private class AwsClientGenerics(private val types: Types) : FluentClientGenerics
 
     /** Bounds for generated `send()` functions */
     override fun sendBounds(input: Symbol, output: Symbol, error: RuntimeType): Writable = writable { }
+
+    override fun toGenericsGenerator(): GenericsGenerator {
+        return GenericsGenerator(mutableListOf())
+    }
 }
 
 class AwsFluentClientDecorator : RustCodegenDecorator<ClientCodegenContext> {
@@ -80,15 +86,20 @@ class AwsFluentClientDecorator : RustCodegenDecorator<ClientCodegenContext> {
     override val order: Byte = (AwsPresigningDecorator.ORDER + 1).toByte()
 
     override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
-        val types = Types(codegenContext.runtimeConfig)
+        val runtimeConfig = codegenContext.runtimeConfig
+        val types = Types(runtimeConfig)
+        val generics = AwsClientGenerics(types)
         FluentClientGenerator(
             codegenContext,
-            generics = AwsClientGenerics(types),
+            generics,
             customizations = listOf(
-                AwsPresignedFluentBuilderMethod(codegenContext.runtimeConfig),
+                AwsPresignedFluentBuilderMethod(runtimeConfig),
                 AwsFluentClientDocs(codegenContext),
             ),
         ).render(rustCrate)
+        rustCrate.withModule(FluentClientGenerator.customizableOperationModule) { writer ->
+            renderCustomizableOperationSendMethod(runtimeConfig, generics, writer)
+        }
         rustCrate.withModule(FluentClientGenerator.clientModule) { writer ->
             AwsFluentClientExtensions(types).render(writer)
         }
@@ -246,3 +257,43 @@ private class AwsFluentClientDocs(private val coreCodegenContext: CoreCodegenCon
         }
     }
 }
+
+private fun renderCustomizableOperationSendMethod(
+    runtimeConfig: RuntimeConfig,
+    generics: FluentClientGenerics,
+    writer: RustWriter,
+) {
+    val smithyHttp = CargoDependency.SmithyHttp(runtimeConfig).asType()
+
+    val operationGenerics = GenericsGenerator(mutableListOf("O" to Optional.empty(), "Retry" to Optional.empty()))
+    val handleGenerics = generics.toGenericsGenerator()
+    val combinedGenerics = operationGenerics + handleGenerics
+
+    val codegenScope = arrayOf(
+        "combined_generics_decl" to combinedGenerics.declaration(),
+        "handle_generics_bounds" to handleGenerics.bounds(),
+        "SdkSuccess" to smithyHttp.member("result::SdkSuccess"),
+        "ClassifyResponse" to smithyHttp.member("retry::ClassifyResponse"),
+        "ParseHttpResponse" to smithyHttp.member("response::ParseHttpResponse"),
+    )
+
+    writer.rustTemplate(
+        """
+        impl#{combined_generics_decl:W} CustomizableOperation#{combined_generics_decl:W}
+        where
+            #{handle_generics_bounds:W}
+        {
+            /// Sends this operation's request
+            pub async fn send<T, E>(self) -> Result<T, SdkError<E>>
+            where
+                E: std::error::Error,
+                O: #{ParseHttpResponse}<Output = Result<T, E>> + Send + Sync + Clone + 'static,
+                Retry: #{ClassifyResponse}<#{SdkSuccess}<T>, SdkError<E>> + Send + Sync + Clone,
+            {
+                self.handle.client.call(self.operation).await
+            }
+        }
+        """,
+        *codegenScope,
+    )
+}

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

@@ -5,7 +5,10 @@
 
 package software.amazon.smithy.rustsdk
 
+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.transform.ModelTransformer
 import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.writable
@@ -16,6 +19,7 @@ import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationCustomization
 import software.amazon.smithy.rust.codegen.smithy.customize.OperationSection
 import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
+import software.amazon.smithy.rust.codegen.smithy.traits.RetryPolicyTrait
 
 class RetryPolicyDecorator : RustCodegenDecorator<ClientCodegenContext> {
     override val name: String = "RetryPolicy"
@@ -31,6 +35,21 @@ class RetryPolicyDecorator : RustCodegenDecorator<ClientCodegenContext> {
 
     override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
         clazz.isAssignableFrom(ClientCodegenContext::class.java)
+
+    override fun transformModel(service: ServiceShape, model: Model): Model {
+        // Allow downstream codegen that needs to know the retry type for a given operation
+        // to recreate it by passing the `runtimeConfig`
+        return ModelTransformer.create().mapShapes(model) { shape ->
+            if (shape is OperationShape) {
+                val retryType = { runtimeConfig: RuntimeConfig ->
+                    runtimeConfig.awsHttp().asType().member("retry::AwsErrorRetryPolicy")
+                }
+                shape.toBuilder().addTrait(RetryPolicyTrait(retryType)).build()
+            } else {
+                shape
+            }
+        }
+    }
 }
 
 class RetryPolicyFeature(private val runtimeConfig: RuntimeConfig) : OperationCustomization() {

aws/sdk/integration-tests/s3/tests/customizable-operation.rs

@@ -0,0 +1,75 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_http::user_agent::AwsUserAgent;
+use aws_sdk_s3::{Credentials, Region};
+use aws_smithy_async::rt::sleep::TokioSleep;
+use aws_smithy_client::test_connection::capture_request;
+
+use std::convert::Infallible;
+use std::sync::Arc;
+use std::time::{Duration, UNIX_EPOCH};
+
+#[tokio::test]
+async fn test_s3_ops_are_customizable() -> Result<(), aws_sdk_s3::Error> {
+    let creds = Credentials::new(
+        "ANOTREAL",
+        "notrealrnrELgWzOk3IfjzDKtFBhDby",
+        Some("notarealsessiontoken".to_string()),
+        None,
+        "test",
+    );
+    let conf = aws_sdk_s3::Config::builder()
+        .credentials_provider(creds)
+        .region(Region::new("us-east-1"))
+        .sleep_impl(Arc::new(TokioSleep::new()))
+        .build();
+    let (conn, rcvr) = capture_request(None);
+
+    let client = aws_sdk_s3::Client::from_conf_conn(conf, conn);
+
+    let op = client
+        .list_buckets()
+        .customize()
+        .await
+        .expect("list_buckets is customizable")
+        .map_operation(|mut op| {
+            op.properties_mut()
+                .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
+            op.properties_mut().insert(AwsUserAgent::for_tests());
+
+            Result::<_, Infallible>::Ok(op)
+        })
+        .expect("inserting into the property bag is infallible");
+
+    // The response from the fake connection won't return the expected XML but we don't care about
+    // that error in this test
+    let _ = op
+        .send()
+        .await
+        .expect_err("this will fail due to not receiving a proper XML response.");
+
+    let expected_req = rcvr.expect_request();
+    let auth_header = expected_req
+        .headers()
+        .get("Authorization")
+        .unwrap()
+        .to_owned();
+
+    // This is a snapshot test taken from a known working test result
+    let snapshot_signature =
+        "Signature=c2028dc806248952fc533ab4b1d9f1bafcdc9b3380ed00482f9935541ae11671";
+    assert!(
+        auth_header
+            .to_str()
+            .unwrap()
+            .contains(snapshot_signature),
+        "authorization header signature did not match expected signature: got {}, expected it to contain {}",
+        auth_header.to_str().unwrap(),
+        snapshot_signature
+    );
+
+    Ok(())
+}

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

@@ -150,6 +150,7 @@ object RustReservedWords : ReservedWords {
         "abstract",
         "become",
         "box",
+        "customize",
         "do",
         "final",
         "macro",

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.smithy.generators
+
+import software.amazon.smithy.rust.codegen.rustlang.rust
+import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.rustlang.writable
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.util.orNull
+import java.util.Optional
+
+class GenericsGenerator(
+    private val types: MutableList<Pair<String, Optional<RuntimeType>>>,
+) {
+    fun add(type: Pair<String, RuntimeType>) {
+        types.add(type.first to Optional.of(type.second))
+    }
+
+    fun declaration() = writable {
+        // Write nothing if this generator is empty
+        if (types.isNotEmpty()) {
+            val typeArgs = types.joinToString(", ") { it.first }
+            rust("<$typeArgs>")
+        }
+    }
+
+    fun bounds() = writable {
+        // Only write bounds for generic type params with a bound
+        types.filter { it.second.isPresent }.map {
+            val (typeArg, runtimeType) = it
+            rustTemplate("$typeArg: #{runtimeType},\n", "runtimeType" to runtimeType.orNull()!!)
+        }
+    }
+
+    operator fun plus(operationGenerics: GenericsGenerator): GenericsGenerator {
+        return GenericsGenerator(listOf(types, operationGenerics.types).flatten().toMutableList())
+    }
+}