"""One-shot migration: copy PDFs from one storage backend to another. Designed for the local -> r2 case but works in both directions. Reads from the source backend, writes to the destination backend, and optionally rewrites books.local_path so subsequent code paths know the file is now hosted somewhere different. Usage: python -m src.lib.storage.migrate # local -> r2 (dry-run) python -m src.lib.storage.migrate --apply # actually copy python -m src.lib.storage.migrate --apply --update-paths Idempotent: skips any book that already exists at the destination (checked via Storage.exists). Re-run after a partial failure to finish. """ from __future__ import annotations import argparse import sys from src.stage1_inventory.db import connect from .factory import _resolve_backend_name # noqa: PLC2701 — internal but stable from .local import LocalStorage from .r2 import R2Storage def _make(name: str): if name == "local": return LocalStorage() if name == "r2": return R2Storage() raise ValueError(f"unknown backend {name!r}") def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser(description="Copy PDFs between storage backends.") p.add_argument("--from", dest="src", default="local", choices=("local", "r2")) p.add_argument("--to", dest="dst", default="r2", choices=("local", "r2")) p.add_argument("--apply", action="store_true", help="Actually copy. Without it, prints the plan only.") p.add_argument("--update-paths", action="store_true", help="Rewrite books.local_path to the destination URI.") p.add_argument("--book-id", default=None, help="Limit to one book; otherwise iterates every row in books.") args = p.parse_args(argv) if args.src == args.dst: print(f"src == dst ({args.src}); nothing to do.", file=sys.stderr) return 1 src = _make(args.src) # Defer destination construction — dry-run shouldn't need R2 creds just # to print a plan. Built lazily on first use below. dst = None # Pull the inventory list (authoritative source of book_ids). sql = "SELECT book_id, local_path FROM books" sql_args: list = [] if args.book_id: sql += " WHERE book_id = ?" sql_args.append(args.book_id) sql += " ORDER BY book_id" with connect() as conn: rows = conn.execute(sql, sql_args).fetchall() if not rows: print("No books in inventory.") return 0 print(f"{args.src} -> {args.dst}: {len(rows)} book(s) in inventory") print(f"backends.storage active = {_resolve_backend_name()!r}") print(f"--apply = {args.apply}") print(f"--update-paths = {args.update_paths}") print() copied = 0 skipped = 0 missing = 0 text_synced = 0 failed = 0 for r in rows: bid = r["book_id"] has_src_pdf = src.exists(bid) # Dry run: never construct the destination, so no creds needed. if not args.apply: verb = "would copy PDF" if has_src_pdf else "no source PDF (text only)" print(f" [plan] {bid}: {verb} {args.src} -> {args.dst}") continue if dst is None: dst = _make(args.dst) try: # --- PDF copy (only if the source actually has it) --- if not has_src_pdf: missing += 1 print(f" [skip-pdf] {bid}: no PDF in {args.src}") elif dst.exists(bid): skipped += 1 print(f" [done] {bid}: PDF already in {args.dst}") if args.update_paths: _update_path_to(bid, dst, args.dst) else: local_path = src.ensure_local(bid) uri = dst.put_pdf(bid, local_path) copied += 1 print(f" [copy] {bid}: -> {uri}") if args.update_paths: _update_path_to(bid, dst, args.dst, uri=uri) # --- Text bundle (independent of the PDF) --- # Always (re)push the OCR/clean text from local disk if present. # Idempotent, and decoupled from the PDF so a book whose source # PDF was already pruned locally still gets its text layer to the # cloud (the whole point of promoting a book is its text, too). for tu in dst.put_text(bid): text_synced += 1 print(f" [text] {bid}: -> {tu}") except Exception as e: # noqa: BLE001 failed += 1 print(f" [FAIL] {bid}: {type(e).__name__}: {e}", file=sys.stderr) print() print(f"copied={copied} skipped(already)={skipped} " f"missing(in src)={missing} text_synced={text_synced} failed={failed}") return 0 if failed == 0 else 2 def _update_path_to(book_id: str, dst, dst_name: str, *, uri: str | None = None) -> None: """Rewrite books.local_path so the inventory reflects the new home.""" if uri is None: # Recompute the URI without re-uploading. if dst_name == "local": from src.config import load_config cfg = load_config() uri = f"file://{(cfg.paths.raw_dir / f'{book_id}.pdf').resolve()}" else: # r2 uri = f"r2://{dst._bucket}/{book_id}.pdf" # noqa: SLF001 with connect() as conn: conn.execute( "UPDATE books SET local_path = ?, updated_at = datetime('now') WHERE book_id = ?", (uri, book_id), ) conn.commit() if __name__ == "__main__": raise SystemExit(main())