""" Database connection layer with connection pooling for VideoDetect. Provides: - Connection pooling via DBUtils + PyMySQL - Automatic reconnection on disconnect (ping + retry on transient errors) - Context manager support - Prepared statements for all queries - Transaction support for atomic state transitions """ import contextlib import logging import time from typing import Optional from dbutils.pooled_db import PooledDB import pymysql logger = logging.getLogger(__name__) class DBConnector: """Thread-safe database connection pool manager.""" # Transient errors that indicate a dropped/stale connection and are safe # to retry with a fresh connection from the pool. _RETRYABLE_ERRORS = (pymysql.err.InterfaceError, pymysql.err.OperationalError) _MAX_RETRIES = 3 _RETRY_BACKOFF_SECONDS = 0.5 def __init__( self, host: str = "mariadb", port: int = 3306, database: str = "videodetect", user: str = "videodetect", password: str = "videodetect123", pool_size: int = 20, pool_min: int = 5, pool_recycle: int = 3600, ): # PooledDB parameters pool_config = { "creator": pymysql, "maxconnections": pool_size, "mincached": pool_min, "maxcached": pool_size, "maxusage": 200, "blocking": True, # ping(0) = check if connection is alive BEFORE returning it from # the pool. This is the correct DBUtils mechanism for detecting # dead/stale connections (idle timeout, network blip, DB restart). "ping": 0, } # PyMySQL connection parameters connection_params = { "host": host, "port": port, "database": database, "user": user, "password": password, "charset": "utf8mb4", "cursorclass": pymysql.cursors.DictCursor, "read_timeout": 30, "write_timeout": 30, } self._pool = PooledDB(**pool_config, **connection_params) self._pool_recycle = pool_recycle logger.info( "DB pool initialized: host=%s db=%s pool_size=%d min=%d ping=0", host, database, pool_size, pool_min, ) def get_connection(self): """Get a connection from the pool.""" return self._pool.connection() def _run_query(self, query: str, params, transaction: bool, fetch: str): """Run a single query on a fresh connection, with retry on transient errors. A pooled connection can be dropped by the server (idle timeout, network blip, DB restart). When that happens the query fails with an InterfaceError/OperationalError. We discard the dead connection and retry with a fresh one from the pool. Args: query: SQL statement. params: Bound parameters (or None). transaction: If True, wrap the query in a transaction. fetch: One of "none", "one", or "all" controlling the return value. Returns: - "none": the affected row count. - "one": a single row (dict) or None. - "all": a list of rows (dicts). """ last_exc = None for attempt in range(1, self._MAX_RETRIES + 1): conn = self.get_connection() try: cursor = conn.cursor() if transaction: conn.begin() cursor.execute(query, params or ()) if fetch == "one": result = cursor.fetchone() elif fetch == "all": result = cursor.fetchall() else: result = cursor.rowcount if transaction: conn.commit() cursor.close() conn.close() return result except self._RETRYABLE_ERRORS as e: last_exc = e logger.warning( "Transient DB error (attempt %d/%d): %s", attempt, self._MAX_RETRIES, e, ) # Discard the dead connection; a fresh one is taken on retry. self._safe_close(conn) if attempt < self._MAX_RETRIES: time.sleep(self._RETRY_BACKOFF_SECONDS * attempt) except Exception: # Non-transient error: roll back if in a transaction and re-raise. if transaction: self._safe_rollback(conn) self._safe_close(conn) raise raise last_exc @staticmethod def _safe_close(conn): try: conn.close() except Exception: pass @staticmethod def _safe_rollback(conn): try: conn.rollback() except Exception: pass @contextlib.contextmanager def get_cursor(self, transaction: bool = False): """Context manager for getting a cursor with optional transaction support. Note: this does not retry on transient errors; use execute/fetchone/ fetchall for retry-safe single-statement queries. """ conn = self.get_connection() cursor = None try: cursor = conn.cursor() if transaction: conn.begin() yield cursor if transaction: conn.commit() except Exception: if transaction: self._safe_rollback(conn) raise finally: if cursor is not None: try: cursor.close() except Exception: pass self._safe_close(conn) @contextlib.contextmanager def transaction(self): """Context manager for a full transaction.""" conn = self.get_connection() try: conn.begin() yield conn conn.commit() except Exception: self._safe_rollback(conn) raise finally: self._safe_close(conn) def execute(self, query: str, params=None, transaction: bool = False): """Execute a query and return affected rows. Retries on transient errors.""" return self._run_query(query, params, transaction, fetch="none") def fetchone(self, query: str, params=None): """Execute a query and return one row. Retries on transient errors.""" return self._run_query(query, params, transaction=False, fetch="one") def fetchall(self, query: str, params=None): """Execute a query and return all rows. Retries on transient errors.""" return self._run_query(query, params, transaction=False, fetch="all") def initialize_schema(self, schema_path: str = "db/schema.sql"): """Initialize the database schema from SQL file.""" with open(schema_path, "r") as f: sql = f.read() # Split on semicolons and execute each statement statements = [s.strip() for s in sql.split(";") if s.strip()] for stmt in statements: if stmt.startswith("--"): continue try: self.execute(stmt) except Exception as e: logger.debug("Schema statement skipped (may already exist): %s", e) logger.info("Schema initialized from %s", schema_path) def health_check(self) -> bool: """Check if the database is reachable.""" try: result = self.fetchone("SELECT 1") return result is not None except Exception as e: logger.error("Health check failed: %s", e) return False