import os from contextvars import ContextVar from typing import Optional logger = __import__('logging').getLogger(__name__) # Context variable for async DB connections _db_connection: ContextVar = ContextVar('db_connection', default=None) # Global connection pool _connection_pool: Optional[any] = None # Read DB config from environment (MySQL only) DB_NAME = os.getenv('DB_NAME', os.getenv('MYSQL_DB', 'eatance_website_builder')) DB_USER = os.getenv('DB_USER', os.getenv('MYSQL_USER', 'root')) DB_PASSWORD = os.getenv('DB_PASSWORD', os.getenv('MYSQL_PASSWORD', '')) DB_HOST = os.getenv('DB_HOST', os.getenv('MYSQL_HOST', 'localhost')) DB_PORT = os.getenv('DB_PORT', os.getenv('MYSQL_PORT', '3306')) # Connection pool configuration POOL_MIN_SIZE = int(os.getenv('DB_POOL_MIN_SIZE', '5')) POOL_MAX_SIZE = int(os.getenv('DB_POOL_MAX_SIZE', '20')) POOL_RECYCLE = int(os.getenv('DB_POOL_RECYCLE', '3600')) class AsyncDBAdapter: """Async adapter that exposes execute(), executemany(), cursor(), commit(), and fetch helpers for underlying async connections from aiomysql. """ def __init__(self, conn, backend: str): self._conn = conn self.backend = backend self._last_cursor = None async def cursor(self): # For aiomysql, get a cursor from the connection if self.backend == 'mysql': return await self._conn.cursor() return await self._conn.cursor() async def execute(self, query, params=None): cur = await self.cursor() q = query p = params # Convert sqlite-style '?' placeholders to '%s' for paramstyle drivers if p is not None and self.backend in ('postgres', 'mysql'): try: if '?' in q: q = q.replace('?', '%s') except Exception: pass # normalize params: avoid passing a bare string/number which would be iterated if p is not None and not isinstance(p, (list, tuple, dict)): p = (p,) try: if p is None: await cur.execute(q) else: await cur.execute(q, p) self._last_cursor = cur return cur except Exception as e: # attempt a tolerant retry for common param-formatting issues msg = str(e) try: logger.debug(f"DB execute failed, attempting tolerant retry: backend={self.backend}; error={msg}") except Exception: pass if p is not None and isinstance(p, (list, tuple)): try: alt = tuple(None if x is None else str(x) for x in p) await cur.execute(q, alt) self._last_cursor = cur return cur except Exception: pass try: logger.error(f"DB execute error (backend={self.backend}): {e}; query={q}; params={p}") except Exception: pass raise async def executemany(self, query, seq_of_params): cur = await self.cursor() res = await cur.executemany(query, seq_of_params) self._last_cursor = cur return res async def fetchone(self): if self._last_cursor is None: cur = await self.cursor() return await cur.fetchone() return await self._last_cursor.fetchone() async def fetchall(self): if self._last_cursor is None: cur = await self.cursor() return await cur.fetchall() return await self._last_cursor.fetchall() async def commit(self): try: await self._conn.commit() except Exception: pass async def close(self): try: await self._conn.close() except Exception: pass async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): # commit on normal exit, rollback on exception if supported try: if exc_type is None: try: await self._conn.commit() except Exception: pass else: try: await self._conn.rollback() except Exception: pass except Exception: pass def __getattr__(self, name): # Proxy other attributes to underlying connection return getattr(self._conn, name) async def _create_connection_pool(): """Create a global connection pool for better performance.""" global _connection_pool if _connection_pool is not None: return _connection_pool try: import aiomysql except Exception as e: # pragma: no cover - optional dependency logger.error('aiomysql is required for mysql backend but is not installed') raise conn_info = { 'host': DB_HOST, 'user': DB_USER, 'password': DB_PASSWORD, 'db': DB_NAME, 'charset': 'utf8mb4', 'minsize': POOL_MIN_SIZE, 'maxsize': POOL_MAX_SIZE, 'pool_recycle': POOL_RECYCLE, 'autocommit': False } if DB_PORT: try: conn_info['port'] = int(DB_PORT) except Exception: pass try: _connection_pool = await aiomysql.create_pool(**conn_info) logger.info(f"Created MySQL connection pool: {DB_NAME}@{DB_HOST}:{DB_PORT} (min={POOL_MIN_SIZE}, max={POOL_MAX_SIZE})") except Exception as e: logger.error(f"Failed to create connection pool: {e}") raise return _connection_pool async def _connect_mysql(): """Get a connection from the pool.""" try: import aiomysql except Exception as e: # pragma: no cover - optional dependency logger.error('aiomysql is required for mysql backend but is not installed') raise # Check if we already have a connection in the context existing_conn = _db_connection.get() if existing_conn is not None: return existing_conn # Create pool if it doesn't exist pool = await _create_connection_pool() # Get connection from pool conn = await pool.acquire() adapter = AsyncDBAdapter(conn, 'mysql') _db_connection.set(adapter) return adapter async def get_connection(): """Return an async DB connection (MySQL only).""" return await _connect_mysql() async def close_connection(): """Release connection back to the pool.""" global _connection_pool try: conn = _db_connection.get() if conn is not None: # Release connection back to pool instead of closing if _connection_pool is not None: _connection_pool.release(conn._conn) _db_connection.set(None) except Exception: pass async def close_pool(): """Close the entire connection pool. Should be called on application shutdown.""" global _connection_pool try: if _connection_pool is not None: _connection_pool.close() await _connection_pool.wait_closed() _connection_pool = None logger.info("Connection pool closed") except Exception as e: logger.error(f"Error closing connection pool: {e}")