From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 10 Nov 2021 01:21:37 +0000 (+0100)
Subject: Don't leave connections in transaction after pool.check()
X-Git-Tag: 3.0.3~2^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=38278987daf426c7c3cb7456c9e5c8f5025bc610;p=thirdparty%2Fpsycopg.git

Don't leave connections in transaction after pool.check()

Close #144.
---

diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 27cc35fb7..74dda9e7a 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -358,6 +358,8 @@ class ConnectionPool(BasePool[Connection[Any]]):
             conn = conns.pop()
             try:
                 conn.execute("SELECT 1")
+                if conn.pgconn.transaction_status == TransactionStatus.INTRANS:
+                    conn.rollback()
             except Exception:
                 self._stats[self._CONNECTIONS_LOST] += 1
                 logger.warning("discarding broken connection: %s", conn)
diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index ae1612978..f1da22240 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -294,6 +294,8 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
             conn = conns.pop()
             try:
                 await conn.execute("SELECT 1")
+                if conn.pgconn.transaction_status == TransactionStatus.INTRANS:
+                    await conn.rollback()
             except Exception:
                 self._stats[self._CONNECTIONS_LOST] += 1
                 logger.warning("discarding broken connection: %s", conn)
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index dfeca1d81..703bf70e0 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -910,6 +910,14 @@ def test_check(dsn, caplog):
         assert pid not in pids2
 
 
+def test_check_idle(dsn):
+    with pool.ConnectionPool(dsn, min_size=2) as p:
+        p.wait(1.0)
+        p.check()
+        with p.connection() as conn:
+            assert conn.info.transaction_status == TransactionStatus.IDLE
+
+
 @pytest.mark.skipif(
     sys.version_info >= (3, 7), reason="async pool supported from Python 3.7"
 )
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index a05fe8283..608c35e0c 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -907,6 +907,14 @@ async def test_check(dsn, caplog):
         assert pid not in pids2
 
 
+async def test_check_idle(dsn):
+    async with pool.AsyncConnectionPool(dsn, min_size=2) as p:
+        await p.wait(1.0)
+        await p.check()
+        async with p.connection() as conn:
+            assert conn.info.transaction_status == TransactionStatus.IDLE
+
+
 @pytest.mark.slow
 @pytest.mark.timing
 async def test_stats_measures(dsn):