Skip to content

Commit

Permalink
Port customizable operation to orchestrator (#2706)
Browse files Browse the repository at this point in the history 
## Motivation and Context
Port [Customizable
Operation](#1647) to
orchestrator

## Description
This PR implements `CustomizableOperation` in the orchestrator. Just
like the counterpart in the middleware, it is created when the
`customize` method (in the orchestrator mode) on a fluent builder is
called. The `customize` method in the orchestrator could technically be
made a synchronous method because there is no need to create an
operation, which requires `async`, therefore making the `customize`
method in the middleware `async`. However, during the transition from
the middleware to the orchestrator, the integration tests
([example](https://github.com/awslabs/smithy-rs/blob/31c152d9af53afb9a5e6edf9df3def57931b9c1e/aws/sdk/integration-tests/s3/tests/signing-it.rs#L36))
need to be able to run in both modes. For this reason, the `customize`
method in the orchestrator is temporarily marked as `async`.

Regarding methods defined on the new `CustomizableOperation`, they
include `mutate_request` and `map_request` from the counterpart in the
middleware. However, it did not port `map_operation` because there is no
operation to map on. Most use cases for `map_operation` is put things in
a property bag. The new `CustomizableOperation` provides an
`interceptor` method to accomplish the same, i.e putting things in a
config bag.

Finally, for integration tests to run in both modes, the code gen emits
the implementation of the `customize` method differently depending on
the active Smithy runtime mode, similar to what the implementation of
`send` method does.

## Testing
Added one `sra-test` for mutating a request.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Co-authored-by: Yuki Saito <awsaito@amazon.com>
Co-authored-by: John DiSanti <jdisanti@amazon.com>
  • Loading branch information
@ysaito1001 @jdisanti @david-perez
3 people authored and david-perez committed May 18, 2023
1 parent edcb765 commit 1543c3d
Show file tree
Hide file tree
Showing 6 changed files with 319 additions and 37 deletions.
58 changes: 52 additions & 6 deletions aws/sra-test/integration-tests/aws-sdk-s3/tests/interceptors.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
);
Expand Down Expand Up @@ -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
);
Expand Down
7 changes: 4 additions & 3 deletions aws/sra-test/integration-tests/aws-sdk-s3/tests/request_information_headers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
);
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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))
Expand Down
143 changes: 142 additions & 1 deletion ...zon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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}<crate::config::Builder>,
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<F, E>(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<F>(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}<crate::config::Builder>,
) -> 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}<dyn #{RuntimePlugin} + #{Send} + #{Sync}>>::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}<impl #{RuntimePlugin} + #{Send} + #{Sync} + 'static>
) -> #{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) {
Expand Down
Loading

0 comments on commit 1543c3d

Please sign in to comment.