Jul 22, 2024 · Jul 22, 2024
diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs
 use crate::codec::{BackendMessages, FrontendMessage};
 use crate::codec::BackendMessages;
 use crate::config::SslMode;
 use crate::connection::{Request, RequestMessages};
 use crate::copy_out::CopyOutStream;
 use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 use postgres_protocol::message::{backend::Message, frontend};
 use postgres_protocol::message::backend::Message;
 use postgres_types::BorrowToSql;
 use std::collections::HashMap;
 use std::fmt;
    ///
    /// The transaction will roll back by default - use the `commit` method to commit it.
    pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
        struct RollbackIfNotDone<'me> {
            client: &'me Client,
            done: bool,
        }

        impl<'a> Drop for RollbackIfNotDone<'a> {
            fn drop(&mut self) {
                if self.done {
                    return;
                }

                let buf = self.client.inner().with_buf(|buf| {
                    frontend::query("ROLLBACK", buf).unwrap();
                    buf.split().freeze()
                });
                let _ = self
                    .client
                    .inner()
                    .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
            }
        }

        // This is done, as `Future` created by this method can be dropped after
        // `RequestMessages` is synchronously send to the `Connection` by
        // `batch_execute()`, but before `Responses` is asynchronously polled to
        // completion. In that case `Transaction` won't be created and thus
        // won't be rolled back.
        {
            let mut cleaner = RollbackIfNotDone {
                client: self,
                done: false,
            };
            self.batch_execute("BEGIN").await?;
            cleaner.done = true;
        }

        Ok(Transaction::new(self))
        self.build_transaction().start().await
    }

    /// Returns a builder for a transaction with custom settings.
diff --git a/tokio-postgres/src/transaction_builder.rs b/tokio-postgres/src/transaction_builder.rs
 use crate::{Client, Error, Transaction};
 use postgres_protocol::message::frontend;

 use crate::{codec::FrontendMessage, connection::RequestMessages, Client, Error, Transaction};

 /// The isolation level of a database transaction.
 #[derive(Debug, Copy, Clone)]
            query.push_str(s);
        }

        self.client.batch_execute(&query).await?;
        struct RollbackIfNotDone<'me> {
            client: &'me Client,
            done: bool,
        }

        impl<'a> Drop for RollbackIfNotDone<'a> {
            fn drop(&mut self) {
                if self.done {
                    return;
                }

                let buf = self.client.inner().with_buf(|buf| {
                    frontend::query("ROLLBACK", buf).unwrap();
                    buf.split().freeze()
                });
                let _ = self
                    .client
                    .inner()
                    .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
            }
        }

        // This is done as `Future` created by this method can be dropped after
        // `RequestMessages` is synchronously send to the `Connection` by
        // `batch_execute()`, but before `Responses` is asynchronously polled to
        // completion. In that case `Transaction` won't be created and thus
        // won't be rolled back.
        {
            let mut cleaner = RollbackIfNotDone {
                client: self.client,
                done: false,
            };
            self.client.batch_execute(&query).await?;
            cleaner.done = true;
        }

        Ok(Transaction::new(self.client))
    }
Original file line number	Diff line number	Diff line change
		@@ -1,4 +1,4 @@
		use crate::codec::{BackendMessages, FrontendMessage};
		use crate::codec::BackendMessages;
		use crate::config::SslMode;
		use crate::connection::{Request, RequestMessages};
		use crate::copy_out::CopyOutStream;
Expand All		@@ -21,7 +21,7 @@ use fallible_iterator::FallibleIterator;
		use futures_channel::mpsc;
		use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
		use parking_lot::Mutex;
		use postgres_protocol::message::{backend::Message, frontend};
		use postgres_protocol::message::backend::Message;
		use postgres_types::BorrowToSql;
		use std::collections::HashMap;
		use std::fmt;
Expand DownExpand Up		@@ -532,43 +532,7 @@ impl Client {
		///
		/// The transaction will roll back by default - use the `commit` method to commit it.
		pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
		struct RollbackIfNotDone<'me> {
		client: &'me Client,
		done: bool,
		}

		impl<'a> Drop for RollbackIfNotDone<'a> {
		fn drop(&mut self) {
		if self.done {
		return;
		}

		let buf = self.client.inner().with_buf(\|buf\| {
		frontend::query("ROLLBACK", buf).unwrap();
		buf.split().freeze()
		});
		let _ = self
		.client
		.inner()
		.send(RequestMessages::Single(FrontendMessage::Raw(buf)));
		}
		}

		// This is done, as `Future` created by this method can be dropped after
		// `RequestMessages` is synchronously send to the `Connection` by
		// `batch_execute()`, but before `Responses` is asynchronously polled to
		// completion. In that case `Transaction` won't be created and thus
		// won't be rolled back.
		{
		let mut cleaner = RollbackIfNotDone {
		client: self,
		done: false,
		};
		self.batch_execute("BEGIN").await?;
		cleaner.done = true;
		}

		Ok(Transaction::new(self))
		self.build_transaction().start().await
		}

		/// Returns a builder for a transaction with custom settings.
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -1,4 +1,6 @@
		use crate::{Client, Error, Transaction};
		use postgres_protocol::message::frontend;

		use crate::{codec::FrontendMessage, connection::RequestMessages, Client, Error, Transaction};

		/// The isolation level of a database transaction.
		#[derive(Debug, Copy, Clone)]
Expand DownExpand Up		@@ -106,7 +108,41 @@ impl<'a> TransactionBuilder<'a> {
		query.push_str(s);
		}

		self.client.batch_execute(&query).await?;
		struct RollbackIfNotDone<'me> {
		client: &'me Client,
		done: bool,
		}

		impl<'a> Drop for RollbackIfNotDone<'a> {
		fn drop(&mut self) {
		if self.done {
		return;
		}

		let buf = self.client.inner().with_buf(\|buf\| {
		frontend::query("ROLLBACK", buf).unwrap();
		buf.split().freeze()
		});
		let _ = self
		.client
		.inner()
		.send(RequestMessages::Single(FrontendMessage::Raw(buf)));
		}
		}

		// This is done as `Future` created by this method can be dropped after
		// `RequestMessages` is synchronously send to the `Connection` by
		// `batch_execute()`, but before `Responses` is asynchronously polled to
		// completion. In that case `Transaction` won't be created and thus
		// won't be rolled back.
		{
		let mut cleaner = RollbackIfNotDone {
		client: self.client,
		done: false,
		};
		self.client.batch_execute(&query).await?;
		cleaner.done = true;
		}

		Ok(Transaction::new(self.client))
		}
Expand Down