Add a Transaction type to simplify dealing with Transactions

sqlx-core/src/pool/executor.rs

@@ -1,3 +1,5 @@
+use std::ops::DerefMut;
+
 use futures_core::{future::BoxFuture, stream::BoxStream};
 use futures_util::StreamExt;
 
@@ -9,6 +11,8 @@ use crate::{
     Database,
 };
 
+use super::PoolConnection;
+
 impl<C> Executor for Pool<C>
 where
     C: Connection + Connect<Connection = C>,
@@ -108,3 +112,45 @@ where
         Box::pin(async move { self.acquire().await?.describe(query).await })
     }
 }
+
+impl<C> Executor for PoolConnection<C>
+where
+    C: Connection + Connect<Connection = C>,
+{
+    type Database = <C as Executor>::Database;
+
+    fn send<'e, 'q: 'e>(&'e mut self, commands: &'q str) -> BoxFuture<'e, crate::Result<()>> {
+        self.deref_mut().send(commands)
+    }
+
+    fn execute<'e, 'q: 'e>(
+        &'e mut self,
+        query: &'q str,
+        args: <<C as Executor>::Database as Database>::Arguments,
+    ) -> BoxFuture<'e, crate::Result<u64>> {
+        self.deref_mut().execute(query, args)
+    }
+
+    fn fetch<'e, 'q: 'e>(
+        &'e mut self,
+        query: &'q str,
+        args: <<C as Executor>::Database as Database>::Arguments,
+    ) -> BoxStream<'e, crate::Result<<<C as Executor>::Database as Database>::Row>> {
+        self.deref_mut().fetch(query, args)
+    }
+
+    fn fetch_optional<'e, 'q: 'e>(
+        &'e mut self,
+        query: &'q str,
+        args: <<C as Executor>::Database as Database>::Arguments,
+    ) -> BoxFuture<'e, crate::Result<Option<<<C as Executor>::Database as Database>::Row>>> {
+        self.deref_mut().fetch_optional(query, args)
+    }
+
+    fn describe<'e, 'q: 'e>(
+        &'e mut self,
+        query: &'q str,
+    ) -> BoxFuture<'e, crate::Result<Describe<Self::Database>>> {
+        self.deref_mut().describe(query)
+    }
+}

sqlx-core/src/pool/mod.rs

@@ -2,21 +2,26 @@
 
 use std::{
     fmt,
+    mem,
     ops::{Deref, DerefMut},
     sync::Arc,
     time::{Duration, Instant},
 };
 
+use futures_core::future::BoxFuture;
+
 use crate::connection::{Connect, Connection};
+use crate::transaction::Transaction;
 
 use self::inner::SharedPool;
-pub use self::options::Builder;
 use self::options::Options;
 
 mod executor;
 mod inner;
 mod options;
 
+pub use self::options::Builder;
+
 /// A pool of database connections.
 pub struct Pool<C>(Arc<SharedPool<C>>);
 
@@ -84,6 +89,11 @@ where
         })
     }
 
+    /// Retrieves a new connection and immediately begins a new transaction.
+    pub async fn begin(&self) -> crate::Result<Transaction<PoolConnection<C>>> {
+        Ok(Transaction::new(0, self.acquire().await?).await?)
+    }
+
     /// Ends the use of a connection pool. Prevents any new connections
     /// and will close all active connections when they are returned to the pool.
     ///
@@ -172,6 +182,27 @@ where
     }
 }
 
+impl<C> Connection for PoolConnection<C>
+where
+    C: Connection + Connect<Connection = C>,
+{
+    fn close(mut self) -> BoxFuture<'static, crate::Result<()>> {
+        Box::pin(async move {
+            if let Some(live) = self.live.take() {
+                let raw = live.raw;
+
+                // Explicitly close the connection
+                raw.close().await?;
+            }
+
+            // Forget ourself so it does not go back to the pool
+            mem::forget(self);
+
+            Ok(())
+        })
+    }
+}
+
 impl<C> Drop for PoolConnection<C>
 where
     C: Connection + Connect<Connection = C>,