# -*- 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, )