107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
"""异步数据库连接与会话。
|
|
|
|
连接按当前配置惰性创建，测试可以安全地切换到临时 SQLite 数据库。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncIterator, Optional
|
|
from weakref import WeakKeyDictionary
|
|
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
|
|
from .config import get_settings
|
|
from .models import Base
|
|
|
|
_engine: Optional[AsyncEngine] = None
|
|
_engine_url: Optional[str] = None
|
|
_session_factory: Optional[async_sessionmaker[AsyncSession]] = None
|
|
_sqlite_session_locks: WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = (
|
|
WeakKeyDictionary()
|
|
)
|
|
|
|
|
|
def _sqlite_session_lock() -> asyncio.Lock:
|
|
"""Return one transaction lock per event loop for the SQLite prototype.
|
|
|
|
SQLite permits only one writer at a time. Serializing the complete
|
|
session_scope transaction prevents the simulator and human control routes
|
|
from reading the same event sequence and then racing on commit. PostgreSQL
|
|
deployments skip this application-level lock.
|
|
"""
|
|
|
|
loop = asyncio.get_running_loop()
|
|
lock = _sqlite_session_locks.get(loop)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
_sqlite_session_locks[loop] = lock
|
|
return lock
|
|
|
|
|
|
def get_engine() -> AsyncEngine:
|
|
global _engine, _engine_url, _session_factory
|
|
url = get_settings().database_url
|
|
if _engine is None or _engine_url != url:
|
|
_engine = create_async_engine(url, future=True)
|
|
_engine_url = url
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
return _engine
|
|
|
|
|
|
def get_session_factory() -> async_sessionmaker[AsyncSession]:
|
|
get_engine()
|
|
assert _session_factory is not None
|
|
return _session_factory
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _transaction_scope() -> AsyncIterator[AsyncSession]:
|
|
session = get_session_factory()()
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def session_scope() -> AsyncIterator[AsyncSession]:
|
|
if get_settings().is_sqlite:
|
|
async with _sqlite_session_lock():
|
|
async with _transaction_scope() as session:
|
|
yield session
|
|
return
|
|
|
|
async with _transaction_scope() as session:
|
|
yield session
|
|
|
|
|
|
async def init_db() -> None:
|
|
engine = get_engine()
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
|
|
|
|
async def drop_db() -> None:
|
|
engine = get_engine()
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.drop_all)
|
|
|
|
|
|
async def dispose_db() -> None:
|
|
global _engine, _engine_url, _session_factory
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
_engine = None
|
|
_engine_url = None
|
|
_session_factory = None
|