70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import threading
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.core.config import settings
|
|
|
|
_engine = None
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
_lock = threading.RLock()
|
|
|
|
|
|
def _get_engine():
|
|
global _engine
|
|
if _engine is None:
|
|
with _lock:
|
|
if _engine is None:
|
|
_engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=settings.debug,
|
|
pool_pre_ping=True,
|
|
pool_recycle=300,
|
|
)
|
|
return _engine
|
|
|
|
|
|
def _get_session_factory():
|
|
global _session_factory
|
|
if _session_factory is None:
|
|
with _lock:
|
|
if _session_factory is None:
|
|
_session_factory = async_sessionmaker(
|
|
_get_engine(), class_=AsyncSession, expire_on_commit=False
|
|
)
|
|
return _session_factory
|
|
|
|
|
|
async def dispose_engine():
|
|
"""Dispose the async engine and reset cached DB handles.
|
|
|
|
Celery tasks run async code through ``asyncio.run()``, which creates and closes
|
|
one event loop per task. Asyncpg connections are bound to the loop that created
|
|
them, so worker tasks must not keep pooled connections across task runs.
|
|
"""
|
|
global _engine, _session_factory
|
|
|
|
engine = _engine
|
|
if engine is not None:
|
|
await engine.dispose()
|
|
|
|
with _lock:
|
|
if _engine is engine:
|
|
_engine = None
|
|
_session_factory = None
|
|
|
|
|
|
async def init_db():
|
|
"""Create tables if they do not exist."""
|
|
from app.db.models import Base # main models
|
|
|
|
engine = _get_engine()
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
|
|
async def get_db():
|
|
"""Yield a DB session with proper cleanup."""
|
|
session_factory = _get_session_factory()
|
|
async with session_factory() as session:
|
|
yield session
|