osdb-1-claude-opusplan/scripts/import_publications.py

#!/usr/bin/env python3
"""
Import publications from 03-data/osws.xml into the OSDB database.

Usage:
    DATABASE_URL=postgres://... python3 scripts/import_publications.py

Idempotent: upserts publication by osdb_pub_id; deletes and re-inserts
linked fruits, descriptions, and fruit images per publication on re-run.
"""

import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

import psycopg2

DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "03-data")


def parse_publications(xml_path: str) -> list:
    """
    Parse osws.xml and return list of publication dicts for <obj> elements with <osw>=1.
    Each dict: id, title, author (None if "."), img_path (relative to data root, or None).
    """
    tree = ET.parse(xml_path)
    root = tree.getroot()
    pubs = []
    for obj in root.findall("obj"):
        if (obj.findtext("osw") or "").strip() != "1":
            continue
        author_raw = (obj.findtext("author") or "").strip()
        img = obj.findtext("img")
        pubs.append({
            "id": (obj.findtext("id") or "").strip(),
            "title": (obj.findtext("name") or "").strip(),
            "author": author_raw if author_raw and author_raw != "." else None,
            "img_path": img.strip() if img else None,
        })
    return pubs


def scan_pub_dir(pub_dir: str, pub_id: str) -> dict:
    """
    Scan pub_dir for files matching {osdb_number}_{pub_id}_s0.jpg and {osdb_number}_{pub_id}.pdf.
    Returns dict mapping osdb_number → {img_path?: str, pdf_path?: str}.
    Ignores _tn.jpg thumbnails and any other files.
    """
    d = Path(pub_dir)
    if not d.is_dir():
        return {}

    result = {}
    img_re = re.compile(rf"^(.+)_{re.escape(pub_id)}_s0\.jpg$")
    pdf_re = re.compile(rf"^(.+)_{re.escape(pub_id)}\.pdf$")

    for f in d.iterdir():
        m = img_re.match(f.name)
        if m:
            result.setdefault(m.group(1), {})["img_path"] = str(f)
            continue
        m = pdf_re.match(f.name)
        if m:
            result.setdefault(m.group(1), {})["pdf_path"] = str(f)

    return result


def import_publications(conn, data_dir: str = None) -> dict:
    """
    Import all publications from osws.xml into the database.
    Returns counts: publications, covers, fruits_linked, pdfs, images, skipped_fruits.
    """
    if data_dir is None:
        data_dir = DATA_DIR

    data_path = Path(data_dir)
    pubs = parse_publications(str(data_path / "osws.xml"))

    counts = {
        "publications": 0,
        "covers": 0,
        "fruits_linked": 0,
        "pdfs": 0,
        "images": 0,
        "skipped_fruits": 0,
    }

    with conn.cursor() as cur:
        for pub in pubs:
            pub_id = pub["id"]

            # Load cover image bytes (skip silently if file missing or unreadable)
            cover_data = None
            if pub["img_path"]:
                cover_path = data_path / pub["img_path"]
                try:
                    cover_data = cover_path.read_bytes()
                    counts["covers"] += 1
                except OSError:
                    pass

            # Upsert publication row
            cur.execute(
                """
                INSERT INTO publications (title, author, osdb_pub_id, image_data, updated_at)
                VALUES (%s, %s, %s, %s, NOW())
                ON CONFLICT (osdb_pub_id) DO UPDATE
                  SET title      = EXCLUDED.title,
                      author     = EXCLUDED.author,
                      image_data = EXCLUDED.image_data,
                      updated_at = NOW()
                RETURNING id
                """,
                (
                    pub["title"],
                    pub["author"],
                    pub_id,
                    psycopg2.Binary(cover_data) if cover_data else None,
                ),
            )
            db_pub_id = cur.fetchone()[0]
            counts["publications"] += 1

            # Scan filesystem for linked fruits
            pub_dir = str(data_path / "osdb" / pub_id)
            fruit_files = scan_pub_dir(pub_dir, pub_id)

            # Resolve osdb_numbers → DB fruit IDs; warn and skip unknowns
            linked = {}
            for osdb_number, files in fruit_files.items():
                cur.execute("SELECT id FROM fruits WHERE osdb_number = %s", (osdb_number,))
                row = cur.fetchone()
                if row is None:
                    print(
                        f"  WARNING: osdb_number not in fruits table, skipping: {osdb_number}",
                        file=sys.stderr,
                    )
                    counts["skipped_fruits"] += 1
                    continue
                linked[osdb_number] = {"db_id": row[0], **files}

            # Idempotent: clear previous linked data for this publication
            cur.execute("DELETE FROM publication_fruit_images WHERE publication_id = %s", (db_pub_id,))
            cur.execute("DELETE FROM publication_descriptions  WHERE publication_id = %s", (db_pub_id,))
            cur.execute("DELETE FROM publication_fruits        WHERE publication_id = %s", (db_pub_id,))

            # Re-insert linked fruits, PDFs, and images
            for osdb_number, info in linked.items():
                fruit_db_id = info["db_id"]

                cur.execute(
                    "INSERT INTO publication_fruits (publication_id, fruit_id) VALUES (%s, %s)",
                    (db_pub_id, fruit_db_id),
                )
                counts["fruits_linked"] += 1

                if "pdf_path" in info:
                    pdf_data = Path(info["pdf_path"]).read_bytes()
                    cur.execute(
                        """
                        INSERT INTO publication_descriptions (publication_id, fruit_id, pdf_data)
                        VALUES (%s, %s, %s)
                        """,
                        (db_pub_id, fruit_db_id, psycopg2.Binary(pdf_data)),
                    )
                    counts["pdfs"] += 1

                if "img_path" in info:
                    img_data = Path(info["img_path"]).read_bytes()
                    filename = Path(info["img_path"]).name
                    cur.execute(
                        """
                        INSERT INTO publication_fruit_images
                          (publication_id, fruit_id, filename, data, image_type)
                        VALUES (%s, %s, %s, %s, 'fruit')
                        """,
                        (db_pub_id, fruit_db_id, filename, psycopg2.Binary(img_data)),
                    )
                    counts["images"] += 1

    conn.commit()
    return counts


def main():
    database_url = os.environ.get("DATABASE_URL")
    if not database_url:
        print("Error: DATABASE_URL environment variable not set", file=sys.stderr)
        sys.exit(1)

    print("Connecting to database...")
    conn = psycopg2.connect(database_url)
    try:
        print("Importing publications...")
        counts = import_publications(conn)
        print("\nDone.")
        print(f"  Publications upserted: {counts['publications']}")
        print(f"  Cover images loaded:   {counts['covers']}")
        print(f"  Fruits linked:         {counts['fruits_linked']}")
        print(f"  PDFs imported:         {counts['pdfs']}")
        print(f"  Fruit images imported: {counts['images']}")
        print(f"  Fruits skipped:        {counts['skipped_fruits']}")
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        conn.rollback()
        sys.exit(1)
    finally:
        conn.close()


if __name__ == "__main__":
    main()