From: Kanav Kalucha <kanav@openai.com>
Date: Tue, 25 Feb 2025 18:28:38 +0000 (-0800)
Subject: fix(pool): reset connection transaction status after failed check
X-Git-Tag: 3.2.12~3^2~2
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=0755481dd9137a408ea832b444c88dcbdae4d73b;p=thirdparty%2Fpsycopg.git

fix(pool): reset connection transaction status after failed check

Close #1014
---

diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 64906e60d..63cd806e0 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -669,6 +669,7 @@ class ConnectionPool(Generic[CT], BasePool):
         """
         Return a connection to the pool after usage.
         """
+        self._reset_connection(conn)
         if from_getconn:
             if conn.pgconn.transaction_status == TransactionStatus.UNKNOWN:
                 self._stats[self._CONNECTIONS_LOST] += 1
@@ -676,14 +677,12 @@ class ConnectionPool(Generic[CT], BasePool):
                 self.run_task(AddConnection(self))
                 logger.info("not serving connection found broken")
                 return
-        else:
-            self._reset_connection(conn)
-            if conn.pgconn.transaction_status == TransactionStatus.UNKNOWN:
-                self._stats[self._RETURNS_BAD] += 1
-                # Connection no more in working state: create a new one.
-                self.run_task(AddConnection(self))
-                logger.warning("discarding closed connection: %s", conn)
-                return
+        elif conn.pgconn.transaction_status == TransactionStatus.UNKNOWN:
+            self._stats[self._RETURNS_BAD] += 1
+            # Connection no more in working state: create a new one.
+            self.run_task(AddConnection(self))
+            logger.warning("discarding closed connection: %s", conn)
+            return
 
         # Check if the connection is past its best before date
         if conn._expire_at <= monotonic():
diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index f39605e91..8c511be49 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -724,6 +724,7 @@ class AsyncConnectionPool(Generic[ACT], BasePool):
         """
         Return a connection to the pool after usage.
         """
+        await self._reset_connection(conn)
         if from_getconn:
             if conn.pgconn.transaction_status == TransactionStatus.UNKNOWN:
                 self._stats[self._CONNECTIONS_LOST] += 1
@@ -733,7 +734,6 @@ class AsyncConnectionPool(Generic[ACT], BasePool):
                 return
 
         else:
-            await self._reset_connection(conn)
             if conn.pgconn.transaction_status == TransactionStatus.UNKNOWN:
                 self._stats[self._RETURNS_BAD] += 1
                 # Connection no more in working state: create a new one.
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index a0e6b2aec..1b7415619 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -1017,6 +1017,28 @@ def test_check_backoff(dsn, caplog, monkeypatch):
         want *= 2
 
 
+@pytest.mark.slow
+@pytest.mark.parametrize("status", ["ERROR", "INTRANS"])
+def test_check_returns_an_ok_connection(dsn, status):
+
+    def check(conn):
+        if status == "ERROR":
+            conn.execute("wat")
+        elif status == "INTRANS":
+            conn.execute("select 1")
+            1 / 0
+        else:
+            assert False
+
+    with pool.ConnectionPool(dsn, min_size=1, check=check) as p:
+        p.wait(1.0)
+        with pytest.raises(pool.PoolTimeout):
+            conn = p.getconn(0.5)
+
+        conn = list(p._pool)[0]
+        assert conn.info.transaction_status == TransactionStatus.IDLE
+
+
 def test_override_close(dsn):
     # Verify that it's possible to override `close()` to act as `putconn()`.
     # which allows to use the psycopg pool in a sqlalchemy NullPool.
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index a301b07e0..ff8509b7d 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -1020,6 +1020,27 @@ async def test_check_backoff(dsn, caplog, monkeypatch):
         want *= 2
 
 
+@pytest.mark.slow
+@pytest.mark.parametrize("status", ["ERROR", "INTRANS"])
+async def test_check_returns_an_ok_connection(dsn, status):
+    async def check(conn):
+        if status == "ERROR":
+            await conn.execute("wat")
+        elif status == "INTRANS":
+            await conn.execute("select 1")
+            1 / 0
+        else:
+            assert False
+
+    async with pool.AsyncConnectionPool(dsn, min_size=1, check=check) as p:
+        await p.wait(1.0)
+        with pytest.raises(pool.PoolTimeout):
+            conn = await p.getconn(0.5)
+
+        conn = list(p._pool)[0]
+        assert conn.info.transaction_status == TransactionStatus.IDLE
+
+
 async def test_override_close(dsn):
     # Verify that it's possible to override `close()` to act as `putconn()`.
     # which allows to use the psycopg pool in a sqlalchemy NullPool.