mirror of
https://github.com/ruvnet/RuView
synced 2026-07-26 18:01:48 +00:00
updates
This commit is contained in:
rUv
@@ -0,0 +1,503 @@
|
||||
"""
|
||||
Database connection management for WiFi-DensePose API
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional, Dict, Any, AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import create_engine, event, pool
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import sessionmaker, Session
|
||||
from sqlalchemy.pool import QueuePool, NullPool
|
||||
from sqlalchemy.exc import SQLAlchemyError, DisconnectionError
|
||||
import redis.asyncio as redis
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
from src.config.settings import Settings
|
||||
from src.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class DatabaseConnectionError(Exception):
|
||||
"""Database connection error."""
|
||||
pass
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
"""Database connection manager."""
|
||||
|
||||
def __init__(self, settings: Settings):
|
||||
self.settings = settings
|
||||
self._async_engine = None
|
||||
self._sync_engine = None
|
||||
self._async_session_factory = None
|
||||
self._sync_session_factory = None
|
||||
self._redis_client = None
|
||||
self._initialized = False
|
||||
self._connection_pool_size = settings.db_pool_size
|
||||
self._max_overflow = settings.db_max_overflow
|
||||
self._pool_timeout = settings.db_pool_timeout
|
||||
self._pool_recycle = settings.db_pool_recycle
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize database connections."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
logger.info("Initializing database connections")
|
||||
|
||||
try:
|
||||
# Initialize PostgreSQL connections
|
||||
await self._initialize_postgresql()
|
||||
|
||||
# Initialize Redis connection
|
||||
await self._initialize_redis()
|
||||
|
||||
self._initialized = True
|
||||
logger.info("Database connections initialized successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize database connections: {e}")
|
||||
raise DatabaseConnectionError(f"Database initialization failed: {e}")
|
||||
|
||||
async def _initialize_postgresql(self):
|
||||
"""Initialize PostgreSQL connections."""
|
||||
# Build database URL
|
||||
if self.settings.database_url:
|
||||
db_url = self.settings.database_url
|
||||
async_db_url = self.settings.database_url.replace("postgresql://", "postgresql+asyncpg://")
|
||||
else:
|
||||
db_url = (
|
||||
f"postgresql://{self.settings.db_user}:{self.settings.db_password}"
|
||||
f"@{self.settings.db_host}:{self.settings.db_port}/{self.settings.db_name}"
|
||||
)
|
||||
async_db_url = (
|
||||
f"postgresql+asyncpg://{self.settings.db_user}:{self.settings.db_password}"
|
||||
f"@{self.settings.db_host}:{self.settings.db_port}/{self.settings.db_name}"
|
||||
)
|
||||
|
||||
# Create async engine
|
||||
self._async_engine = create_async_engine(
|
||||
async_db_url,
|
||||
poolclass=QueuePool,
|
||||
pool_size=self._connection_pool_size,
|
||||
max_overflow=self._max_overflow,
|
||||
pool_timeout=self._pool_timeout,
|
||||
pool_recycle=self._pool_recycle,
|
||||
pool_pre_ping=True,
|
||||
echo=self.settings.db_echo,
|
||||
future=True,
|
||||
)
|
||||
|
||||
# Create sync engine for migrations and admin tasks
|
||||
self._sync_engine = create_engine(
|
||||
db_url,
|
||||
poolclass=QueuePool,
|
||||
pool_size=max(2, self._connection_pool_size // 2),
|
||||
max_overflow=self._max_overflow // 2,
|
||||
pool_timeout=self._pool_timeout,
|
||||
pool_recycle=self._pool_recycle,
|
||||
pool_pre_ping=True,
|
||||
echo=self.settings.db_echo,
|
||||
future=True,
|
||||
)
|
||||
|
||||
# Create session factories
|
||||
self._async_session_factory = async_sessionmaker(
|
||||
self._async_engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
)
|
||||
|
||||
self._sync_session_factory = sessionmaker(
|
||||
self._sync_engine,
|
||||
expire_on_commit=False,
|
||||
)
|
||||
|
||||
# Add connection event listeners
|
||||
self._setup_connection_events()
|
||||
|
||||
# Test connections
|
||||
await self._test_postgresql_connection()
|
||||
|
||||
logger.info("PostgreSQL connections initialized")
|
||||
|
||||
async def _initialize_redis(self):
|
||||
"""Initialize Redis connection."""
|
||||
if not self.settings.redis_enabled:
|
||||
logger.info("Redis disabled, skipping initialization")
|
||||
return
|
||||
|
||||
try:
|
||||
# Build Redis URL
|
||||
if self.settings.redis_url:
|
||||
redis_url = self.settings.redis_url
|
||||
else:
|
||||
redis_url = (
|
||||
f"redis://{self.settings.redis_host}:{self.settings.redis_port}"
|
||||
f"/{self.settings.redis_db}"
|
||||
)
|
||||
|
||||
# Create Redis client
|
||||
self._redis_client = redis.from_url(
|
||||
redis_url,
|
||||
password=self.settings.redis_password,
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
max_connections=self.settings.redis_max_connections,
|
||||
retry_on_timeout=True,
|
||||
socket_timeout=self.settings.redis_socket_timeout,
|
||||
socket_connect_timeout=self.settings.redis_connect_timeout,
|
||||
)
|
||||
|
||||
# Test Redis connection
|
||||
await self._test_redis_connection()
|
||||
|
||||
logger.info("Redis connection initialized")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize Redis: {e}")
|
||||
if self.settings.redis_required:
|
||||
raise
|
||||
else:
|
||||
logger.warning("Redis initialization failed but not required, continuing without Redis")
|
||||
|
||||
def _setup_connection_events(self):
|
||||
"""Setup database connection event listeners."""
|
||||
|
||||
@event.listens_for(self._sync_engine, "connect")
|
||||
def set_sqlite_pragma(dbapi_connection, connection_record):
|
||||
"""Set database-specific settings on connection."""
|
||||
if "sqlite" in str(self._sync_engine.url):
|
||||
cursor = dbapi_connection.cursor()
|
||||
cursor.execute("PRAGMA foreign_keys=ON")
|
||||
cursor.close()
|
||||
|
||||
@event.listens_for(self._sync_engine, "checkout")
|
||||
def receive_checkout(dbapi_connection, connection_record, connection_proxy):
|
||||
"""Log connection checkout."""
|
||||
logger.debug("Database connection checked out")
|
||||
|
||||
@event.listens_for(self._sync_engine, "checkin")
|
||||
def receive_checkin(dbapi_connection, connection_record):
|
||||
"""Log connection checkin."""
|
||||
logger.debug("Database connection checked in")
|
||||
|
||||
@event.listens_for(self._sync_engine, "invalidate")
|
||||
def receive_invalidate(dbapi_connection, connection_record, exception):
|
||||
"""Handle connection invalidation."""
|
||||
logger.warning(f"Database connection invalidated: {exception}")
|
||||
|
||||
async def _test_postgresql_connection(self):
|
||||
"""Test PostgreSQL connection."""
|
||||
try:
|
||||
async with self._async_engine.begin() as conn:
|
||||
result = await conn.execute("SELECT 1")
|
||||
await result.fetchone()
|
||||
logger.debug("PostgreSQL connection test successful")
|
||||
except Exception as e:
|
||||
logger.error(f"PostgreSQL connection test failed: {e}")
|
||||
raise DatabaseConnectionError(f"PostgreSQL connection test failed: {e}")
|
||||
|
||||
async def _test_redis_connection(self):
|
||||
"""Test Redis connection."""
|
||||
if not self._redis_client:
|
||||
return
|
||||
|
||||
try:
|
||||
await self._redis_client.ping()
|
||||
logger.debug("Redis connection test successful")
|
||||
except Exception as e:
|
||||
logger.error(f"Redis connection test failed: {e}")
|
||||
if self.settings.redis_required:
|
||||
raise DatabaseConnectionError(f"Redis connection test failed: {e}")
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
|
||||
"""Get async database session."""
|
||||
if not self._initialized:
|
||||
await self.initialize()
|
||||
|
||||
if not self._async_session_factory:
|
||||
raise DatabaseConnectionError("Async session factory not initialized")
|
||||
|
||||
session = self._async_session_factory()
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
await session.rollback()
|
||||
logger.error(f"Database session error: {e}")
|
||||
raise
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_sync_session(self) -> Session:
|
||||
"""Get sync database session."""
|
||||
if not self._initialized:
|
||||
await self.initialize()
|
||||
|
||||
if not self._sync_session_factory:
|
||||
raise DatabaseConnectionError("Sync session factory not initialized")
|
||||
|
||||
session = self._sync_session_factory()
|
||||
try:
|
||||
yield session
|
||||
session.commit()
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
logger.error(f"Database session error: {e}")
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
async def get_redis_client(self) -> Optional[redis.Redis]:
|
||||
"""Get Redis client."""
|
||||
if not self._initialized:
|
||||
await self.initialize()
|
||||
|
||||
return self._redis_client
|
||||
|
||||
async def health_check(self) -> Dict[str, Any]:
|
||||
"""Perform database health check."""
|
||||
health_status = {
|
||||
"postgresql": {"status": "unknown", "details": {}},
|
||||
"redis": {"status": "unknown", "details": {}},
|
||||
"overall": "unknown"
|
||||
}
|
||||
|
||||
# Check PostgreSQL
|
||||
try:
|
||||
start_time = datetime.utcnow()
|
||||
async with self.get_async_session() as session:
|
||||
result = await session.execute("SELECT 1")
|
||||
await result.fetchone()
|
||||
|
||||
response_time = (datetime.utcnow() - start_time).total_seconds()
|
||||
|
||||
health_status["postgresql"] = {
|
||||
"status": "healthy",
|
||||
"details": {
|
||||
"response_time_ms": round(response_time * 1000, 2),
|
||||
"pool_size": self._async_engine.pool.size(),
|
||||
"checked_out": self._async_engine.pool.checkedout(),
|
||||
"overflow": self._async_engine.pool.overflow(),
|
||||
}
|
||||
}
|
||||
except Exception as e:
|
||||
health_status["postgresql"] = {
|
||||
"status": "unhealthy",
|
||||
"details": {"error": str(e)}
|
||||
}
|
||||
|
||||
# Check Redis
|
||||
if self._redis_client:
|
||||
try:
|
||||
start_time = datetime.utcnow()
|
||||
await self._redis_client.ping()
|
||||
response_time = (datetime.utcnow() - start_time).total_seconds()
|
||||
|
||||
info = await self._redis_client.info()
|
||||
|
||||
health_status["redis"] = {
|
||||
"status": "healthy",
|
||||
"details": {
|
||||
"response_time_ms": round(response_time * 1000, 2),
|
||||
"connected_clients": info.get("connected_clients", 0),
|
||||
"used_memory": info.get("used_memory_human", "unknown"),
|
||||
"uptime": info.get("uptime_in_seconds", 0),
|
||||
}
|
||||
}
|
||||
except Exception as e:
|
||||
health_status["redis"] = {
|
||||
"status": "unhealthy",
|
||||
"details": {"error": str(e)}
|
||||
}
|
||||
else:
|
||||
health_status["redis"] = {
|
||||
"status": "disabled",
|
||||
"details": {"message": "Redis not enabled"}
|
||||
}
|
||||
|
||||
# Determine overall status
|
||||
postgresql_healthy = health_status["postgresql"]["status"] == "healthy"
|
||||
redis_healthy = (
|
||||
health_status["redis"]["status"] in ["healthy", "disabled"] or
|
||||
not self.settings.redis_required
|
||||
)
|
||||
|
||||
if postgresql_healthy and redis_healthy:
|
||||
health_status["overall"] = "healthy"
|
||||
elif postgresql_healthy:
|
||||
health_status["overall"] = "degraded"
|
||||
else:
|
||||
health_status["overall"] = "unhealthy"
|
||||
|
||||
return health_status
|
||||
|
||||
async def get_connection_stats(self) -> Dict[str, Any]:
|
||||
"""Get database connection statistics."""
|
||||
stats = {
|
||||
"postgresql": {},
|
||||
"redis": {}
|
||||
}
|
||||
|
||||
# PostgreSQL stats
|
||||
if self._async_engine:
|
||||
pool = self._async_engine.pool
|
||||
stats["postgresql"] = {
|
||||
"pool_size": pool.size(),
|
||||
"checked_out": pool.checkedout(),
|
||||
"overflow": pool.overflow(),
|
||||
"checked_in": pool.checkedin(),
|
||||
"total_connections": pool.size() + pool.overflow(),
|
||||
"available_connections": pool.size() - pool.checkedout(),
|
||||
}
|
||||
|
||||
# Redis stats
|
||||
if self._redis_client:
|
||||
try:
|
||||
info = await self._redis_client.info()
|
||||
stats["redis"] = {
|
||||
"connected_clients": info.get("connected_clients", 0),
|
||||
"blocked_clients": info.get("blocked_clients", 0),
|
||||
"total_connections_received": info.get("total_connections_received", 0),
|
||||
"rejected_connections": info.get("rejected_connections", 0),
|
||||
}
|
||||
except Exception as e:
|
||||
stats["redis"] = {"error": str(e)}
|
||||
|
||||
return stats
|
||||
|
||||
async def close_connections(self):
|
||||
"""Close all database connections."""
|
||||
logger.info("Closing database connections")
|
||||
|
||||
# Close PostgreSQL connections
|
||||
if self._async_engine:
|
||||
await self._async_engine.dispose()
|
||||
logger.debug("Async PostgreSQL engine disposed")
|
||||
|
||||
if self._sync_engine:
|
||||
self._sync_engine.dispose()
|
||||
logger.debug("Sync PostgreSQL engine disposed")
|
||||
|
||||
# Close Redis connection
|
||||
if self._redis_client:
|
||||
await self._redis_client.close()
|
||||
logger.debug("Redis connection closed")
|
||||
|
||||
self._initialized = False
|
||||
logger.info("Database connections closed")
|
||||
|
||||
async def reset_connections(self):
|
||||
"""Reset all database connections."""
|
||||
logger.info("Resetting database connections")
|
||||
await self.close_connections()
|
||||
await self.initialize()
|
||||
logger.info("Database connections reset")
|
||||
|
||||
|
||||
# Global database manager instance
|
||||
_db_manager: Optional[DatabaseManager] = None
|
||||
|
||||
|
||||
def get_database_manager(settings: Settings) -> DatabaseManager:
|
||||
"""Get database manager instance."""
|
||||
global _db_manager
|
||||
if _db_manager is None:
|
||||
_db_manager = DatabaseManager(settings)
|
||||
return _db_manager
|
||||
|
||||
|
||||
async def get_async_session(settings: Settings) -> AsyncGenerator[AsyncSession, None]:
|
||||
"""Dependency to get async database session."""
|
||||
db_manager = get_database_manager(settings)
|
||||
async with db_manager.get_async_session() as session:
|
||||
yield session
|
||||
|
||||
|
||||
async def get_redis_client(settings: Settings) -> Optional[redis.Redis]:
|
||||
"""Dependency to get Redis client."""
|
||||
db_manager = get_database_manager(settings)
|
||||
return await db_manager.get_redis_client()
|
||||
|
||||
|
||||
class DatabaseHealthCheck:
|
||||
"""Database health check utility."""
|
||||
|
||||
def __init__(self, db_manager: DatabaseManager):
|
||||
self.db_manager = db_manager
|
||||
|
||||
async def check_postgresql(self) -> Dict[str, Any]:
|
||||
"""Check PostgreSQL health."""
|
||||
try:
|
||||
start_time = datetime.utcnow()
|
||||
async with self.db_manager.get_async_session() as session:
|
||||
result = await session.execute("SELECT version()")
|
||||
version = (await result.fetchone())[0]
|
||||
|
||||
response_time = (datetime.utcnow() - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"version": version,
|
||||
"response_time_ms": round(response_time * 1000, 2),
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": str(e),
|
||||
}
|
||||
|
||||
async def check_redis(self) -> Dict[str, Any]:
|
||||
"""Check Redis health."""
|
||||
redis_client = await self.db_manager.get_redis_client()
|
||||
|
||||
if not redis_client:
|
||||
return {
|
||||
"status": "disabled",
|
||||
"message": "Redis not configured"
|
||||
}
|
||||
|
||||
try:
|
||||
start_time = datetime.utcnow()
|
||||
pong = await redis_client.ping()
|
||||
response_time = (datetime.utcnow() - start_time).total_seconds()
|
||||
|
||||
info = await redis_client.info("server")
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"ping": pong,
|
||||
"version": info.get("redis_version", "unknown"),
|
||||
"response_time_ms": round(response_time * 1000, 2),
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"error": str(e),
|
||||
}
|
||||
|
||||
async def full_health_check(self) -> Dict[str, Any]:
|
||||
"""Perform full database health check."""
|
||||
postgresql_health = await self.check_postgresql()
|
||||
redis_health = await self.check_redis()
|
||||
|
||||
overall_status = "healthy"
|
||||
if postgresql_health["status"] != "healthy":
|
||||
overall_status = "unhealthy"
|
||||
elif redis_health["status"] == "unhealthy":
|
||||
overall_status = "degraded"
|
||||
|
||||
return {
|
||||
"overall_status": overall_status,
|
||||
"postgresql": postgresql_health,
|
||||
"redis": redis_health,
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
}
|
||||
Reference in New Issue
Block a user