Transaction support

Author: sfackler

tokio-postgres/src/client.rs

@@ -7,7 +7,7 @@ use crate::tls::TlsConnect;
 use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
-use crate::{cancel_query, cancel_query_raw, query};
+use crate::{cancel_query, cancel_query_raw, query, Transaction};
 use crate::{prepare, SimpleQueryMessage};
 use crate::{simple_query, Row};
 use crate::{Error, Statement};
@@ -274,12 +274,21 @@ impl Client {
         simple_query::batch_execute(self.inner(), query)
     }
 
+    /// Begins a new database transaction.
+    ///
+    /// The transaction will roll back by default - use the `commit` method to commit it.
+    pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
+        self.batch_execute("BEGIN").await?;
+        Ok(Transaction::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
     /// only be returned if the client was unable to connect to the database.
     ///
     /// Requires the `runtime` Cargo feature (enabled by default).
+    #[cfg(feature = "runtime")]
     pub fn cancel_query<T>(&mut self, tls: T) -> impl Future<Output = Result<(), Error>>
     where
         T: MakeTlsConnect<Socket>,

tokio-postgres/src/lib.rs

@@ -108,6 +108,7 @@
 
 pub use crate::client::Client;
 pub use crate::config::Config;
+pub use crate::transaction::Transaction;
 pub use crate::connection::Connection;
 use crate::error::DbError;
 pub use crate::error::Error;
@@ -141,6 +142,7 @@ mod simple_query;
 #[cfg(feature = "runtime")]
 mod socket;
 mod statement;
+mod transaction;
 pub mod tls;
 pub mod types;
 

tokio-postgres/src/transaction.rs

@@ -0,0 +1,157 @@
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+#[cfg(feature = "runtime")]
+use crate::tls::MakeTlsConnect;
+use crate::tls::TlsConnect;
+use crate::types::{ToSql, Type};
+#[cfg(feature = "runtime")]
+use crate::Socket;
+use crate::{query, Client, Error, Row, SimpleQueryMessage, Statement};
+use futures::Stream;
+use postgres_protocol::message::frontend;
+use std::future::Future;
+use tokio::io::{AsyncRead, AsyncWrite};
+
+/// A representation of a PostgreSQL database transaction.
+///
+/// Transactions will implicitly roll back when dropped. Use the `commit` method to commit the changes made in the
+/// transaction. Transactions can be nested, with inner transactions implemented via safepoints.
+pub struct Transaction<'a> {
+    client: &'a mut Client,
+    done: bool,
+}
+
+impl<'a> Drop for Transaction<'a> {
+    fn drop(&mut self) {
+        if self.done {
+            return;
+        }
+
+        let mut buf = vec![];
+        frontend::query("ROLLBACK", &mut buf).unwrap();
+        let _ = self
+            .client
+            .inner()
+            .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+    }
+}
+
+impl<'a> Transaction<'a> {
+    pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> {
+        Transaction {
+            client,
+            done: false,
+        }
+    }
+
+    /// Consumes the transaction, committing all changes made within it.
+    pub async fn commit(mut self) -> Result<(), Error> {
+        self.done = true;
+        self.client.batch_execute("COMMIT").await
+    }
+
+    /// Rolls the transaction back, discarding all changes made within it.
+    ///
+    /// This is equivalent to `Transaction`'s `Drop` implementation, but provides any error encountered to the caller.
+    pub async fn rollback(mut self) -> Result<(), Error> {
+        self.done = true;
+        self.client.batch_execute("ROLLBACK").await
+    }
+
+    /// Like `Client::prepare`.
+    pub fn prepare(&mut self, query: &str) -> impl Future<Output = Result<Statement, Error>> {
+        self.client.prepare(query)
+    }
+
+    /// Like `Client::prepare_typed`.
+    pub fn prepare_typed(
+        &mut self,
+        query: &str,
+        parameter_types: &[Type],
+    ) -> impl Future<Output = Result<Statement, Error>> {
+        self.client.prepare_typed(query, parameter_types)
+    }
+
+    /// Like `Client::query`.
+    pub fn query(
+        &mut self,
+        statement: &Statement,
+        params: &[&dyn ToSql],
+    ) -> impl Stream<Item = Result<Row, Error>> {
+        self.client.query(statement, params)
+    }
+
+    /// Like `Client::query_iter`.
+    pub fn query_iter<'b, I>(
+        &mut self,
+        statement: &Statement,
+        params: I,
+    ) -> impl Stream<Item = Result<Row, Error>> + 'static
+    where
+        I: IntoIterator<Item = &'b dyn ToSql>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        // https://github.com/rust-lang/rust/issues/63032
+        let buf = query::encode(statement, params);
+        query::query(self.client.inner(), statement.clone(), buf)
+    }
+
+    /// Like `Client::execute`.
+    pub fn execute(
+        &mut self,
+        statement: &Statement,
+        params: &[&dyn ToSql],
+    ) -> impl Future<Output = Result<u64, Error>> {
+        self.client.execute(statement, params)
+    }
+
+    /// Like `Client::execute_iter`.
+    pub fn execute_iter<'b, I>(
+        &mut self,
+        statement: &Statement,
+        params: I,
+    ) -> impl Future<Output = Result<u64, Error>>
+    where
+        I: IntoIterator<Item = &'b dyn ToSql>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        // https://github.com/rust-lang/rust/issues/63032
+        let buf = query::encode(statement, params);
+        query::execute(self.client.inner(), buf)
+    }
+
+    /// Like `Client::simple_query`.
+    pub fn simple_query(
+        &mut self,
+        query: &str,
+    ) -> impl Stream<Item = Result<SimpleQueryMessage, Error>> {
+        self.client.simple_query(query)
+    }
+
+    /// Like `Client::batch_execute`.
+    pub fn batch_execute(&mut self, query: &str) -> impl Future<Output = Result<(), Error>> {
+        self.client.batch_execute(query)
+    }
+
+    /// Like `Client::cancel_query`.
+    #[cfg(feature = "runtime")]
+    pub fn cancel_query<T>(&mut self, tls: T) -> impl Future<Output = Result<(), Error>>
+    where
+        T: MakeTlsConnect<Socket>,
+    {
+        self.client.cancel_query(tls)
+    }
+
+    /// Like `Client::cancel_query_raw`.
+    pub fn cancel_query_raw<S, T>(
+        &mut self,
+        stream: S,
+        tls: T,
+    ) -> impl Future<Output = Result<(), Error>>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+        T: TlsConnect<S>,
+    {
+        self.client.cancel_query_raw(stream, tls)
+    }
+}