Skip to content

Commit 80ee06b

Browse files
vcjanavcjana
committed
Refactor endpoint override tracking to use interceptor pattern
1 parent c0e8fdd commit 80ee06b

File tree

6 files changed

+111
-77
lines changed

6 files changed

+111
-77
lines changed

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

Lines changed: 29 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -7,51 +7,42 @@ package software.amazon.smithy.rustsdk
77

88
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
99
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
10-
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
10+
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
11+
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection
12+
import software.amazon.smithy.rust.codegen.core.rustlang.writable
1113
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
12-
import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization
13-
import software.amazon.smithy.rust.codegen.core.smithy.customize.adhocCustomization
14-
import software.amazon.smithy.rust.codegen.core.util.dq
15-
import software.amazon.smithy.rust.codegen.core.util.sdkId
14+
import software.amazon.smithy.rustsdk.InlineAwsDependency
1615

1716
/**
18-
* Decorator that tracks endpoint override business metric when endpoint URL is configured
19-
* via any source: global endpoint (AWS_ENDPOINT_URL), service-specific endpoint
20-
* (AWS_ENDPOINT_URL_<SERVICE>), config file, or SdkConfig builder.
21-
*
22-
* Note: Direct Config::builder().endpoint_url() calls are not tracked due to architecture
23-
* limitations that would cause method name conflicts.
17+
* Decorator that tracks endpoint override business metric when endpoint URL is configured.
2418
*/
2519
class EndpointOverrideMetricDecorator : ClientCodegenDecorator {
2620
override val name: String = "EndpointOverrideMetric"
2721
override val order: Byte = 0
2822

29-
override fun extraSections(codegenContext: ClientCodegenContext): List<AdHocCustomization> {
30-
val runtimeConfig = codegenContext.runtimeConfig
31-
val awsRuntime = AwsRuntimeType.awsRuntime(runtimeConfig)
32-
val serviceId = codegenContext.serviceShape.sdkId()
23+
override fun serviceRuntimePluginCustomizations(
24+
codegenContext: ClientCodegenContext,
25+
baseCustomizations: List<ServiceRuntimePluginCustomization>,
26+
): List<ServiceRuntimePluginCustomization> =
27+
baseCustomizations + listOf(EndpointOverrideFeatureTrackerInterceptor(codegenContext))
28+
}
3329

34-
return listOf(
35-
adhocCustomization<SdkConfigSection.CopySdkConfigToClientConfig> { section ->
36-
rustTemplate(
37-
"""
38-
// Track endpoint override metric if endpoint URL is configured from any source
39-
let has_global_endpoint = ${section.sdkConfig}.endpoint_url().is_some();
40-
let has_service_specific_endpoint = ${section.sdkConfig}
41-
.service_config()
42-
.and_then(|conf| conf.load_config(service_config_key(${serviceId.dq()}, ${"AWS_ENDPOINT_URL".dq()}, ${"endpoint_url".dq()})))
43-
.is_some();
44-
45-
if has_global_endpoint || has_service_specific_endpoint {
46-
${section.serviceConfigBuilder}.push_runtime_plugin(
47-
#{SharedRuntimePlugin}::new(#{EndpointOverrideRuntimePlugin}::new_with_feature_flag())
48-
);
49-
}
50-
""",
51-
"EndpointOverrideRuntimePlugin" to awsRuntime.resolve("endpoint_override::EndpointOverrideRuntimePlugin"),
52-
"SharedRuntimePlugin" to RuntimeType.smithyRuntimeApiClient(runtimeConfig).resolve("client::runtime_plugin::SharedRuntimePlugin"),
53-
)
54-
},
55-
)
56-
}
30+
private class EndpointOverrideFeatureTrackerInterceptor(codegenContext: ClientCodegenContext) :
31+
ServiceRuntimePluginCustomization() {
32+
override fun section(section: ServiceRuntimePluginSection) =
33+
writable {
34+
if (section is ServiceRuntimePluginSection.RegisterRuntimeComponents) {
35+
section.registerInterceptor(this) {
36+
rustTemplate(
37+
"#{Interceptor}",
38+
"Interceptor" to
39+
RuntimeType.forInlineDependency(
40+
InlineAwsDependency.forRustFile(
41+
"endpoint_override",
42+
),
43+
).resolve("EndpointOverrideFeatureTrackerInterceptor"),
44+
)
45+
}
46+
}
47+
}
5748
}

‎aws/codegen-aws-sdk/src/test/kotlin/software/amazon/smithy/rustsdk/EndpointOverrideMetricDecoratorTest.kt‎

Lines changed: 42 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ package software.amazon.smithy.rustsdk
77

88
import org.junit.jupiter.api.Test
99
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
10+
import software.amazon.smithy.rust.codegen.core.rustlang.Feature
1011
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
1112
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
1213
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
@@ -83,15 +84,24 @@ class EndpointOverrideMetricDecoratorTest {
8384
fun `endpoint override metric appears when set via SdkConfig`() {
8485
val testParams = awsIntegrationTestParams()
8586

86-
awsSdkIntegrationTest(model, testParams) { context, rustCrate ->
87+
awsSdkIntegrationTest(
88+
model,
89+
testParams,
90+
environment = mapOf("RUSTUP_TOOLCHAIN" to "1.88.0"),
91+
) { context, rustCrate ->
8792
val rc = context.runtimeConfig
8893
val moduleName = context.moduleUseName()
94+
95+
// Enable test-util feature for aws-runtime
96+
rustCrate.mergeFeature(Feature("test-util", true, listOf("aws-runtime/test-util")))
97+
8998
rustCrate.integrationTest("endpoint_override_via_sdk_config") {
9099
rustTemplate(
91100
"""
92101
use $moduleName::config::Region;
93102
use $moduleName::Client;
94103
use #{capture_request};
104+
use #{assert_ua_contains_metric_values};
95105
96106
##[#{tokio}::test]
97107
async fn metric_tracked_when_endpoint_set_via_sdk_config() {
@@ -121,29 +131,18 @@ class EndpointOverrideMetricDecoratorTest {
121131
uri
122132
);
123133
124-
// Verify metric 'N' is present
125-
let user_agent = std::str::from_utf8(
126-
request
127-
.headers()
128-
.get("x-amz-user-agent")
129-
.expect("x-amz-user-agent header missing")
130-
.as_bytes(),
131-
)
132-
.expect("valid utf8");
133-
134-
let has_metric = user_agent
135-
.split_whitespace()
136-
.any(|part| part.starts_with("m/") && part.contains("N"));
137-
138-
assert!(
139-
has_metric,
140-
"Expected metric 'N' in user agent, got: {}",
141-
user_agent
142-
);
134+
// Verify metric 'N' is present in x-amz-user-agent header
135+
let user_agent = request
136+
.headers()
137+
.get("x-amz-user-agent")
138+
.expect("x-amz-user-agent header missing");
139+
140+
assert_ua_contains_metric_values(user_agent, &["N"]);
143141
}
144142
""",
145143
*preludeScope,
146144
"capture_request" to RuntimeType.captureRequest(rc),
145+
"assert_ua_contains_metric_values" to AwsRuntimeType.awsRuntime(rc).resolve("user_agent::test_util::assert_ua_contains_metric_values"),
147146
"SdkConfig" to AwsRuntimeType.awsTypes(rc).resolve("sdk_config::SdkConfig"),
148147
"tokio" to CargoDependency.Tokio.toType(),
149148
)
@@ -155,15 +154,24 @@ class EndpointOverrideMetricDecoratorTest {
155154
fun `no endpoint override metric when endpoint not set`() {
156155
val testParams = awsIntegrationTestParams()
157156

158-
awsSdkIntegrationTest(model, testParams) { context, rustCrate ->
157+
awsSdkIntegrationTest(
158+
model,
159+
testParams,
160+
environment = mapOf("RUSTUP_TOOLCHAIN" to "1.88.0"),
161+
) { context, rustCrate ->
159162
val rc = context.runtimeConfig
160163
val moduleName = context.moduleUseName()
164+
165+
// Enable test-util feature for aws-runtime
166+
rustCrate.mergeFeature(Feature("test-util", true, listOf("aws-runtime/test-util")))
167+
161168
rustCrate.integrationTest("no_endpoint_override") {
162169
rustTemplate(
163170
"""
164171
use $moduleName::config::{Credentials, Region, SharedCredentialsProvider};
165172
use $moduleName::{Config, Client};
166173
use #{capture_request};
174+
use #{assert_ua_contains_metric_values};
167175
168176
##[#{tokio}::test]
169177
async fn no_metric_when_endpoint_not_overridden() {
@@ -192,28 +200,25 @@ class EndpointOverrideMetricDecoratorTest {
192200
);
193201
194202
// Verify metric 'N' is NOT present
195-
let user_agent = std::str::from_utf8(
196-
request
197-
.headers()
198-
.get("x-amz-user-agent")
199-
.map(|v| v.as_bytes())
200-
.unwrap_or(b""),
201-
)
202-
.unwrap_or("");
203-
204-
let has_metric = user_agent
205-
.split_whitespace()
206-
.any(|part| part.starts_with("m/") && part.contains("N"));
207-
203+
let user_agent = request
204+
.headers()
205+
.get("x-amz-user-agent")
206+
.expect("x-amz-user-agent header should be present");
207+
208+
// This should panic if 'N' is found
209+
let result = std::panic::catch_unwind(|| {
210+
assert_ua_contains_metric_values(user_agent, &["N"]);
211+
});
212+
208213
assert!(
209-
!has_metric,
210-
"Did not expect metric 'N' when endpoint not overridden, got: {}",
211-
user_agent
214+
result.is_err(),
215+
"Metric 'N' should NOT be present when endpoint not overridden"
212216
);
213217
}
214218
""",
215219
*preludeScope,
216220
"capture_request" to RuntimeType.captureRequest(rc),
221+
"assert_ua_contains_metric_values" to AwsRuntimeType.awsRuntime(rc).resolve("user_agent::test_util::assert_ua_contains_metric_values"),
217222
"tokio" to CargoDependency.Tokio.toType(),
218223
)
219224
}

‎aws/rust-runtime/Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎aws/rust-runtime/aws-inlineable/src/endpoint_override.rs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
use aws_runtime::sdk_feature::AwsSdkFeature;
7+
use aws_smithy_runtime_api::{
8+
box_error::BoxError,
9+
client::interceptors::{context::BeforeSerializationInterceptorContextRef, Intercept},
10+
};
11+
use aws_smithy_types::config_bag::ConfigBag;
12+
use aws_types::endpoint_config::EndpointUrl;
13+
14+
/// Interceptor that tracks AWS SDK features for endpoint override.
15+
#[derive(Debug, Default)]
16+
pub(crate) struct EndpointOverrideFeatureTrackerInterceptor;
17+
18+
impl Intercept for EndpointOverrideFeatureTrackerInterceptor {
19+
fn name(&self) -> &'static str {
20+
"EndpointOverrideFeatureTrackerInterceptor"
21+
}
22+
23+
fn read_before_execution(
24+
&self,
25+
_context: &BeforeSerializationInterceptorContextRef<'_>,
26+
cfg: &mut ConfigBag,
27+
) -> Result<(), BoxError> {
28+
if cfg.load::<EndpointUrl>().is_some() {
29+
cfg.interceptor_state()
30+
.store_append(AwsSdkFeature::EndpointOverride);
31+
}
32+
Ok(())
33+
}
34+
}

‎aws/rust-runtime/aws-inlineable/src/lib.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ pub mod account_id_endpoint;
2929
/// Supporting code for the aws-chunked content encoding.
3030
pub mod aws_chunked;
3131

32+
/// Supporting code for tracking endpoint override feature.
33+
#[allow(dead_code)]
34+
pub mod endpoint_override;
35+
3236
/// Supporting code to determine auth scheme options based on the `authSchemes` endpoint list property.
3337
#[allow(dead_code)]
3438
pub mod endpoint_auth;

‎gradle.properties‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,4 @@ allowLocalDeps=false
1717
# Avoid registering dependencies/plugins/tasks that are only used for testing purposes
1818
isTestingEnabled=true
1919
# codegen publication version
20-
codegenVersion=0.1.6
20+
codegenVersion=0.1.7

0 commit comments

Comments
 (0)