transcription/tests/tools/test_export_import_migration.py

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"