add script to fix project metadata

Author: jdkent

compose/backend/manage.py

@@ -14,6 +14,9 @@
 from neurosynth_compose.database import db
 from neurosynth_compose import models
 from neurosynth_compose.ingest import neurostore as ingest_nstore
+from neurosynth_compose.backfill_extraction_metadata import (
+    add_missing_extraction_ids,
+)
 
 
 app.config.from_object(os.environ["APP_SETTINGS"])
@@ -44,3 +47,12 @@ def create_meta_analyses(n_studysets, neurostore_url):
     if n_studysets is not None:
         n_studysets = int(n_studysets)
     ingest_nstore.create_meta_analyses(url=neurostore_url, n_studysets=n_studysets)
+
+
+@app.cli.command("backfill-extraction-metadata")
+def backfill_extraction_metadata():
+    """Add missing extractionMetadata ids to project provenance."""
+    updated, skipped = add_missing_extraction_ids()
+    click.echo(
+        f"Updated {updated} project(s); skipped {skipped} project(s) with no changes."
+    )

compose/backend/neurosynth_compose/scripts/__init__.py

@@ -0,0 +1,57 @@
+import logging
+from typing import Tuple
+
+from sqlalchemy import select
+
+from neurosynth_compose.database import db
+from neurosynth_compose.models.analysis import Project
+
+logger = logging.getLogger(__name__)
+
+
+def add_missing_extraction_ids(session=None) -> Tuple[int, int]:
+    """Add null studysetId/annotationId keys to extractionMetadata when absent."""
+    sess = session or db.session
+    updated = 0
+    skipped = 0
+
+    projects = sess.scalars(select(Project)).all()
+
+    for project in projects:
+        provenance = project.provenance or {}
+        extraction_metadata = provenance.get("extractionMetadata")
+
+        if not isinstance(extraction_metadata, dict):
+            skipped += 1
+            continue
+
+        changed = False
+
+        if "studysetId" not in extraction_metadata:
+            extraction_metadata["studysetId"] = None
+            changed = True
+
+        if "annotationId" not in extraction_metadata:
+            extraction_metadata["annotationId"] = None
+            changed = True
+
+        if changed:
+            provenance["extractionMetadata"] = extraction_metadata
+            project.provenance = provenance
+            updated += 1
+        else:
+            skipped += 1
+
+    if updated:
+        try:
+            sess.commit()
+        except Exception:
+            sess.rollback()
+            logger.exception(
+                "Failed to commit extractionMetadata backfill for projects."
+            )
+            raise
+    else:
+        sess.rollback()
+
+    return updated, skipped