Inline MaybeOwned in ConnectionSource and add another variant to store owned or ref connections

sqlx-core/src/connection.rs

@@ -43,13 +43,14 @@ pub trait Connect: Connection {
         Self: Sized;
 }
 
+#[allow(dead_code)]
 pub(crate) enum ConnectionSource<'c, C>
 where
     C: Connect,
 {
-    Connection(MaybeOwned<PoolConnection<C>, &'c mut C>),
-
-    #[allow(dead_code)]
+    ConnectionRef(&'c mut C),
+    Connection(C),
+    PoolConnection(Pool<C>, PoolConnection<C>),
     Pool(Pool<C>),
 }
 
@@ -60,15 +61,43 @@ where
     #[allow(dead_code)]
     pub(crate) async fn resolve(&mut self) -> crate::Result<&'_ mut C> {
         if let ConnectionSource::Pool(pool) = self {
-            *self = ConnectionSource::Connection(MaybeOwned::Owned(pool.acquire().await?));
+            let conn = pool.acquire().await?;
+
+            *self = ConnectionSource::PoolConnection(pool.clone(), conn);
         }
 
         Ok(match self {
-            ConnectionSource::Connection(conn) => match conn {
-                MaybeOwned::Borrowed(conn) => &mut *conn,
-                MaybeOwned::Owned(ref mut conn) => conn,
-            },
+            ConnectionSource::ConnectionRef(conn) => conn,
+            ConnectionSource::PoolConnection(_, ref mut conn) => conn,
+            ConnectionSource::Connection(ref mut conn) => conn,
             ConnectionSource::Pool(_) => unreachable!(),
         })
     }
 }
+
+impl<'c, C> From<C> for ConnectionSource<'c, C>
+where
+    C: Connect,
+{
+    fn from(connection: C) -> Self {
+        ConnectionSource::Connection(connection)
+    }
+}
+
+impl<'c, C> From<PoolConnection<C>> for ConnectionSource<'c, C>
+where
+    C: Connect,
+{
+    fn from(connection: PoolConnection<C>) -> Self {
+        ConnectionSource::PoolConnection(Pool(connection.pool.clone()), connection)
+    }
+}
+
+impl<'c, C> From<Pool<C>> for ConnectionSource<'c, C>
+where
+    C: Connect,
+{
+    fn from(pool: Pool<C>) -> Self {
+        ConnectionSource::Pool(pool)
+    }
+}

sqlx-core/src/mysql/cursor.rs

@@ -41,7 +41,7 @@ impl<'c, 'q> Cursor<'c, 'q> for MySqlCursor<'c, 'q> {
         E: Execute<'q, MySql>,
     {
         Self {
-            source: ConnectionSource::Connection(conn.into()),
+            source: ConnectionSource::ConnectionRef(conn),
             column_names: Arc::default(),
             column_types: Vec::new(),
             binary: true,

sqlx-core/src/pool/connection.rs

@@ -15,7 +15,7 @@ where
     C: Connect,
 {
     live: Option<Live<C>>,
-    pool: Arc<SharedPool<C>>,
+    pub(crate) pool: Arc<SharedPool<C>>,
 }
 
 pub(super) struct Live<C> {

sqlx-core/src/pool/inner.rs

@@ -19,7 +19,7 @@ use crate::{
 use super::connection::{Floating, Idle, Live};
 use super::Options;
 
-pub(super) struct SharedPool<C> {
+pub(crate) struct SharedPool<C> {
     url: String,
     idle_conns: ArrayQueue<Idle<C>>,
     waiters: SegQueue<Waker>,

sqlx-core/src/pool/mod.rs

@@ -21,7 +21,7 @@ pub use self::connection::PoolConnection;
 pub use self::options::Builder;
 
 /// A pool of database connections.
-pub struct Pool<C>(Arc<SharedPool<C>>);
+pub struct Pool<C>(pub(crate) Arc<SharedPool<C>>);
 
 impl<C> Pool<C>
 where

sqlx-core/src/postgres/cursor.rs

@@ -41,7 +41,7 @@ impl<'c, 'q> Cursor<'c, 'q> for PgCursor<'c, 'q> {
         E: Execute<'q, Postgres>,
     {
         Self {
-            source: ConnectionSource::Connection(conn.into()),
+            source: ConnectionSource::ConnectionRef(conn),
             columns: Arc::default(),
             formats: Arc::new([] as [TypeFormat; 0]),
             query: Some(query.into_parts()),

sqlx-core/src/sqlite/cursor.rs

@@ -40,7 +40,7 @@ impl<'c, 'q> Cursor<'c, 'q> for SqliteCursor<'c, 'q> {
         let (query, arguments) = query.into_parts();
 
         Self {
-            source: ConnectionSource::Connection(conn.into()),
+            source: ConnectionSource::ConnectionRef(conn),
             statement: None,
             query,
             arguments,