[ENH] clone studyset (#1149)

* codex implementation

* fix style and missing import

Author: jdkent

store/backend/neurostore/openapi

@@ -1,5 +1,7 @@
 import string
+from copy import deepcopy
 from sqlalchemy import func, text
+from sqlalchemy.exc import SQLAlchemyError
 
 from pgvector.sqlalchemy import Vector
 from flask import request
@@ -87,6 +89,7 @@ class StudysetsView(ObjectView, ListView):
     _view_fields = {
         **LIST_CLONE_ARGS,
         **LIST_NESTED_ARGS,
+        "copy_annotations": fields.Boolean(load_default=True),
     }
     # reorg int o2m and m2o
     _o2m = {"studies": "StudiesView", "annotations": "AnnotationsView"}
@@ -213,6 +216,157 @@ def serialize_records(self, records, args):
             return content
         return super().serialize_records(records, args)
 
+    def post(self):
+        args = parser.parse(self._user_args, request, location="query")
+        copy_annotations = args.pop("copy_annotations", True)
+        source_id = args.get("source_id")
+
+        if not source_id:
+            return super().post()
+
+        source = args.get("source") or "neurostore"
+        if source != "neurostore":
+            field_err = make_field_error("source", source, valid_options=["neurostore"])
+            abort_unprocessable(
+                "invalid source, choose from: 'neurostore'", [field_err]
+            )
+
+        unknown = self.__class__._schema.opts.unknown
+        data = parser.parse(
+            self.__class__._schema(exclude=("id",)), request, unknown=unknown
+        )
+
+        clone_payload, source_record = self._build_clone_payload(source_id, data)
+
+        # ensure nested serialization when cloning
+        args["nested"] = bool(args.get("nested") or request.args.get("source_id"))
+
+        with db.session.no_autoflush:
+            record = self.__class__.update_or_create(clone_payload)
+
+        unique_ids = self.get_affected_ids([record.id])
+        clear_cache(unique_ids)
+
+        db.session.flush()
+
+        self.update_base_studies(unique_ids.get("base-studies"))
+
+        try:
+            if copy_annotations:
+                self._clone_annotations(source_record, record)
+            self.update_annotations(unique_ids.get("annotations"))
+        except SQLAlchemyError as e:
+            db.session.rollback()
+            abort_validation(str(e))
+
+        response_context = dict(args)
+        response = self.__class__._schema(context=response_context).dump(record)
+
+        db.session.commit()
+
+        return response
+
+    def _build_clone_payload(self, source_id, override_data):
+        source_record = (
+            Studyset.query.options(
+                selectinload(Studyset.studies),
+                selectinload(Studyset.annotations).options(
+                    selectinload(Annotation.annotation_analyses)
+                ),
+            )
+            .filter_by(id=source_id)
+            .first()
+        )
+
+        if source_record is None:
+            abort_not_found(Studyset.__name__, source_id)
+
+        payload = {
+            "name": source_record.name,
+            "description": source_record.description,
+            "publication": source_record.publication,
+            "doi": source_record.doi,
+            "pmid": source_record.pmid,
+            "authors": source_record.authors,
+            "metadata_": (
+                deepcopy(source_record.metadata_)
+                if source_record.metadata_ is not None
+                else None
+            ),
+            "public": source_record.public,
+            "studies": [{"id": study.id} for study in source_record.studies],
+            "source": "neurostore",
+            "source_id": self._resolve_neurostore_origin(source_record),
+            "source_updated_at": source_record.updated_at or source_record.created_at,
+        }
+
+        if payload.get("metadata_") is None:
+            payload.pop("metadata_", None)
+
+        if override_data:
+            payload.update(override_data)
+
+        return payload, source_record
+
+    def _clone_annotations(self, source_record, cloned_record):
+        if not source_record.annotations:
+            return
+
+        owner_id = cloned_record.user_id
+
+        for annotation in source_record.annotations:
+            clone_annotation = Annotation(
+                name=annotation.name,
+                description=annotation.description,
+                source="neurostore",
+                source_id=self._resolve_neurostore_origin(annotation),
+                source_updated_at=annotation.updated_at or annotation.created_at,
+                user_id=owner_id,
+                metadata_=(
+                    deepcopy(annotation.metadata_) if annotation.metadata_ else None
+                ),
+                public=annotation.public,
+                note_keys=(
+                    deepcopy(annotation.note_keys) if annotation.note_keys else {}
+                ),
+            )
+            clone_annotation.studyset = cloned_record
+            db.session.add(clone_annotation)
+            db.session.flush()
+
+            analyses_to_create = []
+            for aa in annotation.annotation_analyses:
+                analyses_to_create.append(
+                    AnnotationAnalysis(
+                        annotation_id=clone_annotation.id,
+                        analysis_id=aa.analysis_id,
+                        note=deepcopy(aa.note) if aa.note else {},
+                        user_id=owner_id,
+                        study_id=aa.study_id,
+                        studyset_id=cloned_record.id,
+                    )
+                )
+
+            if analyses_to_create:
+                db.session.add_all(analyses_to_create)
+
+    @staticmethod
+    def _resolve_neurostore_origin(record):
+        source_id = record.id
+        parent_source_id = record.source_id
+        parent_source = getattr(record, "source", None)
+        Model = type(record)
+
+        while parent_source_id is not None and parent_source == "neurostore":
+            source_id = parent_source_id
+            parent = Model.query.filter_by(id=parent_source_id).first()
+            if parent is None:
+                break
+            parent_source_id = parent.source_id
+            parent_source = getattr(parent, "source", None)
+
+        return source_id
+
 
 @view_maker
 class AnnotationsView(ObjectView, ListView):

store/backend/neurostore/resources/data.py

@@ -569,6 +569,9 @@ class StudysetSchema(BaseDataSchema):
     studies = StringOrNested(
         StudySchema, many=True
     )  # This needs to be nested, but not cloned
+    source = fields.String(dump_only=True, allow_none=True)
+    source_id = fields.String(dump_only=True, allow_none=True)
+    source_updated_at = fields.DateTime(dump_only=True, allow_none=True)
 
     class Meta:
         additional = ("name", "description", "publication", "doi", "pmid")

store/backend/neurostore/schemas/data.py

@@ -148,3 +148,104 @@ def test_hot_swap_study_in_studyset(auth_client, ingest_neurosynth, session):
         == set(s for s in clone_ss_non_nested.json()["studies"])
         == set(s["id"] for s in clone_ss_nested.json()["studies"])
     )
+
+
+def _create_studyset_with_annotation(auth_client, study_ids, name="clone-source"):
+    studyset_resp = auth_client.post(
+        "/api/studysets/",
+        data={
+            "name": name,
+            "description": "clone me",
+            "studies": study_ids,
+        },
+    )
+    assert studyset_resp.status_code == 200
+    studyset_id = studyset_resp.json()["id"]
+
+    annotation_payload = {
+        "studyset": studyset_id,
+        "note_keys": {"include": "boolean"},
+        "name": "annotation for clone",
+    }
+    annotation_resp = auth_client.post("/api/annotations/", data=annotation_payload)
+    assert annotation_resp.status_code == 200
+
+    annotations = auth_client.get(f"/api/annotations/?studyset_id={studyset_id}")
+    assert annotations.status_code == 200
+    assert len(annotations.json()["results"]) >= 1
+
+    return studyset_resp.json(), annotations.json()["results"]
+
+
+def _study_ids_from_payload(studies):
+    ids = []
+    for entry in studies:
+        if isinstance(entry, dict):
+            ids.append(entry.get("id"))
+        else:
+            ids.append(entry)
+    return ids
+
+
+def test_clone_studyset_copies_annotations_by_default(
+    auth_client, ingest_neurosynth, session
+):
+    studies_payload = auth_client.get("/api/studies/?page_size=2")
+    study_ids = [study["id"] for study in studies_payload.json()["results"]]
+
+    source_studyset, source_annotations = _create_studyset_with_annotation(
+        auth_client, study_ids
+    )
+
+    clone_resp = auth_client.post(
+        f"/api/studysets/?source_id={source_studyset['id']}", data={}
+    )
+
+    assert clone_resp.status_code == 200
+    clone_data = clone_resp.json()
+
+    assert clone_data["source"] == "neurostore"
+    assert clone_data["source_id"] == source_studyset["id"]
+    assert set(_study_ids_from_payload(clone_data["studies"])) == set(
+        _study_ids_from_payload(source_studyset["studies"])
+    )
+    assert clone_data["user"] == auth_client.username
+
+    cloned_annotations = auth_client.get(
+        f"/api/annotations/?studyset_id={clone_data['id']}"
+    )
+    assert cloned_annotations.status_code == 200
+    assert len(cloned_annotations.json()["results"]) == len(source_annotations)
+
+
+def test_clone_studyset_without_annotations_when_disabled(
+    auth_client, ingest_neurosynth, session
+):
+    studies_payload = auth_client.get("/api/studies/?page_size=2")
+    study_ids = [study["id"] for study in studies_payload.json()["results"]]
+
+    source_studyset, source_annotations = _create_studyset_with_annotation(
+        auth_client, study_ids, name="clone-source-no-annots"
+    )
+
+    assert len(source_annotations) >= 1
+
+    clone_resp = auth_client.post(
+        f"/api/studysets/?source_id={source_studyset['id']}&copy_annotations=false",
+        data={},
+    )
+
+    assert clone_resp.status_code == 200
+    clone_data = clone_resp.json()
+
+    assert set(_study_ids_from_payload(clone_data["studies"])) == set(
+        _study_ids_from_payload(source_studyset["studies"])
+    )
+    assert clone_data["source"] == "neurostore"
+    assert clone_data["source_id"] == source_studyset["id"]
+
+    cloned_annotations = auth_client.get(
+        f"/api/annotations/?studyset_id={clone_data['id']}"
+    )
+    assert cloned_annotations.status_code == 200
+    assert cloned_annotations.json()["results"] == []