transcription/tools/migrate_v47_to_v48.py

"""One-time migration for V4.8 queue-membership uniqueness.

This migration enforces the V4 requirement that each ``(job_id, source_id)``
pair appears at most once in ``job_source``.

Steps:
1. Detect duplicate ``job_source`` rows per ``(job_id, source_id)``.
2. Keep one row per pair (prefer the row with the latest attempt evidence).
3. Re-point ``execution_attempt.job_source_id`` from removed duplicates to the
   kept row.
4. Delete duplicate ``job_source`` rows.
5. Add uniqueness enforcement for ``(job_id, source_id)``.

Usage::

    python tools/migrate_v47_to_v48.py --dry-run
    python tools/migrate_v47_to_v48.py
"""

from __future__ import annotations

import argparse
from collections import defaultdict
from collections.abc import Sequence
from datetime import UTC
from datetime import datetime
from uuid import UUID

from sqlalchemy import bindparam
from sqlalchemy import create_engine
from sqlalchemy import func
from sqlalchemy import inspect as sqlalchemy_inspect
from sqlalchemy import select
from sqlalchemy import text
from sqlalchemy import update
from sqlalchemy.engine import Connection
from sqlmodel import SQLModel

from transcription.config import Settings
from transcription.config import get_settings
from transcription.db import models as _models  # noqa: F401  (registers metadata tables)
from transcription.db.engine import get_database_url

UNIQUE_NAME = "uq_job_source_job_source"


def _sync_url(settings: Settings) -> str:
    """Return a sync URL for direct SQLAlchemy Core access."""
    return get_database_url(settings).replace("+aiosqlite", "").replace("+asyncpg", "").replace("+psycopg", "")


def _uniqueness_already_enforced(connection: Connection) -> bool:
    inspector = sqlalchemy_inspect(connection)
    unique_constraints = inspector.get_unique_constraints("job_source")
    if any(constraint.get("name") == UNIQUE_NAME for constraint in unique_constraints):
        return True
    indexes = inspector.get_indexes("job_source")
    return any(index.get("name") == UNIQUE_NAME and index.get("unique") is True for index in indexes)


def _choose_keeper(
    candidates: list[dict[str, object]],
) -> dict[str, object]:
    """Keep the row with latest attempt activity, then lexicographically greatest UUID."""

    def key(item: dict[str, object]) -> tuple[datetime, str]:
        latest_attempt = item["latest_attempt_at"]
        attempt_key = latest_attempt if isinstance(latest_attempt, datetime) else datetime.min.replace(tzinfo=UTC)
        return (attempt_key, str(item["job_source_id"]))

    return max(candidates, key=key)


def deduplicate_job_source_membership(connection: Connection, *, dry_run: bool) -> tuple[int, int]:
    """Return ``(pairs_deduplicated, rows_deleted)``."""
    job_source = SQLModel.metadata.tables["job_source"]
    execution_attempt = SQLModel.metadata.tables["execution_attempt"]

    duplicate_pairs = (
        connection.execute(
            select(job_source.c.job_id, job_source.c.source_id)
            .group_by(job_source.c.job_id, job_source.c.source_id)
            .having(func.count(job_source.c.id) > 1)
        )
        .mappings()
        .all()
    )
    if not duplicate_pairs:
        return (0, 0)

    rows_by_pair: dict[tuple[UUID, UUID], list[dict[str, object]]] = defaultdict(list)
    duplicate_memberships = (
        connection.execute(
            select(
                job_source.c.id.label("job_source_id"),
                job_source.c.job_id,
                job_source.c.source_id,
                func.max(execution_attempt.c.created_at).label("latest_attempt_at"),
            )
            .select_from(job_source.outerjoin(execution_attempt, execution_attempt.c.job_source_id == job_source.c.id))
            .group_by(job_source.c.id, job_source.c.job_id, job_source.c.source_id)
        )
        .mappings()
        .all()
    )
    duplicate_key_set = {(row["job_id"], row["source_id"]) for row in duplicate_pairs}
    for row in duplicate_memberships:
        key = (row["job_id"], row["source_id"])
        if key in duplicate_key_set:
            rows_by_pair[key].append(dict(row))

    rows_deleted = 0
    for key, rows in sorted(rows_by_pair.items(), key=lambda item: (str(item[0][0]), str(item[0][1]))):
        keeper = _choose_keeper(rows)
        keeper_id = keeper["job_source_id"]
        duplicate_ids = [row["job_source_id"] for row in rows if row["job_source_id"] != keeper_id]
        rows_deleted += len(duplicate_ids)
        print(
            f"  dedupe pair job_id={key[0]} source_id={key[1]} "
            f"keep={keeper_id} drop={','.join(str(item) for item in duplicate_ids)}"
        )
        if dry_run or not duplicate_ids:
            continue
        connection.execute(
            update(execution_attempt)
            .where(execution_attempt.c.job_source_id.in_(bindparam("duplicate_ids", expanding=True)))
            .values(job_source_id=keeper_id),
            {"duplicate_ids": duplicate_ids},
        )
        connection.execute(
            job_source.delete().where(job_source.c.id.in_(bindparam("duplicate_ids", expanding=True))),
            {"duplicate_ids": duplicate_ids},
        )

    return (len(duplicate_pairs), rows_deleted)


def add_job_source_uniqueness(connection: Connection, *, dry_run: bool) -> None:
    """Create uniqueness enforcement for ``(job_id, source_id)``."""
    if _uniqueness_already_enforced(connection):
        print(f"  uniqueness already enforced ({UNIQUE_NAME})")
        return

    dialect = connection.dialect.name
    if dialect == "postgresql":
        statement = text(f'alter table "job_source" add constraint "{UNIQUE_NAME}" unique ("job_id", "source_id")')
    else:
        statement = text(f'create unique index "{UNIQUE_NAME}" on "job_source" ("job_id", "source_id")')
    print(f"  applying {UNIQUE_NAME}")
    if not dry_run:
        connection.execute(statement)


def migrate(*, settings: Settings, dry_run: bool) -> None:
    engine = create_engine(_sync_url(settings))
    try:
        with engine.begin() as connection:
            print("Step 1: deduplicate job_source membership")
            pair_count, rows_deleted = deduplicate_job_source_membership(connection, dry_run=dry_run)
            print(f"  duplicate pairs={pair_count} rows_deleted={rows_deleted}")

            print("Step 2: enforce unique (job_id, source_id)")
            add_job_source_uniqueness(connection, dry_run=dry_run)
    finally:
        engine.dispose()

    if dry_run:
        print("\nDry run: nothing was written.")
    else:
        print("\nDone.")


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dry-run", action="store_true", help="Report planned changes without writing")
    args = parser.parse_args(argv)

    settings = get_settings()
    print(f"Target: {_sync_url(settings)}")
    migrate(settings=settings, dry_run=args.dry_run)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())