Line data    Source code

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