Add transaction builders

Closes sfackler#543

Author: sfackler

postgres/src/client.rs

@@ -1,6 +1,6 @@
 use crate::{
     CancelToken, Config, CopyInWriter, CopyOutReader, GenericClient, RowIter, Statement,
-    ToStatement, Transaction,
+    ToStatement, Transaction, TransactionBuilder,
 };
 use std::ops::{Deref, DerefMut};
 use tokio::runtime::Runtime;
@@ -486,6 +486,33 @@ impl Client {
         CancelToken::new(self.client.cancel_token())
     }
 
+    /// Returns a builder for a transaction with custom settings.
+    ///
+    /// Unlike the `transaction` method, the builder can be used to control the transaction's isolation level and other
+    /// attributes.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use postgres::{Client, IsolationLevel, NoTls};
+    ///
+    /// # fn main() -> Result<(), postgres::Error> {
+    /// let mut client = Client::connect("host=localhost user=postgres", NoTls)?;
+    ///
+    /// let mut transaction = client.build_transaction()
+    ///     .isolation_level(IsolationLevel::RepeatableRead)
+    ///     .start()?;
+    /// transaction.execute("UPDATE foo SET bar = 10", &[])?;
+    /// // ...
+    ///
+    /// transaction.commit()?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn build_transaction(&mut self) -> TransactionBuilder<'_> {
+        TransactionBuilder::new(&mut self.runtime, self.client.build_transaction())
+    }
+
     /// Determines if the client's connection has already closed.
     ///
     /// If this returns `true`, the client is no longer usable.

postgres/src/lib.rs

@@ -51,7 +51,8 @@
 
 pub use fallible_iterator;
 pub use tokio_postgres::{
-    error, row, tls, types, Column, Portal, SimpleQueryMessage, Socket, Statement, ToStatement,
+    error, row, tls, types, Column, IsolationLevel, Portal, SimpleQueryMessage, Socket, Statement,
+    ToStatement,
 };
 
 pub use crate::cancel_token::CancelToken;
@@ -68,6 +69,7 @@ pub use crate::row_iter::RowIter;
 #[doc(no_inline)]
 pub use crate::tls::NoTls;
 pub use crate::transaction::*;
+pub use crate::transaction_builder::TransactionBuilder;
 
 pub mod binary_copy;
 mod cancel_token;
@@ -79,6 +81,7 @@ mod generic_client;
 mod lazy_pin;
 mod row_iter;
 mod transaction;
+mod transaction_builder;
 
 #[cfg(test)]
 mod test;

postgres/src/transaction_builder.rs

@@ -0,0 +1,47 @@
+use crate::{Error, IsolationLevel, Transaction};
+use tokio::runtime::Runtime;
+
+/// A builder for database transactions.
+pub struct TransactionBuilder<'a> {
+    runtime: &'a mut Runtime,
+    builder: tokio_postgres::TransactionBuilder<'a>,
+}
+
+impl<'a> TransactionBuilder<'a> {
+    pub(crate) fn new(
+        runtime: &'a mut Runtime,
+        builder: tokio_postgres::TransactionBuilder<'a>,
+    ) -> TransactionBuilder<'a> {
+        TransactionBuilder { runtime, builder }
+    }
+
+    /// Sets the isolation level of the transaction.
+    pub fn isolation_level(mut self, isolation_level: IsolationLevel) -> Self {
+        self.builder = self.builder.isolation_level(isolation_level);
+        self
+    }
+
+    /// Sets the transaction to read-only.
+    pub fn read_only(mut self) -> Self {
+        self.builder = self.builder.read_only();
+        self
+    }
+
+    /// Sets the transaction to be deferrable.
+    ///
+    /// If the transaction is also serializable and read only, creation of the transaction may block, but when it
+    /// completes the transaction is able to run with less overhead and a guarantee that it will not be aborted due to
+    /// serialization failure.
+    pub fn deferrable(mut self) -> Self {
+        self.builder = self.builder.deferrable();
+        self
+    }
+
+    /// Begins the transaction.
+    ///
+    /// The transaction will roll back by default - use the `commit` method to commit it.
+    pub fn start(self) -> Result<Transaction<'a>, Error> {
+        let transaction = self.runtime.block_on(self.builder.start())?;
+        Ok(Transaction::new(self.runtime, transaction))
+    }
+}

tokio-postgres/src/client.rs

@@ -13,6 +13,7 @@ use crate::Socket;
 use crate::{
     copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
     GenericClient, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
+    TransactionBuilder,
 };
 use async_trait::async_trait;
 use bytes::{Buf, BytesMut};
@@ -461,6 +462,14 @@ impl Client {
         }
     }
 
+    /// Returns a builder for a transaction with custom settings.
+    ///
+    /// Unlike the `transaction` method, the builder can be used to control the transaction's isolation level and other
+    /// attributes.
+    pub fn build_transaction(&mut self) -> TransactionBuilder<'_> {
+        TransactionBuilder::new(self)
+    }
+
     /// Attempts to cancel an in-progress query.
     ///
     /// The server provides no information about whether a cancellation attempt was successful or not. An error will

tokio-postgres/src/lib.rs

@@ -120,6 +120,7 @@ use crate::tls::MakeTlsConnect;
 pub use crate::tls::NoTls;
 pub use crate::to_statement::ToStatement;
 pub use crate::transaction::Transaction;
+pub use crate::transaction_builder::{IsolationLevel, TransactionBuilder};
 use crate::types::ToSql;
 
 pub mod binary_copy;
@@ -154,6 +155,7 @@ mod statement;
 pub mod tls;
 mod to_statement;
 mod transaction;
+mod transaction_builder;
 pub mod types;
 
 /// A convenience function which parses a connection string and connects to the database.

tokio-postgres/src/transaction_builder.rs

@@ -0,0 +1,103 @@
+use crate::{Client, Error, Transaction};
+
+/// The isolation level of a database transaction.
+#[derive(Debug, Copy, Clone)]
+#[non_exhaustive]
+pub enum IsolationLevel {
+    /// Equivalent to `ReadCommitted`.
+    ReadUncommitted,
+
+    /// An individual statement in the transaction will see rows committed before it began.
+    ReadCommitted,
+
+    /// All statements in the transaction will see the same view of rows committed before the first query in the
+    /// transaction.
+    RepeatableRead,
+
+    /// The reads and writes in this transaction must be able to be committed as an atomic "unit" with respect to reads
+    /// and writes of all other concurrent serializable transactions without interleaving.
+    Serializable,
+}
+
+/// A builder for database transactions.
+pub struct TransactionBuilder<'a> {
+    client: &'a mut Client,
+    isolation_level: Option<IsolationLevel>,
+    read_only: bool,
+    deferrable: bool,
+}
+
+impl<'a> TransactionBuilder<'a> {
+    pub(crate) fn new(client: &'a mut Client) -> TransactionBuilder<'a> {
+        TransactionBuilder {
+            client,
+            isolation_level: None,
+            read_only: false,
+            deferrable: false,
+        }
+    }
+
+    /// Sets the isolation level of the transaction.
+    pub fn isolation_level(mut self, isolation_level: IsolationLevel) -> Self {
+        self.isolation_level = Some(isolation_level);
+        self
+    }
+
+    /// Sets the transaction to read-only.
+    pub fn read_only(mut self) -> Self {
+        self.read_only = true;
+        self
+    }
+
+    /// Sets the transaction to be deferrable.
+    ///
+    /// If the transaction is also serializable and read only, creation of the transaction may block, but when it
+    /// completes the transaction is able to run with less overhead and a guarantee that it will not be aborted due to
+    /// serialization failure.
+    pub fn deferrable(mut self) -> Self {
+        self.deferrable = true;
+        self
+    }
+
+    /// Begins the transaction.
+    ///
+    /// The transaction will roll back by default - use the `commit` method to commit it.
+    pub async fn start(self) -> Result<Transaction<'a>, Error> {
+        let mut query = "START TRANSACTION".to_string();
+        let mut first = true;
+
+        if let Some(level) = self.isolation_level {
+            first = false;
+
+            query.push_str(" ISOLATION LEVEL ");
+            let level = match level {
+                IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
+                IsolationLevel::ReadCommitted => "READ COMMITTED",
+                IsolationLevel::RepeatableRead => "REPEATABLE READ",
+                IsolationLevel::Serializable => "SERIALIZABLE",
+            };
+            query.push_str(level);
+        }
+
+        if self.read_only {
+            if !first {
+                query.push(',');
+            }
+            first = false;
+
+            query.push_str(" READ ONLY");
+        }
+
+        if self.deferrable {
+            if !first {
+                query.push(',');
+            }
+
+            query.push_str(" DEFERRABLE");
+        }
+
+        self.client.batch_execute(&query).await?;
+
+        Ok(Transaction::new(self.client))
+    }
+}

tokio-postgres/tests/test/main.rs

@@ -12,7 +12,9 @@ use tokio::time;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::{NoTls, NoTlsStream};
 use tokio_postgres::types::{Kind, Type};
-use tokio_postgres::{AsyncMessage, Client, Config, Connection, Error, SimpleQueryMessage};
+use tokio_postgres::{
+    AsyncMessage, Client, Config, Connection, Error, IsolationLevel, SimpleQueryMessage,
+};
 
 mod binary_copy;
 mod parse;
@@ -398,6 +400,41 @@ async fn transaction_rollback_drop() {
     assert_eq!(rows.len(), 0);
 }
 
+#[tokio::test]
+async fn transaction_builder() {
+    let mut client = connect("user=postgres").await;
+
+    client
+        .batch_execute(
+            "CREATE TEMPORARY TABLE foo(
+                id SERIAL,
+                name TEXT
+            )",
+        )
+        .await
+        .unwrap();
+
+    let transaction = client
+        .build_transaction()
+        .isolation_level(IsolationLevel::Serializable)
+        .read_only()
+        .deferrable()
+        .start()
+        .await
+        .unwrap();
+    transaction
+        .batch_execute("INSERT INTO foo (name) VALUES ('steven')")
+        .await
+        .unwrap();
+    transaction.commit().await.unwrap();
+
+    let stmt = client.prepare("SELECT name FROM foo").await.unwrap();
+    let rows = client.query(&stmt, &[]).await.unwrap();
+
+    assert_eq!(rows.len(), 1);
+    assert_eq!(rows[0].get::<_, &str>(0), "steven");
+}
+
 #[tokio::test]
 async fn copy_in() {
     let client = connect("user=postgres").await;