Create add_collections_to_body in sfeos_helpers and added search support for large amount of queries to ElasticSearch database logic.

Author: simonvb00

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -27,6 +27,8 @@
     PartialItem,
     PatchOperation,
 )
+
+from stac_fastapi.opensearch.stac_fastapi.opensearch.database_logic import ES_MAX_URL_LENGTH
 from stac_fastapi.sfeos_helpers import filter
 from stac_fastapi.sfeos_helpers.database import (
     apply_free_text_filter_shared,
@@ -60,6 +62,8 @@
 from stac_fastapi.types.rfc3339 import DateTimeType
 from stac_fastapi.types.stac import Collection, Item
 
+from stac_fastapi.sfeos_helpers.stac_fastapi.sfeos_helpers.database.query import add_collections_to_body
+
 logger = logging.getLogger(__name__)
 
 
@@ -520,6 +524,9 @@ async def execute_search(
         query = search.query.to_dict() if search.query else None
 
         index_param = indices(collection_ids)
+        if len(index_param) > ES_MAX_URL_LENGTH - 300:
+            index_param = ITEM_INDICES
+            query = add_collections_to_body(collection_ids, query)
 
         max_result_window = MAX_LIMIT
 

stac_fastapi/opensearch/stac_fastapi/opensearch/database_logic.py

@@ -62,8 +62,10 @@
 from stac_fastapi.types.rfc3339 import DateTimeType
 from stac_fastapi.types.stac import Collection, Item
 
+from stac_fastapi.sfeos_helpers.stac_fastapi.sfeos_helpers.database.query import add_collections_to_body, \
+    ES_MAX_URL_LENGTH
+
 logger = logging.getLogger(__name__)
-ES_MAX_URL_LENGTH = 4096
 
 
 async def create_index_templates() -> None:
@@ -533,6 +535,12 @@ async def execute_search(
         """
         search_body: Dict[str, Any] = {}
         query = search.query.to_dict() if search.query else None
+
+        index_param = indices(collection_ids)
+        if len(index_param) > ES_MAX_URL_LENGTH - 300:
+            index_param = ITEM_INDICES
+            query = add_collections_to_body(collection_ids, query)
+
         if query:
             search_body["query"] = query
 
@@ -545,19 +553,6 @@ async def execute_search(
 
         search_body["sort"] = sort if sort else DEFAULT_SORT
 
-        index_param = indices(collection_ids)
-
-        if len(index_param) > ES_MAX_URL_LENGTH - 300:
-            index_param = ITEM_INDICES
-            index_filter = {"terms": {"collection": collection_ids}}
-            if "bool" not in search_body["query"]:
-                search_body["query"]["bool"] = {}
-            if "filter" not in search_body["query"]["bool"]:
-                search_body["query"]["bool"]["filter"] = [index_filter]
-            filters = search_body["query"]["bool"]["filter"]
-            if index_filter not in filters:
-                filters.append(index_filter)
-
         max_result_window = MAX_LIMIT
 
         size_limit = min(limit + 1, max_result_window)

stac_fastapi/sfeos_helpers/stac_fastapi/sfeos_helpers/database/query.py

@@ -7,6 +7,7 @@
 
 from stac_fastapi.sfeos_helpers.mappings import Geometry
 
+ES_MAX_URL_LENGTH = 4096
 
 def apply_free_text_filter_shared(
     search: Any, free_text_queries: Optional[List[str]]
@@ -83,3 +84,31 @@ def populate_sort_shared(sortby: List) -> Optional[Dict[str, Dict[str, str]]]:
         return {s.field: {"order": s.direction} for s in sortby}
     else:
         return None
+
+
+def add_collections_to_body(collection_ids: List[str], query: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    """Adds a list of collection ids to the body of a query.
+
+    Args:
+        collection_ids (List[str]): A list of collections ids.
+        query (Optional[Dict[str, Any]]): The query to add collections to. If none, create a query that filters
+        the collection ids.
+
+    Returns:
+        Dict[str, Any]: A query that contains a filter on the given collection ids.
+
+    Notes:
+        This function is needed in the execute_search function when the size of the URL path will exceed the maximum of ES.
+    """
+    index_filter = {"terms": {"collection": collection_ids}}
+    if query is None:
+        query = {"query": {}}
+    if "bool" not in query:
+        query["bool"] = {}
+    if "filter" not in query["bool"]:
+        query["bool"]["filter"] = []
+
+    filters = query["bool"]["filter"]
+    if index_filter not in filters:
+        filters.append(index_filter)
+    return query