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/RustTypes.kt

@@ -42,6 +42,10 @@ sealed class RustType {
 
     open val namespace: kotlin.String? = null
 
+    object Unit : RustType() {
+        override val name: kotlin.String = "()"
+    }
+
     object Bool : RustType() {
         override val name: kotlin.String = "bool"
     }
@@ -173,6 +177,7 @@ fun RustType.render(fullyQualified: Boolean = true): String {
         this.namespace?.let { "$it::" } ?: ""
     } else ""
     val base = when (this) {
+        is RustType.Unit -> this.name
         is RustType.Bool -> this.name
         is RustType.Float -> this.name
         is RustType.Integer -> this.name
@@ -282,7 +287,7 @@ data class RustMetadata(
         this.copy(derives = derives.copy(derives = derives.derives + newDerive))
 
     fun withoutDerives(vararg withoutDerives: RuntimeType) =
-        this.copy(derives = derives.copy(derives = derives.derives - withoutDerives))
+        this.copy(derives = derives.copy(derives = derives.derives - withoutDerives.toSet()))
 
     private fun attributes(): List<Attribute> = additionalAttributes + derives