membership-be/alembic/versions/012_fix_remaining_differences.py

"""fix_remaining_differences

Revision ID: 012_fix_remaining
Revises: 011_align_prod_dev
Create Date: 2026-01-05

Fixes the last 5 schema differences found after migration 011:
1-2. import_rollback_audit nullable constraints (PROD)
3-4. role_permissions type and nullable (PROD)
5. UserStatus enum values (DEV - remove deprecated values)
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import ENUM

# revision identifiers, used by Alembic.
revision: str = '012_fix_remaining'
down_revision: Union[str, None] = '011_align_prod_dev'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Fix remaining schema differences"""
    from sqlalchemy import inspect

    conn = op.get_bind()
    inspector = inspect(conn)

    print("Fixing remaining schema differences...")

    # ============================================================
    # 1. FIX IMPORT_ROLLBACK_AUDIT TABLE (PROD only)
    # ============================================================
    print("\n[1/3] Fixing import_rollback_audit nullable constraints...")

    # Check if there are any NULL values first
    try:
        null_count = conn.execute(sa.text("""
            SELECT COUNT(*) FROM import_rollback_audit
            WHERE created_at IS NULL OR rolled_back_at IS NULL
        """)).scalar()

        if null_count > 0:
            # Fill NULLs with current timestamp
            conn.execute(sa.text("""
                UPDATE import_rollback_audit
                SET created_at = NOW() WHERE created_at IS NULL
            """))
            conn.execute(sa.text("""
                UPDATE import_rollback_audit
                SET rolled_back_at = NOW() WHERE rolled_back_at IS NULL
            """))
            print(f"  ✓ Filled {null_count} NULL timestamps")

        # Now set NOT NULL
        op.alter_column('import_rollback_audit', 'created_at', nullable=False)
        op.alter_column('import_rollback_audit', 'rolled_back_at', nullable=False)
        print("  ✓ Set NOT NULL constraints")
    except Exception as e:
        print(f"  ⚠️  Warning: {e}")

    # ============================================================
    # 2. FIX ROLE_PERMISSIONS TABLE (PROD only)
    # ============================================================
    print("\n[2/3] Fixing role_permissions.role type and nullable...")

    try:
        # Change VARCHAR(50) to VARCHAR(10) to match UserRole enum
        op.alter_column('role_permissions', 'role',
                       type_=sa.String(10),
                       postgresql_using='role::varchar(10)')
        print("  ✓ Changed VARCHAR(50) to VARCHAR(10)")

        # Set NOT NULL
        op.alter_column('role_permissions', 'role', nullable=False)
        print("  ✓ Set NOT NULL constraint")
    except Exception as e:
        print(f"  ⚠️  Warning: {e}")

    # ============================================================
    # 3. FIX USERSTATUS ENUM (DEV only - remove deprecated values)
    # ============================================================
    print("\n[3/3] Fixing UserStatus enum values...")

    try:
        # First, check if the enum has deprecated values
        enum_values = conn.execute(sa.text("""
            SELECT enumlabel
            FROM pg_enum
            WHERE enumtypid = (
                SELECT oid FROM pg_type WHERE typname = 'userstatus'
            )
        """)).fetchall()

        enum_values_list = [row[0] for row in enum_values]
        has_deprecated = 'pending_approval' in enum_values_list or 'pre_approved' in enum_values_list

        if not has_deprecated:
            print("  ✓ UserStatus enum already correct (no deprecated values)")
        else:
            print("  ⏳ Found deprecated enum values, migrating...")

            # Check if any users have deprecated status values
            deprecated_count = conn.execute(sa.text("""
                SELECT COUNT(*) FROM users
                WHERE status IN ('pending_approval', 'pre_approved')
            """)).scalar()

            if deprecated_count > 0:
                print(f"  ⏳ Migrating {deprecated_count} users with deprecated status values...")

                # Migrate deprecated values to new equivalents
                conn.execute(sa.text("""
                    UPDATE users
                    SET status = 'pre_validated'
                    WHERE status = 'pre_approved'
                """))

                conn.execute(sa.text("""
                    UPDATE users
                    SET status = 'payment_pending'
                    WHERE status = 'pending_approval'
                """))

                print("  ✓ Migrated deprecated status values")
            else:
                print("  ✓ No users with deprecated status values")

            # Now remove deprecated enum values
            # PostgreSQL doesn't support removing enum values directly,
            # so we need to recreate the enum
            conn.execute(sa.text("""
                -- Create new enum with correct values (matches models.py)
                CREATE TYPE userstatus_new AS ENUM (
                    'pending_email',
                    'pending_validation',
                    'pre_validated',
                    'payment_pending',
                    'active',
                    'inactive',
                    'canceled',
                    'expired',
                    'rejected',
                    'abandoned'
                );

                -- Update column to use new enum
                ALTER TABLE users
                ALTER COLUMN status TYPE userstatus_new
                USING status::text::userstatus_new;

                -- Drop old enum and rename new one
                DROP TYPE userstatus;
                ALTER TYPE userstatus_new RENAME TO userstatus;
            """))

            print("  ✓ Updated UserStatus enum (removed deprecated values)")

    except Exception as e:
        print(f"  ⚠️  Warning: Enum update failed (may already be correct): {e}")

    print("\n✅ All remaining differences fixed!")


def downgrade() -> None:
    """Revert fixes (not recommended)"""
    print("⚠️  Downgrade not supported")
    pass