server-core/app/infrastructure/db.py

# -*- coding: utf-8 -*-
"""基础设施层 · 异步 SQLAlchemy 引擎 / 会话 / 声明基类。

四层架构：基础设施层（含数据模型层）。全异步：
  - 开发/默认：SQLite + aiosqlite
  - 生产：MySQL + asyncmy（经 config.DATABASE_URL 切换）
每请求一个 async session（``get_session`` 依赖），事件循环内无阻塞。
"""
from __future__ import annotations

from collections.abc import AsyncGenerator

from sqlalchemy import Text, create_engine, types
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.orm import DeclarativeBase

from .. import config


class Base(DeclarativeBase):
    """SQLAlchemy 声明式基类（数据模型层继承）。"""


@compiles(types.String, "mysql")
def _mysql_string_ddl(type_: types.String, compiler, **kw) -> str:
    """MySQL 方言适配：String() 未显式给长度时按 VARCHAR(255) 建表（MySQL 要求
    VARCHAR 必须带长度；SQLite 忽略长度不受影响），Text 渲染 MEDIUMTEXT
    （富文本/JSON 大字段，避免 MySQL TEXT 64KB 上限）。"""
    if isinstance(type_, Text):
        return "MEDIUMTEXT"
    if not type_.length:
        return "VARCHAR(255)"
    return f"VARCHAR({int(type_.length)})"


def sync_database_url(url: str | None = None) -> str:
    """DATABASE_URL → 同步驱动 URL（培训同步旁路 / 迁移脚本用）。

    sqlite+aiosqlite → sqlite（pysqlite）；mysql+asyncmy → mysql+pymysql。
    """
    url = url or config.DATABASE_URL
    if url.startswith("sqlite+aiosqlite"):
        return url.replace("sqlite+aiosqlite", "sqlite", 1)
    if url.startswith("mysql+asyncmy"):
        return url.replace("mysql+asyncmy", "mysql+pymysql", 1)
    return url


def make_sync_engine(url: str | None = None):
    """按 DATABASE_URL 创建同步引擎（asyncio.to_thread / 脚本场景）。"""
    u = sync_database_url(url)
    kwargs: dict = {"pool_pre_ping": True, "future": True}
    if u.startswith("mysql"):
        kwargs["pool_recycle"] = 3600
        if "charset=" not in u:
            kwargs["connect_args"] = {"charset": "utf8mb4"}
    elif u.startswith("sqlite"):
        kwargs["connect_args"] = {"check_same_thread": False}
    return create_engine(u, **kwargs)


engine = create_async_engine(
    config.DATABASE_URL,
    pool_pre_ping=True,
)

AsyncSessionLocal = async_sessionmaker(
    bind=engine,
    class_=AsyncSession,
    autoflush=False,
    expire_on_commit=False,
)


async def init_db() -> None:
    """幂等地创建全部表（create_all 对已存在表无操作）。"""
    from . import models  # noqa: F401

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def get_session() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI 依赖：请求级异步 SQLAlchemy 会话，请求结束自动关闭。"""
    async with AsyncSessionLocal() as session:
        yield session


def make_async_engine(db_url: str):
    """按给定 URL 创建独立异步引擎（测试用 tmp 数据库）。"""
    return create_async_engine(db_url, pool_pre_ping=True)


def make_async_session_factory(db_url: str) -> async_sessionmaker:
    return async_sessionmaker(
        bind=make_async_engine(db_url),
        class_=AsyncSession,
        autoflush=False,
        expire_on_commit=False,
    )