Skip to content

Commit

Permalink
update to expose framework metadata instead
Browse files Browse the repository at this point in the history
  • Loading branch information
TingDaoK committed Dec 5, 2024
1 parent 70ab365 commit 9a61f4e
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 35 deletions.
11 changes: 8 additions & 3 deletions aws-s3-transfer-manager/examples/cp.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::borrow::Cow;
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
Expand All @@ -7,13 +8,13 @@ use std::path::PathBuf;
use std::str::FromStr;
use std::time;

use aws_runtime::user_agent::FrameworkMetadata;
use aws_s3_transfer_manager::io::InputStream;
use aws_s3_transfer_manager::metrics::unit::ByteUnit;
use aws_s3_transfer_manager::metrics::Throughput;
use aws_s3_transfer_manager::operation::download::body::Body;
use aws_s3_transfer_manager::types::{ConcurrencySetting, PartSize};
use aws_sdk_s3::error::DisplayErrorContext;
use aws_types::app_name::AppName;
use bytes::Buf;
use clap::{CommandFactory, Parser};
use tokio::fs;
Expand Down Expand Up @@ -152,7 +153,9 @@ async fn do_download(args: Args) -> Result<(), BoxError> {
let tm_config = aws_s3_transfer_manager::from_env()
.concurrency(ConcurrencySetting::Explicit(args.concurrency))
.part_size(PartSize::Target(args.part_size))
.app_name(Some(AppName::new("transfer_manager_example_cp").unwrap()))
.frame_metadata(Some(
FrameworkMetadata::new("some-framework", Some(Cow::Borrowed("1.3"))).unwrap(),
))
.load()
.await;

Expand Down Expand Up @@ -230,7 +233,9 @@ async fn do_upload(args: Args) -> Result<(), BoxError> {
let tm_config = aws_s3_transfer_manager::from_env()
.concurrency(ConcurrencySetting::Explicit(args.concurrency))
.part_size(PartSize::Target(args.part_size))
.app_name(Some(AppName::new("transfer_manager_example_cp").unwrap()))
.frame_metadata(Some(
FrameworkMetadata::new("some-framework", Some(Cow::Borrowed("1.3"))).unwrap(),
))
.load()
.await;

Expand Down
16 changes: 8 additions & 8 deletions aws-s3-transfer-manager/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
* SPDX-License-Identifier: Apache-2.0
*/

use aws_types::app_name::AppName;
use aws_runtime::user_agent::FrameworkMetadata;

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

Expand All @@ -47,8 +47,8 @@ impl Config {
}

/// Returns the name of the app that is using the transfer manager, if it was provided.
pub fn app_name(&self) -> Option<&AppName> {
self.app_name.as_ref()
pub fn frame_metadata(&self) -> Option<&FrameworkMetadata> {
self.frame_metadata.as_ref()
}

/// The Amazon S3 client instance that will be used to send requests to S3.
Expand All @@ -63,7 +63,7 @@ pub struct Builder {
multipart_threshold_part_size: PartSize,
target_part_size: PartSize,
concurrency: ConcurrencySetting,
app_name: Option<AppName>,
frame_metadata: Option<FrameworkMetadata>,
client: Option<aws_sdk_s3::Client>,
}

Expand Down Expand Up @@ -135,8 +135,8 @@ impl Builder {
///
/// This _optional_ name is used to identify the application in the user agent that
/// gets sent along with requests.
pub fn app_name(mut self, app_name: Option<AppName>) -> Self {
self.app_name = app_name;
pub fn frame_metadata(mut self, frame_metadata: Option<FrameworkMetadata>) -> Self {
self.frame_metadata = frame_metadata;
self
}

Expand All @@ -155,7 +155,7 @@ impl Builder {
multipart_threshold: self.multipart_threshold_part_size,
target_part_size: self.target_part_size,
concurrency: self.concurrency,
app_name: self.app_name,
frame_metadata: self.frame_metadata,
client: self.client.expect("client set"),
}
}
Expand Down
59 changes: 35 additions & 24 deletions aws-s3-transfer-manager/src/config/loader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,9 @@

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::app_name::AppName;
use aws_types::os_shim_internal::Env;

use crate::config::Builder;
use crate::{
Expand All @@ -16,20 +17,34 @@ use crate::{
};

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

impl Intercept for TransferManagerFeatureInterceptor {
impl Intercept for S3TransferManagerInterceptor {
fn name(&self) -> &'static str {
"TransferManagerFeature"
"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>(AwsSdkFeature::S3Transfer);
.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(())
}
}
Expand Down Expand Up @@ -80,8 +95,8 @@ impl ConfigLoader {
///
/// This _optional_ name is used to identify the application in the user agent that
/// gets sent along with requests.
pub fn app_name(mut self, app_name: Option<AppName>) -> Self {
self.builder = self.builder.app_name(app_name);
pub fn frame_metadata(mut self, frame_metadata: Option<FrameworkMetadata>) -> Self {
self.builder = self.builder.frame_metadata(frame_metadata);
self
}

Expand All @@ -90,15 +105,17 @@ impl ConfigLoader {
/// 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 mut config_loader =
aws_config::defaults(BehaviorVersion::latest()).http_client(http::default_client());
if let Some(app_name) = self.builder.app_name.as_ref() {
config_loader = config_loader.app_name(app_name.clone());
}
let shared_config = config_loader.load().await;
let shared_config = aws_config::defaults(BehaviorVersion::latest())
.http_client(http::default_client())
.load()
.await;

let mut sdk_client_builder = aws_sdk_s3::config::Builder::from(&shared_config);
sdk_client_builder.push_interceptor(TransferManagerFeatureInterceptor.into_shared());

let interceptor = S3TransferManagerInterceptor {
frame_work_meta_data: self.builder.frame_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()));
Expand All @@ -110,26 +127,20 @@ impl ConfigLoader {
mod tests {
use crate::types::{ConcurrencySetting, PartSize};
use aws_sdk_s3::config::Intercept;
use aws_types::app_name::AppName;

#[tokio::test]
async fn load_with_app_name_and_interceptor() {
let expected_app_name = "bananas";

async fn load_with_frame_metadata_and_interceptor() {
let config = crate::from_env()
.concurrency(ConcurrencySetting::Explicit(123))
.part_size(PartSize::Target(8))
.app_name(Some(AppName::new(expected_app_name).unwrap()))
.load()
.await;
let sdk_s3_config = config.client().config();
assert_eq!(
sdk_s3_config.app_name().unwrap().as_ref(),
expected_app_name
);
let tm_interceptor_exists = sdk_s3_config
.interceptors()
.any(|item| item.name() == "TransferManagerFeature");
.any(|item| item.name() == "S3TransferManager");
assert!(tm_interceptor_exists);
}

/* TODO: test the framework metadata added correctly */
}

0 comments on commit 9a61f4e

Please sign in to comment.