from __future__ import annotations from datetime import UTC from datetime import datetime from pathlib import Path from uuid import uuid4 from sqlalchemy import create_engine from sqlalchemy import text from sqlalchemy import select from sqlmodel import SQLModel from transcription.db.migration import export_bundle from transcription.db.migration import import_bundle from transcription.db.migration import sqlite_url_from_path # Register table metadata. from transcription.db import models as _models # noqa: F401 def test_export_import_migration_round_trips_db_and_uploads(tmp_path): source_db_path = tmp_path / "source.db" target_db_path = tmp_path / "target.db" source_upload_dir = tmp_path / "source_uploads" target_upload_dir = tmp_path / "target_uploads" bundle_dir = tmp_path / "bundle" source_db_url = sqlite_url_from_path(source_db_path) target_db_url = sqlite_url_from_path(target_db_path) document_id = uuid4() job_id = uuid4() source_id = uuid4() job_source_id = uuid4() attempt_id = uuid4() filename = f"{source_id}.jpg" media_path = source_upload_dir / "documents" / str(document_id) / filename media_path.parent.mkdir(parents=True, exist_ok=True) media_path.write_bytes(b"sample-image") engine = create_engine(source_db_url) try: SQLModel.metadata.create_all(engine) with engine.begin() as connection: connection.execute( SQLModel.metadata.tables["document"].insert(), [{"id": document_id, "name": "Export doc"}], ) connection.execute( SQLModel.metadata.tables["job"].insert(), [{"id": job_id, "document_id": document_id, "status": "queued"}], ) connection.execute( SQLModel.metadata.tables["source"].insert(), [ { "id": source_id, "document_id": document_id, "page_number": 1, "upload_name": "upload.jpg", "filename": filename, "file_path": str(media_path), "file_hash": "a" * 64, "file_size_bytes": len(b"sample-image"), } ], ) connection.execute( SQLModel.metadata.tables["job_source"].insert(), [{"id": job_source_id, "job_id": job_id, "source_id": source_id, "status": "pending"}], ) connection.execute( SQLModel.metadata.tables["execution_attempt"].insert(), [ { "id": attempt_id, "job_source_id": job_source_id, "job_id": job_id, "source_id": source_id, "attempt_number": 1, "status": "transcribed", "provider": "fixture", "model": "fixture-model", "transport_body": b"body", "raw_transcription": "hello", "started_at": datetime.now(UTC), "finished_at": datetime.now(UTC), "duration_ms": 10, } ], ) finally: engine.dispose() export_bundle(source_db_url=source_db_url, source_upload_dir=source_upload_dir, bundle_dir=bundle_dir) import_bundle(target_db_url=target_db_url, target_upload_dir=target_upload_dir, bundle_dir=bundle_dir) target_engine = create_engine(target_db_url) try: with target_engine.connect() as connection: source_row = connection.execute( select(SQLModel.metadata.tables["source"].c.file_path).where( SQLModel.metadata.tables["source"].c.id == source_id ) ).one() attempt_row = connection.execute( select(SQLModel.metadata.tables["execution_attempt"].c.transport_body).where( SQLModel.metadata.tables["execution_attempt"].c.id == attempt_id ) ).one() assert source_row[0] == f"documents/{document_id}/{filename}" assert attempt_row[0] == b"body" finally: target_engine.dispose() copied_media_path = target_upload_dir / "documents" / str(document_id) / filename assert copied_media_path.read_bytes() == b"sample-image" def test_export_import_migration_backfills_legacy_portraits_and_homepage_images(tmp_path): source_db_path = tmp_path / "source-legacy.db" target_db_path = tmp_path / "target-legacy.db" source_upload_dir = tmp_path / "source_uploads" target_upload_dir = tmp_path / "target_uploads" bundle_dir = tmp_path / "bundle-legacy" source_db_url = sqlite_url_from_path(source_db_path) target_db_url = sqlite_url_from_path(target_db_path) person_id = "11" * 16 portrait_file = source_upload_dir / "persons" / "legacy" / "portrait.png" portrait_file.parent.mkdir(parents=True, exist_ok=True) portrait_file.write_bytes(b"portrait") homepage_file = source_upload_dir / "homepage" / "banner.jpg" homepage_file.parent.mkdir(parents=True, exist_ok=True) homepage_file.write_bytes(b"homepage") (source_upload_dir / "homepage" / "homepage.md").write_text("# Legacy Home", encoding="utf-8") engine = create_engine(source_db_url) try: with engine.begin() as connection: connection.execute( text( 'create table "person" (' "id char(32) primary key, " "full_name varchar not null, " "portrait_path varchar, " "created_at datetime not null, " "updated_at datetime not null" ")" ) ) connection.execute( text( 'insert into "person" (id, full_name, portrait_path, created_at, updated_at) ' 'values (:id, :full_name, :portrait_path, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)' ), {"id": person_id, "full_name": "Legacy Portrait", "portrait_path": "persons/legacy/portrait.png"}, ) finally: engine.dispose() export_bundle(source_db_url=source_db_url, source_upload_dir=source_upload_dir, bundle_dir=bundle_dir) import_bundle(target_db_url=target_db_url, target_upload_dir=target_upload_dir, bundle_dir=bundle_dir) target_engine = create_engine(target_db_url) try: with target_engine.connect() as connection: photos = connection.execute( text('select person_id, path, is_primary from "photo" order by person_id is not null desc, created_at asc') ).all() assert len(photos) == 2 person_photo = next(row for row in photos if row[0] is not None) homepage_photo = next(row for row in photos if row[0] is None) assert person_photo[2] == 1 assert homepage_photo[2] == 1 assert str(person_photo[1]).startswith("photos/") assert str(homepage_photo[1]).startswith("photos/") finally: target_engine.dispose() assert (target_upload_dir / str(person_photo[1])).read_bytes() == b"portrait" assert (target_upload_dir / str(homepage_photo[1])).read_bytes() == b"homepage" assert (target_upload_dir / "homepage.md").read_text(encoding="utf-8") == "# Legacy Home"