add(scripts): Add script to rebuild rdm_records_state.json

scripts/rebuild_records_state.py

@@ -0,0 +1,398 @@
+"""Regenerate ``rdm_records_state.json`` straight from the CDS-RDM database.
+
+Root cause - previously, the code had a bug:
+
+    for version in versions.keys():
+        draft = self._pre_publish(...)
+        published_record = current_rdm_records_service.publish(...)
+        self._after_publish(...)
+    records.append(published_record._record)   # left outside the loop
+
+It ran exactly once per record, for only the latest version. Now fixed in code already.
+
+``invenio migration stats run`` uses only the ``versions[]`` in that state
+file to attribute downloads, so any multi-version record migrated in that
+window lost every download event belonging to any version except the latest.
+
+This script rebuilds the state file from scratch, entirely from CDS-RDM, for
+a given collection's community id(s) (from that collection's
+``transform.communities_ids`` in ``streams.yaml`` / ``streams_done.yaml`` /
+``streams_shelved.yaml``):
+
+1. Resolve the given community id(s) to every parent record in them using
+   ``RDMParentCommunity``.
+2. Resolve each parent to its legacy recid and its parent recid in one PID
+   query (``pid_type`` ``lrecid`` and ``recid``).
+3. Load every published version of those parents and the file rows. Each
+   query is limited to ``CHUNK_SIZE`` ids so one result set is not the whole
+   community. ``rdm_records_metadata.parent_id`` is not indexed, so this is a
+   handful of table scans instead of one scan per record.
+4. Resolve each record uuid to its recid through the ``recid`` PID. The
+   version query does not read the record JSON. Versions are stored once
+   per parent, so a redirect ``lrecid`` minted onto that parent gets the
+   same versions as the migrated record.
+5. Build the state list in memory. ``latest_version`` is the highest
+   ``index`` that still has a migrated file. A record migrated with no
+   files keeps the highest published index so pageviews have a target.
+   File entries are not stored on the record JSON
+   (``FilesField(store=False)``); they come from ``rdm_records_files``.
+
+Output is written once, next to the original, as
+``<collection>/rdm_records_state.fixed.json``.
+
+Usage:
+    invenio shell
+
+    main(community_ids=["<community-id>"],
+         output_path="/migration/tmp/<name>/rdm_records_state.fixed.json",
+         log_file="/migration/tmp/<name>/records_state.log",
+         dry_run=True)
+"""
+
+import json
+import traceback
+from pathlib import Path
+
+from invenio_db import db
+from invenio_files_rest.models import FileInstance, ObjectVersion
+from invenio_pidstore.models import PersistentIdentifier, PIDStatus
+from invenio_rdm_records.records.models import (
+    RDMFileRecordMetadata,
+    RDMParentCommunity,
+    RDMRecordMetadata,
+)
+
+log_fp = None
+
+
+def log(msg):
+    print(msg)
+    if log_fp is not None:
+        log_fp.write(msg + "\n")
+        log_fp.flush()
+
+
+CHUNK_SIZE = 10000
+
+
+def chunked(items, size=CHUNK_SIZE):
+    """Yield ``items`` in slices of ``size``."""
+    items = list(items)
+    for start in range(0, len(items), size):
+        yield items[start : start + size]
+
+
+def find_parent_pids(community_ids):
+    """Return legacy recids and parent recids for migrated records.
+
+    Community ids map to parent uuids via ``RDMParentCommunity``. One PID
+    query per chunk then splits ``lrecid`` and ``recid`` for those parents.
+
+    :returns: ``({legacy_recid: parent_uuid}, {parent_uuid: parent_recid})``
+    """
+    parent_uuids = list(
+        {
+            row.record_id
+            for row in RDMParentCommunity.query.filter(
+                RDMParentCommunity.community_id.in_(community_ids)
+            )
+        }
+    )
+    legacy_recids = {}
+    parent_recids = {}
+    for batch in chunked(parent_uuids):
+        pids = PersistentIdentifier.query.filter(
+            PersistentIdentifier.pid_type.in_(("lrecid", "recid")),
+            PersistentIdentifier.object_type == "rec",
+            PersistentIdentifier.status == PIDStatus.REGISTERED,
+            PersistentIdentifier.object_uuid.in_(batch),
+        )
+        for pid in pids:
+            object_uuid = str(pid.object_uuid)
+            if pid.pid_type == "lrecid":
+                legacy_recids[pid.pid_value] = object_uuid
+            else:
+                parent_recids[object_uuid] = pid.pid_value
+    db.session.expunge_all()
+    return legacy_recids, parent_recids
+
+
+def load_pids(object_uuids, pid_type):
+    """Return ``{object_uuid: pid_value}`` for registered PIDs.
+
+    ``object_type`` is included so Postgres can use ``idx_object``
+    ``(object_type, object_uuid)`` instead of scanning ``pidstore_pid``.
+    """
+    found = {}
+    for batch in chunked(object_uuids):
+        pids = PersistentIdentifier.query.filter(
+            PersistentIdentifier.pid_type == pid_type,
+            PersistentIdentifier.object_type == "rec",
+            PersistentIdentifier.status == PIDStatus.REGISTERED,
+            PersistentIdentifier.object_uuid.in_(batch),
+        )
+        for pid in pids:
+            found[str(pid.object_uuid)] = pid.pid_value
+    db.session.expunge_all()
+    return found
+
+
+def load_versions(parent_uuids):
+    """Return ``({parent_uuid: {version_index: version dict}}, failed_parents)``.
+
+    Version rows are loaded without the JSON column. The recid is the
+    ``recid`` PID for that row's primary key. Soft-deleted rows (``json``
+    is NULL) are excluded, matching ``RDMRecord.get_record()``.
+
+    Keyed by parent so every ``lrecid`` on that parent (the migrated record
+    and any redirect minted onto it) can share one version list.
+    """
+    versions = {}
+    parents = list(dict.fromkeys(parent_uuids))
+    for n, batch in enumerate(chunked(parents), start=1):
+        rows = (
+            db.session.query(
+                RDMRecordMetadata.id,
+                RDMRecordMetadata.parent_id,
+                RDMRecordMetadata.index,
+                RDMRecordMetadata.bucket_id,
+            )
+            .filter(RDMRecordMetadata.parent_id.in_(batch))
+            .filter(RDMRecordMetadata.json.isnot(None))
+            .all()
+        )
+        for rec_id, parent_id, index, bucket_id in rows:
+            parent_uuid = str(parent_id)
+            versions.setdefault(parent_uuid, {})[index] = {
+                "record_uuid": str(rec_id),
+                "bucket_id": str(bucket_id) if bucket_id else None,
+                "parent_object_uuid": parent_uuid,
+                "files": [],
+            }
+        log(f"loaded version chunk {n}, parents with versions={len(versions)}")
+    db.session.expunge_all()
+
+    record_uuids = [
+        version["record_uuid"]
+        for by_index in versions.values()
+        for version in by_index.values()
+    ]
+    record_recids = load_pids(record_uuids, "recid")
+    failed = set()
+    for parent_uuid, by_index in versions.items():
+        for version in by_index.values():
+            new_recid = record_recids.get(version["record_uuid"])
+            if not new_recid:
+                log(
+                    f"record {version['record_uuid']} has no recid PID, "
+                    f"parent_id={parent_uuid}"
+                )
+                failed.add(parent_uuid)
+                continue
+            version["new_recid"] = new_recid
+    return versions, failed
+
+
+def load_files(record_uuids):
+    """Return {record_uuid: [file dict, ...]} from ``rdm_records_files``.
+
+    ``record_id`` is indexed. ``file_id`` and ``size`` live on the object
+    version and its file instance, not on the record JSON.
+    """
+    files = {}
+    uuids = list(record_uuids)
+    for n, batch in enumerate(chunked(uuids), start=1):
+        rows = (
+            db.session.query(
+                RDMFileRecordMetadata.record_id,
+                RDMFileRecordMetadata.key,
+                RDMFileRecordMetadata.json,
+                ObjectVersion.file_id,
+                FileInstance.size,
+            )
+            .join(
+                ObjectVersion,
+                ObjectVersion.version_id == RDMFileRecordMetadata.object_version_id,
+            )
+            .join(FileInstance, FileInstance.id == ObjectVersion.file_id)
+            .filter(RDMFileRecordMetadata.record_id.in_(batch))
+            .filter(RDMFileRecordMetadata.json.isnot(None))
+            .all()
+        )
+        for record_id, key, file_json, file_id, size in rows:
+            metadata = (file_json or {}).get("metadata") or {}
+            files.setdefault(str(record_id), []).append(
+                {
+                    "legacy_file_id": metadata.get("legacy_file_id"),
+                    "file_key": key,
+                    "file_id": str(file_id),
+                    "size": str(size),
+                }
+            )
+        del rows
+        log(f"loaded file chunk {n}, records with files={len(files)}")
+    db.session.expunge_all()
+    return files
+
+
+def _file_for_state(bucket_id, file_entry):
+    """One file entry in the shape ``_load_record_state`` writes.
+
+    Files added after migration have no ``legacy_file_id``. Those are not
+    part of the state file; the caller logs and drops the ``None``.
+    """
+    legacy_file_id = file_entry.get("legacy_file_id")
+    if legacy_file_id is None:
+        return None
+    return {
+        "legacy_file_id": legacy_file_id,
+        "bucket_id": bucket_id,
+        "file_key": file_entry["file_key"],
+        "file_id": file_entry["file_id"],
+        "size": file_entry["size"],
+    }
+
+
+def _version_files(legacy_recid, version_index, version):
+    """Migrated files for one version. Non-migrated files are logged and skipped."""
+    files = []
+    for file_entry in version["files"]:
+        formatted = _file_for_state(version["bucket_id"], file_entry)
+        if formatted is None:
+            log(
+                "skipping (non-migrated record) file with no legacy_file_id: "
+                f"legacy_recid={legacy_recid} "
+                f"version={version_index} "
+                f"new_recid={version['new_recid']} "
+                f"record_uuid={version['record_uuid']} "
+                f"file_key={file_entry.get('file_key')}"
+            )
+            continue
+        files.append(formatted)
+    return files
+
+
+def build_state_entry(legacy_recid, rdm_versions, parent_recid):
+    """Rebuild one ``rdm_records_state.json`` entry.
+
+    Same fields as ``RecordLoad._load_record_state``
+    (cds_migrator_kit/rdm/records/load/entities/record.py). ``latest_version``
+    is the highest ``index`` that still has a migrated file. A record
+    migrated with no files keeps the highest published index so pageviews
+    have a target.
+    """
+    recid_state = {
+        "legacy_recid": str(legacy_recid),
+        "parent_recid": parent_recid,
+        "versions": [],
+    }
+    latest_version = None
+    for version_index in sorted(rdm_versions):
+        version = rdm_versions[version_index]
+        if "parent_object_uuid" not in recid_state:
+            recid_state["parent_object_uuid"] = version["parent_object_uuid"]
+        files = _version_files(legacy_recid, version_index, version)
+        if not files:
+            # No migrated files: created after migration, or the files are gone.
+            continue
+        recid_state["versions"].append(
+            {
+                "new_recid": version["new_recid"],
+                "version": version_index,
+                "files": files,
+            }
+        )
+        latest_version = version
+    if latest_version is None:
+        latest_version = rdm_versions[max(rdm_versions)]
+    recid_state["latest_version"] = latest_version["new_recid"]
+    recid_state["latest_version_object_uuid"] = latest_version["record_uuid"]
+    return recid_state
+
+
+def write_state_file(filepath, entries):
+    """Write entries in the same JSON-list-of-compact-objects format
+    ``RecordStateLogger.finalise()`` uses (cds_migrator_kit/reports/log.py),
+    which is what ``invenio migration stats run --filepath`` expects."""
+    Path(filepath).parent.mkdir(parents=True, exist_ok=True)
+    with open(filepath, "w", encoding="utf-8") as f:
+        f.write("[\n")
+        for i, entry in enumerate(entries):
+            json_str = json.dumps(entry, ensure_ascii=False, separators=(",", ":"))
+            comma = "," if i < len(entries) - 1 else ""
+            f.write(f"{json_str}{comma}\n")
+        f.write("]\n")
+
+
+def main(community_ids, output_path, log_file, dry_run=True):
+    """Rebuild ``rdm_records_state.json`` for one stream/collection.
+
+    :param community_ids: the community UUIDs for a collection, as set in
+        its ``transform.communities_ids`` in streams.yaml.
+    :param output_path: should be `<CDS_MIGRATOR_KIT_LOGS_PATH>/<collection>/rdm_records_state.fixed.json`
+        next to the original state file (no overwriting).
+    :param log_file: progress log for this run.
+    :param dry_run: pass dry_run=False to actually write entries to disk.
+    """
+    global log_fp
+
+    log_fp = open(log_file, "a")
+    try:
+        legacy_recids, parent_recids = find_parent_pids(community_ids)
+        log(
+            f"starting, community_ids={str(community_ids)}, "
+            f"dry_run={dry_run}, found legacy_recids={len(legacy_recids)}"
+        )
+        versions, failed = load_versions(legacy_recids.values())
+        record_uuids = [
+            version["record_uuid"]
+            for by_index in versions.values()
+            for version in by_index.values()
+        ]
+        files = load_files(record_uuids)
+        for by_index in versions.values():
+            for version in by_index.values():
+                version["files"] = files.get(version["record_uuid"], [])
+
+        stats = {"checked": 0, "no_versions": 0, "fixed": 0, "errors": 0}
+        entries = []
+        for legacy_recid, parent_uuid in legacy_recids.items():
+            stats["checked"] += 1
+            if parent_uuid in failed:
+                stats["errors"] += 1
+                continue
+            rdm_versions = versions.get(parent_uuid)
+            if not rdm_versions:
+                log(
+                    f"legacy_recid={legacy_recid} - no published versions found, skipping"
+                )
+                stats["no_versions"] += 1
+                continue
+            try:
+                entries.append(
+                    build_state_entry(
+                        legacy_recid,
+                        rdm_versions,
+                        parent_recids[parent_uuid],
+                    )
+                )
+                stats["fixed"] += 1
+            except Exception as exc:
+                log(f"unexpected error for legacy_recid={legacy_recid}: {exc}")
+                log(traceback.format_exc())
+                stats["errors"] += 1
+
+        if not dry_run:
+            write_state_file(output_path, entries)
+            log(f"wrote {output_path} ({len(entries)} entries)")
+        else:
+            log(f"dry run — would write {output_path} ({len(entries)} entries)")
+
+        log(
+            f"\nsummary:\nchecked={stats['checked']}\n"
+            f"no_versions={stats['no_versions']}\nfixed={stats['fixed']}\n"
+            f"errors={stats['errors']}"
+        )
+    finally:
+        log_fp.close()
+        log_fp = None