Skip to content

Commit 06999c9

Browse files
committed
Create add_collections_to_body in sfeos_helpers and added search support for large amount of queries to ElasticSearch database logic.
1 parent ab4f56e commit 06999c9

File tree

3 files changed

+45
-14
lines changed

3 files changed

+45
-14
lines changed

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
PartialItem,
2828
PatchOperation,
2929
)
30+
31+
from stac_fastapi.opensearch.stac_fastapi.opensearch.database_logic import ES_MAX_URL_LENGTH
3032
from stac_fastapi.sfeos_helpers import filter
3133
from stac_fastapi.sfeos_helpers.database import (
3234
apply_free_text_filter_shared,
@@ -60,6 +62,8 @@
6062
from stac_fastapi.types.rfc3339 import DateTimeType
6163
from stac_fastapi.types.stac import Collection, Item
6264

65+
from stac_fastapi.sfeos_helpers.stac_fastapi.sfeos_helpers.database.query import add_collections_to_body
66+
6367
logger = logging.getLogger(__name__)
6468

6569

@@ -520,6 +524,9 @@ async def execute_search(
520524
query = search.query.to_dict() if search.query else None
521525

522526
index_param = indices(collection_ids)
527+
if len(index_param) > ES_MAX_URL_LENGTH - 300:
528+
index_param = ITEM_INDICES
529+
query = add_collections_to_body(collection_ids, query)
523530

524531
max_result_window = MAX_LIMIT
525532

stac_fastapi/opensearch/stac_fastapi/opensearch/database_logic.py

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,10 @@
6262
from stac_fastapi.types.rfc3339 import DateTimeType
6363
from stac_fastapi.types.stac import Collection, Item
6464

65+
from stac_fastapi.sfeos_helpers.stac_fastapi.sfeos_helpers.database.query import add_collections_to_body, \
66+
ES_MAX_URL_LENGTH
67+
6568
logger = logging.getLogger(__name__)
66-
ES_MAX_URL_LENGTH = 4096
6769

6870

6971
async def create_index_templates() -> None:
@@ -533,6 +535,12 @@ async def execute_search(
533535
"""
534536
search_body: Dict[str, Any] = {}
535537
query = search.query.to_dict() if search.query else None
538+
539+
index_param = indices(collection_ids)
540+
if len(index_param) > ES_MAX_URL_LENGTH - 300:
541+
index_param = ITEM_INDICES
542+
query = add_collections_to_body(collection_ids, query)
543+
536544
if query:
537545
search_body["query"] = query
538546

@@ -545,19 +553,6 @@ async def execute_search(
545553

546554
search_body["sort"] = sort if sort else DEFAULT_SORT
547555

548-
index_param = indices(collection_ids)
549-
550-
if len(index_param) > ES_MAX_URL_LENGTH - 300:
551-
index_param = ITEM_INDICES
552-
index_filter = {"terms": {"collection": collection_ids}}
553-
if "bool" not in search_body["query"]:
554-
search_body["query"]["bool"] = {}
555-
if "filter" not in search_body["query"]["bool"]:
556-
search_body["query"]["bool"]["filter"] = [index_filter]
557-
filters = search_body["query"]["bool"]["filter"]
558-
if index_filter not in filters:
559-
filters.append(index_filter)
560-
561556
max_result_window = MAX_LIMIT
562557

563558
size_limit = min(limit + 1, max_result_window)

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from stac_fastapi.sfeos_helpers.mappings import Geometry
99

10+
ES_MAX_URL_LENGTH = 4096
1011

1112
def apply_free_text_filter_shared(
1213
search: Any, free_text_queries: Optional[List[str]]
@@ -83,3 +84,31 @@ def populate_sort_shared(sortby: List) -> Optional[Dict[str, Dict[str, str]]]:
8384
return {s.field: {"order": s.direction} for s in sortby}
8485
else:
8586
return None
87+
88+
89+
def add_collections_to_body(collection_ids: List[str], query: Optional[Dict[str, Any]]) -> Dict[str, Any]:
90+
"""Adds a list of collection ids to the body of a query.
91+
92+
Args:
93+
collection_ids (List[str]): A list of collections ids.
94+
query (Optional[Dict[str, Any]]): The query to add collections to. If none, create a query that filters
95+
the collection ids.
96+
97+
Returns:
98+
Dict[str, Any]: A query that contains a filter on the given collection ids.
99+
100+
Notes:
101+
This function is needed in the execute_search function when the size of the URL path will exceed the maximum of ES.
102+
"""
103+
index_filter = {"terms": {"collection": collection_ids}}
104+
if query is None:
105+
query = {"query": {}}
106+
if "bool" not in query:
107+
query["bool"] = {}
108+
if "filter" not in query["bool"]:
109+
query["bool"]["filter"] = []
110+
111+
filters = query["bool"]["filter"]
112+
if index_filter not in filters:
113+
filters.append(index_filter)
114+
return query

0 commit comments

Comments
 (0)