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 208bd8aa8b..f33ef94ff8 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 @@ -40,11 +40,12 @@ async fn operation_interceptor_test() { let resp = dbg!( client .list_objects_v2() - .config_override( - aws_sdk_s3::Config::builder().interceptor(util::TestUserAgentInterceptor) - ) .bucket("test-bucket") .prefix("prefix~") + .customize() + .await + .unwrap() + .interceptor(util::TestUserAgentInterceptor) .send_orchestrator_with_plugin(Some(fixup)) .await ); @@ -106,11 +107,56 @@ async fn interceptor_priority() { let resp = dbg!( client .list_objects_v2() - .config_override(aws_sdk_s3::Config::builder().interceptor( - RequestTimeAdvanceInterceptor(Duration::from_secs(1624036048)) - )) .bucket("test-bucket") .prefix("prefix~") + .customize() + .await + .unwrap() + .interceptor(RequestTimeAdvanceInterceptor(Duration::from_secs( + 1624036048 + ))) + .send_orchestrator_with_plugin(Some(fixup)) + .await + ); + let resp = resp.expect("valid e2e test"); + assert_eq!(resp.name(), Some("test-bucket")); + conn.full_validate(MediaType::Xml).await.expect("success") +} + +#[tokio::test] +async fn set_test_user_agent_through_request_mutation() { + let conn = dvr::ReplayingConnection::from_file(LIST_BUCKETS_PATH).unwrap(); + + let config = aws_sdk_s3::Config::builder() + .credentials_provider(Credentials::for_tests()) + .region(Region::new("us-east-1")) + .http_connector(DynConnector::new(conn.clone())) + .build(); + let client = Client::from_conf(config); + let fixup = util::FixupPlugin { + timestamp: UNIX_EPOCH + Duration::from_secs(1624036048), + }; + + let resp = dbg!( + client + .list_objects_v2() + .bucket("test-bucket") + .prefix("prefix~") + .customize() + .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(), + ); + }) .send_orchestrator_with_plugin(Some(fixup)) .await ); diff --git a/aws/sra-test/integration-tests/aws-sdk-s3/tests/request_information_headers.rs b/aws/sra-test/integration-tests/aws-sdk-s3/tests/request_information_headers.rs index af0ab416dd..a5138f40de 100644 --- a/aws/sra-test/integration-tests/aws-sdk-s3/tests/request_information_headers.rs +++ b/aws/sra-test/integration-tests/aws-sdk-s3/tests/request_information_headers.rs @@ -68,9 +68,12 @@ async fn three_retries_and_then_success() { let resp = dbg!( client .list_objects_v2() - .config_override(aws_sdk_s3::Config::builder().force_path_style(false)) .bucket("test-bucket") .prefix("prefix~") + .customize() + .await + .unwrap() + .config_override(aws_sdk_s3::Config::builder().force_path_style(false)) .send_orchestrator_with_plugin(Some(fixup)) .await ); @@ -152,7 +155,6 @@ async fn three_retries_and_then_success() { // let resp = dbg!( // client // .list_objects_v2() -// .config_override(aws_sdk_s3::Config::builder().force_path_style(false)) // .bucket("test-bucket") // .prefix("prefix~") // .send_v2_with_plugin(Some(fixup)) @@ -244,7 +246,6 @@ async fn three_retries_and_then_success() { // let resp = dbg!( // client // .list_objects_v2() -// .config_override(aws_sdk_s3::Config::builder().force_path_style(false)) // .bucket("test-bucket") // .prefix("prefix~") // .send_v2_with_plugin(Some(fixup)) 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 d16d927d7f..3bfd660a75 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 @@ -5,6 +5,7 @@ package software.amazon.smithy.rust.codegen.client.smithy.generators.client +import software.amazon.smithy.model.shapes.OperationShape import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency @@ -16,13 +17,14 @@ 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.util.outputShape /** * Generates the code required to add the `.customize()` function to the * fluent client builders. */ class CustomizableOperationGenerator( - codegenContext: ClientCodegenContext, + private val codegenContext: ClientCodegenContext, private val generics: FluentClientGenerics, ) { private val runtimeConfig = codegenContext.runtimeConfig @@ -126,6 +128,145 @@ class CustomizableOperationGenerator( *codegenScope, ) } + + fun renderForOrchestrator(writer: RustWriter, operation: OperationShape) { + val symbolProvider = codegenContext.symbolProvider + val model = codegenContext.model + + val builderName = operation.fluentBuilderType(symbolProvider).name + val outputType = symbolProvider.toSymbol(operation.outputShape(model)) + val errorType = symbolProvider.symbolForOperationError(operation) + + val codegenScope = arrayOf( + *preludeScope, + "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig) + .resolve("client::orchestrator::HttpResponse"), + "Interceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig) + .resolve("client::interceptors::Interceptor"), + "MapRequestInterceptor" to RuntimeType.smithyRuntime(runtimeConfig) + .resolve("client::interceptor::MapRequestInterceptor"), + "MutateRequestInterceptor" to RuntimeType.smithyRuntime(runtimeConfig) + .resolve("client::interceptor::MutateRequestInterceptor"), + "OperationError" to errorType, + "OperationOutput" to outputType, + "RuntimePlugin" to RuntimeType.runtimePlugin(runtimeConfig), + "SdkBody" to RuntimeType.sdkBody(runtimeConfig), + "SdkError" to RuntimeType.sdkError(runtimeConfig), + "SharedInterceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig) + .resolve("client::interceptors::SharedInterceptor"), + ) + + writer.rustTemplate( + """ + /// A wrapper type for [`$builderName`]($builderName) that allows for configuring a single + /// operation invocation. + pub struct CustomizableOperation { + pub(crate) fluent_builder: $builderName, + pub(crate) config_override: #{Option}, + pub(crate) interceptors: Vec<#{SharedInterceptor}>, + } + + impl CustomizableOperation { + /// Adds an [`Interceptor`](#{Interceptor}) that runs at specific stages of the request execution pipeline. + /// + /// Note that interceptors can also be added to `CustomizableOperation` by `config_override`, + /// `map_request`, and `mutate_request` (the last two are implemented via interceptors under the hood). + /// The order in which those user-specified operation interceptors are invoked should not be relied upon + /// as it is an implementation detail. + pub fn interceptor(mut self, interceptor: impl #{Interceptor} + #{Send} + #{Sync} + 'static) -> Self { + self.interceptors.push(#{SharedInterceptor}::new(interceptor)); + self + } + + /// Allows for customizing the operation's request. + pub fn map_request(mut self, f: F) -> Self + where + F: #{Fn}(&mut http::Request<#{SdkBody}>) -> #{Result}<(), E> + + #{Send} + + #{Sync} + + 'static, + E: ::std::error::Error + #{Send} + #{Sync} + 'static, + { + self.interceptors.push( + #{SharedInterceptor}::new( + #{MapRequestInterceptor}::new(f), + ), + ); + self + } + + /// Convenience for `map_request` where infallible direct mutation of request is acceptable. + pub fn mutate_request(mut self, f: F) -> Self + where + F: #{Fn}(&mut http::Request<#{SdkBody}>) + #{Send} + #{Sync} + 'static, + { + self.interceptors.push( + #{SharedInterceptor}::new( + #{MutateRequestInterceptor}::new(f), + ), + ); + self + } + + /// Overrides config for a single operation invocation. + /// + /// `config_override` is applied to the operation configuration level. + /// The fields in the builder that are `Some` override those applied to the service + /// configuration level. For instance, + /// + /// Config A overridden by Config B == Config C + /// field_1: None, field_1: Some(v2), field_1: Some(v2), + /// field_2: Some(v1), field_2: Some(v2), field_2: Some(v2), + /// field_3: Some(v1), field_3: None, field_3: Some(v1), + pub fn config_override( + mut self, + config_override: impl #{Into}, + ) -> Self { + self.config_override = Some(config_override.into()); + self + } + + /// Sends the request and returns the response. + pub async fn send( + self + ) -> #{Result}< + #{OperationOutput}, + #{SdkError}< + #{OperationError}, + #{HttpResponse} + > + > { + self.send_orchestrator_with_plugin(#{Option}::<#{Box}>::None) + .await + } + + ##[doc(hidden)] + // TODO(enableNewSmithyRuntime): Delete when unused + /// Equivalent to [`Self::send`] but adds a final runtime plugin to shim missing behavior + pub async fn send_orchestrator_with_plugin( + self, + final_plugin: #{Option} + ) -> #{Result}<#{OperationOutput}, #{SdkError}<#{OperationError}, #{HttpResponse}>> { + let mut config_override = if let Some(config_override) = self.config_override { + config_override + } else { + crate::config::Builder::new() + }; + + self.interceptors.into_iter().for_each(|interceptor| { + config_override.add_interceptor(interceptor); + }); + + self.fluent_builder + .config_override(config_override) + .send_orchestrator_with_plugin(final_plugin) + .await + } + } + """, + *codegenScope, + ) + } } fun renderCustomizableOperationSend(runtimeConfig: RuntimeConfig, generics: FluentClientGenerics, writer: RustWriter) { 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 10a4f0c02f..53c17b8d61 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 @@ -84,13 +84,17 @@ class FluentClientGenerator( fun render(crate: RustCrate) { renderFluentClient(crate) + val customizableOperationGenerator = CustomizableOperationGenerator(codegenContext, generics) operations.forEach { operation -> crate.withModule(symbolProvider.moduleForBuilder(operation)) { renderFluentBuilder(operation) + if (codegenContext.smithyRuntimeMode.generateOrchestrator) { + customizableOperationGenerator.renderForOrchestrator(this, operation) + } } } - CustomizableOperationGenerator(codegenContext, generics).render(crate) + customizableOperationGenerator.render(crate) } private fun renderFluentClient(crate: RustCrate) { @@ -307,9 +311,9 @@ class FluentClientGenerator( ) rustTemplate( """ - /// Consume this builder, creating a customizable operation that can be modified before being - /// sent. The operation's inner [http::Request] can be modified as well. - pub async fn customize(self) -> #{Result}< + // This function will go away in the near future. Do not rely on it. + ##[doc(hidden)] + pub async fn customize_middleware(self) -> #{Result}< #{CustomizableOperation}#{customizable_op_type_params:W}, #{SdkError}<#{OperationError}> > #{send_bounds:W} { @@ -349,6 +353,15 @@ class FluentClientGenerator( #{send_bounds:W} { self.send_middleware().await } + + /// Consumes this builder, creating a customizable operation that can be modified before being + /// sent. The operation's inner [http::Request] can be modified as well. + pub async fn customize(self) -> #{Result}< + #{CustomizableOperation}#{customizable_op_type_params:W}, + #{SdkError}<#{OperationError}> + > #{send_bounds:W} { + self.customize_middleware().await + } """, *middlewareScope, ) @@ -357,6 +370,8 @@ class FluentClientGenerator( if (smithyRuntimeMode.generateOrchestrator) { val orchestratorScope = arrayOf( *preludeScope, + "CustomizableOperation" to symbolProvider.moduleForBuilder(operation).toType() + .resolve("CustomizableOperation"), "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig) .resolve("client::orchestrator::HttpResponse"), "OperationError" to errorType, @@ -382,10 +397,10 @@ class FluentClientGenerator( pub async fn send_orchestrator_with_plugin(self, final_plugin: #{Option}) -> #{Result}<#{OperationOutput}, #{SdkError}<#{OperationError}, #{HttpResponse}>> { let mut runtime_plugins = #{RuntimePlugins}::new() .with_client_plugin(crate::config::ServiceRuntimePlugin::new(self.handle.clone())); + runtime_plugins = runtime_plugins.with_operation_plugin(#{Operation}::new()); if let Some(config_override) = self.config_override { runtime_plugins = runtime_plugins.with_operation_plugin(config_override); } - runtime_plugins = runtime_plugins.with_operation_plugin(#{Operation}::new()); if let Some(final_plugin) = final_plugin { runtime_plugins = runtime_plugins.with_client_plugin(final_plugin); } @@ -402,6 +417,12 @@ class FluentClientGenerator( })?; #{Ok}(#{TypedBox}::<#{OperationOutput}>::assume_from(output).expect("correct output type").unwrap()) } + + ##[doc(hidden)] + // TODO(enableNewSmithyRuntime): Remove `async` once we switch to orchestrator + pub async fn customize_orchestrator(self) -> #{CustomizableOperation} { + #{CustomizableOperation} { fluent_builder: self, config_override: None, interceptors: vec![] } + } """, *orchestratorScope, ) @@ -419,6 +440,16 @@ class FluentClientGenerator( pub async fn send(self) -> #{Result}<#{OperationOutput}, #{SdkError}<#{OperationError}, #{HttpResponse}>> { self.send_orchestrator().await } + + /// Consumes this builder, creating a customizable operation that can be modified before being + /// sent. + // TODO(enableNewSmithyRuntime): Remove `async` and `Result` once we switch to orchestrator + pub async fn customize(self) -> #{Result}< + #{CustomizableOperation}, + #{SdkError}<#{OperationError}> + > { + #{Ok}(self.customize_orchestrator().await) + } """, *orchestratorScope, ) @@ -426,17 +457,7 @@ class FluentClientGenerator( rustTemplate( """ - /// Sets the `config_override` for the builder. - /// - /// `config_override` is applied to the operation configuration level. - /// The fields in the builder that are `Some` override those applied to the service - /// configuration level. For instance, - /// - /// Config A overridden by Config B == Config C - /// field_1: None, field_1: Some(v2), field_1: Some(v2), - /// field_2: Some(v1), field_2: Some(v2), field_2: Some(v2), - /// field_3: Some(v1), field_3: None, field_3: Some(v1), - pub fn config_override( + pub(crate) fn config_override( mut self, config_override: impl Into, ) -> Self { @@ -444,17 +465,7 @@ class FluentClientGenerator( self } - /// Sets the `config_override` for the builder. - /// - /// `config_override` is applied to the operation configuration level. - /// The fields in the builder that are `Some` override those applied to the service - /// configuration level. For instance, - /// - /// Config A overridden by Config B == Config C - /// field_1: None, field_1: Some(v2), field_1: Some(v2), - /// field_2: Some(v1), field_2: Some(v2), field_2: Some(v2), - /// field_3: Some(v1), field_3: None, field_3: Some(v1), - pub fn set_config_override( + pub(crate) fn set_config_override( &mut self, config_override: Option, ) -> &mut Self { diff --git a/rust-runtime/aws-smithy-runtime/src/client.rs b/rust-runtime/aws-smithy-runtime/src/client.rs index c2edaced8b..1a6a106da3 100644 --- a/rust-runtime/aws-smithy-runtime/src/client.rs +++ b/rust-runtime/aws-smithy-runtime/src/client.rs @@ -29,3 +29,6 @@ pub mod runtime_plugin; /// Smithy identity used by auth and signing. pub mod identity; + +/// Interceptors for Smithy clients. +pub mod interceptor; diff --git a/rust-runtime/aws-smithy-runtime/src/client/interceptor.rs b/rust-runtime/aws-smithy-runtime/src/client/interceptor.rs new file mode 100644 index 0000000000..4e5de36b63 --- /dev/null +++ b/rust-runtime/aws-smithy-runtime/src/client/interceptor.rs @@ -0,0 +1,80 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use aws_smithy_http::body::SdkBody; +use aws_smithy_runtime_api::client::interceptors::context::phase::BeforeTransmit; +use aws_smithy_runtime_api::client::interceptors::{BoxError, Interceptor, InterceptorContext}; +use aws_smithy_runtime_api::config_bag::ConfigBag; +use std::fmt; +use std::marker::PhantomData; + +pub struct MapRequestInterceptor { + f: F, + _phantom: PhantomData, +} + +impl fmt::Debug for MapRequestInterceptor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "MapRequestInterceptor") + } +} + +impl MapRequestInterceptor { + pub fn new(f: F) -> Self { + Self { + f, + _phantom: PhantomData, + } + } +} + +impl Interceptor for MapRequestInterceptor +where + F: Fn(&mut http::Request) -> Result<(), E> + Send + Sync + 'static, + E: std::error::Error + Send + Sync + 'static, +{ + fn modify_before_signing( + &self, + context: &mut InterceptorContext, + _cfg: &mut ConfigBag, + ) -> Result<(), BoxError> { + let request = context.request_mut(); + (self.f)(request)?; + + Ok(()) + } +} + +pub struct MutateRequestInterceptor { + f: F, +} + +impl fmt::Debug for MutateRequestInterceptor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "MutateRequestInterceptor") + } +} + +impl MutateRequestInterceptor { + pub fn new(f: F) -> Self { + Self { f } + } +} + +impl Interceptor for MutateRequestInterceptor +where + F: Fn(&mut http::Request) + Send + Sync + 'static, +{ + fn modify_before_signing( + &self, + context: &mut InterceptorContext, + _cfg: &mut ConfigBag, + ) -> Result<(), BoxError> { + let request = context.request_mut(); + (self.f)(request); + + Ok(()) + } +}