headquarter/apps/api/alembic/versions/0014_add_profile_resolver_fields.py

"""add profile resolver fields to config profiles and mounts

Revision ID: 0014_add_profile_resolver_fields
Revises: 0013_add_config_profiles
Create Date: 2026-05-24 14:00:00.000000

"""
from collections.abc import Sequence

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

# revision identifiers, used by Alembic.
revision: str = "0014_add_profile_resolver_fields"
down_revision: str | None = "0013_add_config_profiles"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def _table_exists(table_name: str) -> bool:
    return sa.inspect(op.get_bind()).has_table(table_name)


def _column_exists(table_name: str, column_name: str) -> bool:
    if not _table_exists(table_name):
        return False
    return column_name in {
        column["name"] for column in sa.inspect(op.get_bind()).get_columns(table_name)
    }


def _index_exists(table_name: str, index_name: str) -> bool:
    if not _table_exists(table_name):
        return False
    return index_name in {
        index["name"] for index in sa.inspect(op.get_bind()).get_indexes(table_name)
    }


def _foreign_key_exists(
    table_name: str,
    constrained_columns: list[str],
    referred_table: str,
) -> bool:
    if not _table_exists(table_name):
        return False
    for foreign_key in sa.inspect(op.get_bind()).get_foreign_keys(table_name):
        if (
            foreign_key.get("constrained_columns") == constrained_columns
            and foreign_key.get("referred_table") == referred_table
        ):
            return True
    return False


def _foreign_key_names_for_column(table_name: str, column_name: str) -> list[str]:
    if not _table_exists(table_name):
        return []
    names: list[str] = []
    for foreign_key in sa.inspect(op.get_bind()).get_foreign_keys(table_name):
        if column_name in foreign_key.get("constrained_columns", []):
            name = foreign_key.get("name")
            if name:
                names.append(name)
    return names


def upgrade() -> None:
    if not _column_exists("config_profiles", "project_id"):
        op.add_column(
            "config_profiles",
            sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True),
        )
    if not _column_exists("config_profiles", "tool_type_id"):
        op.add_column(
            "config_profiles",
            sa.Column("tool_type_id", postgresql.UUID(as_uuid=True), nullable=True),
        )
    if not _column_exists("config_profiles", "environment_variables"):
        op.add_column(
            "config_profiles",
            sa.Column("environment_variables", sa.JSON(), nullable=True),
        )
    if not _column_exists("config_profiles", "start_command"):
        op.add_column(
            "config_profiles",
            sa.Column("start_command", sa.Text(), nullable=True),
        )
    if not _column_exists("config_profiles", "working_directory"):
        op.add_column(
            "config_profiles",
            sa.Column("working_directory", sa.Text(), nullable=True),
        )
    if not _column_exists("config_profiles", "port"):
        op.add_column("config_profiles", sa.Column("port", sa.Integer(), nullable=True))
    if not _column_exists("config_profiles", "is_default"):
        op.add_column(
            "config_profiles",
            sa.Column("is_default", sa.Boolean(), nullable=False, server_default="false"),
        )

    if not _foreign_key_exists("config_profiles", ["project_id"], "projects"):
        op.create_foreign_key(
            "fk_config_profiles_project",
            "config_profiles",
            "projects",
            ["project_id"],
            ["id"],
            ondelete="CASCADE",
        )
    if not _foreign_key_exists("config_profiles", ["tool_type_id"], "tool_types"):
        op.create_foreign_key(
            "fk_config_profiles_tool_type",
            "config_profiles",
            "tool_types",
            ["tool_type_id"],
            ["id"],
            ondelete="CASCADE",
        )

    if not _index_exists("config_profiles", "idx_config_profiles_project"):
        op.create_index("idx_config_profiles_project", "config_profiles", ["project_id"])
    if not _index_exists("config_profiles", "idx_config_profiles_tool_type"):
        op.create_index(
            "idx_config_profiles_tool_type", "config_profiles", ["tool_type_id"]
        )

    if _column_exists("config_mounts", "mount_path") and not _column_exists(
        "config_mounts", "target_path"
    ):
        op.alter_column("config_mounts", "mount_path", new_column_name="target_path")
    if not _column_exists("config_mounts", "mode"):
        op.add_column(
            "config_mounts",
            sa.Column("mode", sa.String(length=10), nullable=False, server_default="rw"),
        )
    if not _column_exists("config_mounts", "files"):
        op.add_column(
            "config_mounts",
            sa.Column("files", sa.JSON(), nullable=True),
        )
    for constraint_name in _foreign_key_names_for_column(
        "config_mounts", "source_profile_id"
    ):
        op.drop_constraint(constraint_name, "config_mounts", type_="foreignkey")
    if _column_exists("config_mounts", "content"):
        op.drop_column("config_mounts", "content")
    if _column_exists("config_mounts", "source_profile_id"):
        op.drop_column("config_mounts", "source_profile_id")


def downgrade() -> None:
    # Restore config_mounts
    op.add_column(
        "config_mounts",
        sa.Column("source_profile_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    op.add_column(
        "config_mounts",
        sa.Column("content", sa.Text(), nullable=True),
    )
    op.drop_column("config_mounts", "files")
    op.drop_column("config_mounts", "mode")
    op.alter_column("config_mounts", "target_path", new_column_name="mount_path")

    # Restore config_profiles
    op.drop_index("idx_config_profiles_tool_type", table_name="config_profiles")
    op.drop_index("idx_config_profiles_project", table_name="config_profiles")
    op.drop_constraint(
        "fk_config_profiles_tool_type", "config_profiles", type_="foreignkey"
    )
    op.drop_constraint("fk_config_profiles_project", "config_profiles", type_="foreignkey")
    op.drop_column("config_profiles", "is_default")
    op.drop_column("config_profiles", "port")
    op.drop_column("config_profiles", "working_directory")
    op.drop_column("config_profiles", "start_command")
    op.drop_column("config_profiles", "environment_variables")
    op.drop_column("config_profiles", "tool_type_id")
    op.drop_column("config_profiles", "project_id")