Repository: cancervariants/metakb (language: python)

Issue:
Create a batch get for the metakb API 
Hi there! Wanted to get a conversation started for implementing a batch getter for the metakb API.

In contexts like processing a large number of variant IDs from a VCF, it could be helpful to send multiple IDs at once. For reference, on my 2017 MacBook Pro (2.3 GHz Dual-Core Intel Core i5 cores, 8GB memory) with multiprocessing, I max out at ~25 calls per second. Using this, working with a 16MB VCF (one patient, one chromosome) with 330k variants would take ~3.5 hours. The compute doesn't seem to be a problem, as I'm far from maxing out my CPU or memory. Considering this is a relatively small VCF, scaling this to the multi-patient, larger VCFs, and all chromosomes could end up being lengthy regardless of the machine.

Just wanted to provide some info for one use case, but happy to provide more info or discuss this more in detail.
Streamline graph DB queries
It's been true for a while now that some of our core API endpoints require multiple ordered requests to the graph DB to succeed. For example, searching for a study first retrieves that study and its relationships, but can then spur further calls for certain kinds of variations and therapies. This is totally reasonable for now but as requirements and the data models begin to settle, we should be aiming to condense those protocols into a single DB call eventually.


Discussion / hints:
@quinnwai Thanks for creating an issue for this. MetaKB V2 is still in early stages of development. This is a reasonable ask, I think the question is more on where it lives on the roadmap and which team would implement this. We can discuss this more at the MetaKB meeting on Tuesday and report back here afterwards.
It might also be worth looking at whether some queries should be managed transactions vs the new `driver.execute_query` API https://neo4j.com/docs/python-manual/current/query-simple/
`get_nested_study` makes a LOT of queries. Ideally this would be one. At minimum, rather than iterating through relations, there should be a single structured result.

Fix patch:
diff --git a/src/metakb/main.py b/src/metakb/main.py
index 75b9b6c4..849a0c22 100644
--- a/src/metakb/main.py
+++ b/src/metakb/main.py
@@ -98,3 +98,29 @@ async def get_studies(
     """
     resp = await query.search_studies(variation, disease, therapy, gene, study_id)
     return resp.model_dump(exclude_none=True)
+
+
+_batch_search_studies_descr = {
+    "summary": "Get nested studies for all provided variations.",
+    "description": "Return nested studies associated with any of the provided variations.",
+    "arg_variations": "Variations (subject) to search. Can be free text or VRS variation ID.",
+}
+
+
+@app.get(
+    "/api/v2/batch_search/studies",
+    summary=_batch_search_studies_descr["summary"],
+    description=_batch_search_studies_descr["description"],
+)
+async def batch_get_studies(
+    variations: list[str] | None = Query(  # noqa: B008
+        None, description=_batch_search_studies_descr["arg_variations"]
+    ),
+) -> dict:
+    """Fetch all studies associated with `any` of the provided variations.
+
+    :param variations: variations to match against
+    :return: batch response object
+    """
+    response = await query.batch_search_studies(variations)
+    return response.model_dump(exclude_none=True)
diff --git a/src/metakb/query.py b/src/metakb/query.py
index 7fda72c6..9a746589 100644
--- a/src/metakb/query.py
+++ b/src/metakb/query.py
@@ -20,7 +20,13 @@
 from metakb.database import Graph
 from metakb.normalizers import ViccNormalizers
 from metakb.schemas.annotation import Document, Method
-from metakb.schemas.api import SearchStudiesService, ServiceMeta
+from metakb.schemas.api import (
+    BatchSearchStudiesQuery,
+    BatchSearchStudiesService,
+    NormalizedQuery,
+    SearchStudiesService,
+    ServiceMeta,
+)
 from metakb.schemas.app import SourceName
 from metakb.schemas.categorical_variation import CategoricalVariation
 from metakb.schemas.variation_statement import (
@@ -708,3 +714,66 @@ def _get_therapeutic_agent(in_ta_params: dict) -> TherapeuticAgent:
 
         ta_params["extensions"] = extensions
         return TherapeuticAgent(**ta_params)
+
+    async def batch_search_studies(
+        self,
+        variations: list[str] | None = None,
+    ) -> BatchSearchStudiesService:
+        """Fetch all studies associated with any of the provided variation description
+        strings.
+
+        Because this method could be expanded to include other kinds of search terms,
+        ``variations`` is optionally nullable.
+
+        >>> from metakb.query import QueryHandler
+        >>> qh = QueryHandler()
+        >>> response = await qh.batch_search_studies(["EGFR L858R"])
+        >>> response.study_ids[:3]
+        ['civic.eid:229', 'civic.eid:3811', 'moa.assertion:268']
+
+        All terms are normalized, so redundant terms don't alter search results:
+
+        >>> redundant_response = await qh.batch_search_studies(
+        ...     ["EGFR L858R", "NP_005219.2:p.Leu858Arg"]
+        ... )
+        >>> len(response.study_ids) == len(redundant_response.study_ids)
+        True
+
+        :param variations: a list of variation description strings, e.g.
+            ``["BRAF V600E"]``
+        :return: response object including all matching studies
+        """
+        response = BatchSearchStudiesService(
+            query=BatchSearchStudiesQuery(variations=[]),
+            service_meta_=ServiceMeta(),
+            warnings=[],
+        )
+        if not variations:
+            return response
+
+        for query_variation in set(variations):
+            variation_id = await self._get_normalized_variation(
+                query_variation, response.warnings
+            )
+            response.query.variations.append(
+                NormalizedQuery(term=query_variation, normalized_id=variation_id)
+            )
+        variation_ids = list(
+            {v.normalized_id for v in response.query.variations if v.normalized_id}
+        )
+        if not variation_ids:
+            return response
+
+        with self.driver.session() as session:
+            query = """
+                MATCH (s) -[:HAS_VARIANT] -> (cv:CategoricalVariation)
+                MATCH (cv) -[:HAS_DEFINING_CONTEXT|HAS_MEMBERS] -> (v:Variation)
+                WHERE v.id IN $v_ids
+                RETURN s
+            """
+            result = session.run(query, v_ids=variation_ids)
+            study_nodes = [r[0] for r in result]
+            response.study_ids = list({n["id"] for n in study_nodes})
+            studies = self._get_nested_studies(session, study_nodes)
+            response.studies = [VariantTherapeuticResponseStudy(**s) for s in studies]
+        return response
diff --git a/src/metakb/schemas/api.py b/src/metakb/schemas/api.py
index adea7de2..37e9b4fc 100644
--- a/src/metakb/schemas/api.py
+++ b/src/metakb/schemas/api.py
@@ -45,3 +45,26 @@ class SearchStudiesService(BaseModel):
     study_ids: list[StrictStr] = []
     studies: list[VariantTherapeuticResponseStudy] = []
     service_meta_: ServiceMeta
+
+
+class NormalizedQuery(BaseModel):
+    """Define structure of user-provided query. If possible, add normalized ID."""
+
+    term: StrictStr
+    normalized_id: StrictStr | None = None
+
+
+class BatchSearchStudiesQuery(BaseModel):
+    """Define query as reported in batch search studies endpoint."""
+
+    variations: list[NormalizedQuery] = []
+
+
+class BatchSearchStudiesService(BaseModel):
+    """Define response model for batch search studies endpoint response."""
+
+    query: BatchSearchStudiesQuery
+    warnings: list[StrictStr] = []
+    study_ids: list[StrictStr] = []
+    studies: list[VariantTherapeuticResponseStudy] = []
+    service_meta_: ServiceMeta


Test patch:
diff --git a/tests/unit/test_search_studies.py b/tests/unit/test_search_studies.py
index e8108e02..85667841 100644
--- a/tests/unit/test_search_studies.py
+++ b/tests/unit/test_search_studies.py
@@ -3,7 +3,11 @@
 import pytest
 
 from metakb.query import QueryHandler
-from metakb.schemas.api import SearchStudiesService
+from metakb.schemas.api import (
+    BatchSearchStudiesService,
+    NormalizedQuery,
+    SearchStudiesService,
+)
 
 
 @pytest.fixture(scope="module")
@@ -30,7 +34,7 @@ def assert_no_match(response):
 
 
 def find_and_check_study(
-    resp: SearchStudiesService,
+    resp: SearchStudiesService | BatchSearchStudiesService,
     expected_study: dict,
     assertion_checks: callable,
     should_find_match: bool = True,
@@ -220,3 +224,48 @@ async def test_no_matches(query_handler):
     # valid queries, but no matches with combination
     resp = await query_handler.search_studies(variation="BRAF V600E", gene="EGFR")
     assert_no_match(resp)
+
+
+@pytest.mark.asyncio(scope="module")
+async def test_batch_search(
+    query_handler: QueryHandler,
+    assertion_checks,
+    civic_eid2997_study,
+    civic_eid816_study,
+):
+    """Test batch search studies method."""
+    assert_no_match(await query_handler.batch_search_studies([]))
+    assert_no_match(await query_handler.batch_search_studies(["gibberish variant"]))
+
+    braf_va_id = "ga4gh:VA.Otc5ovrw906Ack087o1fhegB4jDRqCAe"
+    braf_response = await query_handler.batch_search_studies([braf_va_id])
+    assert braf_response.query.variations == [
+        NormalizedQuery(
+            term=braf_va_id,
+            normalized_id=braf_va_id,
+        )
+    ]
+    find_and_check_study(braf_response, civic_eid816_study, assertion_checks)
+
+    redundant_braf_response = await query_handler.batch_search_studies(
+        [braf_va_id, "NC_000007.13:g.140453136A>T"]
+    )
+    assert redundant_braf_response.query.variations == [
+        NormalizedQuery(
+            term=braf_va_id,
+            normalized_id=braf_va_id,
+        ),
+        NormalizedQuery(
+            term="NC_000007.13:g.140453136A>T",
+            normalized_id=braf_va_id,
+        ),
+    ]
+    find_and_check_study(redundant_braf_response, civic_eid816_study, assertion_checks)
+    assert len(braf_response.study_ids) == len(redundant_braf_response.study_ids)
+
+    braf_egfr_response = await query_handler.batch_search_studies(
+        [braf_va_id, "EGFR L858R"]
+    )
+    find_and_check_study(braf_egfr_response, civic_eid816_study, assertion_checks)
+    find_and_check_study(braf_egfr_response, civic_eid2997_study, assertion_checks)
+    assert len(braf_egfr_response.study_ids) > len(braf_response.study_ids)