leocrm/alembic/versions/0037_user_tenant_model.py

"""User-Tenant model cleanup: single source of truth for membership and role.

Revision ID: 0037_user_tenant_model
Revises: 0036_cross_tenant_fk
Create Date: 2026-07-25

Changes:
1. Make users.email globally unique (drop composite uq_users_tenant_email, add UNIQUE on email).
2. Drop tenant_id, role, role_id columns from users table (with data migration to user_tenants).
3. Add role column to user_tenants (built-in role string: admin/editor/viewer).
4. Add status column to user_tenants (active/invited/disabled).
5. Migrate existing data: copy users.tenant_id + users.role_id → user_tenants (if not already present).
"""

from __future__ import annotations

from typing import Union

import sqlalchemy as sa
from alembic import op

# revision identifiers
revision: str = "0037_user_tenant_model"
down_revision: Union[str, None] = "0036_cross_tenant_fk"
branch_labels: Union[str, None] = None
depends_on: Union[str, None] = None


def _constraint_exists(name: str, table: str) -> str:
    """Return SQL that checks if a constraint exists on a table."""
    return (
        f"SELECT 1 FROM information_schema.table_constraints "
        f"WHERE constraint_name = '{name}' AND table_name = '{table}'"
    )


def _column_exists(table: str, column: str) -> str:
    """Return SQL that checks if a column exists on a table."""
    return (
        f"SELECT 1 FROM information_schema.columns "
        f"WHERE table_name = '{table}' AND column_name = '{column}'"
    )


def upgrade() -> None:
    conn = op.get_bind()

    # ── 1. Add UNIQUE constraint on users.email (globally unique) ───────────
    # First check if a unique constraint on email already exists
    email_unique_result = conn.execute(
        sa.text(
            "SELECT 1 FROM information_schema.table_constraints "
            "WHERE constraint_name = 'uq_users_email' AND table_name = 'users'"
        )
    ).fetchone()
    if email_unique_result is None:
        # Check if there's a unique index on email already
        email_index_result = conn.execute(
            sa.text(
                "SELECT 1 FROM pg_index i "
                "JOIN pg_class c ON i.indexrelid = c.oid "
                "WHERE c.relname = 'ix_users_email' AND i.indisunique = true"
            )
        ).fetchone()
        if email_index_result is None:
            op.execute("ALTER TABLE users ADD CONSTRAINT uq_users_email UNIQUE (email)")

    # ── 2. Drop composite unique constraint uq_users_tenant_email ───────────
    result = conn.execute(sa.text(_constraint_exists("uq_users_tenant_email", "users"))).fetchone()
    if result is not None:
        op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS uq_users_tenant_email")

    # ── 3. Add role column to user_tenants ─────────────────────────────────
    role_col_result = conn.execute(sa.text(_column_exists("user_tenants", "role"))).fetchone()
    if role_col_result is None:
        op.add_column("user_tenants", sa.Column("role", sa.String(50), nullable=False, server_default="viewer"))

    # ── 4. Add status column to user_tenants ───────────────────────────────
    status_col_result = conn.execute(sa.text(_column_exists("user_tenants", "status"))).fetchone()
    if status_col_result is None:
        op.add_column("user_tenants", sa.Column("status", sa.String(20), nullable=False, server_default="active"))

    # ── 4b. Add updated_at column to user_tenants ──────────────────────────
    updated_col_result = conn.execute(sa.text(_column_exists("user_tenants", "updated_at"))).fetchone()
    if updated_col_result is None:
        op.add_column("user_tenants", sa.Column("updated_at", sa.DateTime(timezone=True), nullable=True, server_default=sa.func.now()))

    # ── 5. Data migration: copy tenant_id, role, role_id from users to user_tenants ─
    # Only create UserTenant rows that don't already exist
    conn.execute(sa.text("""
        INSERT INTO user_tenants (user_id, tenant_id, is_default, role, role_id, status, created_at)
        SELECT 
            u.id, 
            u.tenant_id, 
            TRUE, 
            COALESCE(u.role, 'viewer'), 
            u.role_id,
            'active',
            NOW()
        FROM users u
        WHERE NOT EXISTS (
            SELECT 1 FROM user_tenants ut 
            WHERE ut.user_id = u.id AND ut.tenant_id = u.tenant_id
        )
        AND u.tenant_id IS NOT NULL
    """))

    # Update existing UserTenant rows with role from users table (if they don't have one set yet)
    conn.execute(sa.text("""
        UPDATE user_tenants ut
        SET role = COALESCE(u.role, 'viewer'),
            role_id = COALESCE(ut.role_id, u.role_id)
        FROM users u
        WHERE ut.user_id = u.id
        AND ut.tenant_id = u.tenant_id
    """))

    # ── 6. Drop role_id FK from users (if it exists) ───────────────────────
    # Find and drop the FK on users.role_id
    role_id_fk_result = conn.execute(
        sa.text(
            "SELECT conname FROM pg_constraint c "
            "JOIN pg_class cls ON c.conrelid = cls.oid "
            "JOIN pg_namespace nsp ON c.connamespace = nsp.oid "
            "WHERE cls.relname = 'users' "
            "AND nsp.nspname = 'public' "
            "AND c.contype = 'f' "
            "AND EXISTS ("
            "  SELECT 1 FROM pg_attribute a "
            "  WHERE a.attrelid = c.conrelid AND a.attname = 'role_id' "
            "  AND a.attnum = ANY(c.conkey)"
            ")"
        )
    ).fetchone()
    if role_id_fk_result is not None:
        fk_name = role_id_fk_result[0]
        op.execute(f"ALTER TABLE users DROP CONSTRAINT IF EXISTS {fk_name}")

    # ── 7. Drop tenant_id, role, role_id columns from users ────────────────
    # Drop tenant_id
    tenant_col_result = conn.execute(sa.text(_column_exists("users", "tenant_id"))).fetchone()
    if tenant_col_result is not None:
        # Drop RLS policy that depends on tenant_id
        op.execute("DROP POLICY IF EXISTS tenant_isolation ON users")
        # Drop any indexes on tenant_id first
        op.execute("DROP INDEX IF EXISTS ix_users_tenant_id")
        op.drop_column("users", "tenant_id")

    # Drop role
    role_col_result = conn.execute(sa.text(_column_exists("users", "role"))).fetchone()
    if role_col_result is not None:
        op.drop_column("users", "role")

    # Drop role_id
    role_id_col_result = conn.execute(sa.text(_column_exists("users", "role_id"))).fetchone()
    if role_id_col_result is not None:
        op.execute("DROP INDEX IF EXISTS ix_users_role_id")
        op.drop_column("users", "role_id")


def downgrade() -> None:
    conn = op.get_bind()

    # ── Re-add tenant_id, role, role_id to users ───────────────────────────
    tenant_col_result = conn.execute(sa.text(_column_exists("users", "tenant_id"))).fetchone()
    if tenant_col_result is None:
        op.add_column("users", sa.Column("tenant_id", sa.dialects.postgresql.UUID(as_uuid=True), nullable=True))
        op.execute('CREATE INDEX IF NOT EXISTS ix_users_tenant_id ON users (tenant_id)')

    role_col_result = conn.execute(sa.text(_column_exists("users", "role"))).fetchone()
    if role_col_result is None:
        op.add_column("users", sa.Column("role", sa.String(50), nullable=False, server_default="viewer"))

    role_id_col_result = conn.execute(sa.text(_column_exists("users", "role_id"))).fetchone()
    if role_id_col_result is None:
        op.add_column("users", sa.Column("role_id", sa.dialects.postgresql.UUID(as_uuid=True), nullable=True))
        op.execute('CREATE INDEX IF NOT EXISTS ix_users_role_id ON users (role_id)')
        # Re-add FK
        op.create_foreign_key("fk_users_role_id", "users", "roles", ["role_id"], ["id"], ondelete="SET NULL")

    # ── Restore data from user_tenants to users (default tenant) ────────────
    conn.execute(sa.text("""
        UPDATE users u
        SET tenant_id = ut.tenant_id,
            role = ut.role,
            role_id = ut.role_id
        FROM user_tenants ut
        WHERE ut.user_id = u.id AND ut.is_default = TRUE
    """))

    # ── Re-add composite unique constraint ─────────────────────────────────
    op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS uq_users_email")
    op.execute("ALTER TABLE users ADD CONSTRAINT uq_users_tenant_email UNIQUE (tenant_id, email)")

    # ── Drop role and status columns from user_tenants ──────────────────────
    op.drop_column("user_tenants", "status")
    op.drop_column("user_tenants", "role")