import logging import time from typing import Optional, Dict, Any, List, Type from contextlib import contextmanager from abc import ABC, abstractmethod from enum import Enum from sqlalchemy import create_engine, text, event from sqlalchemy.engine import Engine from sqlalchemy.pool import QueuePool from sqlalchemy.exc import SQLAlchemyError, OperationalError logger = logging.getLogger(__name__) class DatabaseType(Enum): POSTGRESQL = "postgresql" MYSQL = "mysql" SQLITE = "sqlite" MSSQL = "mssql" class ConnectionConfig: def __init__( self, db_type: DatabaseType, host: str = "localhost", port: int = 5432, database: str = "", username: str = "", password: str = "", pool_min_size: int = 2, pool_max_size: int = 10, pool_timeout: int = 30, pool_recycle: int = 300, connect_timeout: int = 10, query_timeout: int = 10, ): self.db_type = db_type self.host = host self.port = port self.database = database self.username = username self.password = password self.pool_min_size = pool_min_size self.pool_max_size = pool_max_size self.pool_timeout = pool_timeout self.pool_recycle = pool_recycle self.connect_timeout = connect_timeout self.query_timeout = query_timeout def get_connection_string(self) -> str: if self.db_type == DatabaseType.SQLITE: return f"sqlite:///{self.database}" if self.db_type == DatabaseType.POSTGRESQL: driver = "postgresql+psycopg2" default_port = 5432 elif self.db_type == DatabaseType.MYSQL: driver = "mysql+pymysql" default_port = 3306 elif self.db_type == DatabaseType.MSSQL: driver = "mssql+pyodbc" default_port = 1433 else: raise ValueError(f"Unsupported database type: {self.db_type}") port = self.port or default_port return f"{driver}://{self.username}:{self.password}@{self.host}:{port}/{self.database}" class DatabaseAdapter: def __init__(self, config: ConnectionConfig): self._config = config self._engine: Optional[Engine] = None self._is_connected = False self._retry_delays = [1, 2, 4, 8] @property def is_connected(self) -> bool: return self._is_connected @property def engine(self) -> Optional[Engine]: return self._engine def connect(self) -> bool: if self._is_connected and self._engine is not None: return True for attempt, delay in enumerate(self._retry_delays): try: self._engine = self._create_engine() self._verify_connection() self._is_connected = True logger.info(f"Database connection established: {self._config.db_type.value}") return True except OperationalError as e: logger.warning(f"Connection attempt {attempt + 1} failed: {e}") if attempt < len(self._retry_delays) - 1: time.sleep(delay) else: logger.error("All connection attempts exhausted") self._is_connected = False return False except Exception as e: logger.error(f"Unexpected connection error: {e}") self._is_connected = False return False return False def _create_engine(self) -> Engine: connection_string = self._config.get_connection_string() pool_args = {} if self._config.db_type != DatabaseType.SQLITE: pool_args = { "poolclass": QueuePool, "pool_size": self._config.pool_min_size, "max_overflow": self._config.pool_max_size - self._config.pool_min_size, "pool_timeout": self._config.pool_timeout, "pool_recycle": self._config.pool_recycle, "pool_pre_ping": True, } engine = create_engine( connection_string, echo=False, **pool_args ) self._setup_timeout_handlers(engine) return engine def _setup_timeout_handlers(self, engine: Engine) -> None: @event.listens_for(engine, "connect") def set_timeout(dbapi_connection, connection_record): timeout = self._config.query_timeout if self._config.db_type == DatabaseType.POSTGRESQL: cursor = dbapi_connection.cursor() cursor.execute(f"SET statement_timeout = {timeout * 1000}") cursor.close() elif self._config.db_type == DatabaseType.MYSQL: cursor = dbapi_connection.cursor() cursor.execute(f"SET SESSION MAX_EXECUTION_TIME = {timeout * 1000}") cursor.close() def _verify_connection(self) -> None: with self._engine.connect() as conn: conn.execute(text("SELECT 1")) def disconnect(self) -> None: if self._engine is not None: self._engine.dispose() self._engine = None self._is_connected = False logger.info("Database connection closed") @contextmanager def get_connection(self): if not self._is_connected or self._engine is None: raise RuntimeError("Database not connected. Call connect() first.") conn = self._engine.connect() try: yield conn finally: conn.close() def execute_query( self, query: str, params: Optional[Dict[str, Any]] = None, timeout: Optional[int] = None ) -> List[Dict[str, Any]]: if not self._is_connected: raise RuntimeError("Database not connected") start_time = time.time() try: with self.get_connection() as conn: result = conn.execute(text(query), params or {}) rows = result.fetchall() columns = result.keys() formatted_results = [ dict(zip(columns, row)) for row in rows ] execution_time = (time.time() - start_time) * 1000 logger.debug(f"Query executed in {execution_time:.2f}ms, {len(rows)} rows returned") return formatted_results except SQLAlchemyError as e: logger.error(f"Query execution failed: {e}") raise def get_pool_status(self) -> Dict[str, Any]: if self._engine is None: return {"status": "not_initialized"} pool = self._engine.pool return { "status": "active" if self._is_connected else "disconnected", "pool_size": pool.size() if hasattr(pool, "size") else 0, "checked_in": pool.checkedin() if hasattr(pool, "checkedin") else 0, "checked_out": pool.checkedout() if hasattr(pool, "checkedout") else 0, "overflow": pool.overflow() if hasattr(pool, "overflow") else 0, } def health_check(self) -> Dict[str, Any]: try: if not self._is_connected: return {"healthy": False, "error": "Not connected"} start_time = time.time() with self.get_connection() as conn: conn.execute(text("SELECT 1")) latency = (time.time() - start_time) * 1000 return { "healthy": True, "latency_ms": round(latency, 2), "pool": self.get_pool_status() } except Exception as e: return {"healthy": False, "error": str(e)} _adapter_instance: Optional[DatabaseAdapter] = None def get_database_adapter(config: Optional[ConnectionConfig] = None) -> Optional[DatabaseAdapter]: global _adapter_instance if _adapter_instance is None and config is not None: _adapter_instance = DatabaseAdapter(config) return _adapter_instance def reset_database_adapter() -> None: global _adapter_instance if _adapter_instance is not None: _adapter_instance.disconnect() _adapter_instance = None