chore(repo): initialize team collaboration repository

flashops/services/control-plane/flashops_control/db.py

@@ -0,0 +1,106 @@
+"""异步数据库连接与会话。
+
+连接按当前配置惰性创建，测试可以安全地切换到临时 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