portfolio-hub/migrations/versions/20260519_0018_portfolios_asset_migration.py

"""Portfolios asset migration: data migration + asset_id FK + symbol/market DROP (REQ-011, REQ-014)

Revision ID: 20260519_0018
Revises: 20260519_0017
Create Date: 2026-05-19 00:00:00
"""

from __future__ import annotations

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

revision = "20260519_0018"
down_revision = "20260519_0017"
branch_labels = None
depends_on = None


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


def _constraint_exists(inspector: sa.engine.Inspector, table_name: str, constraint_name: str) -> bool:
    fk_constraints = inspector.get_foreign_keys(table_name)
    return any(fk.get("name") == constraint_name for fk in fk_constraints)


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

    # Step 1: Populate assets catalog from existing portfolios (REQ-014)
    # Only if portfolios.symbol column still exists (idempotent guard)
    if _column_exists(inspector, "portfolios", "symbol"):
        bind.execute(text("""
            INSERT INTO assets (symbol, markets)
            SELECT
                symbol,
                COALESCE(
                    ARRAY_AGG(DISTINCT market) FILTER (WHERE market IS NOT NULL),
                    ARRAY[]::text[]
                )
            FROM portfolios
            GROUP BY symbol
            ON CONFLICT (symbol) DO NOTHING
        """))

    # Step 2: Add portfolios.asset_id column (nullable, temporary)
    if not _column_exists(inspector, "portfolios", "asset_id"):
        op.add_column("portfolios", sa.Column("asset_id", sa.UUID(), nullable=True))

    # Step 3: Back-fill asset_id from assets table
    if _column_exists(inspector, "portfolios", "symbol") and _column_exists(inspector, "portfolios", "asset_id"):
        bind.execute(text("""
            UPDATE portfolios p
            SET asset_id = a.asset_id
            FROM assets a
            WHERE a.symbol = p.symbol
              AND p.asset_id IS NULL
        """))

    # Step 4: Set NOT NULL constraint on asset_id
    if _column_exists(inspector, "portfolios", "asset_id"):
        op.alter_column("portfolios", "asset_id", nullable=False)

    # Step 5: Add FK constraint (if not already present)
    if not _constraint_exists(inspector, "portfolios", "fk_portfolios_asset_id"):
        op.create_foreign_key(
            "fk_portfolios_asset_id",
            "portfolios",
            "assets",
            ["asset_id"],
            ["asset_id"],
            ondelete="RESTRICT",
        )

    # Step 6: Drop legacy columns
    if _column_exists(inspector, "portfolios", "symbol"):
        op.drop_column("portfolios", "symbol")

    if _column_exists(inspector, "portfolios", "market"):
        op.drop_column("portfolios", "market")


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

    # Restore symbol and market columns
    if not _column_exists(inspector, "portfolios", "symbol"):
        op.add_column("portfolios", sa.Column("symbol", sa.Text(), nullable=True))

    if not _column_exists(inspector, "portfolios", "market"):
        op.add_column("portfolios", sa.Column("market", sa.Text(), nullable=True))

    # Back-fill symbol from assets table
    if _column_exists(inspector, "portfolios", "asset_id") and _column_exists(inspector, "portfolios", "symbol"):
        bind.execute(text("""
            UPDATE portfolios p
            SET symbol = a.symbol
            FROM assets a
            WHERE a.asset_id = p.asset_id
        """))

    # Remove FK constraint
    if _constraint_exists(inspector, "portfolios", "fk_portfolios_asset_id"):
        op.drop_constraint("fk_portfolios_asset_id", "portfolios", type_="foreignkey")

    # Remove asset_id column
    if _column_exists(inspector, "portfolios", "asset_id"):
        op.drop_column("portfolios", "asset_id")