228 lines
7.6 KiB
Python
228 lines
7.6 KiB
Python
"""
|
|
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
|