219 lines
8.2 KiB
Python
219 lines
8.2 KiB
Python
"""
|
|
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) |