From fdb0076e93bfa29e2506038eb7b96cb60996d8d4 Mon Sep 17 00:00:00 2001 From: Russell Cohen Date: Tue, 23 May 2023 12:50:16 -0400 Subject: [PATCH] Add tests of timestream + docs --- .../src/cache/expiring_cache.rs | 8 - .../aws-inlineable/src/endpoint_discovery.rs | 165 ++++++- aws/rust-runtime/aws-inlineable/src/lib.rs | 1 + .../smithy/rustsdk/AwsCrateDocsDecorator.kt | 9 +- .../software/amazon/smithy/rustsdk/AwsDocs.kt | 25 +- .../rustsdk/IntegrationTestDependencies.kt | 3 + .../customize/ServiceSpecificDecorator.kt | 6 + .../timestream/TimestreamDecorator.kt | 44 +- aws/sdk/gradle.properties | 3 +- aws/sdk/integration-tests/Cargo.toml | 1 + .../timestreamquery/Cargo.toml | 20 + .../timestreamquery/tests/endpoint_disco.rs | 76 ++++ .../timestreamquery/tests/traffic.json | 407 ++++++++++++++++++ .../integration-tests/aws-sdk-s3/Cargo.toml | 1 + .../src/main/kotlin/aws/sdk/ServiceLoader.kt | 2 +- .../rust/codegen/core/rustlang/RustWriter.kt | 13 +- .../core/smithy/customize/Customization.kt | 15 + .../aws-smithy-client/src/dvr/replay.rs | 29 +- 18 files changed, 772 insertions(+), 56 deletions(-) create mode 100644 aws/sdk/integration-tests/timestreamquery/Cargo.toml create mode 100644 aws/sdk/integration-tests/timestreamquery/tests/endpoint_disco.rs create mode 100644 aws/sdk/integration-tests/timestreamquery/tests/traffic.json diff --git a/aws/rust-runtime/aws-credential-types/src/cache/expiring_cache.rs b/aws/rust-runtime/aws-credential-types/src/cache/expiring_cache.rs index 93ea51cf11..67556aad9b 100644 --- a/aws/rust-runtime/aws-credential-types/src/cache/expiring_cache.rs +++ b/aws/rust-runtime/aws-credential-types/src/cache/expiring_cache.rs @@ -59,14 +59,6 @@ where .map(|(creds, _expiry)| creds) } - /// Attempts to load the cached value if it has been set - /// - /// # Panics - /// This function panics if it is called from an asynchronous context - pub fn try_blocking_get(&self) -> Option { - self.value.blocking_read().get().map(|(v, _exp)| v.clone()) - } - /// Attempts to refresh the cached value with the given future. /// If multiple threads attempt to refresh at the same time, one of them will win, /// and the others will await that thread's result rather than multiple refreshes occurring. diff --git a/aws/rust-runtime/aws-inlineable/src/endpoint_discovery.rs b/aws/rust-runtime/aws-inlineable/src/endpoint_discovery.rs index be10bfe1d9..9c1ab6400e 100644 --- a/aws/rust-runtime/aws-inlineable/src/endpoint_discovery.rs +++ b/aws/rust-runtime/aws-inlineable/src/endpoint_discovery.rs @@ -6,6 +6,7 @@ //! Maintain a cache of discovered endpoints use aws_smithy_async::rt::sleep::AsyncSleep; +use aws_smithy_async::time::TimeSource; use aws_smithy_client::erase::boxclone::BoxFuture; use aws_smithy_http::endpoint::{ResolveEndpoint, ResolveEndpointError}; use aws_smithy_types::endpoint::Endpoint; @@ -24,6 +25,7 @@ pub struct ReloadEndpoint { error: Arc>>, rx: Receiver<()>, sleep: Arc, + time: Arc, } impl Debug for ReloadEndpoint { @@ -45,24 +47,29 @@ impl ReloadEndpoint { /// An infinite loop task that will reload the endpoint /// - /// This task will terminate when the corresponding [`EndpointCache`] is dropped. + /// This task will terminate when the corresponding [`Client`](crate::Client) is dropped. pub async fn reload_task(mut self) { loop { match self.rx.try_recv() { Ok(_) | Err(TryRecvError::Closed) => break, _ => {} } - let should_reload = self - .endpoint - .lock() - .unwrap() - .as_ref() - .map(|e| e.is_expired()) - .unwrap_or(true); - if should_reload { - self.reload_once().await; - } - self.sleep.sleep(Duration::from_secs(60)).await + self.reload_increment(self.time.now()).await; + self.sleep.sleep(Duration::from_secs(60)).await; + } + } + + async fn reload_increment(&self, now: SystemTime) { + let should_reload = self + .endpoint + .lock() + .unwrap() + .as_ref() + .map(|e| e.is_expired(now)) + .unwrap_or(true); + if should_reload { + tracing::debug!("reloading endpoint, previous endpoint was expired"); + self.reload_once().await; } } } @@ -88,9 +95,10 @@ struct ExpiringEndpoint { } impl ExpiringEndpoint { - fn is_expired(&self) -> bool { - match SystemTime::now().duration_since(self.expiry) { - Err(e) => true, + fn is_expired(&self, now: SystemTime) -> bool { + tracing::debug!(expiry = ?self.expiry, now = ?now, delta = ?self.expiry.duration_since(now), "checking expiry status of endpoint"); + match self.expiry.duration_since(now) { + Err(_) => true, Ok(t) => t < Duration::from_secs(120), } } @@ -99,6 +107,7 @@ impl ExpiringEndpoint { pub(crate) async fn create_cache( loader_fn: impl Fn() -> F + Send + Sync + 'static, sleep: Arc, + time: Arc, ) -> Result<(EndpointCache, ReloadEndpoint), ResolveEndpointError> where F: Future> + Send + 'static, @@ -117,11 +126,12 @@ where error: error_holder, rx, sleep, + time, }; reloader.reload_once().await; - if let Err(e) = cache.resolve_endpoint() { - return Err(e); - } + // if we didn't successfully get an endpoint, bail out so the client knows + // configuration failed to work + cache.resolve_endpoint()?; Ok((cache, reloader)) } @@ -145,26 +155,135 @@ impl EndpointCache { #[cfg(test)] mod test { use crate::endpoint_discovery::{create_cache, EndpointCache}; + use aws_credential_types::time_source::TimeSource; use aws_smithy_async::rt::sleep::TokioSleep; + use aws_smithy_async::test_util::controlled_time_and_sleep; + use aws_smithy_async::time::SystemTimeSource; use aws_smithy_http::endpoint::ResolveEndpointError; + use aws_smithy_types::endpoint::Endpoint; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; - - fn check_send() {} + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + use tokio::time::timeout; fn check_send_v(t: T) -> T { t } #[tokio::test] + #[allow(unused_must_use)] async fn check_traits() { - // check_send::(); - let (cache, reloader) = create_cache( - || async { Err(ResolveEndpointError::message("stub")) }, + || async { + Ok(( + Endpoint::builder().url("http://foo.com").build(), + SystemTime::now(), + )) + }, Arc::new(TokioSleep::new()), + Arc::new(SystemTimeSource::new()), ) .await .unwrap(); check_send_v(reloader.reload_task()); + check_send_v(cache); + } + + #[tokio::test] + async fn erroring_endpoint_always_reloaded() { + let expiry = UNIX_EPOCH + Duration::from_secs(123456789); + let ct = Arc::new(AtomicUsize::new(0)); + let (cache, reloader) = create_cache( + move || { + let shared_ct = ct.clone(); + shared_ct.fetch_add(1, Ordering::AcqRel); + async move { + Ok(( + Endpoint::builder() + .url(format!("http://foo.com/{shared_ct:?}")) + .build(), + expiry, + )) + } + }, + Arc::new(TokioSleep::new()), + Arc::new(SystemTimeSource::new()), + ) + .await + .expect("returns an endpoint"); + assert_eq!( + cache.resolve_endpoint().expect("ok").url(), + "http://foo.com/1" + ); + // 120 second buffer + reloader + .reload_increment(expiry - Duration::from_secs(240)) + .await; + assert_eq!( + cache.resolve_endpoint().expect("ok").url(), + "http://foo.com/1" + ); + + reloader.reload_increment(expiry).await; + assert_eq!( + cache.resolve_endpoint().expect("ok").url(), + "http://foo.com/2" + ); + } + + #[tokio::test] + async fn test_advance_of_task() { + let expiry = UNIX_EPOCH + Duration::from_secs(123456789); + // expires in 8 minutes + let (time, sleep, mut gate) = controlled_time_and_sleep(expiry - Duration::from_secs(239)); + let ct = Arc::new(AtomicUsize::new(0)); + let (cache, reloader) = create_cache( + move || { + let shared_ct = ct.clone(); + shared_ct.fetch_add(1, Ordering::AcqRel); + async move { + Ok(( + Endpoint::builder() + .url(format!("http://foo.com/{shared_ct:?}")) + .build(), + expiry, + )) + } + }, + Arc::new(sleep.clone()), + Arc::new(time.clone()), + ) + .await + .expect("first load success"); + let reload_task = tokio::spawn(reloader.reload_task()); + assert!(!reload_task.is_finished()); + // expiry occurs after 2 sleeps + // t = 0 + assert_eq!( + gate.expect_sleep().await.duration(), + Duration::from_secs(60) + ); + assert_eq!(cache.resolve_endpoint().unwrap().url(), "http://foo.com/1"); + // t = 60 + + let sleep = gate.expect_sleep().await; + // we're still holding the drop guard, so we haven't expired yet. + assert_eq!(cache.resolve_endpoint().unwrap().url(), "http://foo.com/1"); + assert_eq!(sleep.duration(), Duration::from_secs(60)); + sleep.allow_progress(); + // t = 120 + + let sleep = gate.expect_sleep().await; + assert_eq!(cache.resolve_endpoint().unwrap().url(), "http://foo.com/2"); + sleep.allow_progress(); + + let sleep = gate.expect_sleep().await; + drop(cache); + sleep.allow_progress(); + + timeout(Duration::from_secs(1), reload_task) + .await + .expect("task finishes successfully") + .expect("finishes"); } } diff --git a/aws/rust-runtime/aws-inlineable/src/lib.rs b/aws/rust-runtime/aws-inlineable/src/lib.rs index 19f1a2a2d2..3d8043fdb5 100644 --- a/aws/rust-runtime/aws-inlineable/src/lib.rs +++ b/aws/rust-runtime/aws-inlineable/src/lib.rs @@ -46,4 +46,5 @@ pub mod route53_resource_id_preprocessor; /// Convert a streaming `SdkBody` into an aws-chunked streaming body with checksum trailers pub mod http_body_checksum; +#[allow(dead_code)] pub mod endpoint_discovery; diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCrateDocsDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCrateDocsDecorator.kt index a810b3e8a5..f167015ac2 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCrateDocsDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCrateDocsDecorator.kt @@ -21,6 +21,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rawTemplate import software.amazon.smithy.rust.codegen.core.rustlang.rust import software.amazon.smithy.rust.codegen.core.rustlang.writable import software.amazon.smithy.rust.codegen.core.smithy.RustCrate +import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocSection import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsSection import software.amazon.smithy.rust.codegen.core.smithy.generators.ManifestCustomizations @@ -85,6 +86,10 @@ class AwsCrateDocsDecorator : ClientCodegenDecorator { SdkSettings.from(codegenContext.settings).generateReadme } +sealed class DocSection(name: String) : AdHocSection(name) { + data class CreateClient(val crateName: String, val clientName: String = "client", val indent: String) : DocSection("CustomExample") +} + internal class AwsCrateDocGenerator(private val codegenContext: ClientCodegenContext) { private val logger: Logger = Logger.getLogger(javaClass.name) private val awsConfigVersion by lazy { @@ -154,8 +159,7 @@ internal class AwsCrateDocGenerator(private val codegenContext: ClientCodegenCon ##[#{tokio}::main] async fn main() -> Result<(), $shortModuleName::Error> { - let config = #{aws_config}::load_from_env().await; - let client = $shortModuleName::Client::new(&config); + #{constructClient} // ... make some calls with the client @@ -171,6 +175,7 @@ internal class AwsCrateDocGenerator(private val codegenContext: ClientCodegenCon true -> AwsCargoDependency.awsConfig(codegenContext.runtimeConfig).toDevDependency().toType() else -> writable { rust("aws_config") } }, + "constructClient" to AwsDocs.constructClient(codegenContext, indent = " "), ) template( diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsDocs.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsDocs.kt index a8dcdc33d8..c49d4639f2 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsDocs.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsDocs.kt @@ -9,6 +9,9 @@ import software.amazon.smithy.model.shapes.ShapeId import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext import software.amazon.smithy.rust.codegen.core.rustlang.Writable import software.amazon.smithy.rust.codegen.core.rustlang.docsTemplate +import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate +import software.amazon.smithy.rust.codegen.core.rustlang.writable +import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizationsOrElse import software.amazon.smithy.rust.codegen.core.util.toSnakeCase object AwsDocs { @@ -23,6 +26,24 @@ object AwsDocs { ShapeId.from("com.amazonaws.sso#SWBPortalService"), ).contains(codegenContext.serviceShape.id) + fun constructClient(codegenContext: ClientCodegenContext, indent: String): Writable { + val crateName = codegenContext.moduleName.toSnakeCase() + return writable { + writeCustomizationsOrElse( + codegenContext.rootDecorator.extraSections(codegenContext), + DocSection.CreateClient(crateName = crateName, indent = indent), + ) { + addDependency(AwsCargoDependency.awsConfig(codegenContext.runtimeConfig).toDevDependency()) + rustTemplate( + """ + let config = aws_config::load_from_env().await; + let client = $crateName::Client::new(&config); + """.trimIndent().prependIndent(indent), + ) + } + } + } + fun clientConstructionDocs(codegenContext: ClientCodegenContext): Writable = { if (canRelyOnAwsConfig(codegenContext)) { val crateName = codegenContext.moduleName.toSnakeCase() @@ -40,8 +61,7 @@ object AwsDocs { In the simplest case, creating a client looks as follows: ```rust,no_run ## async fn wrapper() { - let config = #{aws_config}::load_from_env().await; - let client = $crateName::Client::new(&config); + #{constructClient} ## } ``` @@ -76,6 +96,7 @@ object AwsDocs { [builder pattern]: https://rust-lang.github.io/api-guidelines/type-safety.html##builders-enable-construction-of-complex-values-c-builder """.trimIndent(), "aws_config" to AwsCargoDependency.awsConfig(codegenContext.runtimeConfig).toDevDependency().toType(), + "constructClient" to constructClient(codegenContext, indent = ""), ) } } diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt index 556e71d2b0..744b79c441 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt @@ -77,7 +77,10 @@ class IntegrationTestDependencies( if (hasTests) { val smithyClient = CargoDependency.smithyClient(codegenContext.runtimeConfig) .copy(features = setOf("test-util"), scope = DependencyScope.Dev) + val smithyAsync = CargoDependency.smithyAsync(codegenContext.runtimeConfig) + .copy(features = setOf("test-util"), scope = DependencyScope.Dev) addDependency(smithyClient) + addDependency(smithyAsync) addDependency(CargoDependency.smithyProtocolTestHelpers(codegenContext.runtimeConfig)) addDependency(SerdeJson) addDependency(Tokio) diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/ServiceSpecificDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/ServiceSpecificDecorator.kt index cc9c34797d..ce4d3f88ae 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/ServiceSpecificDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/ServiceSpecificDecorator.kt @@ -20,6 +20,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.generators.config.Confi import software.amazon.smithy.rust.codegen.client.smithy.generators.error.ErrorCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.ProtocolTestGenerator import software.amazon.smithy.rust.codegen.core.smithy.RustCrate +import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.BuilderCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization @@ -148,6 +149,11 @@ class ServiceSpecificDecorator( delegateTo.protocolTestGenerator(codegenContext, baseGenerator) } + override fun extraSections(codegenContext: ClientCodegenContext): List = + listOf().maybeApply(codegenContext.serviceShape) { + delegateTo.extraSections(codegenContext) + } + override fun operationRuntimePluginCustomizations( codegenContext: ClientCodegenContext, operation: OperationShape, diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/timestream/TimestreamDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/timestream/TimestreamDecorator.kt index a9d0968a02..efd80567d1 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/timestream/TimestreamDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/timestream/TimestreamDecorator.kt @@ -14,14 +14,37 @@ import software.amazon.smithy.rust.codegen.core.rustlang.Visibility import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate import software.amazon.smithy.rust.codegen.core.rustlang.toType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType +import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope import software.amazon.smithy.rust.codegen.core.smithy.RustCrate +import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization +import software.amazon.smithy.rust.codegen.core.smithy.customize.adhocCustomization +import software.amazon.smithy.rustsdk.AwsCargoDependency +import software.amazon.smithy.rustsdk.DocSection import software.amazon.smithy.rustsdk.InlineAwsDependency -/** Adds Endpoint Discovery Utility to Timestream */ +/** + * This decorator does two things: + * 1. Adds the `endpoint_discovery` inlineable + * 2. Adds a `enable_endpoint_discovery` method on client that returns a wrapped client with endpoint discovery enabled + */ class TimestreamDecorator : ClientCodegenDecorator { override val name: String = "Timestream" - override val order: Byte = 0 + override val order: Byte = -1 + override fun extraSections(codegenContext: ClientCodegenContext): List { + return listOf( + adhocCustomization { + addDependency(AwsCargoDependency.awsConfig(codegenContext.runtimeConfig).toDevDependency()) + rustTemplate( + """ + let config = aws_config::load_from_env().await; + // You MUST call `enable_endpoint_discovery` to produce a working client for this service. + let ${it.clientName} = ${it.crateName}::Client::new(&config).enable_endpoint_discovery(); + """.replaceIndent(it.indent), + ) + }, + ) + } override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) { val endpointDiscovery = InlineAwsDependency.forRustFile( "endpoint_discovery", @@ -29,6 +52,7 @@ class TimestreamDecorator : ClientCodegenDecorator { CargoDependency.Tokio.copy(scope = DependencyScope.Compile, features = setOf("sync")), ) rustCrate.lib { + // helper function to resolve an endpoint given a base client rustTemplate( """ async fn resolve_endpoint(client: &crate::Client) -> Result<(#{Endpoint}, #{SystemTime}), #{ResolveEndpointError}> { @@ -37,8 +61,7 @@ class TimestreamDecorator : ClientCodegenDecorator { #{ResolveEndpointError}::from_source("failed to call describe_endpoints", e) })?; let endpoint = describe_endpoints.endpoints().unwrap().get(0).unwrap(); - let expiry = - #{SystemTime}::now() + #{Duration}::from_secs(endpoint.cache_period_in_minutes() as u64 * 60); + let expiry = client.conf().time_source.now() + #{Duration}::from_secs(endpoint.cache_period_in_minutes() as u64 * 60); Ok(( #{Endpoint}::builder() .url(format!("https://{}", endpoint.address().unwrap())) @@ -48,18 +71,24 @@ class TimestreamDecorator : ClientCodegenDecorator { } impl Client { - pub async fn enable_endpoint_discovery(self) -> Result<(Self, #{endpoint_discovery}::ReloadEndpoint), #{ResolveEndpointError}> { + /// Enable endpoint discovery for this client + /// + /// This method MUST be called to construct a working client. + ##[must_use] + pub async fn enable_endpoint_discovery(self) -> #{Result}<(Self, #{endpoint_discovery}::ReloadEndpoint), #{ResolveEndpointError}> { let mut new_conf = self.conf().clone(); let sleep = self.conf().sleep_impl().expect("sleep impl must be provided"); + let time = ::std::sync::Arc::new(self.conf().time_source.clone()); let (resolver, reloader) = #{endpoint_discovery}::create_cache( move || { let client = self.clone(); async move { resolve_endpoint(&client).await } }, sleep, + time ) .await?; - new_conf.endpoint_resolver = std::sync::Arc::new(resolver); + new_conf.endpoint_resolver = ::std::sync::Arc::new(resolver); Ok((Self::from_conf(new_conf), reloader)) } } @@ -68,7 +97,10 @@ class TimestreamDecorator : ClientCodegenDecorator { "endpoint_discovery" to endpointDiscovery.toType(), "SystemTime" to RuntimeType.std.resolve("time::SystemTime"), "Duration" to RuntimeType.std.resolve("time::Duration"), + "SystemTimeSource" to RuntimeType.smithyAsync(codegenContext.runtimeConfig) + .resolve("time::SystemTimeSource"), *Types(codegenContext.runtimeConfig).toArray(), + *preludeScope, ) } } diff --git a/aws/sdk/gradle.properties b/aws/sdk/gradle.properties index 1cd163ac60..3bd3028606 100644 --- a/aws/sdk/gradle.properties +++ b/aws/sdk/gradle.properties @@ -3,8 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 # -# timestream requires endpoint discovery: https://github.com/awslabs/aws-sdk-rust/issues/114 -aws.services=-timestreamwrite,-timestreamquery +aws.services= # List of services to generate Event Stream operations for: aws.services.eventstream.allowlist=\ diff --git a/aws/sdk/integration-tests/Cargo.toml b/aws/sdk/integration-tests/Cargo.toml index 74a6f261dc..284bc1bcb1 100644 --- a/aws/sdk/integration-tests/Cargo.toml +++ b/aws/sdk/integration-tests/Cargo.toml @@ -15,5 +15,6 @@ members = [ "s3control", "sts", "transcribestreaming", + "timestreamquery", "webassembly", ] diff --git a/aws/sdk/integration-tests/timestreamquery/Cargo.toml b/aws/sdk/integration-tests/timestreamquery/Cargo.toml new file mode 100644 index 0000000000..b81b618c88 --- /dev/null +++ b/aws/sdk/integration-tests/timestreamquery/Cargo.toml @@ -0,0 +1,20 @@ +# This Cargo.toml is unused in generated code. It exists solely to enable these tests to compile in-situ +[package] +name = "timestream-tests" +version = "0.1.0" +authors = ["Russell Cohen "] +edition = "2021" +license = "Apache-2.0" +repository = "https://github.com/awslabs/smithy-rs" +publish = false + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dev-dependencies] +aws-credential-types = { path = "../../build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] } +aws-sdk-timestreamquery = { path = "../../build/aws-sdk/sdk/timestreamquery" } +tokio = { version = "1.23.1", features = ["full", "test-util"] } +aws-smithy-client = { path = "../../build/aws-sdk/sdk/aws-smithy-client", features = ["test-util"] } +aws-smithy-async = { path = "../../build/aws-sdk/sdk/aws-smithy-async", features = ["test-util"] } +aws-types = { path = "../../build/aws-sdk/sdk/aws-types" } +tracing-subscriber = "0.3.17" diff --git a/aws/sdk/integration-tests/timestreamquery/tests/endpoint_disco.rs b/aws/sdk/integration-tests/timestreamquery/tests/endpoint_disco.rs new file mode 100644 index 0000000000..a045840907 --- /dev/null +++ b/aws/sdk/integration-tests/timestreamquery/tests/endpoint_disco.rs @@ -0,0 +1,76 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use aws_credential_types::provider::SharedCredentialsProvider; +use aws_sdk_timestreamquery as query; +use aws_sdk_timestreamquery::config::Credentials; +use aws_smithy_async::test_util::controlled_time_and_sleep; +use aws_smithy_async::time::{SharedTimeSource, TimeSource}; +use aws_smithy_client::dvr::{MediaType, ReplayingConnection}; +use aws_types::region::Region; +use aws_types::SdkConfig; +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; + +#[tokio::test] +async fn do_endpoint_discovery() { + tracing_subscriber::fmt::init(); + let conn = ReplayingConnection::from_file("tests/traffic.json").unwrap(); + //let conn = aws_smithy_client::dvr::RecordingConnection::new(conn); + let start = UNIX_EPOCH + Duration::from_secs(1234567890); + let (ts, sleep, mut gate) = controlled_time_and_sleep(start); + let config = SdkConfig::builder() + .http_connector(conn.clone()) + .region(Region::from_static("us-west-2")) + .sleep_impl(Arc::new(sleep)) + .credentials_provider(SharedCredentialsProvider::new(Credentials::for_tests())) + .time_source(SharedTimeSource::new(ts.clone())) + .build(); + let conf = query::config::Builder::from(&config) + .make_token("0000-0000-0000") + .build(); + let (client, reloader) = query::Client::from_conf(conf) + .enable_endpoint_discovery() + .await + .expect("initial setup of endpoint discovery failed"); + + tokio::spawn(reloader.reload_task()); + + let _resp = client + .query() + .query_string("SELECT now() as time_now") + .send() + .await + .unwrap(); + + // wait 10 minutes for the endpoint to expire + while ts.now() < start + Duration::from_secs(60 * 10) { + assert_eq!( + gate.expect_sleep().await.duration(), + Duration::from_secs(60) + ); + } + + // the recording validates that this request hits another endpoint + let _resp = client + .query() + .query_string("SELECT now() as time_now") + .send() + .await + .unwrap(); + // if you want to update this test: + // conn.dump_to_file("tests/traffic.json").unwrap(); + conn.validate_body_and_headers( + Some(&[ + "x-amz-security-token", + "x-amz-date", + "content-type", + "x-amz-target", + ]), + MediaType::Json, + ) + .await + .unwrap(); +} diff --git a/aws/sdk/integration-tests/timestreamquery/tests/traffic.json b/aws/sdk/integration-tests/timestreamquery/tests/traffic.json new file mode 100644 index 0000000000..6b692ce1dc --- /dev/null +++ b/aws/sdk/integration-tests/timestreamquery/tests/traffic.json @@ -0,0 +1,407 @@ +{ + "events": [ + { + "connection_id": 0, + "action": { + "Request": { + "request": { + "uri": "https://query.timestream.us-west-2.amazonaws.com/", + "headers": { + "content-type": [ + "application/x-amz-json-1.0" + ], + "x-amz-user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head api/timestreamquery/0.0.0-local os/macos lang/rust/1.67.1" + ], + "authorization": [ + "AWS4-HMAC-SHA256 Credential=ANOTREAL/20090213/us-west-2/timestream/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token;x-amz-target;x-amz-user-agent, Signature=1b50e2545f06c8e1ca0e205c20f25a34b6aab82f3a47e4cc370e9a5fea01d08c" + ], + "x-amz-security-token": [ + "notarealsessiontoken" + ], + "x-amz-target": [ + "Timestream_20181101.DescribeEndpoints" + ], + "user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head os/macos lang/rust/1.67.1" + ], + "x-amz-date": [ + "20090213T233130Z" + ] + }, + "method": "POST" + } + } + } + }, + { + "connection_id": 0, + "action": { + "Data": { + "data": { + "Utf8": "{}" + }, + "direction": "Request" + } + } + }, + { + "connection_id": 0, + "action": { + "Eof": { + "ok": true, + "direction": "Request" + } + } + }, + { + "connection_id": 0, + "action": { + "Response": { + "response": { + "Ok": { + "status": 200, + "version": "HTTP/1.1", + "headers": { + "x-amzn-requestid": [ + "fcfdab03-2bb8-45e9-a284-b789cf7efb63" + ], + "content-type": [ + "application/x-amz-json-1.0" + ], + "date": [ + "Wed, 24 May 2023 15:51:07 GMT" + ], + "content-length": [ + "104" + ] + } + } + } + } + } + }, + { + "connection_id": 0, + "action": { + "Data": { + "data": { + "Utf8": "{\"Endpoints\":[{\"Address\":\"query-cell1.timestream.us-west-2.amazonaws.com\",\"CachePeriodInMinutes\":10}]}" + }, + "direction": "Response" + } + } + }, + { + "connection_id": 0, + "action": { + "Eof": { + "ok": true, + "direction": "Response" + } + } + }, + { + "connection_id": 1, + "action": { + "Request": { + "request": { + "uri": "https://query-cell1.timestream.us-west-2.amazonaws.com/", + "headers": { + "content-type": [ + "application/x-amz-json-1.0" + ], + "content-length": [ + "73" + ], + "x-amz-security-token": [ + "notarealsessiontoken" + ], + "user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head os/macos lang/rust/1.67.1" + ], + "x-amz-date": [ + "20090213T233130Z" + ], + "authorization": [ + "AWS4-HMAC-SHA256 Credential=ANOTREAL/20090213/us-west-2/timestream/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-date;x-amz-security-token;x-amz-target;x-amz-user-agent, Signature=8174b6ca0ece22834b562b60785f47ef354b2c1ddf7a541482f255006b5f98c2" + ], + "x-amz-user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head api/timestreamquery/0.0.0-local os/macos lang/rust/1.67.1" + ], + "x-amz-target": [ + "Timestream_20181101.Query" + ] + }, + "method": "POST" + } + } + } + }, + { + "connection_id": 1, + "action": { + "Data": { + "data": { + "Utf8": "{\"QueryString\":\"SELECT now() as time_now\",\"ClientToken\":\"0000-0000-0000\"}" + }, + "direction": "Request" + } + } + }, + { + "connection_id": 1, + "action": { + "Eof": { + "ok": true, + "direction": "Request" + } + } + }, + { + "connection_id": 1, + "action": { + "Response": { + "response": { + "Ok": { + "status": 200, + "version": "HTTP/1.1", + "headers": { + "content-type": [ + "application/x-amz-json-1.0" + ], + "date": [ + "Wed, 24 May 2023 15:51:08 GMT" + ], + "x-amzn-requestid": [ + "OFO47BQ6XAXGTIK3XH4S6GBFLQ" + ], + "content-length": [ + "318" + ] + } + } + } + } + } + }, + { + "connection_id": 1, + "action": { + "Data": { + "data": { + "Utf8": "{\"ColumnInfo\":[{\"Name\":\"time_now\",\"Type\":{\"ScalarType\":\"TIMESTAMP\"}}],\"QueryId\":\"AEDACANMQDLSTQR3SPDV6HEWYACX5IALTEWGK7JM3VSFQQ5J5F3HGKVEMNHBRHY\",\"QueryStatus\":{\"CumulativeBytesMetered\":10000000,\"CumulativeBytesScanned\":0,\"ProgressPercentage\":100.0},\"Rows\":[{\"Data\":[{\"ScalarValue\":\"2023-05-24 15:51:08.760000000\"}]}]}" + }, + "direction": "Response" + } + } + }, + { + "connection_id": 1, + "action": { + "Eof": { + "ok": true, + "direction": "Response" + } + } + }, + { + "connection_id": 2, + "action": { + "Request": { + "request": { + "uri": "https://query.timestream.us-west-2.amazonaws.com/", + "headers": { + "content-type": [ + "application/x-amz-json-1.0" + ], + "user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head os/macos lang/rust/1.67.1" + ], + "x-amz-user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head api/timestreamquery/0.0.0-local os/macos lang/rust/1.67.1" + ], + "x-amz-date": [ + "20090213T234030Z" + ], + "authorization": [ + "AWS4-HMAC-SHA256 Credential=ANOTREAL/20090213/us-west-2/timestream/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token;x-amz-target;x-amz-user-agent, Signature=32e468e574c514ba0fa4a0f0304cef82b72f82965b9e8792fd63f6efb67b297c" + ], + "x-amz-target": [ + "Timestream_20181101.DescribeEndpoints" + ], + "x-amz-security-token": [ + "notarealsessiontoken" + ] + }, + "method": "POST" + } + } + } + }, + { + "connection_id": 2, + "action": { + "Response": { + "response": { + "Ok": { + "status": 200, + "version": "HTTP/1.1", + "headers": { + "content-type": [ + "application/x-amz-json-1.0" + ], + "content-length": [ + "104" + ], + "date": [ + "Wed, 24 May 2023 15:51:07 GMT" + ], + "x-amzn-requestid": [ + "fcfdab03-2bb8-45e9-a284-b789cf7efb63" + ] + } + } + } + } + } + }, + { + "connection_id": 2, + "action": { + "Data": { + "data": { + "Utf8": "{}" + }, + "direction": "Request" + } + } + }, + { + "connection_id": 2, + "action": { + "Eof": { + "ok": true, + "direction": "Request" + } + } + }, + { + "connection_id": 2, + "action": { + "Data": { + "data": { + "Utf8": "{\"Endpoints\":[{\"Address\":\"query-cell2.timestream.us-west-2.amazonaws.com\",\"CachePeriodInMinutes\":10}]}" + }, + "direction": "Response" + } + } + }, + { + "connection_id": 2, + "action": { + "Eof": { + "ok": true, + "direction": "Response" + } + } + }, + { + "connection_id": 3, + "action": { + "Request": { + "request": { + "uri": "https://query-cell2.timestream.us-west-2.amazonaws.com/", + "headers": { + "authorization": [ + "AWS4-HMAC-SHA256 Credential=ANOTREAL/20090213/us-west-2/timestream/aws4_request, SignedHeaders=content-length;content-type;host;x-amz-date;x-amz-security-token;x-amz-target;x-amz-user-agent, Signature=8b2abd7688aefb4004200e5fa087f6c154fde879f2df79d3f6c57934cdc8f62a" + ], + "x-amz-date": [ + "20090213T234130Z" + ], + "user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head os/macos lang/rust/1.67.1" + ], + "x-amz-target": [ + "Timestream_20181101.Query" + ], + "x-amz-user-agent": [ + "aws-sdk-rust/0.0.0-smithy-rs-head api/timestreamquery/0.0.0-local os/macos lang/rust/1.67.1" + ], + "content-type": [ + "application/x-amz-json-1.0" + ], + "content-length": [ + "73" + ], + "x-amz-security-token": [ + "notarealsessiontoken" + ] + }, + "method": "POST" + } + } + } + }, + { + "connection_id": 3, + "action": { + "Data": { + "data": { + "Utf8": "{\"QueryString\":\"SELECT now() as time_now\",\"ClientToken\":\"0000-0000-0000\"}" + }, + "direction": "Request" + } + } + }, + { + "connection_id": 3, + "action": { + "Response": { + "response": { + "Ok": { + "status": 200, + "version": "HTTP/1.1", + "headers": { + "content-type": [ + "application/x-amz-json-1.0" + ], + "date": [ + "Wed, 24 May 2023 15:51:08 GMT" + ], + "x-amzn-requestid": [ + "OFO47BQ6XAXGTIK3XH4S6GBFLQ" + ], + "content-length": [ + "318" + ] + } + } + } + } + } + }, + { + "connection_id": 3, + "action": { + "Data": { + "data": { + "Utf8": "{\"ColumnInfo\":[{\"Name\":\"time_now\",\"Type\":{\"ScalarType\":\"TIMESTAMP\"}}],\"QueryId\":\"AEDACANMQDLSTQR3SPDV6HEWYACX5IALTEWGK7JM3VSFQQ5J5F3HGKVEMNHBRHY\",\"QueryStatus\":{\"CumulativeBytesMetered\":10000000,\"CumulativeBytesScanned\":0,\"ProgressPercentage\":100.0},\"Rows\":[{\"Data\":[{\"ScalarValue\":\"2023-05-24 15:51:08.760000000\"}]}]}" + }, + "direction": "Response" + } + } + }, + { + "connection_id": 3, + "action": { + "Eof": { + "ok": true, + "direction": "Response" + } + } + } + ], + "docs": "todo docs", + "version": "V0" +} diff --git a/aws/sra-test/integration-tests/aws-sdk-s3/Cargo.toml b/aws/sra-test/integration-tests/aws-sdk-s3/Cargo.toml index d93edb0e84..67b619607c 100644 --- a/aws/sra-test/integration-tests/aws-sdk-s3/Cargo.toml +++ b/aws/sra-test/integration-tests/aws-sdk-s3/Cargo.toml @@ -12,6 +12,7 @@ aws-sdk-s3 = { path = "../../build/sdk/aws-sdk-s3", features = ["test-util"] } aws-smithy-client = { path = "../../../../rust-runtime/aws-smithy-client", features = ["test-util", "rustls"] } aws-smithy-runtime = { path = "../../../../rust-runtime/aws-smithy-runtime" } aws-smithy-runtime-api = { path = "../../../../rust-runtime/aws-smithy-runtime-api" } +aws-smithy-async = { path = "../../../../rust-runtime/aws-smithy-async", features = ["test-util"]} aws-types = { path = "../../../rust-runtime/aws-types" } criterion = { version = "0.4", features = ["async_tokio"] } http = "0.2.3" diff --git a/buildSrc/src/main/kotlin/aws/sdk/ServiceLoader.kt b/buildSrc/src/main/kotlin/aws/sdk/ServiceLoader.kt index a1cbe0dc99..c1f76be0c0 100644 --- a/buildSrc/src/main/kotlin/aws/sdk/ServiceLoader.kt +++ b/buildSrc/src/main/kotlin/aws/sdk/ServiceLoader.kt @@ -207,7 +207,7 @@ fun parseMembership(rawList: String): Membership { val inclusions = mutableSetOf() val exclusions = mutableSetOf() - rawList.split(",").map { it.trim() }.forEach { item -> + rawList.split(",").map { it.trim() }.filter { it.isNotEmpty() }.forEach { item -> when { item.startsWith('-') -> exclusions.add(item.substring(1)) item.startsWith('+') -> inclusions.add(item.substring(1)) diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt index 007f5e24d3..a29d12f7b7 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt @@ -466,6 +466,7 @@ class RustWriter private constructor( val devDependenciesOnly: Boolean = false, ) : SymbolWriter(UseDeclarations(namespace)) { + companion object { fun root() = forModule(null) fun forModule(module: String?): RustWriter = if (module == null) { @@ -486,6 +487,7 @@ class RustWriter private constructor( debugMode = debugMode, devDependenciesOnly = true, ) + fileName == "package.json" -> rawWriter(fileName, debugMode = debugMode) fileName == "stubgen.sh" -> rawWriter(fileName, debugMode = debugMode) else -> RustWriter(fileName, namespace, debugMode = debugMode) @@ -515,6 +517,8 @@ class RustWriter private constructor( return super.write(content, *args) } + fun dirty() = super.toString().isNotBlank() || preamble.isNotEmpty() + /** Helper function to determine if a stack frame is relevant for debug purposes */ private fun StackTraceElement.isRelevant(): Boolean { if (this.className.contains("AbstractCodeWriter") || this.className.startsWith("java.lang")) { @@ -711,7 +715,8 @@ class RustWriter private constructor( override fun toString(): String { val contents = super.toString() val preheader = if (preamble.isNotEmpty()) { - val prewriter = RustWriter(filename, namespace, printWarning = false, devDependenciesOnly = devDependenciesOnly) + val prewriter = + RustWriter(filename, namespace, printWarning = false, devDependenciesOnly = devDependenciesOnly) preamble.forEach { it(prewriter) } prewriter.toString() } else { @@ -757,7 +762,8 @@ class RustWriter private constructor( @Suppress("UNCHECKED_CAST") val func = t as? Writable ?: throw CodegenException("RustWriteableInjector.apply choked on non-function t ($t)") - val innerWriter = RustWriter(filename, namespace, printWarning = false, devDependenciesOnly = devDependenciesOnly) + val innerWriter = + RustWriter(filename, namespace, printWarning = false, devDependenciesOnly = devDependenciesOnly) func(innerWriter) innerWriter.dependencies.forEach { addDependencyTestAware(it) } return innerWriter.toString().trimEnd() @@ -790,7 +796,8 @@ class RustWriter private constructor( @Suppress("UNCHECKED_CAST") val func = t as? Writable ?: throw CodegenException("Invalid function type (expected writable) ($t)") - val innerWriter = RustWriter(filename, namespace, printWarning = false, devDependenciesOnly = devDependenciesOnly) + val innerWriter = + RustWriter(filename, namespace, printWarning = false, devDependenciesOnly = devDependenciesOnly) func(innerWriter) innerWriter.dependencies.forEach { addDependencyTestAware(it) } return innerWriter.toString().trimEnd() diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt index d563349867..c174c81d95 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt @@ -60,3 +60,18 @@ fun RustWriter.writeCustomizations(customizations: List RustWriter.writeCustomizationsOrElse( + customizations: List>, + section: T, + orElse: Writable, +) { + val test = RustWriter.root() + test.writeCustomizations(customizations, section) + if (test.dirty()) { + writeCustomizations(customizations, section) + } else { + orElse(this) + } +} diff --git a/rust-runtime/aws-smithy-client/src/dvr/replay.rs b/rust-runtime/aws-smithy-client/src/dvr/replay.rs index fc5d867741..54065e5fc7 100644 --- a/rust-runtime/aws-smithy-client/src/dvr/replay.rs +++ b/rust-runtime/aws-smithy-client/src/dvr/replay.rs @@ -63,15 +63,7 @@ impl ReplayingConnection { /// Validate all headers and bodies pub async fn full_validate(self, media_type: MediaType) -> Result<(), Box> { - self.validate_base(None, |b1, b2| { - aws_smithy_protocol_test::validate_body( - b1, - std::str::from_utf8(b2).unwrap(), - media_type.clone(), - ) - .map_err(|e| Box::new(e) as _) - }) - .await + self.validate_body_and_headers(None, media_type).await } /// Validate actual requests against expected requests @@ -84,6 +76,25 @@ impl ReplayingConnection { .await } + /// Validate that the bodies match, using a given [`MediaType`] for comparison + /// + /// The specified headers are also validated + pub async fn validate_body_and_headers( + self, + checked_headers: Option<&[&str]>, + media_type: MediaType, + ) -> Result<(), Box> { + self.validate_base(checked_headers, |b1, b2| { + aws_smithy_protocol_test::validate_body( + b1, + std::str::from_utf8(b2).unwrap(), + media_type.clone(), + ) + .map_err(|e| Box::new(e) as _) + }) + .await + } + async fn validate_base( self, checked_headers: Option<&[&str]>,