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 fc8782e808a7d30cf9d9c6a5107a46079cc23e58..7fb4da7d801859ec6942c67f560d91651130a9ef 100644 (file)
--- 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 0c38f0298adb2258c3f4ad7ba76f1f8ab0f4cac8..0366619e7b082d058d152236f162d2368b388772 100644 (file)
--- 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