Track features for account based endpoints (#4298) (eaa6135f) · Commits · Public Repositories / Smithy Rust

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AccountIdEndpointParamsDecorator.kt

+32 −0

Original line number	Diff line number	Diff line
		@@ -14,6 +14,8 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ConditionalDe
		import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointCustomization
		import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointTypesGenerator
		import software.amazon.smithy.rust.codegen.client.smithy.endpoint.rustName
		import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
		import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection
		import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
		import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
		import software.amazon.smithy.rust.codegen.core.rustlang.Writable
		@@ -199,6 +201,10 @@ class AccountIdEndpointModeBuiltInParamDecorator : ConditionalDecorator(
		"AccountIdEndpointMode" to
		AwsRuntimeType.awsTypes(codegenContext.runtimeConfig)
		.resolve("endpoint_config::AccountIdEndpointMode"),
		"AwsSdkFeature" to
		AwsRuntimeType.awsRuntime(codegenContext.runtimeConfig)
		.resolve("sdk_feature::AwsSdkFeature"),
		"tracing" to RuntimeType.Tracing,
		)

		override fun loadBuiltInFromServiceConfig(
		@@ -236,5 +242,31 @@ class AccountIdEndpointModeBuiltInParamDecorator : ConditionalDecorator(
		}
		},
		)

		override fun serviceRuntimePluginCustomizations(
		codegenContext: ClientCodegenContext,
		baseCustomizations: List<ServiceRuntimePluginCustomization>,
		): List<ServiceRuntimePluginCustomization> =
		baseCustomizations + listOf(AccountIdEndpointFeatureTrackerInterceptor(codegenContext))
		},
		)

		private class AccountIdEndpointFeatureTrackerInterceptor(codegenContext: ClientCodegenContext) :
		ServiceRuntimePluginCustomization() {
		override fun section(section: ServiceRuntimePluginSection) =
		writable {
		if (section is ServiceRuntimePluginSection.RegisterRuntimeComponents) {
		section.registerInterceptor(this) {
		rustTemplate(
		"#{Interceptor}",
		"Interceptor" to
		RuntimeType.forInlineDependency(
		InlineAwsDependency.forRustFile(
		"account_id_endpoint",
		),
		).resolve("AccountIdEndpointFeatureTrackerInterceptor"),
		)
		}
		}
		}
		}

aws/rust-runtime/Cargo.lock

+2 −2

Original line number	Diff line number	Diff line
		@@ -85,7 +85,7 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"

		[[package]]
		name = "aws-credential-types"
		version = "1.2.6"
		version = "1.2.7"
		dependencies = [
		"async-trait",
		"aws-smithy-async",
		@@ -143,7 +143,7 @@ dependencies = [

		[[package]]
		name = "aws-runtime"
		version = "1.5.10"
		version = "1.5.11"
		dependencies = [
		"arbitrary",
		"aws-credential-types",

aws/rust-runtime/aws-credential-types/Cargo.toml

+1 −1

Original line number	Diff line number	Diff line
		[package]
		name = "aws-credential-types"
		version = "1.2.6"
		version = "1.2.7"
		authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
		description = "Types for AWS SDK credentials."
		edition = "2021"

aws/rust-runtime/aws-credential-types/src/credential_feature.rs

+2 −0

Original line number	Diff line number	Diff line
		@@ -9,6 +9,8 @@ use aws_smithy_types::config_bag::{Storable, StoreAppend};
		#[non_exhaustive]
		#[derive(Clone, Debug, Eq, PartialEq)]
		pub enum AwsCredentialFeature {
		/// An operation where credential resolution resolved an account ID
		ResolvedAccountId,
		/// An operation called using credentials resolved from code, cli parameters, session object, or client instance
		CredentialsCode,
		/// An operation called using credentials resolved from environment variables

aws/rust-runtime/aws-credential-types/src/credentials_impl.rs

+37 −7

Original line number	Diff line number	Diff line
		@@ -396,11 +396,19 @@ impl From<Credentials> for Identity {

		builder.set_expiration(expiry);

		if let Some(features) = val.get_property::<Vec<AwsCredentialFeature>>().cloned() {
		let features = val.get_property::<Vec<AwsCredentialFeature>>().cloned();
		let has_account_id = val.account_id().is_some();

		if features.is_some() \|\| has_account_id {
		let mut layer = Layer::new("IdentityResolutionFeatureIdTracking");
		if let Some(features) = features {
		for feat in features {
		layer.store_append(feat);
		}
		}
		if has_account_id {
		layer.store_append(AwsCredentialFeature::ResolvedAccountId);
		}
		builder.set_property(layer.freeze());
		}

		@@ -413,9 +421,6 @@ mod test {
		use crate::Credentials;
		use std::time::{Duration, UNIX_EPOCH};

		#[cfg(feature = "test-util")]
		use crate::credential_feature::AwsCredentialFeature;

		#[test]
		fn debug_impl() {
		let creds = Credentials::new(
		@@ -451,7 +456,7 @@ mod test {
		#[derive(Clone, Debug)]
		struct Foo;
		let mut creds1 = Credentials::for_tests_with_session_token();
		creds1.set_property(AwsCredentialFeature::CredentialsCode);
		creds1.set_property(crate::credential_feature::AwsCredentialFeature::CredentialsCode);

		let mut creds2 = Credentials::for_tests_with_session_token();
		creds2.set_property(Foo);
		@@ -462,6 +467,7 @@ mod test {
		#[cfg(feature = "test-util")]
		#[test]
		fn identity_inherits_feature_properties() {
		use crate::credential_feature::AwsCredentialFeature;
		use aws_smithy_runtime_api::client::identity::Identity;
		use aws_smithy_types::config_bag::FrozenLayer;

		@@ -485,4 +491,28 @@ mod test {
		feature_props.reverse();
		assert_eq!(maybe_props, feature_props)
		}

		#[cfg(feature = "test-util")]
		#[test]
		fn from_credentials_adds_resolved_account_id_feature() {
		use crate::credential_feature::AwsCredentialFeature;
		use aws_smithy_runtime_api::client::identity::Identity;
		use aws_smithy_types::config_bag::FrozenLayer;

		let creds = Credentials::builder()
		.access_key_id("test")
		.secret_access_key("test")
		.account_id("123456789012")
		.provider_name("test")
		.build();

		let identity = Identity::from(creds);

		let layer = identity.property::<FrozenLayer>().unwrap();
		let features = layer
		.load::<AwsCredentialFeature>()
		.cloned()
		.collect::<Vec<_>>();
		assert!(features.contains(&AwsCredentialFeature::ResolvedAccountId));
		}
		}