From 413fd4f7192a33b9dd692c796dd9b8055aef0124 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Sun, 4 May 2025 20:35:25 +0200 Subject: [PATCH] feat(pool): add connection subclasses to enable putconn-on-close --- psycopg_pool/psycopg_pool/_compat.py | 42 ++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) 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", ] -- 2.47.3