Aether/alembic/versions/20260228_1100_b2c3d4e5f6a7_cleanup_capability_settings.py

"""Add cache_creation columns, clean up capability settings, add user_model_usage_counts,
enforce global_model_id NOT NULL

1. Add cache_creation_input_tokens_5m and cache_creation_input_tokens_1h to usage table.
2. Clean up cache_1h/context_1m/gemini_files from user-configurable settings
   (now auto-detected via REQUEST_PARAM mode).
3. Create user_model_usage_counts table for per-user per-model atomic usage counters.
4. Enforce models.global_model_id NOT NULL (delete orphan models without global model).

Revision ID: b2c3d4e5f6a7
Revises: 9a0b1c2d3e4f
Create Date: 2026-02-28 14:00:00.000000

"""

from __future__ import annotations

import json
import uuid
from collections.abc import Sequence
from datetime import datetime, timezone

import sqlalchemy as sa
from sqlalchemy import inspect

from alembic import op

revision: str = "b2c3d4e5f6a7"
down_revision: str | None = "9a0b1c2d3e4f"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def column_exists(table_name: str, column_name: str) -> bool:
    bind = op.get_bind()
    insp = inspect(bind)
    columns = [c["name"] for c in insp.get_columns(table_name)]
    return column_name in columns


def table_exists(table_name: str) -> bool:
    bind = op.get_bind()
    insp = inspect(bind)
    return table_name in insp.get_table_names()


def index_exists(table_name: str, index_name: str) -> bool:
    bind = op.get_bind()
    insp = inspect(bind)
    return any(idx["name"] == index_name for idx in insp.get_indexes(table_name))


def upgrade() -> None:
    # --- 1. Add cache_creation columns ---
    if not column_exists("usage", "cache_creation_input_tokens_5m"):
        op.add_column(
            "usage",
            sa.Column(
                "cache_creation_input_tokens_5m",
                sa.Integer(),
                nullable=False,
                server_default=sa.text("0"),
                comment="5min TTL cache creation input tokens",
            ),
        )
    if not column_exists("usage", "cache_creation_input_tokens_1h"):
        op.add_column(
            "usage",
            sa.Column(
                "cache_creation_input_tokens_1h",
                sa.Integer(),
                nullable=False,
                server_default=sa.text("0"),
                comment="1h TTL cache creation input tokens",
            ),
        )

    # --- 2. Clean up stale capability settings (pure Python, DB-agnostic) ---
    stale_keys = {"cache_1h", "context_1m", "gemini_files"}
    conn = op.get_bind()

    # ApiKey.force_capabilities: dict-like JSON, remove stale keys
    rows = conn.execute(
        sa.text("SELECT id, force_capabilities FROM api_keys WHERE force_capabilities IS NOT NULL")
    ).fetchall()
    for row in rows:
        raw = row[1]
        if raw is None:
            continue
        data = raw if isinstance(raw, dict) else json.loads(raw)
        cleaned = {k: v for k, v in data.items() if k not in stale_keys}
        new_val = json.dumps(cleaned) if cleaned else None
        conn.execute(
            sa.text("UPDATE api_keys SET force_capabilities = :val WHERE id = :id"),
            {"val": new_val, "id": row[0]},
        )

    # User.model_capability_settings: nested dict {model_key: {cap: val}}, remove stale keys
    rows = conn.execute(
        sa.text(
            "SELECT id, model_capability_settings FROM users"
            " WHERE model_capability_settings IS NOT NULL"
        )
    ).fetchall()
    for row in rows:
        raw = row[1]
        if raw is None:
            continue
        data = raw if isinstance(raw, dict) else json.loads(raw)
        cleaned = {}
        for model_key, caps in data.items():
            cap_cleaned = {k: v for k, v in caps.items() if k not in stale_keys}
            if cap_cleaned:
                cleaned[model_key] = cap_cleaned
        new_val = json.dumps(cleaned) if cleaned else None
        conn.execute(
            sa.text("UPDATE users SET model_capability_settings = :val WHERE id = :id"),
            {"val": new_val, "id": row[0]},
        )

    # GlobalModel.supported_capabilities: JSON array, remove stale entries
    rows = conn.execute(
        sa.text(
            "SELECT id, supported_capabilities FROM global_models"
            " WHERE supported_capabilities IS NOT NULL"
        )
    ).fetchall()
    for row in rows:
        raw = row[1]
        if raw is None:
            continue
        data = raw if isinstance(raw, list) else json.loads(raw)
        cleaned = [c for c in data if c not in stale_keys]
        new_val = json.dumps(cleaned) if cleaned else None
        conn.execute(
            sa.text("UPDATE global_models SET supported_capabilities = :val WHERE id = :id"),
            {"val": new_val, "id": row[0]},
        )

    # --- 3. Create user_model_usage_counts table ---
    if not table_exists("user_model_usage_counts"):
        op.create_table(
            "user_model_usage_counts",
            sa.Column("id", sa.String(36), primary_key=True),
            sa.Column(
                "user_id",
                sa.String(36),
                sa.ForeignKey("users.id", ondelete="CASCADE"),
                nullable=False,
            ),
            sa.Column("model", sa.String(100), nullable=False),
            sa.Column("usage_count", sa.Integer, nullable=False, server_default="0"),
            sa.Column(
                "created_at",
                sa.DateTime(timezone=True),
                nullable=False,
                server_default=sa.func.now(),
            ),
            sa.Column(
                "updated_at",
                sa.DateTime(timezone=True),
                nullable=False,
                server_default=sa.func.now(),
            ),
            sa.UniqueConstraint("user_id", "model", name="uq_user_model_usage_count"),
        )
    if not index_exists("user_model_usage_counts", "idx_user_model_usage_user"):
        op.create_index("idx_user_model_usage_user", "user_model_usage_counts", ["user_id"])
    if not index_exists("user_model_usage_counts", "idx_user_model_usage_model"):
        op.create_index("idx_user_model_usage_model", "user_model_usage_counts", ["model"])

    # Backfill from existing usage records (truncate first for idempotency)
    conn.execute(sa.text("DELETE FROM user_model_usage_counts"))
    rows = conn.execute(
        sa.text(
            "SELECT user_id, model, COUNT(*) AS cnt FROM usage"
            " WHERE user_id IS NOT NULL GROUP BY user_id, model"
        )
    ).fetchall()
    now = datetime.now(timezone.utc)
    for row in rows:
        conn.execute(
            sa.text(
                "INSERT INTO user_model_usage_counts"
                " (id, user_id, model, usage_count, created_at, updated_at)"
                " VALUES (:id, :user_id, :model, :cnt, :now, :now)"
            ),
            {
                "id": str(uuid.uuid4()),
                "user_id": row[0],
                "model": row[1],
                "cnt": row[2],
                "now": now,
            },
        )

    # --- 4. Enforce models.global_model_id NOT NULL ---
    conn = op.get_bind()
    insp = inspect(conn)
    model_cols = {c["name"]: c for c in insp.get_columns("models")}
    if model_cols.get("global_model_id", {}).get("nullable", True):
        op.execute("DELETE FROM models WHERE global_model_id IS NULL")
        op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=False)


def downgrade() -> None:
    # Revert models.global_model_id to nullable
    if column_exists("models", "global_model_id"):
        op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=True)

    # Drop user_model_usage_counts
    if table_exists("user_model_usage_counts"):
        if index_exists("user_model_usage_counts", "idx_user_model_usage_model"):
            op.drop_index("idx_user_model_usage_model", table_name="user_model_usage_counts")
        if index_exists("user_model_usage_counts", "idx_user_model_usage_user"):
            op.drop_index("idx_user_model_usage_user", table_name="user_model_usage_counts")
        op.drop_table("user_model_usage_counts")

    # Drop cache_creation columns
    if column_exists("usage", "cache_creation_input_tokens_1h"):
        op.drop_column("usage", "cache_creation_input_tokens_1h")
    if column_exists("usage", "cache_creation_input_tokens_5m"):
        op.drop_column("usage", "cache_creation_input_tokens_5m")
    # capability settings cleanup is not reversible