Line data    Source code

       1              : use crate::codec::FrontendMessage;
       2              : use crate::connection::RequestMessages;
       3              : use crate::query::RowStream;
       4              : use crate::{CancelToken, Client, Error, ReadyForQueryStatus};
       5              : use postgres_protocol2::message::frontend;
       6              : 
       7              : /// A representation of a PostgreSQL database transaction.
       8              : ///
       9              : /// Transactions will implicitly roll back when dropped. Use the `commit` method to commit the changes made in the
      10              : /// transaction. Transactions can be nested, with inner transactions implemented via safepoints.
      11              : pub struct Transaction<'a> {
      12              :     client: &'a mut Client,
      13              :     done: bool,
      14              : }
      15              : 
      16              : impl Drop for Transaction<'_> {
      17            0 :     fn drop(&mut self) {
      18            0 :         if self.done {
      19            0 :             return;
      20            0 :         }
      21            0 : 
      22            0 :         let buf = self.client.inner().with_buf(|buf| {
      23            0 :             frontend::query("ROLLBACK", buf).unwrap();
      24            0 :             buf.split().freeze()
      25            0 :         });
      26            0 :         let _ = self
      27            0 :             .client
      28            0 :             .inner()
      29            0 :             .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
      30            0 :     }
      31              : }
      32              : 
      33              : impl<'a> Transaction<'a> {
      34            0 :     pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> {
      35            0 :         Transaction {
      36            0 :             client,
      37            0 :             done: false,
      38            0 :         }
      39            0 :     }
      40              : 
      41              :     /// Consumes the transaction, committing all changes made within it.
      42            0 :     pub async fn commit(mut self) -> Result<ReadyForQueryStatus, Error> {
      43            0 :         self.done = true;
      44            0 :         self.client.batch_execute("COMMIT").await
      45            0 :     }
      46              : 
      47              :     /// Rolls the transaction back, discarding all changes made within it.
      48              :     ///
      49              :     /// This is equivalent to `Transaction`'s `Drop` implementation, but provides any error encountered to the caller.
      50            0 :     pub async fn rollback(mut self) -> Result<ReadyForQueryStatus, Error> {
      51            0 :         self.done = true;
      52            0 :         self.client.batch_execute("ROLLBACK").await
      53            0 :     }
      54              : 
      55              :     /// Like `Client::query_raw_txt`.
      56            0 :     pub async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
      57            0 :     where
      58            0 :         S: AsRef<str>,
      59            0 :         I: IntoIterator<Item = Option<S>>,
      60            0 :         I::IntoIter: ExactSizeIterator,
      61            0 :     {
      62            0 :         self.client.query_raw_txt(statement, params).await
      63            0 :     }
      64              : 
      65              :     /// Like `Client::cancel_token`.
      66            0 :     pub fn cancel_token(&self) -> CancelToken {
      67            0 :         self.client.cancel_token()
      68            0 :     }
      69              : 
      70              :     /// Returns a reference to the underlying `Client`.
      71            0 :     pub fn client(&self) -> &Client {
      72            0 :         self.client
      73            0 :     }
      74              : }