From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 15 Jan 2021 22:23:35 +0000 (+0100)
Subject: Don't clobber the exception on rollback error in connection exit
X-Git-Tag: 3.0.dev0~152
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=e220f99d0b0735eb8e9948785b2276e8d3baf76c;p=thirdparty%2Fpsycopg.git

Don't clobber the exception on rollback error in connection exit
---

diff --git a/psycopg3/psycopg3/connection.py b/psycopg3/psycopg3/connection.py
index df221f59c..1916f2ac5 100644
--- a/psycopg3/psycopg3/connection.py
+++ b/psycopg3/psycopg3/connection.py
@@ -430,7 +430,15 @@ class Connection(BaseConnection):
         exc_tb: Optional[TracebackType],
     ) -> None:
         if exc_type:
-            self.rollback()
+            # try to rollback, but if there are problems (connection in a bad
+            # state) just warn without clobbering the exception bubbling up.
+            try:
+                self.rollback()
+            except Exception as exc2:
+                warnings.warn(
+                    f"error rolling back the transaction on {self}: {exc2}",
+                    RuntimeWarning,
+                )
         else:
             self.commit()
 
@@ -558,7 +566,15 @@ class AsyncConnection(BaseConnection):
         exc_tb: Optional[TracebackType],
     ) -> None:
         if exc_type:
-            await self.rollback()
+            # try to rollback, but if there are problems (connection in a bad
+            # state) just warn without clobbering the exception bubbling up.
+            try:
+                await self.rollback()
+            except Exception as exc2:
+                warnings.warn(
+                    f"error rolling back the transaction on {self}: {exc2}",
+                    RuntimeWarning,
+                )
         else:
             await self.commit()
 
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 79c725db1..398900793 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -134,6 +134,19 @@ def test_context_rollback(conn, dsn):
                 cur.execute("select * from textctx")
 
 
+def test_context_rollback_no_clobber(conn, dsn, recwarn):
+    with pytest.raises(ZeroDivisionError):
+        with psycopg3.connect(dsn) as conn2:
+            conn2.execute("select 1")
+            conn.execute(
+                "select pg_terminate_backend(%s::int)",
+                [conn2.pgconn.backend_pid],
+            )
+            1 / 0
+
+    assert "rolling back" in str(recwarn.pop(RuntimeWarning).message)
+
+
 def test_weakref(dsn):
     conn = psycopg3.connect(dsn)
     w = weakref.ref(conn)
diff --git a/tests/test_connection_async.py b/tests/test_connection_async.py
index f24a56907..71bfb6a1d 100644
--- a/tests/test_connection_async.py
+++ b/tests/test_connection_async.py
@@ -139,6 +139,19 @@ async def test_context_rollback(aconn, dsn):
                 await cur.execute("select * from textctx")
 
 
+async def test_context_rollback_no_clobber(conn, dsn, recwarn):
+    with pytest.raises(ZeroDivisionError):
+        async with await psycopg3.AsyncConnection.connect(dsn) as conn2:
+            await conn2.execute("select 1")
+            conn.execute(
+                "select pg_terminate_backend(%s::int)",
+                [conn2.pgconn.backend_pid],
+            )
+            1 / 0
+
+    assert "rolling back" in str(recwarn.pop(RuntimeWarning).message)
+
+
 async def test_weakref(dsn):
     conn = await psycopg3.AsyncConnection.connect(dsn)
     w = weakref.ref(conn)