run count and search concurrently, but don't wait for count to be done

Author: Phil Varner

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -1,4 +1,5 @@
 """Database logic."""
+import asyncio
 import logging
 from base64 import urlsafe_b64decode, urlsafe_b64encode
 from typing import Dict, List, Optional, Tuple, Type, Union
@@ -196,24 +197,28 @@ async def execute_search(
         base_url: str,
     ) -> Tuple[List[Item], Optional[int], Optional[str]]:
         """Database logic to execute search with limit."""
-        body = search.to_dict()
-
-        maybe_count = (
-            await self.client.count(index=ITEMS_INDEX, body=search.to_dict(count=True))
-        ).get("count")
-
         search_after = None
         if token:
             search_after = urlsafe_b64decode(token.encode()).decode().split(",")
 
-        es_response = await self.client.search(
-            index=ITEMS_INDEX,
-            query=body.get("query"),
-            sort=sort or DEFAULT_SORT,
-            search_after=search_after,
-            size=limit,
+        query = search.query.to_dict() if search.query else None
+
+        search_task = asyncio.create_task(
+            self.client.search(
+                index=ITEMS_INDEX,
+                query=query,
+                sort=sort or DEFAULT_SORT,
+                search_after=search_after,
+                size=limit,
+            )
+        )
+
+        count_task = asyncio.create_task(
+            self.client.count(index=ITEMS_INDEX, body=search.to_dict(count=True))
         )
 
+        es_response = await search_task
+
         hits = es_response["hits"]["hits"]
         items = [
             self.item_serializer.db_to_stac(hit["_source"], base_url=base_url)
@@ -226,6 +231,15 @@ async def execute_search(
                 ",".join([str(x) for x in sort_array]).encode()
             ).decode()
 
+        # (1) count should not block returning results, so don't wait for it to be done
+        # (2) don't cancel the task so that it will populate the ES cache for subsequent counts
+        maybe_count = None
+        if count_task.done():
+            try:
+                maybe_count = count_task.result().get("count")
+            except Exception as e:  # type: ignore
+                logger.error(f"Count task failed: {e}")
+
         return items, maybe_count, next_token
 
     """ TRANSACTION LOGIC """