60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
from datetime import datetime
|
||
import os
|
||
|
||
from sqlalchemy import DateTime, func
|
||
from sqlalchemy.ext.asyncio import AsyncAttrs, async_sessionmaker, create_async_engine
|
||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||
|
||
from app.config import settings
|
||
|
||
|
||
def _is_celery_process() -> bool:
|
||
argv = " ".join(os.sys.argv).lower()
|
||
return "celery" in argv
|
||
|
||
|
||
engine_kwargs = {
|
||
"echo": settings.DEBUG,
|
||
"pool_pre_ping": True,
|
||
}
|
||
|
||
# SQLite 本地调试时不要乱塞 pool_size/max_overflow，PostgreSQL/asyncpg 才建议配置
|
||
if _is_celery_process() and not settings.DATABASE_URL.startswith("sqlite"):
|
||
engine_kwargs.update(
|
||
pool_size=settings.CELERY_DB_POOL_SIZE,
|
||
max_overflow=settings.CELERY_DB_MAX_OVERFLOW,
|
||
pool_timeout=settings.CELERY_DB_POOL_TIMEOUT,
|
||
pool_recycle=settings.CELERY_DB_POOL_RECYCLE,
|
||
)
|
||
|
||
engine = create_async_engine(settings.DATABASE_URL, **engine_kwargs)
|
||
|
||
async_session = async_sessionmaker(
|
||
engine,
|
||
expire_on_commit=False,
|
||
autoflush=False,
|
||
)
|
||
|
||
AsyncSessionLocal = async_session
|
||
|
||
|
||
class Base(AsyncAttrs, DeclarativeBase):
|
||
pass
|
||
|
||
|
||
class TimestampMixin:
|
||
created_at: Mapped[datetime] = mapped_column(
|
||
DateTime(timezone=True), server_default=func.now()
|
||
)
|
||
updated_at: Mapped[datetime] = mapped_column(
|
||
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
||
)
|
||
|
||
|
||
async def init_database() -> None:
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
|
||
|
||
async def close_database() -> None:
|
||
await engine.dispose() |