"""Async MySQL connection management for the independent Data Center. The application still uses PostgreSQL for accounts and graph metadata. This module owns a separate MySQL pool whose default database is the protected ``platform_control`` catalog. Business databases are selected explicitly for the duration of a request and reset before the pooled connection is returned. """ from __future__ import annotations from contextlib import asynccontextmanager import logging import re from typing import AsyncIterator from urllib.parse import unquote, urlparse import aiomysql from app.config import settings logger = logging.getLogger(__name__) _DATABASE_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9_]{0,63}$") _pool: aiomysql.Pool | None = None _control_database = "platform_control" _pool_generation = 0 class DataCenterUnavailable(RuntimeError): """Raised when the independent MySQL service is not connected.""" def _connection_options(*, include_database: bool) -> dict[str, object]: parsed = urlparse(settings.data_mysql_url) if parsed.scheme not in {"mysql", "mysql+aiomysql"}: raise ValueError("DATA_MYSQL_URL 必须使用 mysql:// 协议") control_database = unquote(parsed.path.lstrip("/")) or "platform_control" if not _DATABASE_PATTERN.fullmatch(control_database): raise ValueError("DATA_MYSQL_URL 中的控制数据库名称不合法") global _control_database _control_database = control_database options: dict[str, object] = { "host": parsed.hostname or "localhost", "port": parsed.port or 3306, "user": unquote(parsed.username or "root"), "password": unquote(parsed.password or ""), "charset": "utf8mb4", "autocommit": False, "cursorclass": aiomysql.DictCursor, } if include_database: options["db"] = control_database return options def control_database_name() -> str: return _control_database def data_pool_available() -> bool: return _pool is not None def data_pool_generation() -> int: return _pool_generation async def init_data_pool() -> bool: """Create the Data Center catalog and pool. Development can start without MySQL when ``DATA_MYSQL_REQUIRED=false`` so graph-only pages remain available. Data Center endpoints then return 503 instead of silently falling back to PostgreSQL. """ global _pool, _pool_generation if _pool is not None: return True try: bootstrap = _connection_options(include_database=False) bootstrap["autocommit"] = True conn = await aiomysql.connect(**bootstrap) try: async with conn.cursor() as cur: await cur.execute( "SELECT SCHEMA_NAME FROM information_schema.SCHEMATA " "WHERE SCHEMA_NAME=%s", (_control_database,), ) if not await cur.fetchone(): await cur.execute( f"CREATE DATABASE `{_control_database}` " "CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci" ) finally: conn.close() options = _connection_options(include_database=True) _pool = await aiomysql.create_pool( **options, minsize=max(1, settings.data_mysql_pool_min_size), maxsize=max( settings.data_mysql_pool_min_size, settings.data_mysql_pool_max_size, ), pool_recycle=1800, ) _pool_generation += 1 logger.info("Data Center MySQL pool connected") return True except Exception: _pool = None if settings.data_mysql_required: raise logger.warning( "Data Center MySQL is unavailable; Data Center/API Center routes " "will return 503", exc_info=True, ) return False async def close_data_pool() -> None: global _pool, _pool_generation if _pool is None: return _pool.close() await _pool.wait_closed() _pool = None _pool_generation += 1 @asynccontextmanager async def get_data_conn(database: str | None = None) -> AsyncIterator[aiomysql.Connection]: if _pool is None: raise DataCenterUnavailable("数据中心 MySQL 服务尚未连接") if database is not None and not _DATABASE_PATTERN.fullmatch(database): raise ValueError("MySQL 数据库名称不合法") async with _pool.acquire() as conn: selected = database or _control_database await conn.select_db(selected) try: yield conn except Exception: await conn.rollback() raise finally: # Prevent a later request from inheriting a business database. try: await conn.rollback() if selected != _control_database: await conn.select_db(_control_database) except Exception: logger.warning("Failed to reset a Data Center MySQL connection", exc_info=True) @asynccontextmanager async def get_provisioner_conn() -> AsyncIterator[aiomysql.Connection]: """Open the isolated account-provisioning connection. The normal Data Center account must not hold ``CREATE USER`` or ``GRANT OPTION``. This short-lived connection is used only by the admin-only DBeaver issuance and revocation endpoints. """ username = settings.data_mysql_provisioner_user.strip() password = settings.data_mysql_provisioner_password if not username or not password: raise DataCenterUnavailable("DBeaver账号签发服务尚未配置") if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,31}", username): raise ValueError("MySQL账号签发用户名称不合法") parsed = urlparse(settings.data_mysql_url) database = unquote(parsed.path.lstrip("/")) or "platform_control" if not _DATABASE_PATTERN.fullmatch(database): raise ValueError("DATA_MYSQL_URL 中的控制数据库名称不合法") conn = await aiomysql.connect( host=parsed.hostname or "localhost", port=parsed.port or 3306, user=username, password=password, charset="utf8mb4", autocommit=False, cursorclass=aiomysql.DictCursor, ) try: yield conn except Exception: await conn.rollback() raise finally: try: await conn.rollback() finally: conn.close()