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() + ); + } +}