""" Database connection pool management with health monitoring and automatic recovery. """ import asyncio import logging from typing import Optional, Dict, Any from datetime import datetime, timedelta import asyncpg from asyncpg import Pool from ..config import Config from ..utils.exceptions import ConnectionError logger = logging.getLogger(__name__) class ConnectionPoolManager: """Manages database connection pools with health monitoring and recovery.""" def __init__(self, config: Config): self.config = config self.pool: Optional[Pool] = None self._connection_string = self._build_connection_string() self._health_check_interval = 30 # seconds self._health_check_task: Optional[asyncio.Task] = None self._last_health_check = datetime.utcnow() self._connection_failures = 0 self._max_failures = 5 def _build_connection_string(self) -> str: """Build PostgreSQL connection string from config.""" return ( f"postgresql://{self.config.database.user}:{self.config.database.password}" f"@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}" ) async def initialize(self) -> None: """Initialize connection pool with health monitoring.""" try: logger.info("Creating database connection pool...") self.pool = await asyncpg.create_pool( self._connection_string, min_size=5, max_size=self.config.database.pool_size, command_timeout=60, server_settings={ 'jit': 'off', 'timezone': 'UTC', 'statement_timeout': '30s', 'idle_in_transaction_session_timeout': '60s' }, init=self._init_connection ) # Test initial connection await self._test_connection() # Start health monitoring self._health_check_task = asyncio.create_task(self._health_monitor()) logger.info(f"Database connection pool initialized with {self.config.db_min_connections}-{self.config.db_max_connections} connections") except Exception as e: logger.error(f"Failed to initialize connection pool: {e}") raise ConnectionError(f"Connection pool initialization failed: {e}") async def _init_connection(self, conn: asyncpg.Connection) -> None: """Initialize individual database connections.""" try: # Set connection-specific settings await conn.execute("SET timezone = 'UTC'") await conn.execute("SET statement_timeout = '30s'") # Test TimescaleDB extension result = await conn.fetchval("SELECT extname FROM pg_extension WHERE extname = 'timescaledb'") if not result: logger.warning("TimescaleDB extension not found in database") except Exception as e: logger.error(f"Failed to initialize connection: {e}") raise async def _test_connection(self) -> bool: """Test database connection health.""" try: async with self.pool.acquire() as conn: await conn.execute('SELECT 1') self._connection_failures = 0 return True except Exception as e: self._connection_failures += 1 logger.error(f"Connection test failed (attempt {self._connection_failures}): {e}") if self._connection_failures >= self._max_failures: logger.critical("Maximum connection failures reached, attempting pool recreation") await self._recreate_pool() return False async def _recreate_pool(self) -> None: """Recreate connection pool after failures.""" try: if self.pool: await self.pool.close() self.pool = None # Wait before recreating await asyncio.sleep(5) self.pool = await asyncpg.create_pool( self._connection_string, min_size=5, max_size=self.config.database.pool_size, command_timeout=60, server_settings={ 'jit': 'off', 'timezone': 'UTC' }, init=self._init_connection ) self._connection_failures = 0 logger.info("Connection pool recreated successfully") except Exception as e: logger.error(f"Failed to recreate connection pool: {e}") # Will retry on next health check async def _health_monitor(self) -> None: """Background task to monitor connection pool health.""" while True: try: await asyncio.sleep(self._health_check_interval) if self.pool: await self._test_connection() self._last_health_check = datetime.utcnow() # Log pool statistics periodically if datetime.utcnow().minute % 5 == 0: # Every 5 minutes stats = self.get_pool_stats() logger.debug(f"Connection pool stats: {stats}") except asyncio.CancelledError: logger.info("Health monitor task cancelled") break except Exception as e: logger.error(f"Health monitor error: {e}") async def close(self) -> None: """Close connection pool and stop monitoring.""" if self._health_check_task: self._health_check_task.cancel() try: await self._health_check_task except asyncio.CancelledError: pass if self.pool: await self.pool.close() logger.info("Database connection pool closed") def get_pool_stats(self) -> Dict[str, Any]: """Get connection pool statistics.""" if not self.pool: return {"status": "not_initialized"} return { "status": "active", "size": self.pool.get_size(), "max_size": self.pool.get_max_size(), "min_size": self.pool.get_min_size(), "connection_failures": self._connection_failures, "last_health_check": self._last_health_check.isoformat(), "health_check_interval": self._health_check_interval } def is_healthy(self) -> bool: """Check if connection pool is healthy.""" if not self.pool: return False # Check if health check is recent time_since_check = datetime.utcnow() - self._last_health_check if time_since_check > timedelta(seconds=self._health_check_interval * 2): return False # Check failure count return self._connection_failures < self._max_failures async def acquire(self): """Acquire a connection from the pool.""" if not self.pool: raise ConnectionError("Connection pool not initialized") return self.pool.acquire() async def execute(self, query: str, *args) -> None: """Execute a query using a pooled connection.""" async with self.acquire() as conn: return await conn.execute(query, *args) async def fetch(self, query: str, *args) -> list: """Fetch multiple rows using a pooled connection.""" async with self.acquire() as conn: return await conn.fetch(query, *args) async def fetchrow(self, query: str, *args): """Fetch a single row using a pooled connection.""" async with self.acquire() as conn: return await conn.fetchrow(query, *args) async def fetchval(self, query: str, *args): """Fetch a single value using a pooled connection.""" async with self.acquire() as conn: return await conn.fetchval(query, *args)