Add Pool::get_owned

bb8/src/api.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::error;
 use std::fmt;
 use std::marker::PhantomData;
@@ -54,6 +55,18 @@ impl<M: ManageConnection> Pool<M> {
         self.inner.get().await
     }
 
+    /// Retrieves an owned connection from the pool.
+    ///
+    /// [`Pool::get`] should be preferred for situations where an owned
+    /// connection type isn't necessary.
+    ///
+    /// This contains a reference to the pool and will prevent the it from being
+    /// dropped, while an owned connection is alive. So if an owned connection
+    /// is leaked, the pool will never be shutdown.
+    pub async fn get_owned(&self) -> Result<PooledConnection<'static, M>, RunError<M::Error>> {
+        self.inner.get_owned().await
+    }
+
     /// Get a new dedicated connection that will not be managed by the pool.
     /// An application may want a persistent connection (e.g. to do a
     /// postgres LISTEN) that will not be closed or repurposed by the pool.
@@ -285,7 +298,7 @@ pub struct PooledConnection<'a, M>
 where
     M: ManageConnection,
 {
-    pool: &'a PoolInner<M>,
+    pool: Cow<'a, PoolInner<M>>,
     conn: Option<Conn<M::Connection>>,
 }
 
@@ -295,7 +308,7 @@ where
 {
     pub(crate) fn new(pool: &'a PoolInner<M>, conn: Conn<M::Connection>) -> Self {
         Self {
-            pool,
+            pool: Cow::Borrowed(pool),
             conn: Some(conn),
         }
     }
@@ -305,6 +318,18 @@ where
     }
 }
 
+impl<M> PooledConnection<'static, M>
+where
+    M: ManageConnection,
+{
+    pub(crate) fn new_owned(pool: PoolInner<M>, conn: Conn<M::Connection>) -> Self {
+        Self {
+            pool: Cow::Owned(pool),
+            conn: Some(conn),
+        }
+    }
+}
+
 impl<'a, M> Deref for PooledConnection<'a, M>
 where
     M: ManageConnection,
@@ -340,7 +365,7 @@ where
     M: ManageConnection,
 {
     fn drop(&mut self) {
-        self.pool.put_back(self.conn.take());
+        self.pool.as_ref().put_back(self.conn.take())
     }
 }
 

bb8/src/inner.rs

@@ -83,13 +83,36 @@ where
     }
 
     pub(crate) async fn get(&self) -> Result<PooledConnection<'_, M>, RunError<M::Error>> {
+        self.get_with(|this, conn| PooledConnection::new(this, conn))
+            .await
+    }
+
+    pub(crate) async fn get_owned(
+        &self,
+    ) -> Result<PooledConnection<'static, M>, RunError<M::Error>> {
+        self.get_with(|this, conn| {
+            let pool = PoolInner {
+                inner: Arc::clone(&this.inner),
+            };
+            PooledConnection::new_owned(pool, conn)
+        })
+        .await
+    }
+
+    pub(crate) async fn get_with<'a, 'b, F>(
+        &'a self,
+        make_pool: F,
+    ) -> Result<PooledConnection<'b, M>, RunError<M::Error>>
+    where
+        F: Fn(&'a Self, Conn<M::Connection>) -> PooledConnection<'b, M>,
+    {
         loop {
             let mut conn = {
                 let mut locked = self.inner.internals.lock();
                 match locked.pop(&self.inner.statics) {
                     Some((conn, approvals)) => {
                         self.spawn_replenishing_approvals(approvals);
-                        PooledConnection::new(self, conn)
+                        make_pool(self, conn)
                     }
                     None => break,
                 }
@@ -117,7 +140,7 @@ where
         };
 
         match timeout(self.inner.statics.connection_timeout, rx).await {
-            Ok(Ok(mut guard)) => Ok(PooledConnection::new(self, guard.extract())),
+            Ok(Ok(mut guard)) => Ok(make_pool(self, guard.extract())),
             _ => Err(RunError::TimedOut),
         }
     }