wip: speed up PUT

store/neurostore/tests/api/test_studysets.py

@@ -1,3 +1,6 @@
+import random
+import string
+from neurostore.tests.conftest import add_event_listeners
 from neurostore.models import Studyset, Study
 
 
@@ -20,6 +23,40 @@ def test_post_and_get_studysets(auth_client, ingest_neurosynth, session):
         == post_resp.json()
     )
 
+@add_event_listeners
+def test_add_many_studies_to_studyset(auth_client, ingest_neurosynth, session):
+    existing_studies = Study.query.all()
+    existing_study_ids = [s.id for s in existing_studies]
+
+    # Function to generate a random DOI
+    def generate_doi():
+        doi = "10." + "".join(random.choices(string.digits, k=4)) + "/"
+        doi += "".join(random.choices(string.ascii_lowercase, k=4)) + "."
+        doi += "".join(random.choices(string.ascii_lowercase, k=4))
+        return doi
+
+    # List comprehension to generate the desired structure
+    made_up_studies = [
+        {
+            "pmid": random.randint(100000, 999999),
+            "doi": generate_doi(),
+            "name": ''.join(random.choices(string.ascii_letters, k=10)),
+        } for _ in range(1)
+    ]
+    # create empty studyset
+    ss = auth_client.post("/api/studysets/", data={"name": "mixed_studyset"})
+
+    assert ss.status_code == 200
+
+    ss_id = ss.json()['id']
+
+    # combine made_up and created studies
+    all_studies = existing_study_ids# + made_up_studies
+
+
+    ss_update = auth_client.put(f"/api/studysets/{ss_id}", data={"studies": all_studies})
+
+    assert ss_update.status_code == 200
 
 def test_add_study_to_studyset(auth_client, ingest_neurosynth, session):
     payload = auth_client.get("/api/studies/").json()

store/neurostore/tests/conftest.py

@@ -1,5 +1,7 @@
 import pytest
 from os import environ
+import functools
+from datetime import datetime
 from neurostore.models.data import Analysis, Condition
 from ..database import db as _db
 import sqlalchemy as sa
@@ -20,6 +22,30 @@
 from auth0.v3.authentication import GetToken
 import shortuuid
 
+import logging
+
+LOGGER = logging.getLogger(__name__)
+
+
+"""
+Record sql queries
+"""
+def add_event_listeners(f):
+
+    @functools.wraps(f)
+    def event_listeners(*args, **kwargs):
+        @sa.event.listens_for(_db.engine, "before_cursor_execute")
+        def _record_query_start(conn, cursor, statement, parameters, context, executemany):
+            conn.info["query_start"] = datetime.now()
+
+        @sa.event.listens_for(_db.engine, "after_cursor_execute")
+        def _calculate_query_run_time(conn, cursor, statement, parameters, context, executemany):
+            LOGGER.warning(f"\n\n{statement}")
+            LOGGER.warning("this query took {}".format(datetime.now() - conn.info["query_start"]))
+            print("here!")
+        return f(*args, **kwargs)
+
+    return event_listeners
 
 """
 Test selection arguments