database.py

"""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