From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Mon, 8 Mar 2021 01:24:09 +0000 (+0100)
Subject: Don't throw an error on context exit if the connection is closed
X-Git-Tag: 3.0.dev0~87^2~24
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=7084bc6e30289c03f2938b299dfbff375a2c3f30;p=thirdparty%2Fpsycopg.git

Don't throw an error on context exit if the connection is closed
---

diff --git a/psycopg3/psycopg3/connection.py b/psycopg3/psycopg3/connection.py
index 7cdb90f81..d95434b2e 100644
--- a/psycopg3/psycopg3/connection.py
+++ b/psycopg3/psycopg3/connection.py
@@ -465,17 +465,19 @@ class Connection(BaseConnection):
         exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ) -> None:
+        if self.closed:
+            return
+
         if exc_type:
             # try to rollback, but if there are problems (connection in a bad
             # state) just warn without clobbering the exception bubbling up.
-            if not self.closed:
-                try:
-                    self.rollback()
-                except Exception as exc2:
-                    warnings.warn(
-                        f"error rolling back the transaction on {self}: {exc2}",
-                        RuntimeWarning,
-                    )
+            try:
+                self.rollback()
+            except Exception as exc2:
+                warnings.warn(
+                    f"error rolling back the transaction on {self}: {exc2}",
+                    RuntimeWarning,
+                )
         else:
             self.commit()
 
@@ -641,17 +643,19 @@ class AsyncConnection(BaseConnection):
         exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ) -> None:
+        if self.closed:
+            return
+
         if exc_type:
             # try to rollback, but if there are problems (connection in a bad
             # state) just warn without clobbering the exception bubbling up.
-            if not self.closed:
-                try:
-                    await self.rollback()
-                except Exception as exc2:
-                    warnings.warn(
-                        f"error rolling back the transaction on {self}: {exc2}",
-                        RuntimeWarning,
-                    )
+            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/pool/test_pool.py b/tests/pool/test_pool.py
index fec574c3c..1c3177814 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -241,11 +241,10 @@ def test_queue_timeout_override(dsn):
 
 def test_broken_reconnect(dsn):
     with pool.ConnectionPool(dsn, minconn=1) as p:
-        with pytest.raises(psycopg3.OperationalError):
-            with p.connection() as conn:
-                with conn.execute("select pg_backend_pid()") as cur:
-                    (pid1,) = cur.fetchone()
-                conn.close()
+        with p.connection() as conn:
+            with conn.execute("select pg_backend_pid()") as cur:
+                (pid1,) = cur.fetchone()
+            conn.close()
 
         with p.connection() as conn2:
             with conn2.execute("select pg_backend_pid()") as cur:
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index 810643f8d..fa46ba543 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -257,11 +257,10 @@ async def test_queue_timeout_override(dsn):
 
 async def test_broken_reconnect(dsn):
     async with pool.AsyncConnectionPool(dsn, minconn=1) as p:
-        with pytest.raises(psycopg3.OperationalError):
-            async with p.connection() as conn:
-                cur = await conn.execute("select pg_backend_pid()")
-                (pid1,) = await cur.fetchone()
-                await conn.close()
+        async with p.connection() as conn:
+            cur = await conn.execute("select pg_backend_pid()")
+            (pid1,) = await cur.fetchone()
+            await conn.close()
 
         async with p.connection() as conn2:
             cur = await conn2.execute("select pg_backend_pid()")
diff --git a/tests/test_connection.py b/tests/test_connection.py
index cdfbe2e0d..b97570d49 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -152,6 +152,12 @@ def test_context_rollback(conn, dsn):
                 cur.execute("select * from textctx")
 
 
+def test_context_close(conn):
+    with conn:
+        conn.execute("select 1")
+        conn.close()
+
+
 def test_context_rollback_no_clobber(conn, dsn, recwarn):
     with pytest.raises(ZeroDivisionError):
         with psycopg3.connect(dsn) as conn2:
diff --git a/tests/test_connection_async.py b/tests/test_connection_async.py
index 2c303e140..e4406d1a4 100644
--- a/tests/test_connection_async.py
+++ b/tests/test_connection_async.py
@@ -157,6 +157,12 @@ async def test_context_rollback(aconn, dsn):
                 await cur.execute("select * from textctx")
 
 
+async def test_context_close(aconn):
+    async with aconn:
+        await aconn.execute("select 1")
+        await aconn.close()
+
+
 async def test_context_rollback_no_clobber(conn, dsn, recwarn):
     with pytest.raises(ZeroDivisionError):
         async with await psycopg3.AsyncConnection.connect(dsn) as conn2: