44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
"""异步数据库会话。"""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
def create_engine_from_url(database_url: str) -> AsyncEngine:
|
|
connect_args: dict[str, object] = {}
|
|
if database_url.startswith("sqlite"):
|
|
connect_args["check_same_thread"] = False
|
|
return create_async_engine(
|
|
database_url,
|
|
echo=settings.debug,
|
|
pool_pre_ping=True,
|
|
connect_args=connect_args,
|
|
)
|
|
|
|
|
|
engine = create_engine_from_url(settings.resolved_database_url)
|
|
AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
|
|
|
|
|
def configure_database(database_url: str) -> None:
|
|
global AsyncSessionLocal, engine
|
|
engine = create_engine_from_url(database_url)
|
|
AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
|
|
|
|
|
async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
|
async with AsyncSessionLocal() as session:
|
|
yield session
|
|
|
|
|
|
async def close_engine() -> None:
|
|
await engine.dispose()
|