"""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())