Cloud-Tour-to-Libo/app/data_platform/mysql_db.py

"""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:// 协议")
    database = unquote(parsed.path.lstrip("/")) or "platform_control"
    if not _DATABASE_PATTERN.fullmatch(database):
        raise ValueError("DATA_MYSQL_URL 中的控制数据库名称不合法")

    global _control_database
    _control_database = 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"] = 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)