Add ability to map boxed future in new_type_future

Author: ysaito1001

rust-runtime/aws-smithy-async/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "aws-smithy-async"
-version = "1.2.5"
+version = "1.3.0"
 authors = ["AWS Rust SDK Team <[email protected]>", "John DiSanti <[email protected]>"]
 description = "Async runtime agnostic abstractions for smithy-rs."
 edition = "2021"

rust-runtime/aws-smithy-async/src/future/now_or_later.rs

@@ -125,6 +125,66 @@ impl<T, F> NowOrLater<T, F> {
     }
 }
 
+impl<'a, T, F> NowOrLater<T, F>
+where
+    F: Future<Output = T> + Send + 'a,
+{
+    /// Maps the value inside `NowOrLater` and returns a boxed future with lifetime `'a`.
+    ///
+    /// Examples
+    ///
+    /// ```no_run
+    /// # use aws_smithy_async::future::now_or_later::{NowOrLater, BoxFuture};
+    /// # async fn map_boxed_later_variant() {
+    /// let later = NowOrLater::new(async { 10 });
+    /// let mapped = later.map_boxed(|x| x + 5);
+    /// assert_eq!(15, mapped.await);
+    /// # }
+    /// ```
+    pub fn map_boxed<U, M>(self, map_fn: M) -> NowOrLater<U, BoxFuture<'a, U>>
+    where
+        M: FnOnce(T) -> U + Send + 'a,
+    {
+        match self.inner {
+            Inner::Now { value } => {
+                let mapped = value.map(map_fn);
+                NowOrLater {
+                    inner: Inner::Now { value: mapped },
+                }
+            }
+            Inner::Later { future } => {
+                let fut = async move {
+                    let val = future.await;
+                    map_fn(val)
+                };
+                NowOrLater {
+                    inner: Inner::Later {
+                        future: Box::pin(fut),
+                    },
+                }
+            }
+        }
+    }
+}
+
+impl<T> NowOrLater<T, OnlyReady> {
+    /// Like [`NowOrLater::map_boxed`], but specialized for use with [`OnlyReady`]
+    pub fn map_boxed<U, M>(self, map_fn: M) -> NowOrLater<U, OnlyReady>
+    where
+        M: FnOnce(T) -> U,
+    {
+        match self.inner {
+            Inner::Now { value } => {
+                let mapped = value.map(map_fn);
+                NowOrLater {
+                    inner: Inner::Now { value: mapped },
+                }
+            }
+            Inner::Later { .. } => unreachable!(),
+        }
+    }
+}
+
 impl<T, F> Future for NowOrLater<T, F>
 where
     F: Future<Output = T>,
@@ -194,4 +254,18 @@ mod test {
         let wrapped = NowOrLater::new(async { 5 });
         assert_eq!(wrapped.await, 5);
     }
+
+    #[tokio::test]
+    async fn map_boxed_now_variant() {
+        let now: NowOrLater<i32, OnlyReady> = NowOrLater::ready(21);
+        let mapped = now.map_boxed(|x| x * 2);
+        assert_eq!(42, mapped.await);
+    }
+
+    #[tokio::test]
+    async fn map_boxed_later_variant() {
+        let later = NowOrLater::new(async { 10 });
+        let mapped = later.map_boxed(|x| x + 5);
+        assert_eq!(15, mapped.await);
+    }
 }

rust-runtime/aws-smithy-runtime-api/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "aws-smithy-runtime-api"
-version = "1.8.3"
+version = "1.8.4"
 authors = ["AWS Rust SDK Team <[email protected]>", "Zelda Hessler <[email protected]>"]
 description = "Smithy runtime types."
 edition = "2021"

rust-runtime/aws-smithy-runtime-api/src/client/auth.rs

@@ -15,7 +15,7 @@ use aws_smithy_types::config_bag::{ConfigBag, FrozenLayer, Storable, StoreReplac
 use aws_smithy_types::type_erasure::TypeErasedBox;
 use aws_smithy_types::Document;
 use std::borrow::Cow;
-use std::fmt;
+use std::fmt::{self, Debug};
 use std::sync::Arc;
 
 /// Auth schemes for the HTTP `Authorization` header.
@@ -225,6 +225,25 @@ new_type_future! {
     pub struct AuthSchemeOptionsFuture<'a, Vec<AuthSchemeOption>, BoxError>;
 }
 
+// Currently, we don't add `map_ok` to the `new_type_future` macro in general.
+// It's specifically used for `AuthSchemeOptionsFuture`, but we can expand it later if needed.
+impl<'a> AuthSchemeOptionsFuture<'a> {
+    /// Transforms the `Ok` variant inside this `AuthSchemeOptionsFuture` by applying the provided function.
+    ///
+    /// This method maps over the `Ok` variant of the `Result` wrapped by the future,
+    /// applying `map_fn` to the contained `Vec<AuthSchemeOption>`.
+    ///
+    /// The transformation is applied regardless of whether the future's value is already
+    /// available (`Now`) or will be computed asynchronously (`Later`).
+    pub fn map_ok<F>(self, f: F) -> AuthSchemeOptionsFuture<'a>
+    where
+        F: FnOnce(Vec<AuthSchemeOption>) -> Vec<AuthSchemeOption> + Send + 'a,
+    {
+        let inner = self.inner.map_boxed(|result| result.map(f));
+        Self { inner }
+    }
+}
+
 /// Resolver for auth scheme options.
 ///
 /// The orchestrator needs to select an auth scheme to sign requests with, and potentially
@@ -420,3 +439,108 @@ impl<'a> From<&'a Document> for AuthSchemeEndpointConfig<'a> {
         Self(Some(value))
     }
 }
+
+/// An ordered list of [AuthSchemeId]s
+///
+/// Can be used to reorder already-resolved auth schemes by an auth scheme resolver.
+/// This list is intended as a hint rather than a strict override;
+/// any schemes not present in the resolved auth schemes will be ignored.
+#[derive(Clone, Debug, Default, Eq, PartialEq)]
+pub struct AuthSchemePreference {
+    preference_list: Vec<AuthSchemeId>,
+}
+
+impl Storable for AuthSchemePreference {
+    type Storer = StoreReplace<Self>;
+}
+
+impl IntoIterator for AuthSchemePreference {
+    type Item = AuthSchemeId;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.preference_list.into_iter()
+    }
+}
+
+impl<T> From<T> for AuthSchemePreference
+where
+    T: AsRef<[AuthSchemeId]>,
+{
+    fn from(slice: T) -> Self {
+        AuthSchemePreference {
+            preference_list: slice.as_ref().to_vec(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_map_ok_now_variant() {
+        let input = ["sigv4", "http"]
+            .map(|s| AuthSchemeOption::from(AuthSchemeId::from(s)))
+            .to_vec();
+        let fut = AuthSchemeOptionsFuture::ready(Ok(input));
+
+        let mapped = fut.map_ok(|opts| {
+            opts.into_iter()
+                .filter(|opt| opt.scheme_id().inner() != "http")
+                .collect()
+        });
+
+        let result = mapped.await;
+        assert!(result.is_ok());
+        let vec = result.unwrap();
+        assert_eq!(1, vec.len());
+        assert_eq!("sigv4", vec[0].scheme_id().inner());
+    }
+
+    #[tokio::test]
+    async fn test_map_ok_now_variant_error_no_op() {
+        let fut = AuthSchemeOptionsFuture::ready(Err(BoxError::from("oops")));
+
+        let mapped = fut.map_ok(|opts| opts); // no-op
+
+        let result = mapped.await;
+        assert!(result.is_err());
+        assert_eq!(result.unwrap_err().to_string(), "oops");
+    }
+
+    #[tokio::test]
+    async fn test_map_ok_later_variant() {
+        let input = ["foo", "bar"]
+            .map(|s| AuthSchemeOption::from(AuthSchemeId::from(s)))
+            .to_vec();
+        let fut = AuthSchemeOptionsFuture::new(async move { Ok(input) });
+
+        let mapped = fut.map_ok(|opts| {
+            opts.into_iter()
+                .map(|mut opt| {
+                    opt.scheme_id =
+                        AuthSchemeId::from(Cow::Owned(opt.scheme_id().inner().to_uppercase()));
+                    opt
+                })
+                .collect()
+        });
+
+        let result = mapped.await;
+        assert!(result.is_ok());
+        let vec = result.unwrap();
+        assert_eq!(vec[0].scheme_id().inner(), "FOO");
+        assert_eq!(vec[1].scheme_id().inner(), "BAR");
+    }
+
+    #[tokio::test]
+    async fn test_map_ok_later_variant_error_no_op() {
+        let fut = AuthSchemeOptionsFuture::new(async move { Err(BoxError::from("later fail")) });
+
+        let mapped = fut.map_ok(|opts| opts); // no-op
+
+        let result = mapped.await;
+        assert!(result.is_err());
+        assert_eq!(result.unwrap_err().to_string(), "later fail");
+    }
+}