src/database/db_adapter.py

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