customer-portal/customer_portal/migrations/006_drop_legacy_columns.py

#!/usr/bin/env python3
"""Migration 006: Drop legacy customer columns.

Sprint 12 Final Cleanup:
All customer data is now in custom_fields (JSONB).
Remove the deprecated fixed columns.

Usage:
    python -m customer_portal.migrations.006_drop_legacy_columns

IMPORTANT: Run migrations 004 and 005 first!
- 004: Consolidate data into custom_fields
- 005: Convert TEXT to JSONB
- 006: Drop legacy columns (this migration)
"""

import os
import sys

from sqlalchemy import create_engine, text

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))


# Columns to remove
LEGACY_COLUMNS = ["name", "phone", "address_street", "address_city", "address_zip"]


def get_database_url() -> str:
    """Get database URL from environment or default."""
    return os.environ.get(
        "DATABASE_URL", "postgresql://portal:portal@localhost:5433/customer_portal"
    )


def verify_data_migrated(conn) -> bool:
    """Verify all legacy data is in custom_fields before dropping columns."""
    print("Verifying data migration...")

    # Check for customers with legacy data not in custom_fields
    result = conn.execute(
        text(
            """
        SELECT COUNT(*) as cnt
        FROM customers
        WHERE (
            (name IS NOT NULL AND name != '' AND
             (custom_fields->>'name' IS NULL OR custom_fields->>'name' = ''))
            OR
            (phone IS NOT NULL AND phone != '' AND
             (custom_fields->>'phone' IS NULL OR custom_fields->>'phone' = ''))
            OR
            (address_street IS NOT NULL AND address_street != '' AND
             (custom_fields->>'address_street' IS NULL OR custom_fields->>'address_street' = ''))
            OR
            (address_city IS NOT NULL AND address_city != '' AND
             (custom_fields->>'address_city' IS NULL OR custom_fields->>'address_city' = ''))
            OR
            (address_zip IS NOT NULL AND address_zip != '' AND
             (custom_fields->>'address_zip' IS NULL OR custom_fields->>'address_zip' = ''))
        )
    """
        )
    )
    count = result.fetchone()[0]

    if count > 0:
        print(
            f"  WARNING: {count} customers have data in legacy columns not in custom_fields!"
        )
        return False

    print("  All legacy data is in custom_fields")
    return True


def migrate_remaining_data(conn) -> int:
    """Migrate any remaining data from legacy columns to custom_fields."""
    print("Migrating any remaining legacy data...")

    result = conn.execute(
        text(
            """
        UPDATE customers
        SET custom_fields = custom_fields ||
            jsonb_strip_nulls(jsonb_build_object(
                'name', CASE WHEN name != '' AND (custom_fields->>'name' IS NULL OR custom_fields->>'name' = '') THEN name ELSE NULL END,
                'phone', CASE WHEN phone != '' AND (custom_fields->>'phone' IS NULL OR custom_fields->>'phone' = '') THEN phone ELSE NULL END,
                'address_street', CASE WHEN address_street != '' AND (custom_fields->>'address_street' IS NULL OR custom_fields->>'address_street' = '') THEN address_street ELSE NULL END,
                'address_city', CASE WHEN address_city != '' AND (custom_fields->>'address_city' IS NULL OR custom_fields->>'address_city' = '') THEN address_city ELSE NULL END,
                'address_zip', CASE WHEN address_zip != '' AND (custom_fields->>'address_zip' IS NULL OR custom_fields->>'address_zip' = '') THEN address_zip ELSE NULL END
            ))
        WHERE name IS NOT NULL OR phone IS NOT NULL OR address_street IS NOT NULL
              OR address_city IS NOT NULL OR address_zip IS NOT NULL
    """
        )
    )

    count = result.rowcount
    print(f"  Updated {count} customers")
    return count


def check_columns_exist(conn) -> list[str]:
    """Check which legacy columns still exist."""
    result = conn.execute(
        text(
            """
        SELECT column_name
        FROM information_schema.columns
        WHERE table_name = 'customers'
        AND column_name = ANY(:columns)
    """
        ),
        {"columns": LEGACY_COLUMNS},
    )

    return [row[0] for row in result.fetchall()]


def run_migration():
    """Drop legacy columns after verifying data migration."""
    engine = create_engine(get_database_url())

    with engine.connect() as conn:
        # Check which columns exist
        existing = check_columns_exist(conn)
        if not existing:
            print("No legacy columns found - nothing to do")
            return

        print(f"Found legacy columns: {existing}")

        # First, migrate any remaining data
        migrate_remaining_data(conn)
        conn.commit()

        # Verify all data is migrated
        if not verify_data_migrated(conn):
            print("\nERROR: Data verification failed!")
            print("Run migration 004 first to consolidate data.")
            return

        # Drop columns
        print("\nDropping legacy columns...")
        for col in existing:
            print(f"  Dropping: {col}")
            conn.execute(text(f"ALTER TABLE customers DROP COLUMN IF EXISTS {col}"))

        conn.commit()

        print("\nMigration completed successfully!")
        print("\nCustomer table now has only:")
        print("  - id (PK)")
        print("  - email (unique identifier)")
        print("  - custom_fields (JSONB - all customer data)")
        print("  - wp_user_id")
        print("  - email_* preferences")
        print("  - is_admin")
        print("  - created_at, updated_at, last_login_at")


def rollback():
    """Rollback: Re-add legacy columns (data will be empty)."""
    engine = create_engine(get_database_url())

    with engine.connect() as conn:
        print("Rolling back: Re-adding legacy columns...")

        conn.execute(
            text(
                """
            ALTER TABLE customers
            ADD COLUMN IF NOT EXISTS name VARCHAR(255) NOT NULL DEFAULT '',
            ADD COLUMN IF NOT EXISTS phone VARCHAR(50),
            ADD COLUMN IF NOT EXISTS address_street VARCHAR(255),
            ADD COLUMN IF NOT EXISTS address_city VARCHAR(255),
            ADD COLUMN IF NOT EXISTS address_zip VARCHAR(20)
        """
            )
        )

        # Populate from custom_fields
        print("Populating columns from custom_fields...")
        conn.execute(
            text(
                """
            UPDATE customers
            SET
                name = COALESCE(custom_fields->>'name', ''),
                phone = custom_fields->>'phone',
                address_street = custom_fields->>'address_street',
                address_city = custom_fields->>'address_city',
                address_zip = custom_fields->>'address_zip'
        """
            )
        )

        conn.commit()
        print("Rollback completed")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--rollback":
        rollback()
    else:
        run_migration()