calminer/alembic/versions/20251109_02_add_auth_tables.py

"""Add authentication and RBAC tables"""

from __future__ import annotations

from alembic import op
import sqlalchemy as sa
from passlib.context import CryptContext
from sqlalchemy.sql import column, table

# revision identifiers, used by Alembic.
revision = "20251109_02"
down_revision = "20251109_01"
branch_labels = None
depends_on = None

password_context = CryptContext(schemes=["argon2"], deprecated="auto")


def upgrade() -> None:
    op.create_table(
        "users",
        sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("email", sa.String(length=255), nullable=False),
        sa.Column("username", sa.String(length=128), nullable=False),
        sa.Column("password_hash", sa.String(length=255), nullable=False),
        sa.Column(
            "is_active",
            sa.Boolean(),
            nullable=False,
            server_default=sa.true(),
        ),
        sa.Column(
            "is_superuser",
            sa.Boolean(),
            nullable=False,
            server_default=sa.false(),
        ),
        sa.Column("last_login_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.UniqueConstraint("email", name="uq_users_email"),
        sa.UniqueConstraint("username", name="uq_users_username"),
    )
    op.create_index(
        "ix_users_active_superuser",
        "users",
        ["is_active", "is_superuser"],
        unique=False,
    )

    op.create_table(
        "roles",
        sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("name", sa.String(length=64), nullable=False),
        sa.Column("display_name", sa.String(length=128), nullable=False),
        sa.Column("description", sa.Text(), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.UniqueConstraint("name", name="uq_roles_name"),
    )

    op.create_table(
        "user_roles",
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("role_id", sa.Integer(), nullable=False),
        sa.Column(
            "granted_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.Column("granted_by", sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ["user_id"],
            ["users.id"],
            ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["role_id"],
            ["roles.id"],
            ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["granted_by"],
            ["users.id"],
            ondelete="SET NULL",
        ),
        sa.PrimaryKeyConstraint("user_id", "role_id"),
        sa.UniqueConstraint("user_id", "role_id",
                            name="uq_user_roles_user_role"),
    )
    op.create_index(
        "ix_user_roles_role_id",
        "user_roles",
        ["role_id"],
        unique=False,
    )

    # Seed default roles
    roles_table = table(
        "roles",
        column("id", sa.Integer()),
        column("name", sa.String()),
        column("display_name", sa.String()),
        column("description", sa.Text()),
    )

    op.bulk_insert(
        roles_table,
        [
            {
                "id": 1,
                "name": "admin",
                "display_name": "Administrator",
                "description": "Full platform access with user management rights.",
            },
            {
                "id": 2,
                "name": "project_manager",
                "display_name": "Project Manager",
                "description": "Manage projects, scenarios, and associated data.",
            },
            {
                "id": 3,
                "name": "analyst",
                "display_name": "Analyst",
                "description": "Review dashboards and scenario outputs.",
            },
            {
                "id": 4,
                "name": "viewer",
                "display_name": "Viewer",
                "description": "Read-only access to assigned projects and reports.",
            },
        ],
    )

    admin_password_hash = password_context.hash("ChangeMe123!")

    users_table = table(
        "users",
        column("id", sa.Integer()),
        column("email", sa.String()),
        column("username", sa.String()),
        column("password_hash", sa.String()),
        column("is_active", sa.Boolean()),
        column("is_superuser", sa.Boolean()),
    )

    op.bulk_insert(
        users_table,
        [
            {
                "id": 1,
                "email": "admin@calminer.local",
                "username": "admin",
                "password_hash": admin_password_hash,
                "is_active": True,
                "is_superuser": True,
            }
        ],
    )

    user_roles_table = table(
        "user_roles",
        column("user_id", sa.Integer()),
        column("role_id", sa.Integer()),
        column("granted_by", sa.Integer()),
    )

    op.bulk_insert(
        user_roles_table,
        [
            {
                "user_id": 1,
                "role_id": 1,
                "granted_by": 1,
            }
        ],
    )


def downgrade() -> None:
    op.drop_index("ix_user_roles_role_id", table_name="user_roles")
    op.drop_table("user_roles")

    op.drop_table("roles")

    op.drop_index("ix_users_active_superuser", table_name="users")
    op.drop_table("users")