From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 16 Apr 2025 18:57:37 +0000 (+0200)
Subject: test(pool): verify the overriding of `close()` to act as `putconn()`
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=refs%2Fheads%2Fmaint-pool-3.2;p=thirdparty%2Fpsycopg.git

test(pool): verify the overriding of `close()` to act as `putconn()`

We would like to maintain this possibility on a best-effort basis. The
hack has been mentioned in #1046 and upstream in
https://github.com/sqlalchemy/sqlalchemy/discussions/12522 and is useful
to integrate psycopg pool with the SQLAlchemy NullPool.
---

diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index fc8782e80..7fb4da7d8 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -1037,3 +1037,39 @@ def test_check_returns_an_ok_connection(dsn, status):
 
         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.
+    #
+    # We cannot guarantee 100% that we will never break this implementation,
+    # but we can keep awareness that we use it this way, maintain it on a
+    # best-effort basis, and notify upstream if we are forced to break it.
+    #
+    # https://github.com/sqlalchemy/sqlalchemy/discussions/12522
+    # https://github.com/psycopg/psycopg/issues/1046
+
+    class MyConnection(psycopg.Connection[Row]):
+
+        def close(self) -> None:
+            if pool := getattr(self, "_pool", None):
+                # Connection currently checked out from the pool.
+                # Instead of closing it, return it to the pool.
+                pool.putconn(self)
+            else:
+                # Connection not part of any pool, or currently into the pool.
+                # Close the connection for real.
+                super().close()
+
+    with pool.ConnectionPool(dsn, connection_class=MyConnection, min_size=2) as p:
+        p.wait()
+        assert len(p._pool) == 2
+        conn = p.getconn()
+        assert not conn.closed
+        assert len(p._pool) == 1
+        conn.close()
+        assert not conn.closed
+        assert len(p._pool) == 2
+
+    assert conn.closed
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index 0c38f0298..0366619e7 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -1039,3 +1039,40 @@ async def test_check_returns_an_ok_connection(dsn, status):
 
         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.
+    #
+    # We cannot guarantee 100% that we will never break this implementation,
+    # but we can keep awareness that we use it this way, maintain it on a
+    # best-effort basis, and notify upstream if we are forced to break it.
+    #
+    # https://github.com/sqlalchemy/sqlalchemy/discussions/12522
+    # https://github.com/psycopg/psycopg/issues/1046
+
+    class MyConnection(psycopg.AsyncConnection[Row]):
+        async def close(self) -> None:
+            if pool := getattr(self, "_pool", None):
+                # Connection currently checked out from the pool.
+                # Instead of closing it, return it to the pool.
+                await pool.putconn(self)
+            else:
+                # Connection not part of any pool, or currently into the pool.
+                # Close the connection for real.
+                await super().close()
+
+    async with pool.AsyncConnectionPool(
+        dsn, connection_class=MyConnection, min_size=2
+    ) as p:
+        await p.wait()
+        assert len(p._pool) == 2
+        conn = await p.getconn()
+        assert not conn.closed
+        assert len(p._pool) == 1
+        await conn.close()
+        assert not conn.closed
+        assert len(p._pool) == 2
+
+    assert conn.closed