Feature: Customizable Operations

.gitignore

@@ -45,3 +45,6 @@ gradle-app.setting
 
 # Rust build artifacts
 target/
+
+# IDEs
+.idea/

CHANGELOG.next.toml

@@ -72,3 +72,100 @@ wired up by default if none is provided.
 references = ["smithy-rs#1603", "aws-sdk-rust#586"]
 meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "client" }
 author = "jdisanti"
+
+
+[[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"

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

@@ -25,6 +25,8 @@ 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.GenericTypeArg
+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
@@ -73,6 +75,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()
+    }
 }
 
 class AwsFluentClientDecorator : RustCodegenDecorator<ClientCodegenContext> {
@@ -82,15 +88,21 @@ 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),
             ),
+            retryPolicyType = runtimeConfig.awsHttp().asType().member("retry::AwsErrorRetryPolicy"),
         ).render(rustCrate)
+        rustCrate.withModule(FluentClientGenerator.customizableOperationModule) { writer ->
+            renderCustomizableOperationSendMethod(runtimeConfig, generics, writer)
+        }
         rustCrate.withModule(FluentClientGenerator.clientModule) { writer ->
             AwsFluentClientExtensions(types).render(writer)
         }
@@ -254,3 +266,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(GenericTypeArg("O"), GenericTypeArg("Retry"))
+    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/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/rustlang/RustWriter.kt

@@ -136,6 +136,13 @@ fun <T : AbstractCodeWriter<T>> T.rust(
     this.write(contents.trim(), *args)
 }
 
+fun <T : AbstractCodeWriter<T>> T.rustInline(
+    @Language("Rust", prefix = "macro_rules! foo { () =>  {{\n", suffix = "\n}}}") contents: String,
+    vararg args: Any,
+) {
+    this.writeInline(contents.trim(), *args)
+}
+
 /* rewrite #{foo} to #{foo:T} (the smithy template format) */
 private fun transformTemplate(template: String, scope: Array<out Pair<String, Any>>): String {
     check(scope.distinctBy { it.first.lowercase() }.size == scope.size) { "Duplicate cased keys not supported" }
@@ -317,21 +324,6 @@ private fun Element.changeInto(tagName: String) {
  */
 fun RustWriter.raw(text: String) = writeInline(escape(text))
 
-typealias Writable = RustWriter.() -> Unit
-
-/** Helper to allow coercing the Writeable signature
- *  writable { rust("fn foo() { }")
- */
-fun writable(w: Writable): Writable = w
-
-fun writable(w: String): Writable = writable { rust(w) }
-
-fun Writable.isEmpty(): Boolean {
-    val writer = RustWriter.root()
-    this(writer)
-    return writer.toString() == RustWriter.root().toString()
-}
-
 /**
  * Rustdoc doesn't support `r#` for raw identifiers.
  * This function adjusts doc links to refer to raw identifiers directly.
@@ -471,12 +463,15 @@ class RustWriter private constructor(
                     block(derefName)
                 }
             }
+
             shape is NumberShape -> rustBlock("if ${outerField.removePrefix("&")} != 0") {
                 block(outerField)
             }
+
             shape is BooleanShape -> rustBlock("if ${outerField.removePrefix("&")}") {
                 block(outerField)
             }
+
             else -> this.block(outerField)
         }
     }
@@ -555,10 +550,12 @@ class RustWriter private constructor(
                     // for now, use the fully qualified type name
                     t.fullyQualifiedName()
                 }
+
                 is Symbol -> {
                     addDepsRecursively(t)
                     t.rustType().render(fullyQualified = true)
                 }
+
                 else -> throw CodegenException("Invalid type provided to RustSymbolFormatter: $t")
                 // escaping generates `##` sequences for all the common cases where
                 // it will be run through templating, but in this context, we won't be escaped