This commit is contained in:
Ryan Shpeherd
2026-09-08 18:34:31 -04:00
parent 83f980f7f8
commit 78358e0d7b
4 changed files with 118 additions and 25 deletions
+108 -23
View File
@@ -3,7 +3,7 @@ Database connection layer with connection pooling for VideoDetect.
Provides:
- Connection pooling via DBUtils + PyMySQL
- Automatic reconnection on disconnect
- Automatic reconnection on disconnect (ping + retry on transient errors)
- Context manager support
- Prepared statements for all queries
- Transaction support for atomic state transitions
@@ -11,6 +11,7 @@ Provides:
import contextlib
import logging
import time
from typing import Optional
from dbutils.pooled_db import PooledDB
@@ -22,6 +23,12 @@ 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",
@@ -41,9 +48,13 @@ class DBConnector:
"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
# PyMySQL connection parameters
connection_params = {
"host": host,
"port": port,
@@ -55,10 +66,11 @@ class DBConnector:
"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",
"DB pool initialized: host=%s db=%s pool_size=%d min=%d ping=0",
host, database, pool_size, pool_min,
)
@@ -66,10 +78,85 @@ class DBConnector:
"""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."""
"""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:
@@ -79,11 +166,15 @@ class DBConnector:
conn.commit()
except Exception:
if transaction:
conn.rollback()
self._safe_rollback(conn)
raise
finally:
cursor.close()
conn.close()
if cursor is not None:
try:
cursor.close()
except Exception:
pass
self._safe_close(conn)
@contextlib.contextmanager
def transaction(self):
@@ -94,28 +185,22 @@ class DBConnector:
yield conn
conn.commit()
except Exception:
conn.rollback()
self._safe_rollback(conn)
raise
finally:
conn.close()
self._safe_close(conn)
def execute(self, query: str, params=None, transaction: bool = False):
"""Execute a query and return affected rows."""
with self.get_cursor(transaction=transaction) as cursor:
cursor.execute(query, params or ())
return cursor.rowcount
"""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."""
with self.get_cursor() as cursor:
cursor.execute(query, params or ())
return cursor.fetchone()
"""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."""
with self.get_cursor() as cursor:
cursor.execute(query, params or ())
return cursor.fetchall()
"""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."""
+1
View File
@@ -7,6 +7,7 @@ Handles errors gracefully for corrupt or unsupported files.
import json
import logging
import os
import subprocess
import time
from dataclasses import dataclass, field