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 05aa44e4cb3142e7d6ce863545140ddfc742c872..f31aea60450a69ab61c98f6ab5a98dcc7c48034d 100644 (file)
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -1015,3 +1015,39 @@ def test_check_backoff(dsn, caplog, monkeypatch):
     for delta in deltas:
         assert delta == pytest.approx(want, 0.05), deltas
         want *= 2
+
+
+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 f56748d4225a29136e622161c107a90d25413393..7bfb0139aa899a6e66b39d1a1a16c40faace7cb2 100644 (file)
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -1018,3 +1018,40 @@ async def test_check_backoff(dsn, caplog, monkeypatch):
     for delta in deltas:
         assert delta == pytest.approx(want, 0.05), deltas
         want *= 2
+
+
+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