portfolio-hub/migrations/versions/20260511_0015_portfolio_uuid_migration.py

"""portfolios.portfolio_id UUID 전환 및 (owner_user_id, name) UNIQUE 제약 추가

Revision ID: 20260511_0015
Revises: 20260502_0014
Create Date: 2026-05-11 00:00:00
"""

from __future__ import annotations

import uuid as _uuid

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

revision = "20260511_0015"
down_revision = "20260502_0014"
branch_labels = None
depends_on = None


def _column_exists(inspector: sa.engine.Inspector, table_name: str, column_name: str) -> bool:
    if table_name not in inspector.get_table_names():
        return False
    return any(c["name"] == column_name for c in inspector.get_columns(table_name))


def _constraint_exists(inspector: sa.engine.Inspector, table_name: str, constraint_name: str) -> bool:
    if table_name not in inspector.get_table_names():
        return False
    return any(c["name"] == constraint_name for c in inspector.get_unique_constraints(table_name))


def upgrade() -> None:
    bind = op.get_bind()
    inspector = inspect(bind)

    # Step 1: Add name column IF NOT EXISTS (already present via ensure_schema in most envs)
    if not _column_exists(inspector, "portfolios", "name"):
        op.add_column("portfolios", sa.Column("name", sa.Text(), nullable=True))
        bind.execute(sa.text("UPDATE portfolios SET name = portfolio_id WHERE name IS NULL OR name = ''"))
        op.alter_column("portfolios", "name", nullable=False)

    # Step 2: Guard — skip UUID replacement if already migrated (UNIQUE constraint present)
    already_migrated = _constraint_exists(inspector, "portfolios", "uq_portfolios_owner_name")

    if not already_migrated:
        # Collect all existing portfolio_ids
        result = bind.execute(sa.text("SELECT portfolio_id FROM portfolios"))
        old_pids = [row[0] for row in result.fetchall()]

        for old_pid in old_pids:
            new_uuid = str(_uuid.uuid4())
            # Copy old portfolio_id to name (overwrites display name with identifier)
            bind.execute(
                sa.text("UPDATE portfolios SET name = :old_pid WHERE portfolio_id = :old_pid"),
                {"old_pid": old_pid},
            )
            # Update FK reference tables before changing portfolios PK
            bind.execute(
                sa.text("UPDATE portfolio_ledger SET portfolio_id = :new_uuid WHERE portfolio_id = :old_pid"),
                {"new_uuid": new_uuid, "old_pid": old_pid},
            )
            bind.execute(
                sa.text("UPDATE portfolio_comments SET portfolio_id = :new_uuid WHERE portfolio_id = :old_pid"),
                {"new_uuid": new_uuid, "old_pid": old_pid},
            )
            bind.execute(
                sa.text("UPDATE portfolio_tags SET portfolio_id = :new_uuid WHERE portfolio_id = :old_pid"),
                {"new_uuid": new_uuid, "old_pid": old_pid},
            )
            bind.execute(
                sa.text("UPDATE ledger_trades SET asset_id = :new_uuid WHERE asset_id = :old_pid"),
                {"new_uuid": new_uuid, "old_pid": old_pid},
            )
            # Update portfolios PK last
            bind.execute(
                sa.text("UPDATE portfolios SET portfolio_id = :new_uuid WHERE portfolio_id = :old_pid"),
                {"new_uuid": new_uuid, "old_pid": old_pid},
            )

        # Ensure no name nulls remain before adding UNIQUE constraint
        bind.execute(
            sa.text("UPDATE portfolios SET name = portfolio_id WHERE name IS NULL OR name = ''")
        )

        # Step 3: Add (owner_user_id, name) UNIQUE constraint
        op.create_unique_constraint("uq_portfolios_owner_name", "portfolios", ["owner_user_id", "name"])

        # Step 4: Verify no orphan records in FK reference tables
        for table, col in [
            ("portfolio_ledger", "portfolio_id"),
            ("portfolio_comments", "portfolio_id"),
            ("portfolio_tags", "portfolio_id"),
        ]:
            result = bind.execute(
                sa.text(
                    f"SELECT COUNT(*) FROM {table}"  # noqa: S608
                    f" WHERE {col} NOT IN (SELECT portfolio_id FROM portfolios)"
                )
            )
            count = result.scalar()
            if count and count > 0:
                raise RuntimeError(
                    f"Migration 20260511_0015: {count} orphan record(s) found in {table}.{col}"
                )


def downgrade() -> None:
    bind = op.get_bind()
    inspector = inspect(bind)

    if _constraint_exists(inspector, "portfolios", "uq_portfolios_owner_name"):
        op.drop_constraint("uq_portfolios_owner_name", "portfolios", type_="unique")