scripts/export_snapshot.py

"""Export the local DB to a portable snapshot for cloud import.

Writes:
    seed_data/snapshot/
    ├── metadata.jsonl        # one JSON object per sighting (with embedding)
    └── images/<key>          # original + cropped JPEGs, mirroring the bucket layout

The snapshot can later be loaded by `scripts.import_snapshot` against any
Postgres+pgvector + S3-compatible storage (e.g. Supabase).

Run inside the backend container:
    docker compose exec backend python -m scripts.export_snapshot
"""
import argparse
import io
import json
import logging
from pathlib import Path

from sqlalchemy import select

from app.config import get_settings
from app.db import SessionLocal
from app.models import Sighting
from app.services.storage import _client


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--root",
        type=Path,
        default=Path("/seed_data/snapshot"),
        help="Output directory (mounted from host).",
    )
    parser.add_argument(
        "--include-user",
        action="store_true",
        help="Also export user-uploaded sightings (default: filler + target only).",
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    log = logging.getLogger("export")

    settings = get_settings()
    out_root: Path = args.root
    images_dir = out_root / "images"
    images_dir.mkdir(parents=True, exist_ok=True)
    metadata_path = out_root / "metadata.jsonl"

    sources = ["filler", "target"]
    if args.include_user:
        sources.append("user")

    session = SessionLocal()
    client = _client()
    bucket = settings.minio_bucket

    total_bytes = 0
    counts: dict[str, int] = {}

    try:
        rows = session.scalars(
            select(Sighting)
            .where(Sighting.source.in_(sources))
            .order_by(Sighting.source, Sighting.identity, Sighting.id)
        ).all()
        log.info("Exporting %d sightings to %s...", len(rows), out_root)

        with metadata_path.open("w", encoding="utf-8") as out:
            for row in rows:
                counts[row.source] = counts.get(row.source, 0) + 1

                # Pull both image variants out of MinIO into the snapshot tree.
                for key in (row.image_url, row.cropped_url):
                    dst = images_dir / key
                    dst.parent.mkdir(parents=True, exist_ok=True)
                    if dst.exists():
                        continue
                    obj = client.get_object(bucket, key)
                    try:
                        data = obj.read()
                    finally:
                        obj.close()
                        obj.release_conn()
                    dst.write_bytes(data)
                    total_bytes += len(data)

                # Embedding may come back as a numpy array; coerce to plain list.
                embedding = row.embedding
                if hasattr(embedding, "tolist"):
                    embedding = embedding.tolist()
                else:
                    embedding = list(embedding)

                record = {
                    "id": str(row.id),
                    "source": row.source,
                    "identity": row.identity,
                    "image_key": row.image_url,
                    "cropped_key": row.cropped_url,
                    "embedding": embedding,
                    "latitude": row.latitude,
                    "longitude": row.longitude,
                    "sighted_at": row.sighted_at.isoformat(),
                    "notes": row.notes,
                }
                out.write(json.dumps(record) + "\n")

        log.info("---")
        log.info("Wrote %s", metadata_path)
        for src, n in sorted(counts.items()):
            log.info("  source=%-7s  rows=%d", src, n)
        log.info("Total image bytes: %.1f MB", total_bytes / (1024 * 1024))
    finally:
        session.close()


if __name__ == "__main__":
    main()