File size: 4,334 Bytes
bb56382 971fa87 bb56382 971fa87 067db8e 971fa87 067db8e f598ecc 067db8e 971fa87 bb56382 971fa87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """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
|