storage manager
This commit is contained in:
Dobromir Popov
@ -1,140 +1,219 @@
|
||||
"""
|
||||
Database connection pool management for TimescaleDB.
|
||||
Database connection pool management with health monitoring and automatic recovery.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import asyncpg
|
||||
import logging
|
||||
from typing import Optional, Dict, Any
|
||||
from contextlib import asynccontextmanager
|
||||
from ..config import config
|
||||
from ..utils.logging import get_logger
|
||||
from ..utils.exceptions import StorageError
|
||||
from datetime import datetime, timedelta
|
||||
import asyncpg
|
||||
from asyncpg import Pool
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from ..config import Config
|
||||
from ..utils.exceptions import ConnectionError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatabaseConnectionPool:
|
||||
"""Manages database connection pool for TimescaleDB"""
|
||||
class ConnectionPoolManager:
|
||||
"""Manages database connection pools with health monitoring and recovery."""
|
||||
|
||||
def __init__(self):
|
||||
self._pool: Optional[asyncpg.Pool] = None
|
||||
self._is_initialized = False
|
||||
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 the connection pool"""
|
||||
if self._is_initialized:
|
||||
return
|
||||
|
||||
"""Initialize connection pool with health monitoring."""
|
||||
try:
|
||||
# Build connection string
|
||||
dsn = (
|
||||
f"postgresql://{config.database.user}:{config.database.password}"
|
||||
f"@{config.database.host}:{config.database.port}/{config.database.name}"
|
||||
)
|
||||
logger.info("Creating database connection pool...")
|
||||
|
||||
# Create connection pool
|
||||
self._pool = await asyncpg.create_pool(
|
||||
dsn,
|
||||
self.pool = await asyncpg.create_pool(
|
||||
self._connection_string,
|
||||
min_size=5,
|
||||
max_size=config.database.pool_size,
|
||||
max_queries=50000,
|
||||
max_inactive_connection_lifetime=300,
|
||||
command_timeout=config.database.pool_timeout,
|
||||
max_size=self.config.database.pool_size,
|
||||
command_timeout=60,
|
||||
server_settings={
|
||||
'search_path': config.database.schema,
|
||||
'timezone': 'UTC'
|
||||
}
|
||||
'jit': 'off',
|
||||
'timezone': 'UTC',
|
||||
'statement_timeout': '30s',
|
||||
'idle_in_transaction_session_timeout': '60s'
|
||||
},
|
||||
init=self._init_connection
|
||||
)
|
||||
|
||||
self._is_initialized = True
|
||||
logger.info(f"Database connection pool initialized with {config.database.pool_size} connections")
|
||||
# Test initial connection
|
||||
await self._test_connection()
|
||||
|
||||
# Test connection
|
||||
await self.health_check()
|
||||
# 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 database connection pool: {e}")
|
||||
raise StorageError(f"Database connection failed: {e}", "DB_INIT_ERROR")
|
||||
logger.error(f"Failed to initialize connection pool: {e}")
|
||||
raise ConnectionError(f"Connection pool initialization failed: {e}")
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the connection pool"""
|
||||
if self._pool:
|
||||
await self._pool.close()
|
||||
self._pool = None
|
||||
self._is_initialized = False
|
||||
logger.info("Database connection pool closed")
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_connection(self):
|
||||
"""Get a database connection from the pool"""
|
||||
if not self._is_initialized:
|
||||
await self.initialize()
|
||||
|
||||
if not self._pool:
|
||||
raise StorageError("Connection pool not initialized", "POOL_NOT_READY")
|
||||
|
||||
async with self._pool.acquire() as connection:
|
||||
try:
|
||||
yield connection
|
||||
except Exception as e:
|
||||
logger.error(f"Database operation failed: {e}")
|
||||
raise
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_transaction(self):
|
||||
"""Get a database transaction"""
|
||||
async with self.get_connection() as conn:
|
||||
async with conn.transaction():
|
||||
yield conn
|
||||
|
||||
async def execute_query(self, query: str, *args) -> Any:
|
||||
"""Execute a query and return results"""
|
||||
async with self.get_connection() as conn:
|
||||
return await conn.fetch(query, *args)
|
||||
|
||||
async def execute_command(self, command: str, *args) -> str:
|
||||
"""Execute a command and return status"""
|
||||
async with self.get_connection() as conn:
|
||||
return await conn.execute(command, *args)
|
||||
|
||||
async def execute_many(self, command: str, args_list) -> None:
|
||||
"""Execute a command multiple times with different arguments"""
|
||||
async with self.get_connection() as conn:
|
||||
await conn.executemany(command, args_list)
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check database health"""
|
||||
async def _init_connection(self, conn: asyncpg.Connection) -> None:
|
||||
"""Initialize individual database connections."""
|
||||
try:
|
||||
async with self.get_connection() as conn:
|
||||
result = await conn.fetchval("SELECT 1")
|
||||
if result == 1:
|
||||
logger.debug("Database health check passed")
|
||||
return True
|
||||
else:
|
||||
logger.warning("Database health check returned unexpected result")
|
||||
return False
|
||||
# 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"Database health check failed: {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 get_pool_stats(self) -> Dict[str, Any]:
|
||||
"""Get connection pool statistics"""
|
||||
if not self._pool:
|
||||
return {}
|
||||
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 {
|
||||
'size': self._pool.get_size(),
|
||||
'min_size': self._pool.get_min_size(),
|
||||
'max_size': self._pool.get_max_size(),
|
||||
'idle_size': self._pool.get_idle_size(),
|
||||
'is_closing': self._pool.is_closing()
|
||||
"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
|
||||
}
|
||||
|
||||
@property
|
||||
def is_initialized(self) -> bool:
|
||||
"""Check if pool is initialized"""
|
||||
return self._is_initialized
|
||||
|
||||
|
||||
# Global connection pool instance
|
||||
db_pool = DatabaseConnectionPool()
|
||||
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)
|
||||
Reference in New Issue
Block a user