Fix type filter pagination without messing up type count aggregation

helxplatform · Aug 10, 2023 · 063af0e · 063af0e
1 parent 6b6c2a6
commit 063af0e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/src/dug/core/async_search.py b/src/dug/core/async_search.py
@@ -101,7 +101,7 @@ async def agg_data_type(self):
         return data_type_list
 
     @staticmethod
-    def _build_concepts_query(query, types=None, fuzziness=1, prefix_length=3):
+    def _build_concepts_query(query, fuzziness=1, prefix_length=3):
         "Static data structure populator, pulled for easier testing"
         query_object = {
             "bool": {
@@ -111,15 +111,7 @@ def _build_concepts_query(query, types=None, fuzziness=1, prefix_length=3):
                             {"wildcard": {"description": "?*"}},
                             {"wildcard": {"name": "?*"}}
                         ]
-                    },
-                    **({
-                        "bool": {
-                            "should": [
-                                {'term': {'type': {'value': t}}} for t in types
-                            ],
-                            "minimum_should_match": 1
-                        }
-                    } if isinstance(types, list) else {})
+                    }
                 },
                 "should": [
                     {
@@ -219,10 +211,19 @@ async def search_concepts(self, query, offset=0, size=None, types=None,
         """
         Changed to a long boolean match query to optimize search results
         """
-        query_dict = self._build_concepts_query(query, types, **kwargs)
+        query_dict = self._build_concepts_query(query, **kwargs)
         # Get aggregated counts of biolink types
         search_body = {"query": query_dict}
         search_body['aggs'] = {'type-count': {'terms': {'field': 'type'}}}
+        if isinstance(types, list):
+            search_body['post_filter'] = {
+                "bool": {
+                    "should": [
+                        {'term': {'type': {'value': t}}} for t in types
+                    ],
+                    "minimum_should_match": 1
+                }
+            }
         search_results = await self.es.search(
             index="concepts_index",
             body=search_body,
@@ -233,7 +234,14 @@ async def search_concepts(self, query, offset=0, size=None, types=None,
             size=size,
             explain=True
         )
+        # Aggs/post_filter aren't supported by count
         del search_body["aggs"]
+        if "post_filter" in search_body:
+            # We'll move the post_filter into the actual filter
+            search_body["query"]["bool"]["filter"]["bool"].update(
+                search_body["post_filter"]["bool"]
+            )
+            del search_body["post_filter"]
         total_items = await self.es.count(
             body=search_body,
             index="concepts_index"

diff --git a/tests/unit/test_async_search.py b/tests/unit/test_async_search.py
@@ -28,7 +28,7 @@ def setUp(self):
         "Build mock elasticsearch responses"
         search_result = _brain_search_result()
         self.search = async_search.Search(Config.from_env())
-        self.query_body = self.search._build_concepts_query("brain", types=None)
+        self.query_body = self.search._build_concepts_query("brain")
         self.search.es = es_mock
 
     def test_concepts_search(self):