WonderQ-Project/WonderQ-Admin/alembic/versions/0026_admin_ownership.py

"""Add department and creator ownership to operator-managed resources."""

from alembic import context, op
import sqlalchemy as sa
from sqlalchemy import inspect


revision = "0026_admin_ownership"
down_revision = "0025_admin_rbac"
branch_labels = None
depends_on = None

DEFAULT_DEPT_ID = "00000000-0000-0000-0000-000000000001"
OWNED_TABLES = (
    "MediaAsset",
    "HeroSlide",
    "DestinationHero",
    "VehicleOption",
    "VehicleServiceConfig",
    "HomeExperience",
    "HomeTeamBuilding",
    "HomeWildArchive",
    "HomeWanfaRecommendation",
    "WanfaCategory",
    "WanfaRoute",
    "DetailRecord",
    "ConciergeAdvisor",
    "DemandHero",
    "DemandFeatureCard",
    "DemandForm",
    "Lead",
)


def is_offline_mode() -> bool:
    try:
        return context.is_offline_mode()
    except (NameError, AttributeError):
        return False


def _schema_snapshot() -> tuple[set[str], dict[str, dict[str, dict]], dict[str, set[str]], dict[str, list[dict]]]:
    if is_offline_mode():
        return set(), {}, {}, {}

    inspector = inspect(op.get_bind())
    table_names = set(inspector.get_table_names())
    columns = {
        table_name: {column["name"]: column for column in inspector.get_columns(table_name)}
        for table_name in table_names
    }
    indexes = {
        table_name: {
            index["name"]
            for index in inspector.get_indexes(table_name)
            if index.get("name")
        }
        for table_name in table_names
    }
    foreign_keys = {
        table_name: inspector.get_foreign_keys(table_name)
        for table_name in table_names
    }
    return table_names, columns, indexes, foreign_keys


def _has_foreign_key(
    foreign_keys: list[dict],
    constrained_column: str,
    referred_table: str,
) -> bool:
    return any(
        foreign_key.get("constrained_columns") == [constrained_column]
        and foreign_key.get("referred_table") == referred_table
        and foreign_key.get("referred_columns") == ["id"]
        for foreign_key in foreign_keys
    )


def upgrade() -> None:
    table_names, columns, indexes, foreign_keys = _schema_snapshot()

    for table_name in OWNED_TABLES:
        if table_name not in table_names:
            continue

        table_columns = columns[table_name]
        added_dept_column = "deptId" not in table_columns
        added_creator_column = "createdById" not in table_columns
        if added_dept_column:
            op.add_column(
                table_name,
                sa.Column(
                    "deptId",
                    sa.String(),
                    nullable=True,
                    server_default=sa.text(f"'{DEFAULT_DEPT_ID}'"),
                ),
            )
            table_columns["deptId"] = {"name": "deptId", "nullable": True}
        if added_creator_column:
            op.add_column(table_name, sa.Column("createdById", sa.String(), nullable=True))
            table_columns["createdById"] = {"name": "createdById", "nullable": True}

        table_indexes = indexes[table_name]
        if f"ix_{table_name}_deptId" not in table_indexes:
            op.create_index(f"ix_{table_name}_deptId", table_name, ["deptId"])
            table_indexes.add(f"ix_{table_name}_deptId")
        if f"ix_{table_name}_createdById" not in table_indexes:
            op.create_index(f"ix_{table_name}_createdById", table_name, ["createdById"])
            table_indexes.add(f"ix_{table_name}_createdById")

        table_foreign_keys = foreign_keys[table_name]
        if not _has_foreign_key(table_foreign_keys, "deptId", "AdminDepartment"):
            op.create_foreign_key(
                f"fk_{table_name}_deptId",
                table_name,
                "AdminDepartment",
                ["deptId"],
                ["id"],
                ondelete="RESTRICT",
            )
            table_foreign_keys.append(
                {
                    "constrained_columns": ["deptId"],
                    "referred_table": "AdminDepartment",
                    "referred_columns": ["id"],
                }
            )
        if not _has_foreign_key(table_foreign_keys, "createdById", "AdminUser"):
            op.create_foreign_key(
                f"fk_{table_name}_createdById",
                table_name,
                "AdminUser",
                ["createdById"],
                ["id"],
                ondelete="SET NULL",
            )
            table_foreign_keys.append(
                {
                    "constrained_columns": ["createdById"],
                    "referred_table": "AdminUser",
                    "referred_columns": ["id"],
                }
            )
        quoted_table = f'"{table_name}"'
        needs_dept_backfill = added_dept_column or table_columns["deptId"].get("nullable") is True
        if needs_dept_backfill:
            op.execute(
                sa.text(f'UPDATE {quoted_table} SET "deptId" = :dept_id WHERE "deptId" IS NULL').bindparams(
                    dept_id=DEFAULT_DEPT_ID
                )
            )
            op.alter_column(
                table_name,
                "deptId",
                existing_type=sa.String(),
                nullable=False,
                server_default=sa.text(f"'{DEFAULT_DEPT_ID}'"),
            )
        if added_creator_column:
            op.execute(
                sa.text(
                    f'UPDATE {quoted_table} SET "createdById" = '
                    '(SELECT "id" FROM "AdminUser" ORDER BY "createdAt" ASC LIMIT 1) '
                    'WHERE "createdById" IS NULL'
                )
            )


def downgrade() -> None:
    for table_name in reversed(OWNED_TABLES):
        op.drop_constraint(f"fk_{table_name}_createdById", table_name, type_="foreignkey")
        op.drop_constraint(f"fk_{table_name}_deptId", table_name, type_="foreignkey")
        op.drop_index(f"ix_{table_name}_createdById", table_name=table_name)
        op.drop_index(f"ix_{table_name}_deptId", table_name=table_name)
        op.drop_column(table_name, "createdById")
        op.drop_column(table_name, "deptId")