eryao/backend/alembic/versions/20260402_0001_initial_llm_schema.py

"""initial llm/factory/system_agents schema

Revision ID: 202604020001
Revises:
Create Date: 2026-04-02 18:25:00
"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision: str = "202604020001"
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.create_table(
        "llm_factory",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(length=50), nullable=False),
        sa.Column("request_url", sa.String(length=255), nullable=False),
        sa.Column("avatar", sa.Text(), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("name"),
    )
    op.create_index("ix_llm_factory_name", "llm_factory", ["name"], unique=True)
    _enable_rls("llm_factory")

    op.create_table(
        "llms",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("factory_id", sa.UUID(), nullable=False),
        sa.Column("model_code", sa.String(length=50), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("model_code"),
    )
    op.create_index("ix_llms_factory_id", "llms", ["factory_id"], unique=False)
    op.create_index("ix_llms_model_code", "llms", ["model_code"], unique=True)
    op.create_foreign_key(
        "fk_llms_factory_id",
        "llms",
        "llm_factory",
        ["factory_id"],
        ["id"],
        ondelete="RESTRICT",
    )
    _enable_rls("llms")

    op.create_table(
        "system_agents",
        sa.Column("agent_type", sa.String(length=20), nullable=False),
        sa.Column("llm_id", sa.UUID(), nullable=False),
        sa.Column("status", sa.String(length=20), nullable=False),
        sa.Column(
            "config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default="{}",
            nullable=False,
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("agent_type"),
    )
    op.create_foreign_key(
        "fk_system_agents_llm_id",
        "system_agents",
        "llms",
        ["llm_id"],
        ["id"],
        ondelete="RESTRICT",
    )
    _enable_rls("system_agents")

    op.execute("REVOKE ALL ON TABLE public.alembic_version FROM anon")
    op.execute("REVOKE ALL ON TABLE public.alembic_version FROM authenticated")


def downgrade() -> None:
    _drop_rls("system_agents")
    op.drop_constraint("fk_system_agents_llm_id", "system_agents", type_="foreignkey")
    op.drop_table("system_agents")

    _drop_rls("llms")
    op.drop_constraint("fk_llms_factory_id", "llms", type_="foreignkey")
    op.drop_index("ix_llms_model_code", table_name="llms")
    op.drop_index("ix_llms_factory_id", table_name="llms")
    op.drop_table("llms")

    _drop_rls("llm_factory")
    op.drop_index("ix_llm_factory_name", table_name="llm_factory")
    op.drop_table("llm_factory")


def _enable_rls(table_name: str) -> None:
    for role in ["anon", "authenticated"]:
        for action in ["select", "insert", "update", "delete"]:
            op.execute(
                f"DROP POLICY IF EXISTS {role}_{action}_{table_name} ON {table_name}"
            )
    op.execute(f"ALTER TABLE {table_name} ENABLE ROW LEVEL SECURITY")
    for role in ["anon", "authenticated"]:
        op.execute(
            f"CREATE POLICY {role}_select_{table_name} ON {table_name} FOR SELECT TO {role} USING (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_insert_{table_name} ON {table_name} FOR INSERT TO {role} WITH CHECK (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_update_{table_name} ON {table_name} FOR UPDATE TO {role} USING (false) WITH CHECK (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_delete_{table_name} ON {table_name} FOR DELETE TO {role} USING (false)"
        )


def _drop_rls(table_name: str) -> None:
    for role in ["anon", "authenticated"]:
        op.execute(f"DROP POLICY IF EXISTS {role}_delete_{table_name} ON {table_name}")
        op.execute(f"DROP POLICY IF EXISTS {role}_update_{table_name} ON {table_name}")
        op.execute(f"DROP POLICY IF EXISTS {role}_insert_{table_name} ON {table_name}")
        op.execute(f"DROP POLICY IF EXISTS {role}_select_{table_name} ON {table_name}")
    op.execute(f"ALTER TABLE {table_name} DISABLE ROW LEVEL SECURITY")