diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt
new file mode 100644
index 0000000000..70cc99e52c
--- /dev/null
+++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt
@@ -0,0 +1,103 @@
+/*
+ * 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.client.smithy.generators.client.CustomizableOperationCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.client.CustomizableOperationSection
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.rustlang.writable
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+
+class CustomizableOperationTestHelpers(runtimeConfig: RuntimeConfig) :
+ CustomizableOperationCustomization() {
+ private val codegenScope = arrayOf(
+ *RuntimeType.preludeScope,
+ "AwsUserAgent" to AwsRuntimeType.awsHttp(runtimeConfig)
+ .resolve("user_agent::AwsUserAgent"),
+ "BeforeTransmitInterceptorContextMut" to RuntimeType.smithyRuntimeApi(runtimeConfig)
+ .resolve("client::interceptors::BeforeTransmitInterceptorContextMut"),
+ "ConfigBag" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("config_bag::ConfigBag"),
+ "ConfigBagAccessors" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::orchestrator::ConfigBagAccessors"),
+ "http" to CargoDependency.Http.toType(),
+ "InterceptorContext" to RuntimeType.smithyRuntimeApi(runtimeConfig)
+ .resolve("client::interceptors::InterceptorContext"),
+ "RequestTime" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::orchestrator::RequestTime"),
+ "SharedInterceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig)
+ .resolve("client::interceptors::SharedInterceptor"),
+ "TestParamsSetterInterceptor" to CargoDependency.smithyRuntime(runtimeConfig).withFeature("test-util")
+ .toType().resolve("client::test_util::interceptor::TestParamsSetterInterceptor"),
+ )
+
+ override fun section(section: CustomizableOperationSection): Writable =
+ writable {
+ if (section is CustomizableOperationSection.CustomizableOperationImpl) {
+ if (section.operationShape == null) {
+ // TODO(enableNewSmithyRuntime): Delete this branch when middleware is no longer used
+ // This branch customizes CustomizableOperation in the middleware. section.operationShape being
+ // null means that this customization is rendered in a place where we don't need to figure out
+ // the module for an operation (which is the case for CustomizableOperation in the middleware
+ // that is rendered in the customize module).
+ rustTemplate(
+ """
+ ##[doc(hidden)]
+ // This is a temporary method for testing. NEVER use it in production
+ pub fn request_time_for_tests(mut self, request_time: ::std::time::SystemTime) -> Self {
+ self.operation.properties_mut().insert(request_time);
+ self
+ }
+
+ ##[doc(hidden)]
+ // This is a temporary method for testing. NEVER use it in production
+ pub fn user_agent_for_tests(mut self) -> Self {
+ self.operation.properties_mut().insert(#{AwsUserAgent}::for_tests());
+ self
+ }
+ """,
+ *codegenScope,
+ )
+ } else {
+ // The else branch is for rendering customization for the orchestrator.
+ rustTemplate(
+ """
+ ##[doc(hidden)]
+ // This is a temporary method for testing. NEVER use it in production
+ pub fn request_time_for_tests(mut self, request_time: ::std::time::SystemTime) -> Self {
+ use #{ConfigBagAccessors};
+ let interceptor = #{TestParamsSetterInterceptor}::new(move |_: &mut #{BeforeTransmitInterceptorContextMut}<'_>, cfg: &mut #{ConfigBag}| {
+ cfg.set_request_time(#{RequestTime}::new(request_time));
+ });
+ self.interceptors.push(#{SharedInterceptor}::new(interceptor));
+ self
+ }
+
+ ##[doc(hidden)]
+ // This is a temporary method for testing. NEVER use it in production
+ pub fn user_agent_for_tests(mut self) -> Self {
+ let interceptor = #{TestParamsSetterInterceptor}::new(|context: &mut #{BeforeTransmitInterceptorContextMut}<'_>, _: &mut #{ConfigBag}| {
+ let headers = context.request_mut().headers_mut();
+ let user_agent = #{AwsUserAgent}::for_tests();
+ headers.insert(
+ #{http}::header::USER_AGENT,
+ #{http}::HeaderValue::try_from(user_agent.ua_header()).unwrap(),
+ );
+ headers.insert(
+ #{http}::HeaderName::from_static("x-amz-user-agent"),
+ #{http}::HeaderValue::try_from(user_agent.aws_ua_header()).unwrap(),
+ );
+ });
+ self.interceptors.push(#{SharedInterceptor}::new(interceptor));
+ self
+ }
+ """,
+ *codegenScope,
+ )
+ }
+ }
+ }
+}
diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt
index b0963eaea9..050ba65be0 100644
--- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt
+++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt
@@ -71,7 +71,7 @@ class AwsFluentClientDecorator : ClientCodegenDecorator {
AwsFluentClientDocs(codegenContext),
),
retryClassifier = AwsRuntimeType.awsHttp(runtimeConfig).resolve("retry::AwsResponseRetryClassifier"),
- ).render(rustCrate)
+ ).render(rustCrate, listOf(CustomizableOperationTestHelpers(runtimeConfig)))
rustCrate.withModule(ClientRustModule.Client.customize) {
renderCustomizableOperationSendMethod(runtimeConfig, generics, this)
}
diff --git a/aws/sdk/integration-tests/s3/tests/checksums.rs b/aws/sdk/integration-tests/s3/tests/checksums.rs
index a82aac5e0a..ae533964a1 100644
--- a/aws/sdk/integration-tests/s3/tests/checksums.rs
+++ b/aws/sdk/integration-tests/s3/tests/checksums.rs
@@ -5,7 +5,6 @@
use aws_config::SdkConfig;
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::config::{Credentials, Region};
use aws_sdk_s3::types::ChecksumMode;
use aws_sdk_s3::Client;
@@ -14,10 +13,7 @@ use aws_smithy_client::test_connection::{capture_request, TestConnection};
use aws_smithy_http::body::SdkBody;
use http::header::AUTHORIZATION;
use http::{HeaderValue, Uri};
-use std::{
- convert::Infallible,
- time::{Duration, UNIX_EPOCH},
-};
+use std::time::{Duration, UNIX_EPOCH};
use tracing_test::traced_test;
/// Test connection for the movies IT
@@ -77,14 +73,8 @@ async fn test_checksum_on_streaming_response(
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests()
.send()
.await
.unwrap();
@@ -191,14 +181,8 @@ async fn test_checksum_on_streaming_request<'a>(
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests()
.send()
.await
.unwrap();
diff --git a/aws/sdk/integration-tests/s3/tests/customizable-operation.rs b/aws/sdk/integration-tests/s3/tests/customizable-operation.rs
index 24b53523a9..09dd42fb8f 100644
--- a/aws/sdk/integration-tests/s3/tests/customizable-operation.rs
+++ b/aws/sdk/integration-tests/s3/tests/customizable-operation.rs
@@ -5,12 +5,10 @@
use aws_config::SdkConfig;
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::config::{Credentials, Region};
use aws_sdk_s3::Client;
use aws_smithy_client::test_connection::capture_request;
-use std::convert::Infallible;
use std::time::{Duration, UNIX_EPOCH};
#[tokio::test]
@@ -29,14 +27,8 @@ async fn test_s3_ops_are_customizable() {
.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");
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests();
// The response from the fake connection won't return the expected XML but we don't care about
// that error in this test
diff --git a/aws/sdk/integration-tests/s3/tests/endpoints.rs b/aws/sdk/integration-tests/s3/tests/endpoints.rs
index e74d1a0a83..e6428ccc2b 100644
--- a/aws/sdk/integration-tests/s3/tests/endpoints.rs
+++ b/aws/sdk/integration-tests/s3/tests/endpoints.rs
@@ -9,7 +9,6 @@ use aws_sdk_s3::config::Builder;
use aws_sdk_s3::config::{Credentials, Region};
use aws_sdk_s3::Client;
use aws_smithy_client::test_connection::{capture_request, CaptureRequestReceiver};
-use std::convert::Infallible;
use std::time::{Duration, UNIX_EPOCH};
fn test_client(update_builder: fn(Builder) -> Builder) -> (CaptureRequestReceiver, Client) {
@@ -90,12 +89,7 @@ async fn s3_object_lambda() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1234567890));
- Result::<_, Infallible>::Ok(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1234567890))
.send()
.await
.unwrap();
diff --git a/aws/sdk/integration-tests/s3/tests/ignore-invalid-xml-body-root.rs b/aws/sdk/integration-tests/s3/tests/ignore-invalid-xml-body-root.rs
index f4e8704360..3e55ffe9c3 100644
--- a/aws/sdk/integration-tests/s3/tests/ignore-invalid-xml-body-root.rs
+++ b/aws/sdk/integration-tests/s3/tests/ignore-invalid-xml-body-root.rs
@@ -4,16 +4,12 @@
*/
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::{config::Credentials, config::Region, types::ObjectAttributes, Client};
use aws_smithy_client::test_connection::TestConnection;
use aws_smithy_http::body::SdkBody;
use aws_types::SdkConfig;
use http::header::AUTHORIZATION;
-use std::{
- convert::Infallible,
- time::{Duration, UNIX_EPOCH},
-};
+use std::time::{Duration, UNIX_EPOCH};
const RESPONSE_BODY_XML: &[u8] = b"\ne1AsOh9IyGCa4hLN+2Od7jlnP14=";
@@ -60,14 +56,8 @@ async fn ignore_invalid_xml_body_root() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests()
.send()
.await
.unwrap();
diff --git a/aws/sdk/integration-tests/s3/tests/naughty-string-metadata.rs b/aws/sdk/integration-tests/s3/tests/naughty-string-metadata.rs
index eef5b0b14b..96142f2b43 100644
--- a/aws/sdk/integration-tests/s3/tests/naughty-string-metadata.rs
+++ b/aws/sdk/integration-tests/s3/tests/naughty-string-metadata.rs
@@ -4,15 +4,11 @@
*/
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::{config::Credentials, config::Region, primitives::ByteStream, Client};
use aws_smithy_client::test_connection::capture_request;
use aws_types::SdkConfig;
use http::HeaderValue;
-use std::{
- convert::Infallible,
- time::{Duration, UNIX_EPOCH},
-};
+use std::time::{Duration, UNIX_EPOCH};
const NAUGHTY_STRINGS: &str = include_str!("blns/blns.txt");
@@ -83,14 +79,8 @@ async fn test_s3_signer_with_naughty_string_metadata() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests()
.send()
.await
.unwrap();
diff --git a/aws/sdk/integration-tests/s3/tests/normalize-uri-path.rs b/aws/sdk/integration-tests/s3/tests/normalize-uri-path.rs
index 1f903abb7a..a50867ccd0 100644
--- a/aws/sdk/integration-tests/s3/tests/normalize-uri-path.rs
+++ b/aws/sdk/integration-tests/s3/tests/normalize-uri-path.rs
@@ -5,11 +5,9 @@
use aws_config::SdkConfig;
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::{config::Credentials, config::Region, Client};
use aws_smithy_client::test_connection::capture_request;
-use std::convert::Infallible;
use std::time::{Duration, UNIX_EPOCH};
#[tokio::test]
@@ -33,14 +31,8 @@ async fn test_operation_should_not_normalize_uri_path() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1669257290));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1669257290))
+ .user_agent_for_tests()
.send()
.await
.unwrap();
diff --git a/aws/sdk/integration-tests/s3/tests/query-strings-are-correctly-encoded.rs b/aws/sdk/integration-tests/s3/tests/query-strings-are-correctly-encoded.rs
index 03393a5a74..7d7880fa81 100644
--- a/aws/sdk/integration-tests/s3/tests/query-strings-are-correctly-encoded.rs
+++ b/aws/sdk/integration-tests/s3/tests/query-strings-are-correctly-encoded.rs
@@ -5,11 +5,9 @@
use aws_config::SdkConfig;
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::config::{Credentials, Region};
use aws_sdk_s3::Client;
use aws_smithy_client::test_connection::capture_request;
-use std::convert::Infallible;
use std::time::{Duration, UNIX_EPOCH};
#[tokio::test]
@@ -35,14 +33,8 @@ async fn test_s3_signer_query_string_with_all_valid_chars() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests()
.send()
.await;
diff --git a/aws/sdk/integration-tests/s3/tests/signing-it.rs b/aws/sdk/integration-tests/s3/tests/signing-it.rs
index 6e20e187e2..a02b6b670b 100644
--- a/aws/sdk/integration-tests/s3/tests/signing-it.rs
+++ b/aws/sdk/integration-tests/s3/tests/signing-it.rs
@@ -5,12 +5,10 @@
use aws_config::SdkConfig;
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::config::{Credentials, Region};
use aws_sdk_s3::Client;
use aws_smithy_client::test_connection::TestConnection;
use aws_smithy_http::body::SdkBody;
-use std::convert::Infallible;
use std::time::{Duration, UNIX_EPOCH};
#[tokio::test]
@@ -36,14 +34,8 @@ async fn test_signer() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1624036048))
+ .user_agent_for_tests()
.send()
.await;
diff --git a/aws/sdk/integration-tests/s3control/tests/signing-it.rs b/aws/sdk/integration-tests/s3control/tests/signing-it.rs
index f5c583bbcd..88dc2debfb 100644
--- a/aws/sdk/integration-tests/s3control/tests/signing-it.rs
+++ b/aws/sdk/integration-tests/s3control/tests/signing-it.rs
@@ -4,13 +4,11 @@
*/
use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3control::config::{Credentials, Region};
use aws_sdk_s3control::Client;
use aws_smithy_client::test_connection::TestConnection;
use aws_smithy_http::body::SdkBody;
use aws_types::SdkConfig;
-use std::convert::Infallible;
use std::time::{Duration, UNIX_EPOCH};
#[tokio::test]
@@ -39,14 +37,8 @@ async fn test_signer() {
.customize()
.await
.unwrap()
- .map_operation(|mut op| {
- op.properties_mut()
- .insert(UNIX_EPOCH + Duration::from_secs(1636751225));
- op.properties_mut().insert(AwsUserAgent::for_tests());
-
- Result::Ok::<_, Infallible>(op)
- })
- .unwrap()
+ .request_time_for_tests(UNIX_EPOCH + Duration::from_secs(1636751225))
+ .user_agent_for_tests()
.send()
.await
.expect_err("empty response");
diff --git a/aws/sra-test/integration-tests/aws-sdk-s3/tests/interceptors.rs b/aws/sra-test/integration-tests/aws-sdk-s3/tests/interceptors.rs
index 1f061f7f9d..b1e708aad1 100644
--- a/aws/sra-test/integration-tests/aws-sdk-s3/tests/interceptors.rs
+++ b/aws/sra-test/integration-tests/aws-sdk-s3/tests/interceptors.rs
@@ -5,18 +5,20 @@
mod util;
+use aws_http::user_agent::AwsUserAgent;
use aws_sdk_s3::config::{Credentials, Region};
use aws_sdk_s3::Client;
use aws_smithy_client::dvr;
use aws_smithy_client::dvr::MediaType;
use aws_smithy_client::erase::DynConnector;
-use aws_smithy_runtime_api::client::interceptors::context::{Error, Input, Output};
use aws_smithy_runtime_api::client::interceptors::{
BeforeTransmitInterceptorContextMut, Interceptor,
};
use aws_smithy_runtime_api::client::orchestrator::ConfigBagAccessors;
use aws_smithy_runtime_api::client::orchestrator::RequestTime;
use aws_smithy_runtime_api::config_bag::ConfigBag;
+use http::header::USER_AGENT;
+use http::HeaderValue;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
const LIST_BUCKETS_PATH: &str = "test-data/list-objects-v2.json";
@@ -148,16 +150,16 @@ async fn set_test_user_agent_through_request_mutation() {
.await
.unwrap()
.mutate_request(|request| {
- request.headers_mut()
- .insert(
- http::HeaderName::from_static("user-agent"),
- http::HeaderValue::from_str("aws-sdk-rust/0.123.test os/windows/XPSP3 lang/rust/1.50.0").unwrap(),
- );
- request.headers_mut()
- .insert(
- http::HeaderName::from_static("x-amz-user-agent"),
- http::HeaderValue::from_str("aws-sdk-rust/0.123.test api/test-service/0.123 os/windows/XPSP3 lang/rust/1.50.0").unwrap(),
- );
+ let headers = request.headers_mut();
+ let user_agent = AwsUserAgent::for_tests();
+ headers.insert(
+ USER_AGENT,
+ HeaderValue::try_from(user_agent.ua_header()).unwrap(),
+ );
+ headers.insert(
+ util::X_AMZ_USER_AGENT,
+ HeaderValue::try_from(user_agent.aws_ua_header()).unwrap(),
+ );
})
.send_orchestrator_with_plugin(Some(fixup))
.await
diff --git a/aws/sra-test/integration-tests/aws-sdk-s3/tests/util.rs b/aws/sra-test/integration-tests/aws-sdk-s3/tests/util.rs
index 4f2997f5c6..7ce99e3af3 100644
--- a/aws/sra-test/integration-tests/aws-sdk-s3/tests/util.rs
+++ b/aws/sra-test/integration-tests/aws-sdk-s3/tests/util.rs
@@ -5,7 +5,6 @@
use aws_http::user_agent::AwsUserAgent;
use aws_runtime::invocation_id::InvocationId;
-use aws_smithy_runtime_api::client::interceptors::context::{Error, Input, Output};
use aws_smithy_runtime_api::client::interceptors::{
BeforeTransmitInterceptorContextMut, Interceptor, InterceptorRegistrar,
};
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt
new file mode 100644
index 0000000000..93a01e41bb
--- /dev/null
+++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt
@@ -0,0 +1,19 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.generators.client
+
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.rust.codegen.core.smithy.customize.NamedCustomization
+import software.amazon.smithy.rust.codegen.core.smithy.customize.Section
+
+sealed class CustomizableOperationSection(name: String) : Section(name) {
+ /** Write custom code into a customizable operation's impl block */
+ data class CustomizableOperationImpl(
+ val operationShape: OperationShape?,
+ ) : CustomizableOperationSection("CustomizableOperationImpl")
+}
+
+abstract class CustomizableOperationCustomization : NamedCustomization()
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt
index 3bfd660a75..c3108a2f57 100644
--- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt
+++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt
@@ -13,10 +13,12 @@ import software.amazon.smithy.rust.codegen.core.rustlang.GenericTypeArg
import software.amazon.smithy.rust.codegen.core.rustlang.RustGenerics
import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.rustlang.writable
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
+import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations
import software.amazon.smithy.rust.codegen.core.util.outputShape
/**
@@ -26,6 +28,7 @@ import software.amazon.smithy.rust.codegen.core.util.outputShape
class CustomizableOperationGenerator(
private val codegenContext: ClientCodegenContext,
private val generics: FluentClientGenerics,
+ private val customizations: List,
) {
private val runtimeConfig = codegenContext.runtimeConfig
private val smithyHttp = CargoDependency.smithyHttp(runtimeConfig).toType()
@@ -123,9 +126,14 @@ class CustomizableOperationGenerator(
pub fn request_mut(&mut self) -> &mut #{HttpRequest}<#{SdkBody}> {
self.operation.request_mut()
}
+
+ #{additional_methods}
}
""",
*codegenScope,
+ "additional_methods" to writable {
+ writeCustomizations(customizations, CustomizableOperationSection.CustomizableOperationImpl(null))
+ },
)
}
@@ -262,9 +270,14 @@ class CustomizableOperationGenerator(
.send_orchestrator_with_plugin(final_plugin)
.await
}
+
+ #{additional_methods}
}
""",
*codegenScope,
+ "additional_methods" to writable {
+ writeCustomizations(customizations, CustomizableOperationSection.CustomizableOperationImpl(operation))
+ },
)
}
}
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt
index 53c17b8d61..d81cfa32d1 100644
--- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt
+++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt
@@ -81,10 +81,10 @@ class FluentClientGenerator(
private val core = FluentClientCore(model)
private val smithyRuntimeMode = codegenContext.smithyRuntimeMode
- fun render(crate: RustCrate) {
+ fun render(crate: RustCrate, customizableOperationCustomizations: List = emptyList()) {
renderFluentClient(crate)
- val customizableOperationGenerator = CustomizableOperationGenerator(codegenContext, generics)
+ val customizableOperationGenerator = CustomizableOperationGenerator(codegenContext, generics, customizableOperationCustomizations)
operations.forEach { operation ->
crate.withModule(symbolProvider.moduleForBuilder(operation)) {
renderFluentBuilder(operation)
diff --git a/rust-runtime/aws-smithy-runtime/src/client/test_util.rs b/rust-runtime/aws-smithy-runtime/src/client/test_util.rs
index c85ab8f245..71afb382e4 100644
--- a/rust-runtime/aws-smithy-runtime/src/client/test_util.rs
+++ b/rust-runtime/aws-smithy-runtime/src/client/test_util.rs
@@ -5,4 +5,5 @@
pub mod connector;
pub mod deserializer;
+pub mod interceptor;
pub mod serializer;
diff --git a/rust-runtime/aws-smithy-runtime/src/client/test_util/interceptor.rs b/rust-runtime/aws-smithy-runtime/src/client/test_util/interceptor.rs
new file mode 100644
index 0000000000..74fa60bd70
--- /dev/null
+++ b/rust-runtime/aws-smithy-runtime/src/client/test_util/interceptor.rs
@@ -0,0 +1,78 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// TODO(enableNewSmithyRuntime): Delete this file once test helpers on `CustomizableOperation` have been removed
+
+use aws_smithy_runtime_api::client::interceptors::{
+ BeforeTransmitInterceptorContextMut, BoxError, Interceptor,
+};
+use aws_smithy_runtime_api::config_bag::ConfigBag;
+use std::fmt;
+
+pub struct TestParamsSetterInterceptor {
+ f: F,
+}
+
+impl fmt::Debug for TestParamsSetterInterceptor {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ write!(f, "TestParamsSetterInterceptor")
+ }
+}
+
+impl TestParamsSetterInterceptor {
+ pub fn new(f: F) -> Self {
+ Self { f }
+ }
+}
+
+impl Interceptor for TestParamsSetterInterceptor
+where
+ F: Fn(&mut BeforeTransmitInterceptorContextMut<'_>, &mut ConfigBag) + Send + Sync + 'static,
+{
+ fn modify_before_signing(
+ &self,
+ context: &mut BeforeTransmitInterceptorContextMut<'_>,
+ cfg: &mut ConfigBag,
+ ) -> Result<(), BoxError> {
+ (self.f)(context, cfg);
+
+ Ok(())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use aws_smithy_http::body::SdkBody;
+ use aws_smithy_runtime_api::client::interceptors::InterceptorContext;
+ use aws_smithy_runtime_api::client::orchestrator::{ConfigBagAccessors, RequestTime};
+ use aws_smithy_runtime_api::type_erasure::TypedBox;
+ use std::time::{Duration, UNIX_EPOCH};
+
+ #[test]
+ fn set_test_request_time() {
+ let mut cfg = ConfigBag::base();
+ let mut ctx = InterceptorContext::new(TypedBox::new("anything").erase());
+ ctx.enter_serialization_phase();
+ ctx.set_request(http::Request::builder().body(SdkBody::empty()).unwrap());
+ let _ = ctx.take_input();
+ ctx.enter_before_transmit_phase();
+ let mut ctx = Into::into(&mut ctx);
+ let request_time = UNIX_EPOCH + Duration::from_secs(1624036048);
+ let interceptor = TestParamsSetterInterceptor::new({
+ let request_time = request_time.clone();
+ move |_: &mut BeforeTransmitInterceptorContextMut<'_>, cfg: &mut ConfigBag| {
+ cfg.set_request_time(RequestTime::new(request_time));
+ }
+ });
+ interceptor
+ .modify_before_signing(&mut ctx, &mut cfg)
+ .unwrap();
+ assert_eq!(
+ request_time,
+ cfg.get::().unwrap().system_time()
+ );
+ }
+}