Track features for account based endpoints (#4298)

ysaito1001 · web-flow · commit eaa6135f7e99 · 2025-09-18T13:25:27.000-05:00
diff --git a/aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AccountIdEndpointParamsDecorator.kt b/aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AccountIdEndpointParamsDecorator.kt
@@ -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"),
+                    )
+                }
+            }
+        }
+}
diff --git a/aws/rust-runtime/Cargo.lock b/aws/rust-runtime/Cargo.lock
diff --git a/aws/rust-runtime/aws-credential-types/Cargo.toml b/aws/rust-runtime/aws-credential-types/Cargo.toml
@@ -1,6 +1,6 @@
 [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"
diff --git a/aws/rust-runtime/aws-credential-types/src/credential_feature.rs b/aws/rust-runtime/aws-credential-types/src/credential_feature.rs
@@ -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
diff --git a/aws/rust-runtime/aws-credential-types/src/credentials_impl.rs b/aws/rust-runtime/aws-credential-types/src/credentials_impl.rs
@@ -396,10 +396,18 @@ 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");
-            for feat in features {
-                layer.store_append(feat);
+            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));
+    }
 }
diff --git a/aws/rust-runtime/aws-inlineable/src/account_id_endpoint.rs b/aws/rust-runtime/aws-inlineable/src/account_id_endpoint.rs
@@ -0,0 +1,55 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_runtime::sdk_feature::AwsSdkFeature;
+use aws_smithy_runtime_api::{
+    box_error::BoxError,
+    client::interceptors::{context::BeforeSerializationInterceptorContextRef, Intercept},
+};
+use aws_smithy_types::config_bag::ConfigBag;
+use aws_types::endpoint_config::AccountIdEndpointMode;
+
+// Interceptor that tracks AWS SDK features for the account based endpoints.
+#[derive(Debug, Default)]
+pub(crate) struct AccountIdEndpointFeatureTrackerInterceptor;
+
+impl Intercept for AccountIdEndpointFeatureTrackerInterceptor {
+    fn name(&self) -> &'static str {
+        "AccountIdEndpointFeatureTrackerInterceptor"
+    }
+
+    fn read_before_execution(
+        &self,
+        _context: &BeforeSerializationInterceptorContextRef<'_>,
+        cfg: &mut ConfigBag,
+    ) -> Result<(), BoxError> {
+        match cfg
+            .load::<AccountIdEndpointMode>()
+            .cloned()
+            .unwrap_or_default()
+        {
+            AccountIdEndpointMode::Preferred => {
+                cfg.interceptor_state()
+                    .store_append(AwsSdkFeature::AccountIdModePreferred);
+            }
+            AccountIdEndpointMode::Required => {
+                cfg.interceptor_state()
+                    .store_append(AwsSdkFeature::AccountIdModeRequired);
+            }
+            AccountIdEndpointMode::Disabled => {
+                cfg.interceptor_state()
+                    .store_append(AwsSdkFeature::AccountIdModeDisabled);
+            }
+            otherwise => {
+                ::tracing::warn!(
+                    "Attempted to track an SDK feature for `{otherwise:?}`, which is not recognized in the current version of the SDK. \
+                    Consider upgrading to the latest version to ensure that it is properly tracked."
+                );
+            }
+        }
+
+        Ok(())
+    }
+}
diff --git a/aws/rust-runtime/aws-inlineable/src/lib.rs b/aws/rust-runtime/aws-inlineable/src/lib.rs
@@ -22,6 +22,10 @@
     unreachable_pub
 )]
 
+/// Supporting code for the account based endpoints.
+#[allow(dead_code)]
+pub mod account_id_endpoint;
+
 /// Supporting code to determine auth scheme options based on the `authSchemes` endpoint list property.
 #[allow(dead_code)]
 pub mod endpoint_auth;
diff --git a/aws/rust-runtime/aws-runtime/Cargo.toml b/aws/rust-runtime/aws-runtime/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "aws-runtime"
-version = "1.5.10"
+version = "1.5.11"
 authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
 description = "Runtime support code for the AWS SDK. This crate isn't intended to be used directly."
 edition = "2021"
diff --git a/aws/rust-runtime/aws-runtime/src/sdk_feature.rs b/aws/rust-runtime/aws-runtime/src/sdk_feature.rs
@@ -12,6 +12,12 @@ use aws_smithy_types::config_bag::{Storable, StoreAppend};
 #[non_exhaustive]
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum AwsSdkFeature {
+    /// An operation called with account ID mode set to preferred
+    AccountIdModePreferred,
+    /// An operation called with account ID mode set to disabled
+    AccountIdModeDisabled,
+    /// An operation called with account ID mode set to required
+    AccountIdModeRequired,
     /// Indicates that an operation was called by the S3 Transfer Manager
     S3Transfer,
     /// Calling an SSO-OIDC operation as part of the SSO login flow, when using the OAuth2.0 device code grant
diff --git a/aws/rust-runtime/aws-runtime/src/user_agent/metrics.rs b/aws/rust-runtime/aws-runtime/src/user_agent/metrics.rs
@@ -216,6 +216,9 @@ impl ProvideBusinessMetric for AwsSdkFeature {
     fn provide_business_metric(&self) -> Option<BusinessMetric> {
         use AwsSdkFeature::*;
         match self {
+            AccountIdModePreferred => Some(BusinessMetric::AccountIdModePreferred),
+            AccountIdModeDisabled => Some(BusinessMetric::AccountIdModeDisabled),
+            AccountIdModeRequired => Some(BusinessMetric::AccountIdModeRequired),
             S3Transfer => Some(BusinessMetric::S3Transfer),
             SsoLoginDevice => Some(BusinessMetric::SsoLoginDevice),
             SsoLoginAuth => Some(BusinessMetric::SsoLoginAuth),
@@ -227,6 +230,7 @@ impl ProvideBusinessMetric for AwsCredentialFeature {
     fn provide_business_metric(&self) -> Option<BusinessMetric> {
         use AwsCredentialFeature::*;
         match self {
+            ResolvedAccountId => Some(BusinessMetric::ResolvedAccountId),
             CredentialsCode => Some(BusinessMetric::CredentialsCode),
             CredentialsEnvVars => Some(BusinessMetric::CredentialsEnvVars),
             CredentialsEnvVarsStsWebIdToken => {
diff --git a/aws/sdk/integration-tests/dynamodb/Cargo.toml b/aws/sdk/integration-tests/dynamodb/Cargo.toml
@@ -13,7 +13,7 @@ publish = false
 [dependencies]
 approx = "0.5.1"
 aws-config = { path = "../../build/aws-sdk/sdk/aws-config" }
-aws-runtime = { path = "../../build/aws-sdk/sdk/aws-runtime" }
+aws-runtime = { path = "../../build/aws-sdk/sdk/aws-runtime", features = ["test-util"] }
 aws-credential-types = { path = "../../build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }
 aws-sdk-dynamodb = { path = "../../build/aws-sdk/sdk/dynamodb", features = ["behavior-version-latest"] }
 aws-smithy-async = { path = "../../build/aws-sdk/sdk/aws-smithy-async", features = ["test-util"] }
diff --git a/aws/sdk/integration-tests/dynamodb/tests/account-based-endpoints.rs b/aws/sdk/integration-tests/dynamodb/tests/account-based-endpoints.rs
@@ -5,6 +5,7 @@
 
 use aws_config::Region;
 use aws_credential_types::Credentials;
+use aws_runtime::user_agent::test_util::assert_ua_contains_metric_values;
 use aws_sdk_dynamodb::{
     config::Builder,
     error::{DisplayErrorContext, SdkError},
@@ -53,40 +54,48 @@ async fn call_operation(
 
 #[tokio::test]
 async fn basic_positive_cases() {
-    let test_cases: &[(fn(Builder) -> Builder, &str, &str)] = &[
+    let test_cases: &[(fn(Builder) -> Builder, &str, &str, &[&'static str])] = &[
         (
             std::convert::identity,
             "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
             "https://333333333333.ddb.us-east-1.amazonaws.com/",
+            &["P", "T"],
         ),
         (
             std::convert::identity,
             "table_name", // doesn't specify ARN for the table name
             "https://123456789012.ddb.us-east-1.amazonaws.com/", // the account ID should come from credentials
+            &["P", "T"],
         ),
         (
             |b: Builder| b.credentials_provider(Credentials::for_tests()), // credentials do not provide an account ID
             "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
             "https://333333333333.ddb.us-east-1.amazonaws.com/",
+            &["P"],
         ),
         (
             |b: Builder| b.account_id_endpoint_mode(AccountIdEndpointMode::Preferred), // sets the default mode `Preferred` explicitly
             "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
             "https://333333333333.ddb.us-east-1.amazonaws.com/",
+            &["P", "T"],
         ),
         (
             |b: Builder| b.account_id_endpoint_mode(AccountIdEndpointMode::Disabled),
             "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
             "https://dynamodb.us-east-1.amazonaws.com/",
+            &["Q", "T"],
         ),
         (
             |b: Builder| b.account_id_endpoint_mode(AccountIdEndpointMode::Required),
             "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
             "https://333333333333.ddb.us-east-1.amazonaws.com/",
+            &["R", "T"],
         ),
     ];
 
-    for (i, (update_builder, table_name, expected_uri)) in test_cases.into_iter().enumerate() {
+    for (i, (update_builder, table_name, expected_uri, expected_metrics)) in
+        test_cases.into_iter().enumerate()
+    {
         let (client, rx) = test_client(*update_builder);
         let _ = call_operation(client, table_name).await;
         let req = rx.expect_request();
@@ -95,6 +104,10 @@ async fn basic_positive_cases() {
             req.uri(),
             "on the {i}th test case where table name is `{table_name}`"
         );
+
+        // Test User-Agent metrics for account ID tracking
+        let user_agent = req.headers().get("x-amz-user-agent").unwrap();
+        assert_ua_contains_metric_values(user_agent, expected_metrics);
     }
 }
 
diff --git a/tools/ci-build/changelogger/smithy-rs-maintainers.txt b/tools/ci-build/changelogger/smithy-rs-maintainers.txt
@@ -9,3 +9,4 @@ rschmitt
 rhernandez35
 aajtodd
 landonxjames
+vcjana
