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 142896ad42b6ef39aaa9231034b363a65bb28204..c9ea05e63d640fb8b974ba35c15485277430d6f2 100644 (file)
--- 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",
 ]