update fixtures and solr implementation

backend/src/monarch_py/api/association.py

@@ -5,7 +5,7 @@
 from monarch_py.api.additional_models import OutputFormat, PaginationParams
 from monarch_py.api.config import solr
 from monarch_py.datamodels.model import AssociationResults, MultiEntityAssociationResults
-from monarch_py.datamodels.category_enums import AssociationCategory, AssociationPredicate
+from monarch_py.datamodels.category_enums import AssociationCategory, AssociationPredicate, EntityCategory
 from monarch_py.utils.format_utils import to_tsv
 
 router = APIRouter(
@@ -17,18 +17,18 @@
 @router.get("")
 @router.get("/all", include_in_schema=False)  # We can remove this once the chatgpt plugin & oak aren't using it
 async def _get_associations(
-    category: Union[List[AssociationCategory], None] = Query(default=None),
+    category: List[AssociationCategory] = Query(default_factory=list),
     subject: Union[List[str], None] = Query(default=None),
-    predicate: Union[List[AssociationPredicate], None] = Query(default=None),
-    object: Union[List[str], None] = Query(default=None),
-    entity: Union[List[str], None] = Query(default=None),
-    subject_category: Union[List[str], None] = Query(default=None),
+    subject_category: List[EntityCategory] = Query(default_factory=list),
     subject_namespace: Union[List[str], None] = Query(default=None),
     subject_taxon: Union[List[str], None] = Query(default=None),
-    object_taxon: Union[List[str], None] = Query(default=None),
-    object_category: Union[List[str], None] = Query(default=None),
+    predicate: List[AssociationPredicate] = Query(default_factory=list),
+    object: Union[List[str], None] = Query(default=None),
+    object_category: List[EntityCategory] = Query(default_factory=list),
     object_namespace: Union[List[str], None] = Query(default=None),
-    direct: Union[bool, None] = Query(default=None),
+    object_taxon: Union[List[str], None] = Query(default=None),
+    entity: Union[List[str], None] = Query(default=None),
+    direct: bool = Query(default=False),
     pagination: PaginationParams = Depends(),
     format: OutputFormat = Query(
         default=OutputFormat.json,
@@ -37,21 +37,17 @@ async def _get_associations(
     ),
 ) -> Union[AssociationResults, str]:
     """Retrieves all associations for a given entity, or between two entities."""
-    if category:
-        category = [c.value if isinstance(c, AssociationCategory) else c for c in category]
-    if predicate:
-        predicate = [p.value if isinstance(p, AssociationPredicate) else p for p in predicate]
     response = solr().get_associations(
-        category=category,
-        predicate=predicate,
+        category=[category.value for category in category],
+        predicate=[predicate.value for predicate in predicate],
         subject=subject,
         object=object,
         entity=entity,
-        subject_category=subject_category,
+        subject_category=[subject_category.value for subject_category in subject_category],
         subject_namespace=subject_namespace,
         subject_taxon=subject_taxon,
         object_taxon=object_taxon,
-        object_category=object_category,
+        object_category=[object_category.value for object_category in object_category],
         object_namespace=object_namespace,
         direct=direct,
         offset=pagination.offset,

backend/src/monarch_py/datamodels/solr.py

@@ -1,6 +1,6 @@
 import urllib
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from monarch_py.utils.utils import escape
 from pydantic import BaseModel, Field
@@ -51,18 +51,20 @@ class SolrQuery(BaseModel):
     boost: Optional[str] = None
     sort: Optional[str] = None
 
-    def add_field_filter_query(self, field, value):
-        if field is not None and value is not None:
-            self.filter_queries.append(f"{field}:{escape(value)}")
-        else:
-            raise ValueError("Can't add a field filter query without a field and value")
+    def add_field_filter_query(self, field: str, value: Union[list, str]):
+        if not value:
+            return self
+        if isinstance(value, list) and not len(value) == 0:
+            value = " OR ".join(value)
+        self.filter_queries.append(f"{field}:{escape(value)}")
         return self
 
-    def add_filter_query(self, filter_query):
-        if filter_query is not None:
-            self.filter_queries.append(filter_query)
-        else:
-            raise ValueError("Can't append an empty filter query")
+    def add_filter_query(self, filter_query: Union[list, str]):
+        if not filter_query:
+            return self
+        if isinstance(filter_query, list):
+            filter_query = " OR ".join(escape(filter_query))
+        self.filter_queries.append(filter_query)
         return self
 
     def query_string(self):

backend/src/monarch_py/implementations/solr/solr_implementation.py

@@ -219,12 +219,18 @@ def get_associations(
         self,
         category: List[str] = None,
         subject: List[str] = None,
+        subject_closure: str = None,
+        subject_category: List[str] = None,
+        subject_namespace: List[str] = None,
+        subject_taxon: List[str] = None,
         predicate: List[str] = None,
         object: List[str] = None,
-        subject_closure: str = None,
         object_closure: str = None,
+        object_category: List[str] = None,
+        object_namespace: List[str] = None,
+        object_taxon: List[str] = None,
         entity: List[str] = None,
-        direct: bool = None,
+        direct: bool = False,
         q: str = None,
         offset: int = 0,
         limit: int = 20,
@@ -256,6 +262,12 @@ def get_associations(
             entity=[entity] if isinstance(entity, str) else entity,
             subject_closure=subject_closure,
             object_closure=object_closure,
+            subject_category=[subject_category] if isinstance(subject_category, str) else subject_category,
+            subject_namespace=[subject_namespace] if isinstance(subject_namespace, str) else subject_namespace,
+            subject_taxon=[subject_taxon] if isinstance(subject_taxon, str) else subject_taxon,
+            object_category=[object_category] if isinstance(object_category, str) else object_category,
+            object_taxon=[object_taxon] if isinstance(object_taxon, str) else object_taxon,
+            object_namespace=[object_namespace] if isinstance(object_namespace, str) else object_namespace,
             direct=direct,
             q=q,
             offset=offset,

backend/src/monarch_py/implementations/solr/solr_query_utils.py

@@ -7,13 +7,19 @@
 
 def build_association_query(
     category: List[str] = None,
-    predicate: List[str] = None,
     subject: List[str] = None,
-    object: List[str] = None,
     subject_closure: str = None,
+    subject_category: List[str] = None,
+    subject_namespace: List[str] = None,
+    subject_taxon: List[str] = None,
+    predicate: List[str] = None,
+    object: List[str] = None,
     object_closure: str = None,
+    object_category: List[str] = None,
+    object_namespace: List[str] = None,
+    object_taxon: List[str] = None,
     entity: List[str] = None,
-    direct: bool = None,
+    direct: bool = False,
     q: str = None,
     sort: List[str] = None,
     facet_fields: List[str] = None,
@@ -23,24 +29,26 @@ def build_association_query(
 ) -> SolrQuery:
     """Populate a SolrQuery object with association filters"""
     query = SolrQuery(start=offset, rows=limit)
-    if category:
-        query.add_filter_query(" OR ".join([f"category:{escape(cat)}" for cat in category]))
-    if predicate:
-        query.add_filter_query(" OR ".join([f"predicate:{escape(pred)}" for pred in predicate]))
+    query.add_field_filter_query("category", category)
+    query.add_field_filter_query("predicate", predicate)
+    query.add_field_filter_query("subject_closure", subject_closure)
+    query.add_field_filter_query("subject_category", subject_category)
+    query.add_field_filter_query("subject_namespace", subject_namespace)
+    query.add_field_filter_query("subject_taxon", subject_taxon)
+    query.add_field_filter_query("object_closure", object_closure)
+    query.add_field_filter_query("object_category", object_category)
+    query.add_field_filter_query("object_namespace", object_namespace)
+    query.add_field_filter_query("object_taxon", object_taxon)
     if subject:
         if direct:
             query.add_field_filter_query("subject", " OR ".join(subject))
         else:
             query.add_filter_query(" OR ".join([f'subject:"{s}" OR subject_closure:"{s}"' for s in subject]))
-    if subject_closure:
-        query.add_field_filter_query("subject_closure", subject_closure)
     if object:
         if direct:
             query.add_field_filter_query("object", " OR ".join(object))
         else:
             query.add_filter_query(" OR ".join([f'object:"{o}" OR object_closure:"{o}"' for o in object]))
-    if object_closure:
-        query.add_field_filter_query("object_closure", object_closure)
     if entity:
         if direct:
             query.add_filter_query(" OR ".join([f'subject:"{escape(e)}" OR object:"{escape(e)}"' for e in entity]))
@@ -65,6 +73,7 @@ def build_association_query(
         query.facet_fields = facet_fields
     if facet_queries:
         query.facet_queries = facet_queries
+    print(query)
     return query
 
 

backend/tests/api/test_association_router.py

@@ -17,3 +17,29 @@ def test_associations(associations):
     assert response.status_code == 200
     assert response.headers["content-type"] == "application/json"
     assert response.json() == associations
+
+
+@patch("monarch_py.implementations.solr.solr_implementation.SolrImplementation.get_associations")
+def test_associations_params(mock_get_assoc):
+    """Test that the correct parameters are passed to the monarch api"""
+    mock_get_assoc.return_value = MagicMock()
+    client.get(
+        "/association?category=biolink%3AAssociation&subject=HGNC%3A0000001&predicate=biolink%3Ainteracts_with&object=MONDO%3A0000001&entity=MONDO%3A0000002&subject_category=biolink%3ADisease&subject_namespace=HGNC&subject_taxon=NCBITaxon%3A1234&object_taxon=NCBITaxon%3A9876&object_category=biolink%3AGene&object_namespace=MONDO&direct=true&format=json&limit=20&offset=0"
+    )
+    mock_get_assoc.assert_called_with(
+        category=["biolink:Association"],
+        subject=["HGNC:0000001"],
+        predicate=["biolink:interacts_with"],
+        object=["MONDO:0000001"],
+        entity=["MONDO:0000002"],
+        subject_category=["biolink:Disease"],
+        subject_namespace=["HGNC"],
+        subject_taxon=["NCBITaxon:1234"],
+        object_taxon=["NCBITaxon:9876"],
+        object_category=["biolink:Gene"],
+        object_namespace=["MONDO"],
+        direct=True,
+        format="json",
+        offset=0,
+        limit=20,
+    )

backend/tests/fixtures/association_counts.py

@@ -5,7 +5,7 @@
 def association_counts():
     return {
         "items": [
-            {"label": "Phenotypes", "count": 3879, "category": "biolink:DiseaseToPhenotypicFeatureAssociation"},
+            {"label": "Phenotypes", "count": 3875, "category": "biolink:DiseaseToPhenotypicFeatureAssociation"},
             {"label": "Causal Genes", "count": 124, "category": "biolink:CausalGeneToDiseaseAssociation"},
             {"label": "Correlated Genes", "count": 139, "category": "biolink:CorrelatedGeneToDiseaseAssociation"},
         ]