diff --git a/src/database.py b/src/database.py index 11e8728..e088f92 100644 --- a/src/database.py +++ b/src/database.py @@ -3,7 +3,7 @@ import psycopg2 from psycopg2 import pool from datetime import datetime, timezone -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) @@ -12,6 +12,24 @@ DEFAULT_POOL_MAX_CONN = 50 +class _PooledConnection: + """Wrapper that returns connection to pool on close() instead of closing it.""" + + def __init__(self, conn, pool_instance): + self._conn = conn + self._pool = pool_instance + + def close(self): + try: + self._conn.rollback() + except Exception: + pass + self._pool.putconn(self._conn) + + def __getattr__(self, name: str) -> Any: + return getattr(self._conn, name) + + class Database: """PostgreSQL database manager for Hell Divers 2 API data with connection pooling""" @@ -55,16 +73,7 @@ def _get_pool(self) -> pool.ThreadedConnectionPool: def _get_connection(self): """Get a database connection from the pool (returns to pool on close).""" conn = self._get_pool().getconn() - # Wrap so conn.close() returns to pool instead of closing the underlying connection - def _close(): - try: - conn.rollback() # Reset connection state for reuse - except Exception: - pass - self._get_pool().putconn(conn) - - conn.close = _close - return conn + return _PooledConnection(conn, self._get_pool()) def close_pool(self): """Close the connection pool. Call on application shutdown."""