Add a Transaction type to simplify dealing with Transactions

Author: mehcode

examples/realworld-postgres/src/main.rs

@@ -50,7 +50,9 @@ async fn register(mut req: Request<PgPool>) -> Response {
     let body: RegisterRequestBody = req.body_json().await.unwrap();
     let hash = hash_password(&body.password).unwrap();
 
-    let mut pool = req.state();
+    // Make a new transaction
+    let pool = req.state();
+    let mut tx = pool.begin().await.unwrap();
 
     let rec = sqlx::query!(
         r#"
@@ -62,12 +64,15 @@ RETURNING id, username, email
         body.email,
         hash,
     )
-    .fetch_one(&mut pool)
+    .fetch_one(&mut tx)
     .await
     .unwrap();
 
     let token = generate_token(rec.id).unwrap();
 
+    // Explicitly commit
+    tx.commit().await.unwrap();
+
     #[derive(serde::Serialize)]
     struct RegisterResponseBody {
         user: User,

sqlx-core/src/lib.rs

@@ -16,6 +16,7 @@ mod database;
 mod executor;
 mod query;
 mod query_as;
+mod transaction;
 mod url;
 
 #[macro_use]
@@ -47,6 +48,7 @@ pub use connection::{Connect, Connection};
 pub use executor::Executor;
 pub use query::{query, Query};
 pub use query_as::{query_as, QueryAs};
+pub use transaction::Transaction;
 
 #[doc(hidden)]
 pub use query_as::query_as_mapped;

sqlx-core/src/mysql/connection.rs

@@ -7,7 +7,7 @@ use futures_core::future::BoxFuture;
 use sha1::Sha1;
 
 use crate::cache::StatementCache;
-use crate::connection::Connection;
+use crate::connection::{Connect, Connection};
 use crate::io::{Buf, BufMut, BufStream, MaybeTlsStream};
 use crate::mysql::error::MySqlError;
 use crate::mysql::protocol::{
@@ -475,7 +475,7 @@ impl MySqlConnection {
 }
 
 impl MySqlConnection {
-    pub(super) async fn open(url: crate::Result<Url>) -> crate::Result<Self> {
+    pub(super) async fn establish(url: crate::Result<Url>) -> crate::Result<Self> {
         let url = url?;
         let mut self_ = Self::new(&url).await?;
 
@@ -598,19 +598,19 @@ impl MySqlConnection {
         T: TryInto<Url, Error = crate::Error>,
         Self: Sized,
     {
-        Box::pin(MySqlConnection::open(url.try_into()))
+        Box::pin(MySqlConnection::establish(url.try_into()))
     }
 }
 
 impl Connect for MySqlConnection {
     type Connection = MySqlConnection;
 
-    fn connect<T>(url: T) -> BoxFuture<'static, Result<MySqlConnection>>
+    fn connect<T>(url: T) -> BoxFuture<'static, crate::Result<MySqlConnection>>
     where
         T: TryInto<Url, Error = crate::Error>,
         Self: Sized,
     {
-        Box::pin(PgConnection::open(url.try_into()))
+        Box::pin(MySqlConnection::establish(url.try_into()))
     }
 }
 

sqlx-core/src/mysql/mod.rs

@@ -26,16 +26,3 @@ pub use row::MySqlRow;
 
 /// An alias for [`Pool`], specialized for **MySQL**.
 pub type MySqlPool = super::Pool<MySql>;
-
-use std::convert::TryInto;
-
-use crate::url::Url;
-
-// used in tests and hidden code in examples
-#[doc(hidden)]
-pub async fn connect<T>(url: T) -> crate::Result<MySqlConnection>
-where
-    T: TryInto<Url, Error = crate::Error>,
-{
-    MySqlConnection::open(url.try_into()).await
-}

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>,

sqlx-core/src/postgres/connection.rs

@@ -8,7 +8,7 @@ use rand::Rng;
 use sha2::{Digest, Sha256};
 
 use crate::cache::StatementCache;
-use crate::connection::Connection;
+use crate::connection::{Connect, Connection};
 use crate::io::{Buf, BufStream, MaybeTlsStream};
 use crate::postgres::protocol::{
     self, hi, Authentication, Decode, Encode, Message, SaslInitialResponse, SaslResponse,
@@ -334,7 +334,7 @@ impl PgConnection {
 }
 
 impl PgConnection {
-    pub(super) async fn open(url: Result<Url>) -> Result<Self> {
+    pub(super) async fn establish(url: Result<Url>) -> Result<Self> {
         let url = url?;
 
         let stream = MaybeTlsStream::connect(&url, 5432).await?;
@@ -402,7 +402,7 @@ impl PgConnection {
         T: TryInto<Url, Error = crate::Error>,
         Self: Sized,
     {
-        Box::pin(PgConnection::open(url.try_into()))
+        Box::pin(PgConnection::establish(url.try_into()))
     }
 }
 
@@ -414,7 +414,7 @@ impl Connect for PgConnection {
         T: TryInto<Url, Error = crate::Error>,
         Self: Sized,
     {
-        Box::pin(PgConnection::open(url.try_into()))
+        Box::pin(PgConnection::establish(url.try_into()))
     }
 }
 

sqlx-core/src/postgres/mod.rs

@@ -18,16 +18,3 @@ mod types;
 
 /// An alias for [`Pool`], specialized for **Postgres**.
 pub type PgPool = super::Pool<Postgres>;
-
-use std::convert::TryInto;
-
-use crate::url::Url;
-
-// used in tests and hidden code in examples
-#[doc(hidden)]
-pub async fn connect<T>(url: T) -> crate::Result<PgConnection>
-where
-    T: TryInto<Url, Error = crate::Error>,
-{
-    PgConnection::open(url.try_into()).await
-}