wasi_pg_client/transaction/
mod.rs

1//! Transaction management: BEGIN, COMMIT, ROLLBACK, savepoints.
2//!
3//! This module provides the [`Transaction`] guard type and supporting types
4//! ([`TransactionOptions`], [`IsolationLevel`], [`Savepoint`]).
5
6use crate::connection::Connection;
7use crate::error::Result;
8use crate::query::result::{ExecuteResult, QueryResult};
9
10#[cfg(feature = "tracing")]
11use crate::tracing_ext::TARGET_TRANSACTION;
12
13pub mod options;
14pub mod savepoint;
15
16pub use options::{IsolationLevel, TransactionOptions};
17pub use savepoint::Savepoint;
18
19// ---------------------------------------------------------------------------
20// Transaction guard
21// ---------------------------------------------------------------------------
22
23/// An active transaction guard.
24///
25/// Created via [`Connection::transaction`] or [`Connection::transaction_with`].
26/// The guard provides methods to execute queries within the transaction and
27/// to [`commit`](Self::commit) or [`rollback`](Self::rollback) it.
28///
29/// # Drop behaviour
30///
31/// `Drop` cannot perform async I/O.  If the transaction is not explicitly
32/// committed or rolled back before it goes out of scope, the connection may
33/// be left in an idle-in-transaction state.  **Always** call `.commit().await`
34/// or `.rollback().await` explicitly.
35#[non_exhaustive]
36pub struct Transaction<'a> {
37    pub(crate) conn: &'a mut Connection,
38    pub(crate) committed: bool,
39    pub(crate) savepoint_depth: u32,
40}
41
42impl<'a> Transaction<'a> {
43    pub(crate) fn new(conn: &'a mut Connection) -> Self {
44        Self {
45            conn,
46            committed: false,
47            savepoint_depth: 0,
48        }
49    }
50
51    /// Commit the transaction.
52    #[must_use = "commit errors should be checked"]
53    pub async fn commit(mut self) -> Result<()> {
54        #[cfg(feature = "tracing")]
55        tracing::info!(target: TARGET_TRANSACTION, "COMMIT transaction");
56        self.conn.execute("COMMIT").await?;
57        self.committed = true;
58        Ok(())
59    }
60
61    /// Roll back the transaction.
62    #[must_use = "rollback errors should be checked"]
63    pub async fn rollback(mut self) -> Result<()> {
64        #[cfg(feature = "tracing")]
65        tracing::warn!(target: TARGET_TRANSACTION, "ROLLBACK transaction");
66        self.conn.execute("ROLLBACK").await?;
67        self.committed = true;
68        Ok(())
69    }
70
71    /// Returns `true` if the transaction is in a failed state.
72    pub fn is_failed(&self) -> bool {
73        self.conn.transaction_status() == crate::protocol::TransactionStatus::Failed
74    }
75
76    /// Execute a query that returns rows, within the transaction.
77    #[must_use = "query errors should be checked"]
78    pub async fn query(&mut self, sql: &str) -> Result<QueryResult> {
79        self.conn.query(sql).await
80    }
81
82    /// Execute a statement that does not return rows.
83    #[must_use = "execute errors should be checked"]
84    pub async fn execute(&mut self, sql: &str) -> Result<ExecuteResult> {
85        self.conn.execute(sql).await
86    }
87
88    /// Execute a query and return at most one row.
89    #[must_use = "query errors should be checked"]
90    pub async fn query_one(&mut self, sql: &str) -> Result<Option<crate::Row>> {
91        self.conn.query_one(sql).await
92    }
93
94    /// Execute a parameterized query that returns rows.
95    #[must_use = "query errors should be checked"]
96    pub async fn query_params(
97        &mut self,
98        sql: &str,
99        params: &[&dyn crate::types::ToSql],
100    ) -> Result<QueryResult> {
101        self.conn.query_params(sql, params).await
102    }
103
104    /// Execute a parameterized statement that does not return rows.
105    #[must_use = "execute errors should be checked"]
106    pub async fn execute_params(
107        &mut self,
108        sql: &str,
109        params: &[&dyn crate::types::ToSql],
110    ) -> Result<ExecuteResult> {
111        self.conn.execute_params(sql, params).await
112    }
113
114    /// Prepare a statement within the transaction.
115    #[must_use = "prepare errors should be checked"]
116    pub async fn prepare(&mut self, sql: &str) -> Result<crate::query::PreparedStatement> {
117        self.conn.prepare(sql).await
118    }
119
120    /// Create a savepoint (nested transaction scope).
121    ///
122    /// Only one `Savepoint` guard can be active at a time for a given
123    /// `Transaction` because it holds a mutable borrow.
124    #[must_use = "savepoint errors should be checked"]
125    pub async fn savepoint(&mut self, name: &str) -> Result<Savepoint<'_, 'a>> {
126        let sql = format!("SAVEPOINT {}", quote_identifier(name));
127        self.conn.execute(&sql).await?;
128        self.savepoint_depth += 1;
129        Ok(Savepoint {
130            transaction: self,
131            name: name.to_string(),
132            released: false,
133        })
134    }
135
136    /// Start a COPY IN operation within this transaction.
137    #[must_use = "copy errors should be checked"]
138    pub async fn copy_in(&mut self, sql: &str) -> Result<crate::CopyIn<'_>> {
139        self.conn.copy_in(sql).await
140    }
141
142    /// Start a COPY OUT operation within this transaction.
143    #[must_use = "copy errors should be checked"]
144    pub async fn copy_out(&mut self, sql: &str) -> Result<crate::CopyOut<'_>> {
145        self.conn.copy_out(sql).await
146    }
147}
148
149impl<'a> Drop for Transaction<'a> {
150    #[allow(clippy::needless_return)]
151    fn drop(&mut self) {
152        if self.committed || std::thread::panicking() {
153            return;
154        }
155        #[cfg(feature = "tracing")]
156        tracing::warn!(target: TARGET_TRANSACTION, "Transaction dropped without explicit commit/rollback");
157        // Drop cannot be async.  We cannot send ROLLBACK here.
158        // Users must call .commit().await or .rollback().await explicitly.
159    }
160}
161
162// ---------------------------------------------------------------------------
163// Utility
164// ---------------------------------------------------------------------------
165
166/// Quote a PostgreSQL identifier to prevent SQL injection.
167pub(crate) fn quote_identifier(name: &str) -> String {
168    format!("\"{}\"", name.replace('"', "\"\""))
169}
170
171// ---------------------------------------------------------------------------
172// Connection extensions
173// ---------------------------------------------------------------------------
174
175impl Connection {
176    /// Begin a new transaction.
177    ///
178    /// # Example
179    /// ```ignore
180    /// let mut txn = conn.transaction().await?;
181    /// txn.execute("INSERT INTO users (name) VALUES ('alice')").await?;
182    /// txn.commit().await?;
183    /// ```
184    #[must_use = "transaction errors should be checked"]
185    pub async fn transaction(&mut self) -> Result<Transaction<'_>> {
186        self.execute("BEGIN").await?;
187        #[cfg(feature = "tracing")]
188        tracing::info!(target: TARGET_TRANSACTION, "BEGIN transaction");
189        Ok(Transaction::new(self))
190    }
191
192    /// Begin a new transaction with the given options.
193    ///
194    /// # Example
195    /// ```ignore
196    /// let mut txn = conn.transaction_with(
197    ///     TransactionOptions::new()
198    ///         .isolation_level(IsolationLevel::Serializable)
199    ///         .read_only(true)
200    /// ).await?;
201    /// ```
202    #[must_use = "transaction errors should be checked"]
203    pub async fn transaction_with(
204        &mut self,
205        options: &TransactionOptions,
206    ) -> Result<Transaction<'_>> {
207        let sql = options.to_begin_sql();
208        self.execute(&sql).await?;
209        Ok(Transaction::new(self))
210    }
211
212    /// Execute an async closure within a transaction.
213    ///
214    /// Commits on `Ok`, rolls back on `Err`. This requires Rust 1.85+ (async
215    /// closures are stable).
216    ///
217    /// # Example
218    /// ```ignore
219    /// let rows: Vec<i32> = conn.with_transaction(async |txn| {
220    ///     txn.execute("INSERT INTO nums (v) VALUES (1)").await?;
221    ///     let result = txn.query("SELECT v FROM nums").await?;
222    ///     let vals: Vec<i32> = result.iter().map(|r| r.get(0).unwrap()).collect();
223    ///     Ok(vals)
224    /// }).await?;
225    /// ```
226    #[must_use = "transaction errors should be checked"]
227    pub async fn with_transaction<T, F>(&mut self, f: F) -> Result<T>
228    where
229        F: AsyncFnOnce(&mut Transaction<'_>) -> Result<T>,
230    {
231        let mut txn = self.transaction().await?;
232        match f(&mut txn).await {
233            Ok(val) => {
234                txn.commit().await?;
235                Ok(val)
236            }
237            Err(e) => {
238                let _ = txn.rollback().await;
239                Err(e)
240            }
241        }
242    }
243}
244
245// ---------------------------------------------------------------------------
246// Tests
247// ---------------------------------------------------------------------------
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use crate::auth::{Codec, ServerParams};
253    use crate::config::Config;
254    use crate::connection::ConnectionState;
255    use crate::error::Error;
256    use crate::protocol::TransactionStatus;
257    use crate::transport::{BufferedTransport, ClientTransport, MockTransport, PgTransport};
258    use std::collections::VecDeque;
259
260    fn make_connection(read_data: Vec<u8>) -> Connection {
261        let transport = PgTransport::Plain(BufferedTransport::new(ClientTransport::Mock(
262            MockTransport::new(read_data),
263        )));
264        Connection {
265            transport,
266            codec: Codec::new(),
267            server_params: ServerParams::default(),
268            state: ConnectionState::Idle,
269            config: Config::new(),
270            transaction_status: TransactionStatus::Idle,
271            notification_queue: VecDeque::new(),
272            notice_handler: None,
273            statement_counter: 0,
274            needs_recovery: false,
275            health: crate::reconnect::session::ConnectionHealth::new(),
276            session_state: crate::reconnect::session::SessionState::new(),
277        }
278    }
279
280    fn build_command_complete_msg(tag: &str) -> Vec<u8> {
281        let mut buf = vec![b'C'];
282        let mut body = Vec::new();
283        body.extend_from_slice(tag.as_bytes());
284        body.push(0);
285        let len = (body.len() + 4) as i32;
286        buf.extend_from_slice(&len.to_be_bytes());
287        buf.extend_from_slice(&body);
288        buf
289    }
290
291    fn build_ready_for_query(status: u8) -> Vec<u8> {
292        vec![b'Z', 0, 0, 0, 5, status]
293    }
294
295    fn build_error_response(msg: &str) -> Vec<u8> {
296        let mut buf = vec![b'E'];
297        let mut body = Vec::new();
298        body.push(b'S');
299        body.extend_from_slice(b"ERROR\0");
300        body.push(b'M');
301        body.extend_from_slice(msg.as_bytes());
302        body.push(0);
303        body.push(0);
304        let len = (body.len() + 4) as i32;
305        buf.extend_from_slice(&len.to_be_bytes());
306        buf.extend_from_slice(&body);
307        buf
308    }
309
310    fn build_row_description_msg(fields: &[(&str, u32)]) -> Vec<u8> {
311        let mut buf = vec![b'T'];
312        let mut body = Vec::new();
313        body.extend_from_slice(&(fields.len() as i16).to_be_bytes());
314        for (name, type_oid) in fields {
315            body.extend_from_slice(name.as_bytes());
316            body.push(0);
317            body.extend_from_slice(&0u32.to_be_bytes()); // table_oid
318            body.extend_from_slice(&0i16.to_be_bytes()); // column_id
319            body.extend_from_slice(&type_oid.to_be_bytes()); // type_oid
320            body.extend_from_slice(&(-1i16).to_be_bytes()); // type_size
321            body.extend_from_slice(&(-1i32).to_be_bytes()); // type_modifier
322            body.extend_from_slice(&0i16.to_be_bytes()); // format
323        }
324        let len = (body.len() + 4) as i32;
325        buf.extend_from_slice(&len.to_be_bytes());
326        buf.extend_from_slice(&body);
327        buf
328    }
329
330    fn build_data_row_msg(values: &[Option<&str>]) -> Vec<u8> {
331        let mut buf = vec![b'D'];
332        let mut body = Vec::new();
333        body.extend_from_slice(&(values.len() as i16).to_be_bytes());
334        for val in values {
335            match val {
336                Some(v) => {
337                    let bytes = v.as_bytes();
338                    body.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
339                    body.extend_from_slice(bytes);
340                }
341                None => {
342                    body.extend_from_slice(&(-1i32).to_be_bytes());
343                }
344            }
345        }
346        let len = (body.len() + 4) as i32;
347        buf.extend_from_slice(&len.to_be_bytes());
348        buf.extend_from_slice(&body);
349        buf
350    }
351
352    // -----------------------------------------------------------------------
353    // Identifier quoting
354    // -----------------------------------------------------------------------
355
356    #[test]
357    fn test_quote_identifier_basic() {
358        assert_eq!(quote_identifier("foo"), "\"foo\"");
359    }
360
361    #[test]
362    fn test_quote_identifier_with_quotes() {
363        assert_eq!(quote_identifier("foo\"bar"), "\"foo\"\"bar\"");
364    }
365
366    #[test]
367    fn test_quote_identifier_empty() {
368        assert_eq!(quote_identifier(""), "\"\"");
369    }
370
371    // -----------------------------------------------------------------------
372    // Transaction options
373    // -----------------------------------------------------------------------
374
375    #[test]
376    fn test_transaction_options_default() {
377        let opts = TransactionOptions::new();
378        assert_eq!(opts.to_begin_sql(), "BEGIN");
379    }
380
381    #[test]
382    fn test_transaction_options_isolation() {
383        let opts = TransactionOptions::new().isolation_level(IsolationLevel::Serializable);
384        assert_eq!(opts.to_begin_sql(), "BEGIN ISOLATION LEVEL SERIALIZABLE");
385    }
386
387    #[test]
388    fn test_transaction_options_all() {
389        let opts = TransactionOptions::new()
390            .isolation_level(IsolationLevel::RepeatableRead)
391            .read_only(true)
392            .deferrable(true);
393        assert_eq!(
394            opts.to_begin_sql(),
395            "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY DEFERRABLE"
396        );
397    }
398
399    #[test]
400    fn test_transaction_options_read_write() {
401        let opts = TransactionOptions::new().read_only(false);
402        assert_eq!(opts.to_begin_sql(), "BEGIN READ WRITE");
403    }
404
405    // -----------------------------------------------------------------------
406    // Transaction lifecycle (mock transport)
407    // -----------------------------------------------------------------------
408
409    #[tokio::test]
410    async fn test_transaction_commit_mock() {
411        let mut data = Vec::new();
412        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
413        data.extend_from_slice(&build_ready_for_query(b'T'));
414        data.extend_from_slice(&build_command_complete_msg("COMMIT"));
415        data.extend_from_slice(&build_ready_for_query(b'I'));
416
417        let mut conn = make_connection(data);
418        let txn = conn.transaction().await.unwrap();
419        assert!(!txn.committed);
420        txn.commit().await.unwrap();
421        // After commit txn is consumed; verify connection state
422        assert_eq!(conn.transaction_status(), TransactionStatus::Idle);
423    }
424
425    #[tokio::test]
426    async fn test_transaction_rollback_mock() {
427        let mut data = Vec::new();
428        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
429        data.extend_from_slice(&build_ready_for_query(b'T'));
430        data.extend_from_slice(&build_command_complete_msg("ROLLBACK"));
431        data.extend_from_slice(&build_ready_for_query(b'I'));
432
433        let mut conn = make_connection(data);
434        let txn = conn.transaction().await.unwrap();
435        assert!(!txn.committed);
436        txn.rollback().await.unwrap();
437        assert_eq!(conn.transaction_status(), TransactionStatus::Idle);
438    }
439
440    #[tokio::test]
441    async fn test_transaction_is_failed_mock() {
442        let mut data = Vec::new();
443        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
444        data.extend_from_slice(&build_ready_for_query(b'T'));
445        data.extend_from_slice(&build_error_response("syntax error"));
446        data.extend_from_slice(&build_ready_for_query(b'E'));
447        data.extend_from_slice(&build_command_complete_msg("ROLLBACK"));
448        data.extend_from_slice(&build_ready_for_query(b'I'));
449
450        let mut conn = make_connection(data);
451        let mut txn = conn.transaction().await.unwrap();
452        assert!(!txn.is_failed());
453
454        // Bad query puts the transaction in failed state
455        let err = txn.execute("BAD SQL").await;
456        assert!(err.is_err());
457        assert!(txn.is_failed());
458
459        // Rolling back should clear the failed state
460        txn.rollback().await.unwrap();
461        assert_eq!(conn.transaction_status(), TransactionStatus::Idle);
462    }
463
464    #[tokio::test]
465    async fn test_transaction_query_delegation_mock() {
466        let mut data = Vec::new();
467        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
468        data.extend_from_slice(&build_ready_for_query(b'T'));
469        // RowDescription + DataRow + CommandComplete + ReadyForQuery for SELECT
470        data.extend_from_slice(&build_row_description_msg(&[(
471            "val",
472            crate::types::INT4_OID,
473        )]));
474        data.extend_from_slice(&build_data_row_msg(&[Some("42")]));
475        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
476        data.extend_from_slice(&build_ready_for_query(b'T'));
477        data.extend_from_slice(&build_command_complete_msg("COMMIT"));
478        data.extend_from_slice(&build_ready_for_query(b'I'));
479
480        let mut conn = make_connection(data);
481        let mut txn = conn.transaction().await.unwrap();
482        let result = txn.query("SELECT 42").await.unwrap();
483        assert_eq!(result.len(), 1);
484        let v: i32 = result.rows()[0].get(0).unwrap();
485        assert_eq!(v, 42);
486        txn.commit().await.unwrap();
487    }
488
489    #[tokio::test]
490    async fn test_with_transaction_success_mock() {
491        let mut data = Vec::new();
492        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
493        data.extend_from_slice(&build_ready_for_query(b'T'));
494        data.extend_from_slice(&build_row_description_msg(&[(
495            "val",
496            crate::types::INT4_OID,
497        )]));
498        data.extend_from_slice(&build_data_row_msg(&[Some("42")]));
499        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
500        data.extend_from_slice(&build_ready_for_query(b'T'));
501        data.extend_from_slice(&build_command_complete_msg("COMMIT"));
502        data.extend_from_slice(&build_ready_for_query(b'I'));
503
504        let mut conn = make_connection(data);
505        let result = conn
506            .with_transaction(async |txn| {
507                let qr = txn.query("SELECT 42").await?;
508                let v: i32 = qr.rows()[0].get(0)?;
509                Ok(v)
510            })
511            .await
512            .unwrap();
513        assert_eq!(result, 42);
514        assert_eq!(conn.transaction_status(), TransactionStatus::Idle);
515    }
516
517    #[tokio::test]
518    async fn test_with_transaction_error_rolls_back_mock() {
519        let mut data = Vec::new();
520        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
521        data.extend_from_slice(&build_ready_for_query(b'T'));
522        data.extend_from_slice(&build_row_description_msg(&[(
523            "val",
524            crate::types::INT4_OID,
525        )]));
526        data.extend_from_slice(&build_data_row_msg(&[Some("42")]));
527        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
528        data.extend_from_slice(&build_ready_for_query(b'T'));
529        data.extend_from_slice(&build_command_complete_msg("ROLLBACK"));
530        data.extend_from_slice(&build_ready_for_query(b'I'));
531
532        let mut conn = make_connection(data);
533        let result = conn
534            .with_transaction(async |txn| {
535                let qr = txn.query("SELECT 42").await?;
536                let _v: i32 = qr.rows()[0].get(0)?;
537                // Force an error to trigger rollback
538                Err::<i32, Error>(Error::Config("intentional failure".into()))
539            })
540            .await;
541        assert!(result.is_err());
542        assert_eq!(conn.transaction_status(), TransactionStatus::Idle);
543    }
544
545    #[tokio::test]
546    async fn test_transaction_savepoint_mock() {
547        let mut data = Vec::new();
548        data.extend_from_slice(&build_command_complete_msg("BEGIN"));
549        data.extend_from_slice(&build_ready_for_query(b'T'));
550        data.extend_from_slice(&build_command_complete_msg("SAVEPOINT sp1"));
551        data.extend_from_slice(&build_ready_for_query(b'T'));
552        data.extend_from_slice(&build_command_complete_msg("RELEASE SAVEPOINT sp1"));
553        data.extend_from_slice(&build_ready_for_query(b'T'));
554        data.extend_from_slice(&build_command_complete_msg("COMMIT"));
555        data.extend_from_slice(&build_ready_for_query(b'I'));
556
557        let mut conn = make_connection(data);
558        let mut txn = conn.transaction().await.unwrap();
559        assert_eq!(txn.savepoint_depth, 0);
560        let sp = txn.savepoint("sp1").await.unwrap();
561        sp.release().await.unwrap();
562        assert_eq!(txn.savepoint_depth, 0);
563        txn.commit().await.unwrap();
564    }
565
566    #[test]
567    fn test_transaction_drop_without_commit_does_not_panic() {
568        let mut conn = make_connection(Vec::new());
569        let txn = Transaction::new(&mut conn);
570        assert!(!txn.committed);
571        drop(txn); // must not panic
572    }
573
574    #[test]
575    fn test_transaction_drop_after_commit_is_noop() {
576        let mut conn = make_connection(Vec::new());
577        let txn = Transaction::new(&mut conn);
578        // Simulate committed state (normally set by commit())
579        let mut txn = txn;
580        txn.committed = true;
581        drop(txn); // must not panic
582    }
583}
wasi_pg_client/transaction/mod.rs

wasi_pg_client/transaction/
mod.rs
