From 9d4142eb8346054db5dc73abcf03c5d0558f4cdc Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 16 Apr 2025 20:57:37 +0200
Subject: [PATCH] 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.
---
 tests/pool/test_pool.py       | 36 ++++++++++++++++++++++++++++++++++
 tests/pool/test_pool_async.py | 37 +++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index 05aa44e4c..f31aea604 100644
--- 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 f56748d42..7bfb0139a 100644
--- 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
-- 
2.47.2