from contextlib import asynccontextmanager
from typing import AsyncGenerator, Generator

from sqlmodel import Session

from .engine import engine


def get_session() -> Generator[Session, None, None]:
    """
    用于以依赖注入的方式在 路由端点函数 中获取数据库会话。
    `session: Annotated[Session, Depends(get_session)],`

    Yields:
        数据库会话对象 Session。
    """
    with Session(engine) as session:
        yield session


@asynccontextmanager
async def async_get_session() -> AsyncGenerator[Session, None]:
    with Session(engine) as session:
        yield session
