From 214413d9dc03e892364bc043bae722d57edb3259 Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Wed, 8 Jan 2020 17:19:27 -0800 Subject: [PATCH] Add transaction builders Closes #543 --- postgres/src/client.rs | 29 +++++- postgres/src/lib.rs | 5 +- postgres/src/transaction_builder.rs | 47 ++++++++++ tokio-postgres/src/client.rs | 9 ++ tokio-postgres/src/lib.rs | 2 + tokio-postgres/src/transaction_builder.rs | 103 ++++++++++++++++++++++ tokio-postgres/tests/test/main.rs | 39 +++++++- 7 files changed, 231 insertions(+), 3 deletions(-) create mode 100644 postgres/src/transaction_builder.rs create mode 100644 tokio-postgres/src/transaction_builder.rs diff --git a/postgres/src/client.rs b/postgres/src/client.rs index e3499792..ec973229 100644 --- a/postgres/src/client.rs +++ b/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. diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 1bc8a883..ef669f7d 100644 --- a/postgres/src/lib.rs +++ b/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; diff --git a/postgres/src/transaction_builder.rs b/postgres/src/transaction_builder.rs new file mode 100644 index 00000000..6e0c525b --- /dev/null +++ b/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, Error> { + let transaction = self.runtime.block_on(self.builder.start())?; + Ok(Transaction::new(self.runtime, transaction)) + } +} diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs index 50cca9dd..955435cd 100644 --- a/tokio-postgres/src/client.rs +++ b/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 diff --git a/tokio-postgres/src/lib.rs b/tokio-postgres/src/lib.rs index 98428424..0d8aa843 100644 --- a/tokio-postgres/src/lib.rs +++ b/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. diff --git a/tokio-postgres/src/transaction_builder.rs b/tokio-postgres/src/transaction_builder.rs new file mode 100644 index 00000000..6ed7e0dd --- /dev/null +++ b/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, + 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, 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)) + } +} diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs index 37ee0351..805a76a9 100644 --- a/tokio-postgres/tests/test/main.rs +++ b/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;