UMBRA/backend/alembic/versions/023_auth_migration_users_sessions.py

"""Auth migration: create users + user_sessions tables, migrate pin_hash to User row,
add user_id FK to settings, drop pin_hash from settings.

Revision ID: 023
Revises: 022
Create Date: 2026-02-25

Data migration strategy (handles both fresh DB and existing single-user DB):
  1. Create users table
  2. Create user_sessions table
  3. If settings row exists with a pin_hash, create a User row with
     username='admin' and password_hash = pin_hash (bcrypt hash is valid —
     user will be prompted to change password on first login, hash upgrades
     transparently to Argon2id on first successful login).
  4. Add user_id FK column to settings (nullable initially)
  5. Backfill settings.user_id to point at the migrated admin user (if any)
  6. Drop pin_hash from settings
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql


# revision identifiers
revision = '023'
down_revision = '022'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ------------------------------------------------------------------
    # 1. Create users table
    # ------------------------------------------------------------------
    op.create_table(
        'users',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('username', sa.String(length=50), nullable=False),
        sa.Column('password_hash', sa.String(length=255), nullable=False),
        sa.Column('totp_secret', sa.String(length=500), nullable=True),
        sa.Column('totp_enabled', sa.Boolean(), nullable=False, server_default='false'),
        sa.Column('failed_login_count', sa.Integer(), nullable=False, server_default='0'),
        sa.Column('locked_until', sa.DateTime(), nullable=True),
        sa.Column('is_active', sa.Boolean(), nullable=False, server_default='true'),
        sa.Column('created_at', sa.DateTime(), nullable=False, server_default=sa.text('NOW()')),
        sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.text('NOW()')),
        sa.Column('last_login_at', sa.DateTime(), nullable=True),
        sa.PrimaryKeyConstraint('id'),
    )
    op.create_index('ix_users_id', 'users', ['id'], unique=False)
    op.create_index('ix_users_username', 'users', ['username'], unique=True)

    # ------------------------------------------------------------------
    # 2. Create user_sessions table
    # ------------------------------------------------------------------
    op.create_table(
        'user_sessions',
        sa.Column('id', sa.String(length=64), nullable=False),
        sa.Column('user_id', sa.Integer(), nullable=False),
        sa.Column('created_at', sa.DateTime(), nullable=False, server_default=sa.text('NOW()')),
        sa.Column('expires_at', sa.DateTime(), nullable=False),
        sa.Column('revoked', sa.Boolean(), nullable=False, server_default='false'),
        sa.Column('ip_address', sa.String(length=45), nullable=True),
        sa.Column('user_agent', sa.String(length=255), nullable=True),
        sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE'),
        sa.PrimaryKeyConstraint('id'),
    )
    op.create_index('ix_user_sessions_user_id', 'user_sessions', ['user_id'], unique=False)

    # ------------------------------------------------------------------
    # 3. Data migration: create admin User from existing pin_hash (if any)
    #    Uses raw SQL to avoid any ORM dependency issues in the migration.
    # ------------------------------------------------------------------
    bind = op.get_bind()

    # Check whether settings table has a row with a pin_hash
    result = bind.execute(
        sa.text("SELECT id, pin_hash FROM settings WHERE pin_hash IS NOT NULL LIMIT 1")
    )
    settings_row = result.fetchone()

    admin_user_id = None
    if settings_row is not None:
        settings_id, existing_pin_hash = settings_row[0], settings_row[1]

        # Insert the migrated user — username defaults to 'admin', password_hash
        # retains the existing bcrypt hash (will upgrade to Argon2id on first login)
        insert_result = bind.execute(
            sa.text(
                "INSERT INTO users (username, password_hash, is_active, totp_enabled, "
                "failed_login_count, created_at, updated_at) "
                "VALUES ('admin', :ph, true, false, 0, NOW(), NOW()) RETURNING id"
            ),
            {"ph": existing_pin_hash},
        )
        admin_user_id = insert_result.fetchone()[0]

    # ------------------------------------------------------------------
    # 4. Add user_id FK column to settings (nullable during migration)
    # ------------------------------------------------------------------
    op.add_column(
        'settings',
        sa.Column(
            'user_id',
            sa.Integer(),
            sa.ForeignKey('users.id', ondelete='CASCADE'),
            nullable=True,
        )
    )
    op.create_index('ix_settings_user_id', 'settings', ['user_id'], unique=False)

    # ------------------------------------------------------------------
    # 5. Backfill settings.user_id for the migrated admin user
    # ------------------------------------------------------------------
    if admin_user_id is not None:
        bind.execute(
            sa.text("UPDATE settings SET user_id = :uid"),
            {"uid": admin_user_id},
        )

    # ------------------------------------------------------------------
    # 6. Enforce NOT NULL on user_id now that backfill is complete
    # ------------------------------------------------------------------
    op.alter_column('settings', 'user_id', nullable=False)

    # ------------------------------------------------------------------
    # 7. Drop pin_hash from settings — data now lives in users.password_hash
    # ------------------------------------------------------------------
    op.drop_column('settings', 'pin_hash')


def downgrade() -> None:
    # Restore pin_hash column (empty — data cannot be recovered from users table
    # because the column was dropped, not just migrated; acceptable for downgrade path)
    op.add_column(
        'settings',
        sa.Column('pin_hash', sa.String(length=255), nullable=True)
    )

    op.drop_index('ix_settings_user_id', table_name='settings')
    op.drop_column('settings', 'user_id')

    op.drop_index('ix_user_sessions_user_id', table_name='user_sessions')
    op.drop_table('user_sessions')

    op.drop_index('ix_users_username', table_name='users')
    op.drop_index('ix_users_id', table_name='users')
    op.drop_table('users')