"""Async database engine and session management. Includes a SQLAlchemy ``after_commit`` event listener that pushes the panel SQLite to the configured HF Dataset via :class:`PanelStorageCoordinator` -- so every route that mutates DB state already gets persistence for free. """ from __future__ import annotations from contextlib import asynccontextmanager from typing import AsyncGenerator from sqlalchemy import event from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy import select from config import PANEL_CONFIG from models import Base engine = create_async_engine( PANEL_CONFIG.database_url, echo=PANEL_CONFIG.debug, future=True, ) async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) async def init_db() -> None: """Create all tables and apply lightweight inline migrations. SQLAlchemy's ``create_all`` is idempotent for *tables* but doesn't add new *columns* to an existing table — so we run a tiny additive-migration step for each panel-level schema bump. Keep this list short and only use it for non-destructive ADD COLUMN statements; for anything more involved ship a real Alembic migration. """ async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # Inline additive migrations (SQLite-friendly). Wrapped in try / # except because ``ALTER TABLE ADD COLUMN`` raises if the column # already exists — that's exactly the case we want to ignore. inline_migrations: list[str] = [ "ALTER TABLE users ADD COLUMN plan_max_bots INTEGER NOT NULL DEFAULT 4", # 2026-07-01: invite-code redemptions audit table. The model # already declares it in Base.metadata; this statement exists # so a panel deployed before the model was added still gets # the table on the next start. """ CREATE TABLE IF NOT EXISTS invite_redemptions ( id INTEGER PRIMARY KEY AUTOINCREMENT, invite_id INTEGER NOT NULL REFERENCES invite_codes(id) ON DELETE CASCADE, user_id INTEGER NOT NULL, username VARCHAR(64), redeemed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ) """, "CREATE INDEX IF NOT EXISTS idx_invite_redemptions_invite ON invite_redemptions(invite_id)", "CREATE INDEX IF NOT EXISTS idx_invite_redemptions_user ON invite_redemptions(user_id)", ] for stmt in inline_migrations: try: async with engine.begin() as conn: await conn.exec_driver_sql(stmt) except Exception: # Column already exists (or SQLite is mid-migration) — fine. pass async def get_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency: yield an async session.""" async with async_session() as session: yield session @asynccontextmanager async def db_session() -> AsyncGenerator[AsyncSession, None]: """Context manager for manual DB sessions.""" async with async_session() as session: yield session # -- Write-through event hook ---------------------------------------------- # Fires after EVERY successful commit on any Session / AsyncSession; pushes # to HF Dataset via the coordinator. Sync listener -> fire-and-forget task. # We attach to the underlying sync ``Session`` class because AsyncSession's # events are routed through it. from sqlalchemy.orm import Session as _SyncSession @event.listens_for(_SyncSession, "after_commit", propagate=True) def _panel_after_commit(session) -> None: # Lazy import to avoid circular dep at module-load. from storage import PANEL_STORAGE_COORD try: info = session.info except Exception: info = {} if info.get("_panel_storage_skip"): return # internal / bulk operations may opt out reason = info.get("_panel_storage_reason") or "db commit" PANEL_STORAGE_COORD.schedule_push_sync(reason) @event.listens_for(_SyncSession, "after_rollback", propagate=True) def _panel_after_rollback(session) -> None: # Reset any custom reason so a future commit on the same session # doesn't carry over a stale label. try: session.info.pop("_panel_storage_reason", None) except Exception: pass