66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
|
# -*- 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.ext.asyncio import (
|
|||
|
|
AsyncSession,
|
|||
|
|
async_sessionmaker,
|
|||
|
|
create_async_engine,
|
|||
|
|
)
|
|||
|
|
from sqlalchemy.orm import DeclarativeBase
|
|||
|
|
|
|||
|
|
from .. import config
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Base(DeclarativeBase):
|
|||
|
|
"""SQLAlchemy 声明式基类（数据模型层继承）。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
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,
|
|||
|
|
)
|