orion/alembic/versions/fa7d4d10e358_add_rbac_enhancements.py

"""add_rbac_enhancements

Revision ID: fa7d4d10e358
Revises: 4951b2e50581
Create Date: 2025-11-13 16:51:25.010057

SQLite-compatible version
"""

from typing import Sequence, Union

import sqlalchemy as sa

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "fa7d4d10e358"
down_revision: Union[str, None] = "4951b2e50581"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade():
    """Upgrade database schema for RBAC enhancements.

    SQLite-compatible version using batch operations for table modifications.
    """

    # ========================================================================
    # User table changes
    # ========================================================================
    with op.batch_alter_table("users", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "is_email_verified",
                sa.Boolean(),
                nullable=False,
                server_default="false",
            )
        )

    # Set existing active users as verified
    op.execute("UPDATE users SET is_email_verified = TRUE WHERE is_active = TRUE")

    # ========================================================================
    # VendorUser table changes (requires table recreation for SQLite)
    # ========================================================================
    with op.batch_alter_table("vendor_users", schema=None) as batch_op:
        # Add new columns
        batch_op.add_column(
            sa.Column(
                "user_type",
                sa.String(length=20),
                nullable=False,
                server_default="member",
            )
        )
        batch_op.add_column(
            sa.Column("invitation_token", sa.String(length=100), nullable=True)
        )
        batch_op.add_column(
            sa.Column("invitation_sent_at", sa.DateTime(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("invitation_accepted_at", sa.DateTime(), nullable=True)
        )

        # Create index on invitation_token
        batch_op.create_index("idx_vendor_users_invitation_token", ["invitation_token"])

        # Modify role_id to be nullable (this recreates the table in SQLite)
        batch_op.alter_column("role_id", existing_type=sa.Integer(), nullable=True)

        # Change is_active default (this recreates the table in SQLite)
        batch_op.alter_column(
            "is_active", existing_type=sa.Boolean(), server_default="false"
        )

    # Set owners correctly (after table modifications)
    # SQLite-compatible UPDATE with subquery
    op.execute(
        """
        UPDATE vendor_users 
        SET user_type = 'owner' 
        WHERE (vendor_id, user_id) IN (
            SELECT id, owner_user_id 
            FROM vendors
        )
    """
    )

    # Set existing owners as active
    op.execute(
        """
        UPDATE vendor_users 
        SET is_active = TRUE 
        WHERE user_type = 'owner'
    """
    )

    # ========================================================================
    # Role table changes
    # ========================================================================
    with op.batch_alter_table("roles", schema=None) as batch_op:
        # Create index on vendor_id and name
        batch_op.create_index("idx_roles_vendor_name", ["vendor_id", "name"])

        # Note: JSONB conversion only for PostgreSQL
        # SQLite stores JSON as TEXT by default, no conversion needed


def downgrade():
    """Downgrade database schema.

    SQLite-compatible version using batch operations.
    """

    # ========================================================================
    # Role table changes
    # ========================================================================
    with op.batch_alter_table("roles", schema=None) as batch_op:
        batch_op.drop_index("idx_roles_vendor_name")

    # ========================================================================
    # VendorUser table changes
    # ========================================================================
    with op.batch_alter_table("vendor_users", schema=None) as batch_op:
        # Revert is_active default
        batch_op.alter_column(
            "is_active", existing_type=sa.Boolean(), server_default="true"
        )

        # Revert role_id to NOT NULL
        # Note: This might fail if there are NULL values
        batch_op.alter_column("role_id", existing_type=sa.Integer(), nullable=False)

        # Drop indexes and columns
        batch_op.drop_index("idx_vendor_users_invitation_token")
        batch_op.drop_column("invitation_accepted_at")
        batch_op.drop_column("invitation_sent_at")
        batch_op.drop_column("invitation_token")
        batch_op.drop_column("user_type")

    # ========================================================================
    # User table changes
    # ========================================================================
    with op.batch_alter_table("users", schema=None) as batch_op:
        batch_op.drop_column("is_email_verified")