yunke/yunke-server/app/db/base.py

"""Database core: declarative base, async engine, session factory.

Production targets MySQL via `asyncmy` (configured by `Settings.db_url`).
Unit tests build their own in-memory SQLite engines (see `tests/test_db_base.py`)
and only import `Base` / `async_session_maker` from here.
"""

from collections.abc import AsyncIterator

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase

from app.config import get_settings
from app.db import compat  # noqa: F401,I100 — 注册 MySQL NULLS LAST 编译补丁(导入即生效)


class Base(DeclarativeBase):
    """Common declarative base for all ORM models in the project."""


settings = get_settings()

# SQLite: 写锁竞争时忙等 60s 而不是立刻报 "database is locked"
# （多进程并发写的场景：页面抓取器 × 2 + 后端 API 同时跑，2026-08-19 实测撞锁崩任务）。
_connect_args = {"timeout": 60} if settings.db_url.startswith("sqlite") else {}

async_engine = create_async_engine(
    settings.db_url,
    pool_pre_ping=True,
    pool_recycle=3600,  # recycle connections before MySQL wait_timeout (default 8h) closes them
    connect_args=_connect_args,
    future=True,
)

async_session_maker: async_sessionmaker[AsyncSession] = async_sessionmaker(
    async_engine,
    expire_on_commit=False,
)


async def get_session() -> AsyncIterator[AsyncSession]:
    """FastAPI dependency that yields an `AsyncSession`.

    Usage::

        @app.get("/")
        async def index(session: AsyncSession = Depends(get_session)):
            ...
    """
    async with async_session_maker() as session:
        yield session