V6 Phase 2 complete

src/transcription/db/migration.py

@@ -16,9 +16,12 @@ from uuid import uuid4
 from sqlalchemy import URL
 from sqlalchemy import MetaData
 from sqlalchemy import Table
+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.engine import RowMapping
 from sqlalchemy.engine import make_url
 from sqlmodel import SQLModel
@@ -47,6 +50,7 @@ EXPORT_TABLE_ORDER = (
 )
 
 BYTES_FIELDS = {"transport_body"}
+VERIFICATION_TABLES = EXPORT_TABLE_ORDER
 
 
 @dataclass(frozen=True)
@@ -142,21 +146,74 @@ def import_bundle(*, target_db_url: str, target_upload_dir: Path, bundle_dir: Pa
     engine = create_engine(target_db_url)
     try:
         SQLModel.metadata.create_all(engine)
+        execution_attempt_ids = _collect_execution_attempt_ids(payload)
         with engine.begin() as connection:
             for table_name in reversed(EXPORT_TABLE_ORDER):
                 table = SQLModel.metadata.tables[table_name]
                 connection.execute(table.delete())
 
+            deferred_source_preferred_attempt_updates: list[dict[str, Any]] = []
             for table_name in EXPORT_TABLE_ORDER:
                 rows = payload.get("tables", {}).get(table_name, [])
                 if not rows:
                     continue
                 table = SQLModel.metadata.tables[table_name]
+                if table_name == "source":
+                    prepared_source_rows, updates = _prepare_source_rows_for_import(
+                        rows=rows,
+                        source_table=table,
+                        execution_attempt_ids=execution_attempt_ids,
+                    )
+                    deferred_source_preferred_attempt_updates.extend(updates)
+                    connection.execute(table.insert(), prepared_source_rows)
+                    continue
                 connection.execute(table.insert(), [_deserialize_row(row, table) for row in rows])
+
+            if deferred_source_preferred_attempt_updates:
+                source_table = SQLModel.metadata.tables["source"]
+                connection.execute(
+                    source_table.update()
+                    .where(source_table.c.id == bindparam("source_id"))
+                    .values(preferred_execution_attempt_id=bindparam("preferred_execution_attempt_id")),
+                    deferred_source_preferred_attempt_updates,
+                )
     finally:
         engine.dispose()
 
 
+def _collect_execution_attempt_ids(payload: dict[str, Any]) -> set[str]:
+    execution_attempt_rows = payload.get("tables", {}).get("execution_attempt", [])
+    return {_normalize_uuid_like(row.get("id")) for row in execution_attempt_rows if row.get("id") is not None}
+
+
+def _prepare_source_rows_for_import(
+    *,
+    rows: list[dict[str, Any]],
+    source_table: Table,
+    execution_attempt_ids: set[str],
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    prepared_source_rows: list[dict[str, Any]] = []
+    updates: list[dict[str, Any]] = []
+    for row in rows:
+        source_row = _deserialize_row(row, source_table)
+        source_id = source_row.get("id")
+        preferred_attempt_id = source_row.get("preferred_execution_attempt_id")
+        if (
+            source_id is not None
+            and preferred_attempt_id is not None
+            and _normalize_uuid_like(preferred_attempt_id) in execution_attempt_ids
+        ):
+            updates.append(
+                {
+                    "source_id": source_id,
+                    "preferred_execution_attempt_id": preferred_attempt_id,
+                }
+            )
+        source_row["preferred_execution_attempt_id"] = None
+        prepared_source_rows.append(source_row)
+    return prepared_source_rows, updates
+
+
 def _ensure_sqlite_target_parent_exists(target_db_url: str) -> None:
     parsed = make_url(target_db_url)
     if not parsed.drivername.startswith("sqlite"):
@@ -192,6 +249,24 @@ def migrate_via_bundle(paths: MigrationPaths) -> None:
     )
 
 
+@dataclass(frozen=True)
+class MigrationVerificationReport:
+    source_counts: dict[str, int]
+    target_counts: dict[str, int]
+    mismatched_tables: dict[str, dict[str, int]]
+    integrity_violations: dict[str, int]
+    success: bool
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "success": self.success,
+            "source_counts": self.source_counts,
+            "target_counts": self.target_counts,
+            "mismatched_tables": self.mismatched_tables,
+            "integrity_violations": self.integrity_violations,
+        }
+
+
 def sqlite_url_from_path(path: Path) -> str:
     return URL.create(drivername="sqlite", database=str(path)).render_as_string(hide_password=False)
 
@@ -201,6 +276,87 @@ def default_sync_db_url(settings: Settings | None = None) -> str:
     return get_database_url(runtime_settings).replace("+aiosqlite", "").replace("+asyncpg", "").replace("+psycopg", "")
 
 
+def verify_migration(*, source_db_url: str, target_db_url: str) -> MigrationVerificationReport:
+    source_counts = _table_counts(source_db_url)
+    target_counts = _table_counts(target_db_url)
+    mismatched_tables = {
+        table_name: {"source": source_counts[table_name], "target": target_counts[table_name]}
+        for table_name in VERIFICATION_TABLES
+        if source_counts[table_name] != target_counts[table_name]
+    }
+    integrity_violations = _integrity_violations(target_db_url)
+    success = not mismatched_tables and all(count == 0 for count in integrity_violations.values())
+    return MigrationVerificationReport(
+        source_counts=source_counts,
+        target_counts=target_counts,
+        mismatched_tables=mismatched_tables,
+        integrity_violations=integrity_violations,
+        success=success,
+    )
+
+
+def _table_counts(db_url: str) -> dict[str, int]:
+    engine = create_engine(db_url)
+    try:
+        metadata = MetaData()
+        metadata.reflect(bind=engine)
+        counts: dict[str, int] = {}
+        with engine.connect() as connection:
+            for table_name in VERIFICATION_TABLES:
+                table = metadata.tables.get(table_name)
+                if table is None:
+                    counts[table_name] = 0
+                    continue
+                counts[table_name] = int(connection.execute(select(func.count()).select_from(table)).scalar_one())
+        return counts
+    finally:
+        engine.dispose()
+
+
+def _integrity_violations(db_url: str) -> dict[str, int]:
+    checks = {
+        "orphan_source_document": (
+            "select count(*) from source s left join document d on d.id = s.document_id where d.id is null"
+        ),
+        "orphan_job_document": (
+            "select count(*) from job j left join document d on d.id = j.document_id where d.id is null"
+        ),
+        "orphan_job_source_job": (
+            "select count(*) from job_source js left join job j on j.id = js.job_id where j.id is null"
+        ),
+        "orphan_job_source_source": (
+            "select count(*) from job_source js left join source s on s.id = js.source_id where s.id is null"
+        ),
+        "orphan_attempt_job_source": (
+            "select count(*) from execution_attempt ea "
+            "left join job_source js on js.id = ea.job_source_id "
+            "where js.id is null"
+        ),
+        "orphan_attempt_job": (
+            "select count(*) from execution_attempt ea left join job j on j.id = ea.job_id where j.id is null"
+        ),
+        "orphan_attempt_source": (
+            "select count(*) from execution_attempt ea left join source s on s.id = ea.source_id where s.id is null"
+        ),
+        "duplicate_attempt_numbers": (
+            "select count(*) from ("
+            "  select job_id, source_id, attempt_number, count(*) as c"
+            "  from execution_attempt"
+            "  group by job_id, source_id, attempt_number"
+            "  having count(*) > 1"
+            ") x"
+        ),
+    }
+    engine = create_engine(db_url)
+    try:
+        with engine.connect() as connection:
+            return {
+                check_name: int(connection.execute(text(query)).scalar_one()) for check_name, query in checks.items()
+            }
+    finally:
+        engine.dispose()
+
+
 def _serialize_row(row: RowMapping, *, table_name: str, source_upload_dir: Path) -> dict[str, Any]:
     serialized: dict[str, Any] = {}
     for raw_key, value in row.items():
@@ -283,6 +439,18 @@ def _deserialize_value(python_type: type[Any], value: Any) -> Any:
     return value
 
 
+def _normalize_uuid_like(value: Any) -> str:
+    if isinstance(value, UUID):
+        return str(value)
+    if isinstance(value, str):
+        text_value = value.strip()
+        try:
+            return str(UUID(text_value))
+        except ValueError:
+            return text_value
+    return str(value)
+
+
 def _canonical_media_relative_path(value: str, *, source_upload_dir: Path, preferred_prefix: str) -> str:
     normalized = value.strip().replace("\\", "/")
     lowered = normalized.casefold()