postgresql/
transaction.rs

1//! Transactions.
2
3use std::{future::Future, sync::Arc};
4
5use crate::{
6    connection::{Connection, RowStream},
7    encode::ToSql,
8    error::Error,
9    pool::Checkout,
10};
11
12/// An open transaction.
13///
14/// Owns its connection (no lifetime), so it can be passed to generic code that
15/// accepts any [`crate::Executor`]. Dropping without committing enqueues a `ROLLBACK`
16/// and, for pooled transactions, returns the connection to its pool once any
17/// row stream taken from the transaction has also been dropped.
18pub struct Transaction {
19    checkout: Option<Arc<Checkout>>,
20    done: bool,
21}
22
23impl Transaction {
24    /// Begins a transaction on a connection.
25    ///
26    /// # Errors
27    ///
28    /// Returns an error if `BEGIN` fails.
29    pub async fn begin(conn: Connection) -> Result<Self, Error> {
30        Self::begin_checkout(Arc::new(Checkout::detached(conn))).await
31    }
32
33    /// Runs `BEGIN` on a freshly built checkout.
34    ///
35    /// On failure the checkout is dropped, which returns the connection to its
36    /// pool (or closes it).
37    pub(crate) async fn begin_checkout(checkout: Arc<Checkout>) -> Result<Self, Error> {
38        match checkout.connection().execute("BEGIN", &[]).await {
39            Ok(_) => Ok(Transaction {
40                checkout: Some(checkout),
41                done: false,
42            }),
43            Err(err) => {
44                drop(checkout);
45                Err(err)
46            }
47        }
48    }
49
50    /// Commits the transaction.
51    ///
52    /// # Errors
53    ///
54    /// Returns an error if `COMMIT` fails. The transaction is considered
55    /// finished either way.
56    pub async fn commit(mut self) -> Result<(), Error> {
57        let conn = self.connection();
58        let result = conn.execute("COMMIT", &[]).await.map(|_| ());
59        self.done = true;
60        result
61    }
62
63    /// Rolls the transaction back.
64    ///
65    /// # Errors
66    ///
67    /// Returns an error if `ROLLBACK` fails.
68    pub async fn rollback(mut self) -> Result<(), Error> {
69        let conn = self.connection();
70        let result = conn.execute("ROLLBACK", &[]).await.map(|_| ());
71        self.done = true;
72        result
73    }
74
75    fn checkout(&self) -> &Checkout {
76        self.checkout.as_deref().expect("transaction already completed")
77    }
78
79    fn connection(&self) -> &Connection {
80        self.checkout().connection()
81    }
82}
83
84impl Drop for Transaction {
85    fn drop(&mut self) {
86        if let Some(checkout) = self.checkout.take() {
87            if !self.done {
88                // Enqueue a rollback; the actor serialises it before any later
89                // query on this connection, so no `tokio::spawn` is needed.
90                checkout.connection().enqueue("ROLLBACK");
91            }
92            // Dropping the checkout returns the connection to the pool once any
93            // stream taken from this transaction is also gone.
94            drop(checkout);
95        }
96    }
97}
98
99impl crate::query::sealed::Sealed for Transaction {}
100
101impl crate::query::Capable for Transaction {
102    fn execute<'a>(
103        &'a self,
104        sql: &'a str,
105        params: &'a [&'a dyn ToSql],
106    ) -> impl Future<Output = Result<u64, Error>> + Send + 'a {
107        async move { self.connection().execute(sql, params).await }
108    }
109
110    fn fetch_rows<'a>(
111        &'a self,
112        sql: &'a str,
113        params: &'a [&'a dyn ToSql],
114        max_rows: i32,
115    ) -> impl Future<Output = Result<RowStream, Error>> + Send + 'a {
116        async move {
117            // Tie a clone of the checkout to the stream: the connection (and the
118            // rollback-on-drop) is only settled once the stream is gone.
119            let checkout = self.checkout.as_ref().expect("transaction already completed").clone();
120            let stream = checkout.connection().fetch_rows(sql, params, max_rows).await?;
121            Ok(stream.with_keepalive(checkout))
122        }
123    }
124}
125
126impl crate::query::ScriptExecutor for Transaction {
127    fn execute_script<'a>(&'a self, sql: &'a str) -> impl Future<Output = Result<u64, Error>> + Send + 'a {
128        async move { self.connection().execute_script(sql).await }
129    }
130}
postgresql/transaction.rs

postgresql/
transaction.rs
