mule-image/backend/alembic/versions/0009_users_and_auth.py

"""users table and user_id foreign keys

Revision ID: 0009_users_and_auth
Revises: 0008_photos_date_warning
Create Date: 2026-04-12

Introduces multi-user support:
  1. Creates the `users` table.
  2. Adds `user_id` FK columns to photos, folders, source_roots, heaps, tags.
  3. For existing installs: creates a default admin user and assigns all
     existing rows to that user. The generated password is printed to the
     backend logs — the admin should change it on first login.
  4. Replaces the unique constraint on tags (name, kind) with
     (name, kind, user_id) so each user can have their own tags.
"""
from typing import Sequence, Union
import uuid
import secrets

from alembic import op
import sqlalchemy as sa

revision: str = "0009_users_and_auth"
down_revision: Union[str, None] = "0008_photos_date_warning"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


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

    # 1. Create users table (IF NOT EXISTS — safe on fresh installs where
    #    init_db's create_all has already laid down the schema).
    conn.execute(sa.text("""
        CREATE TABLE IF NOT EXISTS users (
            id VARCHAR NOT NULL PRIMARY KEY,
            username VARCHAR(50) NOT NULL UNIQUE,
            email VARCHAR UNIQUE,
            hashed_password VARCHAR NOT NULL,
            role VARCHAR NOT NULL DEFAULT 'user',
            is_active BOOLEAN NOT NULL DEFAULT true,
            created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now(),
            media_path VARCHAR NOT NULL
        )
    """))
    conn.execute(sa.text(
        "CREATE INDEX IF NOT EXISTS ix_users_username ON users (username)"
    ))

    # 2. Add user_id columns (nullable initially for the data migration)
    for table in ("photos", "folders", "source_roots", "heaps", "tags"):
        conn.execute(sa.text(
            f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS user_id VARCHAR"
        ))
        conn.execute(sa.text(
            f"CREATE INDEX IF NOT EXISTS ix_{table}_user_id ON {table} (user_id)"
        ))
        # FK — check if it already exists before adding
        fk_name = f"fk_{table}_user_id"
        fk_exists = conn.execute(sa.text(
            "SELECT 1 FROM information_schema.table_constraints "
            "WHERE constraint_name = :name AND table_name = :tbl"
        ), {"name": fk_name, "tbl": table}).scalar()
        if not fk_exists:
            conn.execute(sa.text(
                f"ALTER TABLE {table} ADD CONSTRAINT {fk_name} "
                f"FOREIGN KEY (user_id) REFERENCES users(id)"
            ))

    # 3. Data migration: if rows exist, create a default admin and assign
    conn = op.get_bind()
    photo_count = conn.execute(sa.text("SELECT COUNT(*) FROM photos")).scalar()

    if photo_count > 0:
        admin_id = str(uuid.uuid4())
        generated_password = secrets.token_urlsafe(16)

        # Hash the password using passlib at migration time
        from passlib.context import CryptContext
        pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
        hashed = pwd_context.hash(generated_password)

        # Every user gets a subfolder — including the migrated admin.
        conn.execute(
            sa.text(
                "INSERT INTO users (id, username, hashed_password, role, media_path) "
                "VALUES (:id, :username, :hashed, :role, :media_path)"
            ),
            {
                "id": admin_id,
                "username": "admin",
                "hashed": hashed,
                "role": "admin",
                "media_path": "/photos/admin",
            },
        )

        # Assign all existing rows to the default admin
        for table in ("photos", "folders", "source_roots", "heaps", "tags"):
            conn.execute(
                sa.text(f"UPDATE {table} SET user_id = :uid WHERE user_id IS NULL"),
                {"uid": admin_id},
            )

        import logging
        logger = logging.getLogger("alembic.migration")
        logger.warning(
            f"=== MIGRATION 0009 === Default admin created. "
            f"Username: admin | Password: {generated_password} | "
            f"Change this password on first login!"
        )

    # 4. Replace tag unique constraint to include user_id
    # Check whether the old constraint exists before trying to drop it
    # (on fresh installs create_all creates the new constraint directly).
    old_uq_exists = conn.execute(sa.text(
        "SELECT 1 FROM information_schema.table_constraints "
        "WHERE constraint_name = 'uq_tags_name_kind' AND table_name = 'tags'"
    )).scalar()
    if old_uq_exists:
        op.drop_constraint("uq_tags_name_kind", "tags", type_="unique")

    new_uq_exists = conn.execute(sa.text(
        "SELECT 1 FROM information_schema.table_constraints "
        "WHERE constraint_name = 'uq_tags_name_kind_user' AND table_name = 'tags'"
    )).scalar()
    if not new_uq_exists:
        op.create_unique_constraint("uq_tags_name_kind_user", "tags", ["name", "kind", "user_id"])


def downgrade() -> None:
    # Reverse the tag constraint
    op.drop_constraint("uq_tags_name_kind_user", "tags", type_="unique")
    op.create_unique_constraint("uq_tags_name_kind", "tags", ["name", "kind"])

    # Drop user_id columns and FKs
    for table in ("photos", "folders", "source_roots", "heaps", "tags"):
        op.drop_constraint(f"fk_{table}_user_id", table, type_="foreignkey")
        op.drop_index(f"ix_{table}_user_id", table_name=table)
        op.drop_column(table, "user_id")

    # Drop users table
    op.drop_index("ix_users_username", table_name="users")
    op.drop_table("users")