From 9e671fe44fac55d78bc255efbac1f2b844401daf Mon Sep 17 00:00:00 2001
From: Denis Laxalde <denis@laxalde.org>
Date: Mon, 9 May 2022 21:14:59 +0200
Subject: [PATCH] fix: wrap transaction in pipelines if the connection has one

When entering a transaction on a connection in pipeline mode, we open an
inner pipeline to ensure that a Sync is emitted at the end of
transaction thus restoring the connection in its expected state (i.e.
the same as in non-pipeline mode).
---
 psycopg/psycopg/connection.py       | 9 +++++++--
 psycopg/psycopg/connection_async.py | 8 ++++++--
 tests/test_pipeline.py              | 1 -
 tests/test_pipeline_async.py        | 1 -
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/psycopg/psycopg/connection.py b/psycopg/psycopg/connection.py
index 4440c2285..f8a43bb92 100644
--- a/psycopg/psycopg/connection.py
+++ b/psycopg/psycopg/connection.py
@@ -869,8 +869,13 @@ class Connection(BaseConnection[Row]):
             block even if there were no error (e.g. to try a no-op process).
         :rtype: Transaction
         """
-        with Transaction(self, savepoint_name, force_rollback) as tx:
-            yield tx
+        tx = Transaction(self, savepoint_name, force_rollback)
+        if self._pipeline:
+            with tx, self.pipeline():
+                yield tx
+        else:
+            with tx:
+                yield tx
 
     def notifies(self) -> Generator[Notify, None, None]:
         """
diff --git a/psycopg/psycopg/connection_async.py b/psycopg/psycopg/connection_async.py
index af37e1618..606c68386 100644
--- a/psycopg/psycopg/connection_async.py
+++ b/psycopg/psycopg/connection_async.py
@@ -282,8 +282,12 @@ class AsyncConnection(BaseConnection[Row]):
         :rtype: AsyncTransaction
         """
         tx = AsyncTransaction(self, savepoint_name, force_rollback)
-        async with tx:
-            yield tx
+        if self._pipeline:
+            async with tx, self.pipeline():
+                yield tx
+        else:
+            async with tx:
+                yield tx
 
     async def notifies(self) -> AsyncGenerator[Notify, None]:
         while 1:
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index d8783f35b..29816cf23 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -237,7 +237,6 @@ def test_errors_raised_on_transaction_exit(conn):
             with conn.transaction():
                 conn.execute("select 1 from nosuchtable")
                 here = True
-        conn.rollback()  # TODO: inconsistent with non-pipeline.
         cur1 = conn.execute("select 1")
     assert here
     cur2 = conn.execute("select 2")
diff --git a/tests/test_pipeline_async.py b/tests/test_pipeline_async.py
index c69fef788..b3dfef755 100644
--- a/tests/test_pipeline_async.py
+++ b/tests/test_pipeline_async.py
@@ -240,7 +240,6 @@ async def test_errors_raised_on_transaction_exit(aconn):
             async with aconn.transaction():
                 await aconn.execute("select 1 from nosuchtable")
                 here = True
-        await aconn.rollback()  # TODO: inconsistent with non-pipeline.
         cur1 = await aconn.execute("select 1")
     assert here
     cur2 = await aconn.execute("select 2")
-- 
2.47.2