db.py

import os
import logging
import time
from typing import Optional
from contextlib import contextmanager
from mysql.connector import pooling, errors as mysql_errors
import redis
import mysql.connector
from config import get_config

logger = logging.getLogger(__name__)

# Get configuration
config = get_config()

# Globals (per-process)
_mysql_pool = None
_mysql_pool_pid = None
_connection_failures = 0
_last_failure_time = 0
_circuit_breaker_open = False

# Circuit breaker settings
MAX_FAILURES = 5
CIRCUIT_BREAKER_TIMEOUT = 60  # seconds

def _check_circuit_breaker():
    """Check if circuit breaker should be opened or closed"""
    global _circuit_breaker_open, _connection_failures, _last_failure_time
    
    current_time = time.time()
    
    # If circuit breaker is open, check if timeout has passed
    if _circuit_breaker_open:
        if current_time - _last_failure_time >= CIRCUIT_BREAKER_TIMEOUT:
            logger.info("Circuit breaker timeout passed, attempting to close")
            _circuit_breaker_open = False
            _connection_failures = 0
        else:
            raise mysql_errors.DatabaseError("Circuit breaker is open, refusing connection")
    
    # Check if we should open circuit breaker
    if _connection_failures >= MAX_FAILURES:
        logger.error(f"Opening circuit breaker after {_connection_failures} failures")
        _circuit_breaker_open = True
        _last_failure_time = current_time
        raise mysql_errors.DatabaseError("Circuit breaker opened due to too many failures")


def _record_connection_failure():
    """Record a connection failure for circuit breaker"""
    global _connection_failures, _last_failure_time
    _connection_failures += 1
    _last_failure_time = time.time()
    logger.warning(f"Connection failure recorded. Total failures: {_connection_failures}")


def _record_connection_success():
    """Record a successful connection"""
    global _connection_failures
    if _connection_failures > 0:
        logger.info("Connection successful, resetting failure count")
        _connection_failures = 0


def _ensure_pool():
    """Ensure connection pool exists for current process"""
    global _mysql_pool, _mysql_pool_pid
    
    _check_circuit_breaker()
    
    pid = os.getpid()
    if _mysql_pool is None or _mysql_pool_pid != pid:
        pool_name = f"mysql_pool_{pid}"
        logger.info(
            "Creating MySQL pool for pid=%s (pool_name=%s, size=%s)",
            pid, pool_name, config.DB_POOL_SIZE
        )
        
        try:
            _mysql_pool = pooling.MySQLConnectionPool(
                pool_name=pool_name,
                pool_size=config.DB_POOL_SIZE,
                pool_reset_session=True,
                host=config.DB_HOST,
                port=config.DB_PORT,
                user=config.DB_USER,
                password=config.DB_PASSWORD,
                database=config.DB_DATABASE,
                connection_timeout=config.DB_POOL_TIMEOUT,
                autocommit=False,
                charset='utf8mb4',
                collation='utf8mb4_unicode_ci',
            )
            _mysql_pool_pid = pid
            _record_connection_success()
        except Exception as e:
            _record_connection_failure()
            logger.exception("Failed to create MySQL connection pool")
            raise

def get_connection(max_retries: int = 3, retry_delay: float = 0.5):
    """
    Returns a connection from the pool with retry logic.
    
    Args:
        max_retries: Maximum number of retry attempts
        retry_delay: Delay between retries in seconds
        
    Returns:
        MySQL connection object
        
    Raises:
        DatabaseError: If connection cannot be established after retries
    """
    last_exception = None
    
    for attempt in range(max_retries):
        try:
            _ensure_pool()
            conn = _mysql_pool.get_connection()  # type: ignore
            
            # Verify connection is alive
            try:
                conn.ping(reconnect=True, attempts=2, delay=1)
                _record_connection_success()
                return conn
            except Exception as ping_error:
                logger.warning(f"Connection ping failed: {ping_error}")
                safe_close(conn)
                raise
                
        except mysql_errors.PoolError as e:
            last_exception = e
            logger.warning(
                f"Pool connection failed (attempt {attempt + 1}/{max_retries}): {e}"
            )
            
            # Try direct connection as fallback
            if attempt == max_retries - 1:
                try:
                    logger.info("Attempting direct connection as fallback")
                    conn = mysql.connector.connect(
                        host=config.DB_HOST,
                        port=config.DB_PORT,
                        user=config.DB_USER,
                        password=config.DB_PASSWORD,
                        database=config.DB_DATABASE,
                        connection_timeout=config.DB_POOL_TIMEOUT,
                        charset='utf8mb4',
                        collation='utf8mb4_unicode_ci',
                    )
                    _record_connection_success()
                    return conn
                except Exception as direct_error:
                    logger.exception("Direct connection also failed")
                    _record_connection_failure()
                    raise
            
            time.sleep(retry_delay * (attempt + 1))  # Exponential backoff
            
        except Exception as e:
            last_exception = e
            logger.exception(f"Unexpected error getting connection (attempt {attempt + 1}/{max_retries})")
            _record_connection_failure()
            
            if attempt < max_retries - 1:
                time.sleep(retry_delay * (attempt + 1))
    
    # All retries failed
    error_msg = f"Failed to get database connection after {max_retries} attempts"
    logger.error(error_msg)
    raise mysql_errors.DatabaseError(error_msg) from last_exception

def safe_close(conn):
    """Safely close a database connection"""
    if not conn:
        return
    try:
        # Rollback any pending transaction
        if conn.in_transaction:
            conn.rollback()
        conn.close()
    except Exception:
        logger.debug("Error closing connection (may already be closed)", exc_info=True)


@contextmanager
def get_db_cursor(dictionary: bool = True, buffered: bool = False):
    """
    Context manager for database operations.
    Automatically handles connection and cursor lifecycle.
    
    Usage:
        with get_db_cursor() as cursor:
            cursor.execute("SELECT * FROM users")
            results = cursor.fetchall()
    """
    conn = None
    cursor = None
    try:
        conn = get_connection()
        cursor = conn.cursor(dictionary=dictionary, buffered=buffered)
        yield cursor
        conn.commit()
    except Exception:
        if conn:
            conn.rollback()
        raise
    finally:
        if cursor:
            try:
                cursor.close()
            except Exception:
                pass
        if conn:
            safe_close(conn)


# --- Redis Connection ---
try:
    redis_pool = redis.ConnectionPool(
        host=config.REDIS_HOST,
        port=config.REDIS_PORT,
        db=config.REDIS_DB,
        password=config.REDIS_PASSWORD,
        decode_responses=True,
        max_connections=config.REDIS_MAX_CONNECTIONS,
        socket_timeout=5,
        socket_connect_timeout=5,
        socket_keepalive=True,
        health_check_interval=30,
    )
    
    redis_client = redis.Redis(connection_pool=redis_pool)
    
    # Test connection
    redis_client.ping()
    logger.info("Redis connection established successfully")
    
except redis.ConnectionError as e:
    logger.error(f"Failed to connect to Redis: {e}")
    raise
except Exception as e:
    logger.exception("Unexpected error initializing Redis")
    raise


def get_redis_client() -> redis.Redis:
    """Get Redis client with connection verification"""
    try:
        redis_client.ping()
        return redis_client
    except redis.ConnectionError:
        logger.error("Redis connection lost, attempting to reconnect")
        raise