140 lines
4.8 KiB
Python
140 lines
4.8 KiB
Python
"""
|
|
Database connection pool management for TimescaleDB.
|
|
"""
|
|
|
|
import asyncio
|
|
import asyncpg
|
|
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
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class DatabaseConnectionPool:
|
|
"""Manages database connection pool for TimescaleDB"""
|
|
|
|
def __init__(self):
|
|
self._pool: Optional[asyncpg.Pool] = None
|
|
self._is_initialized = False
|
|
|
|
async def initialize(self) -> None:
|
|
"""Initialize the connection pool"""
|
|
if self._is_initialized:
|
|
return
|
|
|
|
try:
|
|
# Build connection string
|
|
dsn = (
|
|
f"postgresql://{config.database.user}:{config.database.password}"
|
|
f"@{config.database.host}:{config.database.port}/{config.database.name}"
|
|
)
|
|
|
|
# Create connection pool
|
|
self._pool = await asyncpg.create_pool(
|
|
dsn,
|
|
min_size=5,
|
|
max_size=config.database.pool_size,
|
|
max_queries=50000,
|
|
max_inactive_connection_lifetime=300,
|
|
command_timeout=config.database.pool_timeout,
|
|
server_settings={
|
|
'search_path': config.database.schema,
|
|
'timezone': 'UTC'
|
|
}
|
|
)
|
|
|
|
self._is_initialized = True
|
|
logger.info(f"Database connection pool initialized with {config.database.pool_size} connections")
|
|
|
|
# Test connection
|
|
await self.health_check()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize database connection pool: {e}")
|
|
raise StorageError(f"Database connection failed: {e}", "DB_INIT_ERROR")
|
|
|
|
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"""
|
|
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
|
|
except Exception as e:
|
|
logger.error(f"Database health check failed: {e}")
|
|
return False
|
|
|
|
async def get_pool_stats(self) -> Dict[str, Any]:
|
|
"""Get connection pool statistics"""
|
|
if not self._pool:
|
|
return {}
|
|
|
|
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()
|
|
}
|
|
|
|
@property
|
|
def is_initialized(self) -> bool:
|
|
"""Check if pool is initialized"""
|
|
return self._is_initialized
|
|
|
|
|
|
# Global connection pool instance
|
|
db_pool = DatabaseConnectionPool() |