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 873daff4ae0b789136efe5d1cecb62aab3cf0b60..c5a9ae3469a371876a6a8d5ee420762d915c3186 100644 (file)
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -1038,3 +1038,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 939c3f707ecd6743789e5109f980ed5e7a362ac8..e288a27e96ec8efd0bc2be1173899900151b8dab 100644 (file)
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -1040,3 +1040,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