69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
|
"""Async SQLAlchemy engine/session setup.
|
||
|
|
|
||
|
|
The engine is created lazily so importing the app (e.g. in tests that override
|
||
|
|
`get_db`) never requires a reachable MySQL server or its driver.
|
||
|
|
"""
|
||
|
|
from collections.abc import AsyncIterator
|
||
|
|
|
||
|
|
from sqlalchemy.ext.asyncio import (
|
||
|
|
AsyncEngine,
|
||
|
|
AsyncSession,
|
||
|
|
async_sessionmaker,
|
||
|
|
create_async_engine,
|
||
|
|
)
|
||
|
|
|
||
|
|
from app.config import get_settings
|
||
|
|
|
||
|
|
_engine: AsyncEngine | None = None
|
||
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_engine() -> AsyncEngine:
|
||
|
|
global _engine
|
||
|
|
if _engine is None:
|
||
|
|
settings = get_settings()
|
||
|
|
url = settings.effective_database_url
|
||
|
|
if url.startswith("sqlite"):
|
||
|
|
# Test runs: fresh connection per checkout (no cross-event-loop
|
||
|
|
# reuse) and a generous busy timeout for concurrent writers.
|
||
|
|
from sqlalchemy import event
|
||
|
|
from sqlalchemy.pool import NullPool
|
||
|
|
|
||
|
|
_engine = create_async_engine(
|
||
|
|
url, poolclass=NullPool, connect_args={"timeout": 30}
|
||
|
|
)
|
||
|
|
|
||
|
|
# SQLite (rollback-journal) deadlocks when a transaction upgrades
|
||
|
|
# from read to write while another writer waits. Taking the write
|
||
|
|
# lock up front (BEGIN IMMEDIATE) serializes transactions cleanly,
|
||
|
|
# mirroring the row-lock semantics InnoDB gives us in production.
|
||
|
|
@event.listens_for(_engine.sync_engine, "connect")
|
||
|
|
def _sqlite_autocommit(dbapi_conn, _record):
|
||
|
|
dbapi_conn.isolation_level = None
|
||
|
|
|
||
|
|
@event.listens_for(_engine.sync_engine, "begin")
|
||
|
|
def _sqlite_begin_immediate(conn):
|
||
|
|
conn.exec_driver_sql("BEGIN IMMEDIATE")
|
||
|
|
else:
|
||
|
|
_engine = create_async_engine(
|
||
|
|
url,
|
||
|
|
pool_pre_ping=True,
|
||
|
|
pool_recycle=1800,
|
||
|
|
echo=False,
|
||
|
|
)
|
||
|
|
return _engine
|
||
|
|
|
||
|
|
|
||
|
|
def get_session_factory() -> async_sessionmaker[AsyncSession]:
|
||
|
|
global _session_factory
|
||
|
|
if _session_factory is None:
|
||
|
|
_session_factory = async_sessionmaker(
|
||
|
|
get_engine(), expire_on_commit=False, autoflush=False
|
||
|
|
)
|
||
|
|
return _session_factory
|
||
|
|
|
||
|
|
|
||
|
|
async def get_db() -> AsyncIterator[AsyncSession]:
|
||
|
|
async with get_session_factory()() as session:
|
||
|
|
yield session
|