feat(mono-2840): Allow injecting fresh credentials per-connection

src/conn/mod.rs

@@ -1021,9 +1021,18 @@ impl Conn {
     }
 
     /// Returns a future that resolves to [`Conn`].
+    ///
+    /// [`Opts::credentials_provider`] is used to supply freshly minted credentials
+    /// (e.g. an IAM auth token) for this specific physical connection.
+    /// See [`OptsBuilder::credentials_provider`].
     pub fn new<T: Into<Opts>>(opts: T) -> crate::BoxFuture<'static, Conn> {
         let opts = opts.into();
         async move {
+            let opts = match opts.credentials_provider() {
+                Some(provider) => provider(opts).await?,
+                None => opts,
+            };
+
             let mut conn = Conn::empty(opts.clone());
 
             let stream = if let Some(_path) = opts.socket() {

src/conn/pool/mod.rs

@@ -589,6 +589,54 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn credentials_provider_runs_once_per_physical_connection() -> super::Result<()> {
+        use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
+
+        let calls = Arc::new(AtomicUsize::new(0));
+        let calls_clone = calls.clone();
+
+        // min == max so returned connections are kept idle (not closed) for reuse below.
+        let constraints = PoolConstraints::new(3, 3).unwrap();
+        let pool_opts = PoolOpts::default().with_constraints(constraints);
+        let opts = get_opts()
+            .pool_opts(pool_opts)
+            .prefer_socket(false)
+            .credentials_provider(move |opts| {
+                let calls = calls_clone.clone();
+                async move {
+                    calls.fetch_add(1, AtomicOrdering::SeqCst);
+                    Ok(opts)
+                }
+                .boxed()
+            });
+
+        let pool = Pool::new(opts);
+
+        let conns = try_join_all((0..3).map(|_| pool.get_conn())).await?;
+        assert_eq!(calls.load(AtomicOrdering::SeqCst), 3);
+
+        drop(conns);
+        sleep(Duration::from_millis(200)).await;
+
+        // Reusing an idling connection must not invoke the provider again.
+        let conn = pool.get_conn().await?;
+        assert_eq!(calls.load(AtomicOrdering::SeqCst), 3);
+
+        drop(conn);
+        pool.disconnect().await
+    }
+
+    #[tokio::test]
+    async fn credentials_provider_error_fails_the_connection() {
+        let opts = get_opts().credentials_provider(|_opts| {
+            async move { Err(crate::Error::Driver(crate::DriverError::PoolDisconnected)) }.boxed()
+        });
+
+        let pool = Pool::new(opts);
+        assert!(pool.get_conn().await.is_err());
+    }
+
     #[tokio::test]
     #[ignore]
     async fn can_handle_the_pressure() {

src/opts/mod.rs

@@ -611,6 +611,27 @@ impl fmt::Debug for AfterConnectCallbackWrapper {
     }
 }
 
+/// Clippy shorthand
+type CredentialsProviderCallback =
+    Arc<dyn Fn(Opts) -> crate::BoxFuture<'static, Opts> + Send + Sync + 'static>;
+
+#[derive(Clone)]
+pub(crate) struct CredentialsProviderCallbackWrapper(CredentialsProviderCallback);
+
+impl Eq for CredentialsProviderCallbackWrapper {}
+
+impl PartialEq for CredentialsProviderCallbackWrapper {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl fmt::Debug for CredentialsProviderCallbackWrapper {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("CredentialsProviderCallbackWrapper").finish()
+    }
+}
+
 /// Mysql connection options.
 ///
 /// Build one with [`OptsBuilder`].
@@ -647,6 +668,9 @@ pub(crate) struct MysqlOpts {
     /// Callback to execute once a new connection is established.
     after_connect: Option<AfterConnectCallbackWrapper>,
 
+    /// Callback invoked immediately before every new physical connection is established.
+    credentials_provider: Option<CredentialsProviderCallbackWrapper>,
+
     /// Commands to execute once new connection is established.
     init: Vec<String>,
 
@@ -862,6 +886,16 @@ impl Opts {
             .map(|cb| cb.0.clone())
     }
 
+    /// Callback invoked immediately before establishing every new physical connection. See
+    /// [`OptsBuilder::credentials_provider`].
+    pub fn credentials_provider(&self) -> Option<CredentialsProviderCallback> {
+        self.inner
+            .mysql_opts
+            .credentials_provider
+            .as_ref()
+            .map(|cb| cb.0.clone())
+    }
+
     /// Commands to execute once new a connection is established.
     pub fn init(&self) -> &[String] {
         self.inner.mysql_opts.init.as_ref()
@@ -1223,6 +1257,7 @@ impl Default for MysqlOpts {
             pass: None,
             db_name: None,
             after_connect: None,
+            credentials_provider: None,
             init: vec![],
             setup: vec![],
             tcp_keepalive: None,
@@ -1447,6 +1482,44 @@ impl OptsBuilder {
         self
     }
 
+    /// A callback, invoked for every new *physical* connection, i.e. every time
+    /// [`Conn::new`][crate::Conn::new] actually dials the server.
+    /// This happens for every connection [`Pool`][crate::Pool] creates to grow
+    /// itself, but *not* connections the pool hands out from its idle set.
+    ///
+    /// The callback receives the connection's current [`Opts`] and returns the
+    /// (possibly updated) `Opts` that will actually be used to connect.
+    /// 
+    /// This is the hook to use for injecting refreshed credentials, such as in 
+    /// cloud provider IAM database authentication tokens with
+    /// a few minutes validation.
+    ///
+    /// If this returns an error, the connection attempt will also fail.
+    ///
+    /// ```no_run
+    /// use futures_util::FutureExt;
+    /// use mysql_async::OptsBuilder;
+    ///
+    /// let opts_builder: OptsBuilder = todo!();
+    ///
+    /// opts_builder.credentials_provider(|opts| {
+    ///     async move {
+    ///         let token: String =
+    ///             todo!("mint a fresh auth token for `opts.user()` / `opts.ip_or_hostname()`");
+    ///         Ok(OptsBuilder::from_opts(opts).pass(Some(token)).into())
+    ///     }
+    ///     .boxed()
+    /// });
+    /// ```
+    pub fn credentials_provider<F>(mut self, callback: F) -> Self
+    where
+        F: Fn(Opts) -> crate::BoxFuture<'static, Opts> + Send + Sync + 'static,
+    {
+        self.opts.credentials_provider =
+            Some(CredentialsProviderCallbackWrapper(Arc::new(callback)));
+        self
+    }
+
     /// Defines initial queries. See [`Opts::init`].
     pub fn init<T: Into<String>>(mut self, init: Vec<T>) -> Self {
         self.opts.init = init.into_iter().map(Into::into).collect();