From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 4 May 2025 18:35:25 +0000 (+0200)
Subject: feat(pool): add connection subclasses to enable putconn-on-close
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=dbab58361b212ea34061a78bbea261d442ea36b7;p=thirdparty%2Fpsycopg.git

feat(pool): add connection subclasses to enable putconn-on-close
---

diff --git a/psycopg_pool/psycopg_pool/_compat.py b/psycopg_pool/psycopg_pool/_compat.py
index 142896ad4..c9ea05e63 100644
--- a/psycopg_pool/psycopg_pool/_compat.py
+++ b/psycopg_pool/psycopg_pool/_compat.py
@@ -18,9 +18,51 @@ if sys.version_info >= (3, 13):
 else:
     from typing_extensions import TypeVar
 
+import psycopg
 import psycopg.errors as e
 
+PSYCOPG_VERSION = tuple(map(int, psycopg.__version__.split(".", 2)[:2]))
+
+if PSYCOPG_VERSION >= (3, 3):
+    AsyncPoolConnection = psycopg.AsyncConnection
+    PoolConnection = psycopg.Connection
+
+else:
+
+    class AsyncPoolConnection(psycopg.AsyncConnection):  # type: ignore[no-redef]
+        """
+        Thin wrapper around the psycopg async connection to improve pool integration.
+        """
+
+        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()
+
+    class PoolConnection(psycopg.Connection):  # type: ignore[no-redef]
+        """
+        Thin wrapper around the psycopg connection to improve pool integration.
+        """
+
+        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()
+
+
 __all__ = [
+    "AsyncPoolConnection",
+    "PoolConnection",
     "Self",
     "TypeVar",
 ]