From 9e399aa93f3260c3fa8f15b5573cbd122fbfef85 Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Sat, 14 Jul 2018 14:59:37 -0700 Subject: [PATCH] Basic transaction support --- tokio-postgres/src/lib.rs | 27 ++++++ tokio-postgres/src/proto/mod.rs | 2 + tokio-postgres/src/proto/transaction.rs | 104 ++++++++++++++++++++++++ tokio-postgres/tests/test.rs | 79 ++++++++++++++++++ 4 files changed, 212 insertions(+) create mode 100644 tokio-postgres/src/proto/transaction.rs diff --git a/tokio-postgres/src/lib.rs b/tokio-postgres/src/lib.rs index 697264c9..a5089d3d 100644 --- a/tokio-postgres/src/lib.rs +++ b/tokio-postgres/src/lib.rs @@ -95,6 +95,14 @@ impl Client { Query(self.0.query(&statement.0, params)) } + pub fn transaction(&mut self, future: T) -> Transaction + where + T: Future, + T::Error: From, + { + Transaction(proto::TransactionFuture::new(self.0.clone(), future)) + } + pub fn batch_execute(&mut self, query: &str) -> BatchExecute { BatchExecute(self.0.batch_execute(query)) } @@ -242,6 +250,25 @@ impl Row { } } +#[must_use = "futures do nothing unless polled"] +pub struct Transaction(proto::TransactionFuture) +where + T: Future, + T::Error: From; + +impl Future for Transaction +where + T: Future, + T::Error: From, +{ + type Item = T::Item; + type Error = T::Error; + + fn poll(&mut self) -> Poll { + self.0.poll() + } +} + #[must_use = "futures do nothing unless polled"] pub struct BatchExecute(proto::SimpleQueryFuture); diff --git a/tokio-postgres/src/proto/mod.rs b/tokio-postgres/src/proto/mod.rs index dc95a9aa..4e9d6006 100644 --- a/tokio-postgres/src/proto/mod.rs +++ b/tokio-postgres/src/proto/mod.rs @@ -21,6 +21,7 @@ mod row; mod simple_query; mod socket; mod statement; +mod transaction; mod typeinfo; mod typeinfo_composite; mod typeinfo_enum; @@ -37,3 +38,4 @@ pub use proto::row::Row; pub use proto::simple_query::SimpleQueryFuture; pub use proto::socket::Socket; pub use proto::statement::Statement; +pub use proto::transaction::TransactionFuture; diff --git a/tokio-postgres/src/proto/transaction.rs b/tokio-postgres/src/proto/transaction.rs new file mode 100644 index 00000000..2a7b0826 --- /dev/null +++ b/tokio-postgres/src/proto/transaction.rs @@ -0,0 +1,104 @@ +use futures::{Async, Future, Poll}; +use proto::client::Client; +use proto::simple_query::SimpleQueryFuture; +use state_machine_future::RentToOwn; + +use Error; + +#[derive(StateMachineFuture)] +pub enum Transaction +where + F: Future, + E: From, +{ + #[state_machine_future(start, transitions(Beginning))] + Start { client: Client, future: F }, + #[state_machine_future(transitions(Running))] + Beginning { + client: Client, + begin: SimpleQueryFuture, + future: F, + }, + #[state_machine_future(transitions(Finishing))] + Running { client: Client, future: F }, + #[state_machine_future(transitions(Finished))] + Finishing { + future: SimpleQueryFuture, + result: Result, + }, + #[state_machine_future(ready)] + Finished(T), + #[state_machine_future(error)] + Failed(E), +} + +impl PollTransaction for Transaction +where + F: Future, + E: From, +{ + fn poll_start<'a>( + state: &'a mut RentToOwn<'a, Start>, + ) -> Poll, E> { + let state = state.take(); + transition!(Beginning { + begin: state.client.batch_execute("BEGIN"), + client: state.client, + future: state.future, + }) + } + + fn poll_beginning<'a>( + state: &'a mut RentToOwn<'a, Beginning>, + ) -> Poll, E> { + try_ready!(state.begin.poll()); + let state = state.take(); + transition!(Running { + client: state.client, + future: state.future, + }) + } + + fn poll_running<'a>( + state: &'a mut RentToOwn<'a, Running>, + ) -> Poll, E> { + match state.future.poll() { + Ok(Async::NotReady) => return Ok(Async::NotReady), + Ok(Async::Ready(t)) => transition!(Finishing { + future: state.client.batch_execute("COMMIT"), + result: Ok(t), + }), + Err(e) => transition!(Finishing { + future: state.client.batch_execute("ROLLBACK"), + result: Err(e), + }), + } + } + + fn poll_finishing<'a>( + state: &'a mut RentToOwn<'a, Finishing>, + ) -> Poll, E> { + match state.future.poll() { + Ok(Async::NotReady) => return Ok(Async::NotReady), + Ok(Async::Ready(())) => { + let t = state.take().result?; + transition!(Finished(t)) + } + Err(e) => match state.take().result { + Ok(_) => Err(e.into()), + // prioritize the future's error over the rollback error + Err(e) => Err(e), + }, + } + } +} + +impl TransactionFuture +where + F: Future, + E: From, +{ + pub fn new(client: Client, future: F) -> TransactionFuture { + Transaction::start(client, future) + } +} diff --git a/tokio-postgres/tests/test.rs b/tokio-postgres/tests/test.rs index b48761d7..bd946c6c 100644 --- a/tokio-postgres/tests/test.rs +++ b/tokio-postgres/tests/test.rs @@ -9,6 +9,7 @@ extern crate log; use futures::future; use futures::sync::mpsc; +use std::error::Error; use std::time::{Duration, Instant}; use tokio::prelude::*; use tokio::runtime::current_thread::Runtime; @@ -477,3 +478,81 @@ fn notifications() { assert_eq!(notifications[1].channel, "test_notifications"); assert_eq!(notifications[1].payload, "world"); } + +#[test] +fn test_transaction_commit() { + let _ = env_logger::try_init(); + let mut runtime = Runtime::new().unwrap(); + + let (mut client, connection) = runtime + .block_on(tokio_postgres::connect( + "postgres://postgres@localhost:5433".parse().unwrap(), + TlsMode::None, + )) + .unwrap(); + let connection = connection.map_err(|e| panic!("{}", e)); + runtime.handle().spawn(connection).unwrap(); + + runtime + .block_on(client.batch_execute( + "CREATE TEMPORARY TABLE foo ( + id SERIAL, + name TEXT + )", + )) + .unwrap(); + + let f = client.batch_execute("INSERT INTO foo (name) VALUES ('steven')"); + runtime.block_on(client.transaction(f)).unwrap(); + + let rows = runtime + .block_on( + client + .prepare("SELECT name FROM foo") + .and_then(|s| client.query(&s, &[]).collect()), + ) + .unwrap(); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::<_, &str>(0), "steven"); +} + +#[test] +fn test_transaction_abort() { + let _ = env_logger::try_init(); + let mut runtime = Runtime::new().unwrap(); + + let (mut client, connection) = runtime + .block_on(tokio_postgres::connect( + "postgres://postgres@localhost:5433".parse().unwrap(), + TlsMode::None, + )) + .unwrap(); + let connection = connection.map_err(|e| panic!("{}", e)); + runtime.handle().spawn(connection).unwrap(); + + runtime + .block_on(client.batch_execute( + "CREATE TEMPORARY TABLE foo ( + id SERIAL, + name TEXT + )", + )) + .unwrap(); + + let f = client + .batch_execute("INSERT INTO foo (name) VALUES ('steven')") + .map_err(|e| Box::new(e) as Box) + .and_then(|_| Err::<(), _>(Box::::from(""))); + runtime.block_on(client.transaction(f)).unwrap_err(); + + let rows = runtime + .block_on( + client + .prepare("SELECT name FROM foo") + .and_then(|s| client.query(&s, &[]).collect()), + ) + .unwrap(); + + assert_eq!(rows.len(), 0); +}