Skip to content

Commit

Permalink
Add user agent (#79)
Browse files Browse the repository at this point in the history
- Add an interceptor when we load the config to add the transfer manager metric for user agent
- Add a framework metadata for the transfer manager config and pass it down to the sdk s3 client as the framework name in the user agent.
- Only added for the `ConfigLoader`, we can decide what to do when user pass in their own client via config directly later.
  • Loading branch information
TingDaoK authored Dec 10, 2024
1 parent 540285a commit d18bc7a
Show file tree
Hide file tree
Showing 4 changed files with 156 additions and 6 deletions.
5 changes: 3 additions & 2 deletions aws-s3-transfer-manager/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@ aws-config = { version = "1.5.6", features = ["behavior-version-latest"] }
aws-sdk-s3 = { version = "1.51.0", features = ["behavior-version-latest"] }
aws-smithy-async = "1.2.1"
aws-smithy-experimental = { version = "0.1.3", features = ["crypto-aws-lc"] }
aws-smithy-runtime-api = "1.7.1"
aws-smithy-runtime-api = "1.7.3"
aws-runtime = "1.4.4"
aws-smithy-types = "1.2.6"
aws-types = "1.3.3"
blocking = "1.6.0"
Expand All @@ -32,7 +33,7 @@ walkdir = "2"
[dev-dependencies]
aws-sdk-s3 = { version = "1.51.0", features = ["behavior-version-latest", "test-util"] }
aws-smithy-mocks-experimental = "0.2.1"
aws-smithy-runtime = { version = "1.7.1", features = ["client", "connector-hyper-0-14-x", "test-util", "wire-mock"] }
aws-smithy-runtime = { version = "1.7.4", features = ["client", "connector-hyper-0-14-x", "test-util", "wire-mock"] }
clap = { version = "4.5.7", default-features = false, features = ["derive", "std", "help"] }
console-subscriber = "0.4.0"
http-02x = { package = "http", version = "0.2.9" }
Expand Down
2 changes: 1 addition & 1 deletion aws-s3-transfer-manager/external-types.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,5 +8,5 @@ allowed_external_types = [
"bytes::bytes::Bytes",
"bytes::buf::buf_impl::Buf",
"aws_types::request_id::RequestId",
"aws_types::request_id::RequestIdExt"
"aws_types::request_id::RequestIdExt",
]
24 changes: 24 additions & 0 deletions aws-s3-transfer-manager/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
* SPDX-License-Identifier: Apache-2.0
*/

use aws_runtime::user_agent::FrameworkMetadata;

use crate::metrics::unit::ByteUnit;
use crate::types::{ConcurrencySetting, PartSize};
use std::cmp;
Expand All @@ -18,6 +20,7 @@ pub struct Config {
multipart_threshold: PartSize,
target_part_size: PartSize,
concurrency: ConcurrencySetting,
framework_metadata: Option<FrameworkMetadata>,
client: aws_sdk_s3::client::Client,
}

Expand All @@ -43,6 +46,12 @@ impl Config {
&self.concurrency
}

/// Returns the framework metadata setting when using transfer manager.
#[doc(hidden)]
pub fn framework_metadata(&self) -> Option<&FrameworkMetadata> {
self.framework_metadata.as_ref()
}

/// The Amazon S3 client instance that will be used to send requests to S3.
pub fn client(&self) -> &aws_sdk_s3::Client {
&self.client
Expand All @@ -55,6 +64,7 @@ pub struct Builder {
multipart_threshold_part_size: PartSize,
target_part_size: PartSize,
concurrency: ConcurrencySetting,
framework_metadata: Option<FrameworkMetadata>,
client: Option<aws_sdk_s3::Client>,
}

Expand Down Expand Up @@ -122,8 +132,21 @@ impl Builder {
self
}

/// Sets the framework metadata for the transfer manager.
///
/// This _optional_ name is used to identify the framework using transfer manager in the user agent that
/// gets sent along with requests.
#[doc(hidden)]
pub fn framework_metadata(mut self, framework_metadata: Option<FrameworkMetadata>) -> Self {
self.framework_metadata = framework_metadata;
self
}

/// Set an explicit S3 client to use.
pub fn client(mut self, client: aws_sdk_s3::Client) -> Self {
// TODO - decide the approach here:
// - Convert the client to build to modify it based on other configs for transfer manager
// - Instead of taking the client, take sdk-config/s3-config/builder?
self.client = Some(client);
self
}
Expand All @@ -134,6 +157,7 @@ impl Builder {
multipart_threshold: self.multipart_threshold_part_size,
target_part_size: self.target_part_size,
concurrency: self.concurrency,
framework_metadata: self.framework_metadata,
client: self.client.expect("client set"),
}
}
Expand Down
131 changes: 128 additions & 3 deletions aws-s3-transfer-manager/src/config/loader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,52 @@
* SPDX-License-Identifier: Apache-2.0
*/

use aws_config::BehaviorVersion;
use aws_runtime::sdk_feature::AwsSdkFeature;
use aws_runtime::user_agent::{ApiMetadata, AwsUserAgent, FrameworkMetadata};
use aws_sdk_s3::config::{Intercept, IntoShared};
use aws_types::os_shim_internal::Env;

use crate::config::Builder;
use crate::{
http,
types::{ConcurrencySetting, PartSize},
Config,
};

#[derive(Debug)]
struct S3TransferManagerInterceptor {
frame_work_meta_data: Option<FrameworkMetadata>,
}

impl Intercept for S3TransferManagerInterceptor {
fn name(&self) -> &'static str {
"S3TransferManager"
}

fn read_before_execution(
&self,
_ctx: &aws_sdk_s3::config::interceptors::BeforeSerializationInterceptorContextRef<'_>,
cfg: &mut aws_sdk_s3::config::ConfigBag,
) -> Result<(), aws_sdk_s3::error::BoxError> {
// Assume the interceptor only be added to the client constructed by the loader.
// In this case, there should not be any user agent was sent before this interceptor starts.
// Create our own user agent with S3Transfer feature and user passed-in framework_meta_data if any.
cfg.interceptor_state()
.store_append(AwsSdkFeature::S3Transfer);
let api_metadata = cfg.load::<ApiMetadata>().unwrap();
// TODO: maybe APP Name someday
let mut ua = AwsUserAgent::new_from_environment(Env::real(), api_metadata.clone());
if let Some(framework_metadata) = self.frame_work_meta_data.clone() {
ua = ua.with_framework_metadata(framework_metadata);
}

cfg.interceptor_state().store_put(ua);

Ok(())
}
}

/// Load transfer manager [`Config`] from the environment.
#[derive(Default, Debug)]
pub struct ConfigLoader {
Expand Down Expand Up @@ -52,17 +91,103 @@ impl ConfigLoader {
self
}

/// Sets the framework metadata for the transfer manager.
///
/// This _optional_ name is used to identify the framework using transfer manager in the user agent that
/// gets sent along with requests.
#[doc(hidden)]
pub fn framework_metadata(mut self, framework_metadata: Option<FrameworkMetadata>) -> Self {
self.builder = self.builder.framework_metadata(framework_metadata);
self
}

/// Load the default configuration
///
/// If fields have been overridden during builder construction, the override values will be
/// used. Otherwise, the default values for each field will be provided.
pub async fn load(self) -> Config {
let shared_config = aws_config::from_env()
let shared_config = aws_config::defaults(BehaviorVersion::latest())
.http_client(http::default_client())
.load()
.await;
let s3_client = aws_sdk_s3::Client::new(&shared_config);
let builder = self.builder.client(s3_client);

let mut sdk_client_builder = aws_sdk_s3::config::Builder::from(&shared_config);

let interceptor = S3TransferManagerInterceptor {
frame_work_meta_data: self.builder.framework_metadata.clone(),
};
sdk_client_builder.push_interceptor(S3TransferManagerInterceptor::into_shared(interceptor));
let builder = self
.builder
.client(aws_sdk_s3::Client::from_conf(sdk_client_builder.build()));
builder.build()
}
}

#[cfg(test)]
mod tests {
use std::borrow::Cow;

use crate::types::{ConcurrencySetting, PartSize};
use aws_config::Region;
use aws_runtime::user_agent::FrameworkMetadata;
use aws_sdk_s3::config::Intercept;
use aws_smithy_runtime::client::http::test_util::capture_request;

#[tokio::test]
async fn load_with_interceptor() {
let config = crate::from_env()
.concurrency(ConcurrencySetting::Explicit(123))
.part_size(PartSize::Target(8))
.load()
.await;
let sdk_s3_config = config.client().config();
let tm_interceptor_exists = sdk_s3_config
.interceptors()
.any(|item| item.name() == "S3TransferManager");
assert!(tm_interceptor_exists);
}

#[tokio::test]
async fn load_with_interceptor_and_framework_metadata() {
let (http_client, captured_request) = capture_request(None);
let config = crate::from_env()
.concurrency(ConcurrencySetting::Explicit(123))
.part_size(PartSize::Target(8))
.framework_metadata(Some(
FrameworkMetadata::new("some-framework", Some(Cow::Borrowed("1.3"))).unwrap(),
))
.load()
.await;
// Inject the captured request to the http client to capture the request made from transfer manager.
let sdk_s3_config = config
.client()
.config()
.to_builder()
.http_client(http_client)
.region(Region::from_static("us-west-2"))
.with_test_defaults()
.build();

let capture_request_config = crate::Config::builder()
.client(aws_sdk_s3::Client::from_conf(sdk_s3_config))
.concurrency(ConcurrencySetting::Explicit(123))
.part_size(PartSize::Target(8))
.build();

let transfer_manager = crate::Client::new(capture_request_config);

let mut handle = transfer_manager
.download()
.bucket("foo")
.key("bar")
.initiate()
.unwrap();
// Expect to fail
let _ = handle.body_mut().next().await;
// Check the request made contains the expected framework meta data in user agent.
let expected_req = captured_request.expect_request();
let user_agent = expected_req.headers().get("x-amz-user-agent").unwrap();
assert!(user_agent.contains("lib/some-framework/1.3"));
}
}

0 comments on commit d18bc7a

Please sign in to comment.