Merge branch 'fix/project_metadata' into staging

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

compose/scripts/add_usernames.py renamed to compose/backend/neurosynth_compose/scripts/add_usernames.py

@@ -364,7 +364,8 @@ def ingest_neurosynth(max_rows=None):
 
         # add notes to annotation
         annot.note_keys = {
-            k: _check_type(v) for k, v in annotation_row._asdict().items()
+            k: {"type": _check_type(v) or "string", "order": idx}
+            for idx, (k, v) in enumerate(annotation_row._asdict().items())
         }
         annot.annotation_analyses = notes
         for note in notes:

compose/backend/neurosynth_compose/scripts/backfill_extraction_metadata.py

@@ -393,6 +393,84 @@ def get_affected_ids(self, ids):
         }
         return unique_ids
 
+    @staticmethod
+    def _ordered_note_keys(note_keys):
+        if not note_keys:
+            return []
+        keys = list(note_keys.keys())
+        alphabetical = sorted(keys)
+        if keys == alphabetical:
+            return alphabetical
+        return keys
+
+    @classmethod
+    def _normalize_note_keys(cls, note_keys):
+        if note_keys is None:
+            return None
+        if not isinstance(note_keys, dict):
+            abort_validation("`note_keys` must be an object.")
+
+        ordered_keys = cls._ordered_note_keys(note_keys)
+        normalized = OrderedDict()
+        used_orders = set()
+        next_order = 0
+
+        for key in ordered_keys:
+            descriptor = note_keys.get(key) or {}
+            note_type = descriptor.get("type")
+            if note_type not in {"string", "number", "boolean"}:
+                abort_validation(
+                    "Invalid `type` for note_keys entry "
+                    f"'{key}', choose from: ['boolean', 'number', 'string']."
+                )
+
+            order = descriptor.get("order")
+            if isinstance(order, bool) or (
+                order is not None and not isinstance(order, int)
+            ):
+                order = None
+
+            if isinstance(order, int) and order not in used_orders:
+                used_orders.add(order)
+                if order >= next_order:
+                    next_order = order + 1
+            else:
+                while next_order in used_orders:
+                    next_order += 1
+                order = next_order
+                used_orders.add(order)
+                next_order += 1
+
+            normalized[key] = {"type": note_type, "order": order}
+
+        return normalized
+
+    @classmethod
+    def _merge_note_keys(cls, existing, additions):
+        """
+        additions is a mapping of key -> type
+        """
+        base = cls._normalize_note_keys(existing or {}) or OrderedDict()
+        used_orders = {v.get("order") for v in base.values() if isinstance(v, dict)}
+        used_orders = {o for o in used_orders if isinstance(o, int)}
+        next_order = max(used_orders, default=-1) + 1
+
+        for key, value_type in additions.items():
+            if key in base:
+                descriptor = base[key] or {}
+                descriptor["type"] = value_type or descriptor.get("type") or "string"
+                base[key] = descriptor
+                continue
+
+            descriptor = {
+                "type": value_type or "string",
+                "order": next_order,
+            }
+            base[key] = descriptor
+            next_order += 1
+
+        return base
+
     @classmethod
     def load_nested_records(cls, data, record=None):
         if not data:
@@ -554,6 +632,9 @@ def put(self, id):
         schema = self._schema()
         data = schema.load(request_data)
 
+        if "note_keys" in data:
+            data["note_keys"] = self._normalize_note_keys(data["note_keys"])
+
         pipeline_payload = data.pop("pipelines", [])
 
         args = {}
@@ -942,12 +1023,10 @@ def _apply_pipeline_columns(self, annotation, data, specs, column_counter):
 
         if column_types:
             if data.get("note_keys") is None:
-                note_keys = dict(annotation.note_keys or {})
+                note_keys = self._normalize_note_keys(annotation.note_keys or {})
             else:
-                note_keys = dict(data["note_keys"])
-            for key, value_type in column_types.items():
-                note_keys[key] = value_type or "string"
-            data["note_keys"] = note_keys
+                note_keys = self._normalize_note_keys(data["note_keys"])
+            data["note_keys"] = self._merge_note_keys(note_keys, column_types)
 
         data["annotation_analyses"] = list(note_map.values())
 

store/backend/neurostore/ingest/__init__.py

@@ -663,6 +663,70 @@ class AnnotationPipelineSchema(BaseSchema):
     columns = fields.List(fields.String(), required=True)
 
 
+class NoteKeysField(fields.Field):
+    allowed_types = {"string", "number", "boolean"}
+
+    def _serialize(self, value, attr, obj, **kwargs):
+        if not value:
+            return {}
+        serialized = {}
+        for key, descriptor in value.items():
+            if not isinstance(descriptor, dict):
+                continue
+            serialized[key] = {
+                "type": descriptor.get("type"),
+                "order": descriptor.get("order"),
+            }
+        return serialized
+
+    def _deserialize(self, value, attr, data, **kwargs):
+        if value is None:
+            return {}
+        if not isinstance(value, dict):
+            raise ValidationError("`note_keys` must be an object.")
+
+        normalized = {}
+        used_orders = set()
+        explicit_orders = []
+        for descriptor in value.values():
+            if isinstance(descriptor, dict) and isinstance(
+                descriptor.get("order"), int
+            ):
+                explicit_orders.append(descriptor["order"])
+        next_order = max(explicit_orders, default=-1) + 1
+
+        for key, descriptor in value.items():
+            if not isinstance(descriptor, dict):
+                raise ValidationError("Each note key must map to an object.")
+
+            note_type = descriptor.get("type")
+            if note_type not in self.allowed_types:
+                raise ValidationError(
+                    f"Invalid note type for '{key}', choose from: {sorted(self.allowed_types)}"
+                )
+
+            order = descriptor.get("order")
+            if isinstance(order, bool) or (
+                order is not None and not isinstance(order, int)
+            ):
+                order = None
+
+            if isinstance(order, int) and order not in used_orders:
+                used_orders.add(order)
+                if order >= next_order:
+                    next_order = order + 1
+            else:
+                while next_order in used_orders:
+                    next_order += 1
+                order = next_order
+                used_orders.add(order)
+                next_order += 1
+
+            normalized[key] = {"type": note_type, "order": order}
+
+        return normalized
+
+
 class AnnotationSchema(BaseDataSchema):
     # serialization
     studyset_id = fields.String(data_key="studyset")
@@ -675,7 +739,7 @@ class AnnotationSchema(BaseDataSchema):
     source_id = fields.String(dump_only=True, allow_none=True)
     source_updated_at = fields.DateTime(dump_only=True, allow_none=True)
 
-    note_keys = fields.Dict()
+    note_keys = NoteKeysField()
     metadata = fields.Dict(attribute="metadata_", dump_only=True)
     # deserialization
     metadata_ = fields.Dict(data_key="metadata", load_only=True, allow_none=True)