46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
|
# app/api/dependencies/database.py
|
|||
|
|
from typing import AsyncIterator, Callable, Type
|
|||
|
|
|
|||
|
|
from asyncpg import Connection, Pool
|
|||
|
|
from fastapi import Depends
|
|||
|
|
from starlette.requests import Request
|
|||
|
|
|
|||
|
|
from app.db.repositories.base import BaseRepository
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_db_pool(request: Request) -> Pool:
|
|||
|
|
"""
|
|||
|
|
从 app.state.pool 取得连接池；若未初始化给出清晰报错。
|
|||
|
|
"""
|
|||
|
|
pool = getattr(request.app.state, "pool", None)
|
|||
|
|
if pool is None:
|
|||
|
|
raise RuntimeError("Database pool not initialized on app.state.pool")
|
|||
|
|
return pool
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def _get_connection_from_pool(
|
|||
|
|
pool: Pool = Depends(_get_db_pool),
|
|||
|
|
) -> AsyncIterator[Connection]:
|
|||
|
|
"""
|
|||
|
|
私有实现：从连接池借出一个连接，使用后自动归还。
|
|||
|
|
"""
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
yield conn
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ✅ 公共别名：供路由里直接使用 Depends(get_connection)
|
|||
|
|
get_connection = _get_connection_from_pool
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_repository(
|
|||
|
|
repo_type: Type[BaseRepository],
|
|||
|
|
) -> Callable[[Connection], BaseRepository]:
|
|||
|
|
"""
|
|||
|
|
兼容旧用法：Depends(get_repository(UserRepo))
|
|||
|
|
内部依赖 get_connection，因此两种写法都能共存。
|
|||
|
|
"""
|
|||
|
|
def _get_repo(conn: Connection = Depends(get_connection)) -> BaseRepository:
|
|||
|
|
return repo_type(conn)
|
|||
|
|
|
|||
|
|
return _get_repo
|