stack/dev/scripts/zotero_archive_migrate.py

"""Migrate host Zotero library into container Zotero under an 'archive' collection.

Creates new keys for all items, remaps all IDs, copies storage directories.
Must be run with Zotero process stopped.
"""

import shutil
import sqlite3
import sys
from pathlib import Path

from zot.db import generate_key

HOST_DB = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/tmp/host-zotero.sqlite")
CONTAINER_DB = (
    Path(sys.argv[2])
    if len(sys.argv) > 2
    else Path("/home/ubuntu/Zotero/zotero.sqlite")
)
HOST_STORAGE = (
    Path(sys.argv[3]) if len(sys.argv) > 3 else Path("/home/ubuntu/data/Zotero/storage")
)
CONTAINER_STORAGE = (
    Path(sys.argv[4]) if len(sys.argv) > 4 else Path("/home/ubuntu/Zotero/storage")
)

LIBRARY_ID = 1  # user library


def gen_key(existing: set[str]) -> str:
    while True:
        k = generate_key()
        if k not in existing:
            existing.add(k)
            return k


def main():
    if not HOST_DB.exists():
        sys.exit(f"Host DB not found: {HOST_DB}")
    if not CONTAINER_DB.exists():
        sys.exit(f"Container DB not found: {CONTAINER_DB}")

    src = sqlite3.connect(str(HOST_DB))
    src.row_factory = sqlite3.Row
    dst = sqlite3.connect(str(CONTAINER_DB))
    dst.row_factory = sqlite3.Row
    dst.execute("PRAGMA journal_mode=WAL")
    dst.execute("PRAGMA foreign_keys=OFF")  # we handle ordering ourselves

    # Collect existing keys in destination
    existing_keys = {
        r[0]
        for r in dst.execute("SELECT key FROM items WHERE libraryID=?", (LIBRARY_ID,))
    }
    existing_keys |= {
        r[0]
        for r in dst.execute(
            "SELECT key FROM collections WHERE libraryID=?", (LIBRARY_ID,)
        )
    }

    # --- Max IDs in destination ---
    max_item_id = dst.execute("SELECT COALESCE(MAX(itemID),0) FROM items").fetchone()[0]
    max_coll_id = dst.execute(
        "SELECT COALESCE(MAX(collectionID),0) FROM collections"
    ).fetchone()[0]
    max_value_id = dst.execute(
        "SELECT COALESCE(MAX(valueID),0) FROM itemDataValues"
    ).fetchone()[0]
    max_creator_id = dst.execute(
        "SELECT COALESCE(MAX(creatorID),0) FROM creators"
    ).fetchone()[0]
    max_tag_id = dst.execute("SELECT COALESCE(MAX(tagID),0) FROM tags").fetchone()[0]
    max_word_id = dst.execute(
        "SELECT COALESCE(MAX(wordID),0) FROM fulltextWords"
    ).fetchone()[0]

    # ============================================================
    # 1. Create "archive" collection + mirror host collection tree
    # ============================================================
    archive_key = gen_key(existing_keys)
    max_coll_id += 1
    archive_coll_id = max_coll_id
    dst.execute(
        "INSERT INTO collections (collectionID, collectionName, parentCollectionID, libraryID, key, version, synced) "
        "VALUES (?,?,NULL,?,?,0,0)",
        (archive_coll_id, "archive", LIBRARY_ID, archive_key),
    )
    print(f"Created 'archive' collection: ID={archive_coll_id}, key={archive_key}")

    # Mirror host collections under archive
    host_colls = src.execute(
        "SELECT collectionID, collectionName, parentCollectionID FROM collections ORDER BY collectionID"
    ).fetchall()

    coll_id_map: dict[int, int] = {}  # host collectionID -> dest collectionID
    for hc in host_colls:
        max_coll_id += 1
        new_id = max_coll_id
        coll_id_map[hc["collectionID"]] = new_id
        # parent: if host has a parent, map it; otherwise parent is archive
        if hc["parentCollectionID"] is not None:
            parent = coll_id_map.get(hc["parentCollectionID"], archive_coll_id)
        else:
            parent = archive_coll_id
        coll_key = gen_key(existing_keys)
        dst.execute(
            "INSERT INTO collections (collectionID, collectionName, parentCollectionID, libraryID, key, version, synced) "
            "VALUES (?,?,?,?,?,0,0)",
            (new_id, hc["collectionName"], parent, LIBRARY_ID, coll_key),
        )
    print(f"Mirrored {len(host_colls)} host collections under archive")

    # ============================================================
    # 2. Migrate itemDataValues (shared, deduplicated by value)
    # ============================================================
    # Build lookup of existing values in dest
    dst_values: dict[str, int] = {}
    for r in dst.execute("SELECT valueID, value FROM itemDataValues"):
        dst_values[str(r["value"])] = r["valueID"]

    src_value_map: dict[int, int] = {}  # host valueID -> dest valueID
    for r in src.execute("SELECT valueID, value FROM itemDataValues"):
        val = str(r["value"])
        if val in dst_values:
            src_value_map[r["valueID"]] = dst_values[val]
        else:
            max_value_id += 1
            dst.execute(
                "INSERT INTO itemDataValues (valueID, value) VALUES (?,?)",
                (max_value_id, r["value"]),
            )
            dst_values[val] = max_value_id
            src_value_map[r["valueID"]] = max_value_id
    print(f"Mapped {len(src_value_map)} itemDataValues")

    # ============================================================
    # 3. Migrate creators (deduplicated by lastName+firstName+fieldMode)
    # ============================================================
    dst_creators: dict[tuple, int] = {}
    for r in dst.execute(
        "SELECT creatorID, firstName, lastName, fieldMode FROM creators"
    ):
        dst_creators[(r["lastName"], r["firstName"], r["fieldMode"])] = r["creatorID"]

    src_creator_map: dict[int, int] = {}
    for r in src.execute(
        "SELECT creatorID, firstName, lastName, fieldMode FROM creators"
    ):
        ck = (r["lastName"], r["firstName"], r["fieldMode"])
        if ck in dst_creators:
            src_creator_map[r["creatorID"]] = dst_creators[ck]
        else:
            max_creator_id += 1
            dst.execute(
                "INSERT INTO creators (creatorID, firstName, lastName, fieldMode) VALUES (?,?,?,?)",
                (max_creator_id, r["firstName"], r["lastName"], r["fieldMode"]),
            )
            dst_creators[ck] = max_creator_id
            src_creator_map[r["creatorID"]] = max_creator_id
    print(f"Mapped {len(src_creator_map)} creators")

    # ============================================================
    # 4. Migrate tags (deduplicated by name)
    # ============================================================
    dst_tags: dict[str, int] = {}
    for r in dst.execute("SELECT tagID, name FROM tags"):
        dst_tags[r["name"]] = r["tagID"]

    src_tag_map: dict[int, int] = {}
    for r in src.execute("SELECT tagID, name FROM tags"):
        if r["name"] in dst_tags:
            src_tag_map[r["tagID"]] = dst_tags[r["name"]]
        else:
            max_tag_id += 1
            dst.execute(
                "INSERT INTO tags (tagID, name) VALUES (?,?)", (max_tag_id, r["name"])
            )
            dst_tags[r["name"]] = max_tag_id
            src_tag_map[r["tagID"]] = max_tag_id
    print(f"Mapped {len(src_tag_map)} tags")

    # ============================================================
    # 5. Migrate fulltextWords (deduplicated by word)
    # ============================================================
    dst_words: dict[str, int] = {}
    for r in dst.execute("SELECT wordID, word FROM fulltextWords"):
        dst_words[r["word"]] = r["wordID"]

    src_word_map: dict[int, int] = {}
    for r in src.execute("SELECT wordID, word FROM fulltextWords"):
        if r["word"] in dst_words:
            src_word_map[r["wordID"]] = dst_words[r["word"]]
        else:
            max_word_id += 1
            dst.execute(
                "INSERT INTO fulltextWords (wordID, word) VALUES (?,?)",
                (max_word_id, r["word"]),
            )
            dst_words[r["word"]] = max_word_id
            src_word_map[r["wordID"]] = max_word_id
    print(f"Mapped {len(src_word_map)} fulltextWords")

    # ============================================================
    # 6. Migrate items (new keys, new IDs)
    # ============================================================
    # Process items in order so parents come before children
    host_items = src.execute(
        "SELECT itemID, itemTypeID, dateAdded, dateModified, clientDateModified, libraryID, key "
        "FROM items WHERE libraryID=? ORDER BY itemID",
        (LIBRARY_ID,),
    ).fetchall()

    item_id_map: dict[int, int] = {}  # host itemID -> dest itemID
    item_key_map: dict[str, str] = {}  # host key -> dest key (for storage)

    for hi in host_items:
        max_item_id += 1
        new_key = gen_key(existing_keys)
        item_id_map[hi["itemID"]] = max_item_id
        item_key_map[hi["key"]] = new_key
        dst.execute(
            "INSERT INTO items (itemID, itemTypeID, dateAdded, dateModified, clientDateModified, libraryID, key, version, synced) "
            "VALUES (?,?,?,?,?,?,?,0,0)",
            (
                max_item_id,
                hi["itemTypeID"],
                hi["dateAdded"],
                hi["dateModified"],
                hi["clientDateModified"],
                LIBRARY_ID,
                new_key,
            ),
        )
    print(f"Migrated {len(item_id_map)} items")

    # ============================================================
    # 7. Migrate itemData
    # ============================================================
    count = 0
    for r in src.execute("SELECT itemID, fieldID, valueID FROM itemData"):
        if r["itemID"] not in item_id_map:
            continue
        dst.execute(
            "INSERT INTO itemData (itemID, fieldID, valueID) VALUES (?,?,?)",
            (item_id_map[r["itemID"]], r["fieldID"], src_value_map[r["valueID"]]),
        )
        count += 1
    print(f"Migrated {count} itemData rows")

    # ============================================================
    # 8. Migrate itemCreators
    # ============================================================
    count = 0
    for r in src.execute(
        "SELECT itemID, creatorID, creatorTypeID, orderIndex FROM itemCreators"
    ):
        if r["itemID"] not in item_id_map:
            continue
        dst.execute(
            "INSERT INTO itemCreators (itemID, creatorID, creatorTypeID, orderIndex) VALUES (?,?,?,?)",
            (
                item_id_map[r["itemID"]],
                src_creator_map[r["creatorID"]],
                r["creatorTypeID"],
                r["orderIndex"],
            ),
        )
        count += 1
    print(f"Migrated {count} itemCreators rows")

    # ============================================================
    # 9. Migrate itemAttachments
    # ============================================================
    count = 0
    for r in src.execute(
        "SELECT itemID, parentItemID, linkMode, contentType, charsetID, path, syncState, storageModTime, storageHash FROM itemAttachments"
    ):
        if r["itemID"] not in item_id_map:
            continue
        parent = (
            item_id_map.get(r["parentItemID"])
            if r["parentItemID"] is not None
            else None
        )
        dst.execute(
            "INSERT INTO itemAttachments (itemID, parentItemID, linkMode, contentType, charsetID, path, syncState, storageModTime, storageHash) "
            "VALUES (?,?,?,?,?,?,?,?,?)",
            (
                item_id_map[r["itemID"]],
                parent,
                r["linkMode"],
                r["contentType"],
                r["charsetID"],
                r["path"],
                r["syncState"],
                r["storageModTime"],
                r["storageHash"],
            ),
        )
        count += 1
    print(f"Migrated {count} itemAttachments rows")

    # ============================================================
    # 10. Migrate itemNotes
    # ============================================================
    count = 0
    for r in src.execute("SELECT itemID, parentItemID, note, title FROM itemNotes"):
        if r["itemID"] not in item_id_map:
            continue
        parent = (
            item_id_map.get(r["parentItemID"])
            if r["parentItemID"] is not None
            else None
        )
        dst.execute(
            "INSERT INTO itemNotes (itemID, parentItemID, note, title) VALUES (?,?,?,?)",
            (item_id_map[r["itemID"]], parent, r["note"], r["title"]),
        )
        count += 1
    print(f"Migrated {count} itemNotes rows")

    # ============================================================
    # 11. Migrate itemTags
    # ============================================================
    count = 0
    for r in src.execute("SELECT itemID, tagID, type FROM itemTags"):
        if r["itemID"] not in item_id_map:
            continue
        dst.execute(
            "INSERT INTO itemTags (itemID, tagID, type) VALUES (?,?,?)",
            (item_id_map[r["itemID"]], src_tag_map[r["tagID"]], r["type"]),
        )
        count += 1
    print(f"Migrated {count} itemTags rows")

    # ============================================================
    # 12. Migrate itemRelations
    # ============================================================
    # Get relationPredicates from both DBs — they should match (standard set)
    count = 0
    for r in src.execute("SELECT itemID, predicateID, object FROM itemRelations"):
        if r["itemID"] not in item_id_map:
            continue
        dst.execute(
            "INSERT OR IGNORE INTO itemRelations (itemID, predicateID, object) VALUES (?,?,?)",
            (item_id_map[r["itemID"]], r["predicateID"], r["object"]),
        )
        count += 1
    print(f"Migrated {count} itemRelations rows")

    # ============================================================
    # 13. Migrate fulltextItems + fulltextItemWords
    # ============================================================
    count = 0
    for r in src.execute(
        "SELECT itemID, indexedPages, totalPages, indexedChars, totalChars FROM fulltextItems"
    ):
        if r["itemID"] not in item_id_map:
            continue
        dst.execute(
            "INSERT INTO fulltextItems (itemID, indexedPages, totalPages, indexedChars, totalChars, version, synced) "
            "VALUES (?,?,?,?,?,0,0)",
            (
                item_id_map[r["itemID"]],
                r["indexedPages"],
                r["totalPages"],
                r["indexedChars"],
                r["totalChars"],
            ),
        )
        count += 1
    print(f"Migrated {count} fulltextItems rows")

    count = 0
    for r in src.execute("SELECT wordID, itemID FROM fulltextItemWords"):
        if r["itemID"] not in item_id_map:
            continue
        if r["wordID"] not in src_word_map:
            continue
        dst.execute(
            "INSERT OR IGNORE INTO fulltextItemWords (wordID, itemID) VALUES (?,?)",
            (src_word_map[r["wordID"]], item_id_map[r["itemID"]]),
        )
        count += 1
    print(f"Migrated {count} fulltextItemWords rows")

    # ============================================================
    # 14. Migrate collectionItems (only top-level items, not child attachments/notes)
    # ============================================================
    # Zotero triggers prevent adding child attachments/notes to collections
    child_item_ids = set()
    for r in src.execute(
        "SELECT itemID FROM itemAttachments WHERE parentItemID IS NOT NULL"
    ):
        child_item_ids.add(r["itemID"])
    for r in src.execute("SELECT itemID FROM itemNotes WHERE parentItemID IS NOT NULL"):
        child_item_ids.add(r["itemID"])

    count = 0
    for r in src.execute(
        "SELECT collectionID, itemID, orderIndex FROM collectionItems"
    ):
        if r["itemID"] not in item_id_map:
            continue
        if r["collectionID"] not in coll_id_map:
            continue
        dst.execute(
            "INSERT OR IGNORE INTO collectionItems (collectionID, itemID, orderIndex) VALUES (?,?,?)",
            (coll_id_map[r["collectionID"]], item_id_map[r["itemID"]], r["orderIndex"]),
        )
        count += 1
    print(f"Migrated {count} collectionItems rows")

    # Also add all top-level items to the archive collection itself
    top_count = 0
    for host_id, dest_id in item_id_map.items():
        if host_id not in child_item_ids:
            dst.execute(
                "INSERT OR IGNORE INTO collectionItems (collectionID, itemID, orderIndex) VALUES (?,?,0)",
                (archive_coll_id, dest_id),
            )
            top_count += 1
    print(f"Added {top_count} top-level items to archive collection")

    # ============================================================
    # 15. Commit
    # ============================================================
    dst.commit()
    print("Database committed successfully")

    # ============================================================
    # 16. Copy storage directories
    # ============================================================
    copied = 0
    skipped = 0
    for old_key, new_key in item_key_map.items():
        src_dir = HOST_STORAGE / old_key
        dst_dir = CONTAINER_STORAGE / new_key
        if src_dir.is_dir():
            if dst_dir.exists():
                skipped += 1
                continue
            shutil.copytree(str(src_dir), str(dst_dir))
            copied += 1
    print(f"Copied {copied} storage directories, skipped {skipped}")

    src.close()
    dst.close()
    print("Migration complete!")


if __name__ == "__main__":
    main()