feat: add retry for datetime search queries

- Add retry decorator for execute_search on datetime queries
- Retry NotFoundError/ConnectionError with exponential backoff
- Configurable retry count via STAC_SEARCH_MAX_RETRY env var

Author: Yuri Zmytrakov

stac_fastapi/core/stac_fastapi/core/utilities.py

@@ -4,9 +4,12 @@
 such as converting bounding boxes to polygon representations.
 """
 
+import asyncio
 import logging
 import os
-from typing import Any, Dict, List, Optional, Set, Union
+import random
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, Set, Union
 
 from stac_fastapi.types.stac import Item
 
@@ -195,3 +198,44 @@ def get_excluded_from_items(obj: dict, field_path: str) -> None:
             return
 
     current.pop(final, None)
+
+
+def datetime_search_retry(func: Callable) -> Callable:
+    """Retry decorator for datetime search queries for NotFoundError/ConnectionError."""
+
+    @wraps(func)
+    async def wrapper(self, *args, **kwargs):
+        datetime_search = kwargs.get("datetime_search")
+
+        is_datetime_query = bool(datetime_search and isinstance(datetime_search, dict))
+
+        if not is_datetime_query:
+            return await func(self, *args, **kwargs)
+
+        base_delay = 0.5
+        max_retries = int(os.getenv("MAX_RETRY", "3"))
+
+        # Validate max_retries is between 1 and 10
+        if max_retries < 1 or max_retries > 10:
+            max_retries = 3
+
+        for attempt in range(max_retries):
+            try:
+                return await func(self, *args, **kwargs)
+            except Exception as e:
+                error_name = type(e).__name__.lower()
+                error_msg = str(e).lower()
+
+                retry_error = (
+                    "notfound" in error_name
+                    or "connection" in error_name
+                    or "not found" in error_msg
+                )
+
+                if not retry_error or attempt == max_retries - 1:
+                    raise
+                delay = base_delay * (2**attempt) + random.uniform(0, 0.1)
+                await asyncio.sleep(delay)
+        raise Exception("Search failed after retries")
+
+    return wrapper

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -22,7 +22,12 @@
     CollectionSerializer,
     ItemSerializer,
 )
-from stac_fastapi.core.utilities import MAX_LIMIT, bbox2polygon, get_bool_env
+from stac_fastapi.core.utilities import (
+    MAX_LIMIT,
+    bbox2polygon,
+    datetime_search_retry,
+    get_bool_env,
+)
 from stac_fastapi.elasticsearch.config import AsyncElasticsearchSettings
 from stac_fastapi.elasticsearch.config import (
     ElasticsearchSettings as SyncElasticsearchSettings,
@@ -804,6 +809,7 @@ def populate_sort(sortby: List) -> Optional[Dict[str, Dict[str, str]]]:
         """
         return populate_sort_shared(sortby=sortby)
 
+    @datetime_search_retry
     async def execute_search(
         self,
         search: Search,

stac_fastapi/opensearch/stac_fastapi/opensearch/database_logic.py

@@ -19,7 +19,12 @@
 from stac_fastapi.core.base_database_logic import BaseDatabaseLogic
 from stac_fastapi.core.datetime_utils import format_datetime_range
 from stac_fastapi.core.serializers import CollectionSerializer, ItemSerializer
-from stac_fastapi.core.utilities import MAX_LIMIT, bbox2polygon, get_bool_env
+from stac_fastapi.core.utilities import (
+    MAX_LIMIT,
+    bbox2polygon,
+    datetime_search_retry,
+    get_bool_env,
+)
 from stac_fastapi.extensions.core.transaction.request import (
     PartialCollection,
     PartialItem,
@@ -831,6 +836,7 @@ def populate_sort(sortby: List) -> Optional[Dict[str, Dict[str, str]]]:
         """
         return populate_sort_shared(sortby=sortby)
 
+    @datetime_search_retry
     async def execute_search(
         self,
         search: Search,

stac_fastapi/tests/api/test_api.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+from stac_fastapi.core.utilities import datetime_search_retry
 from stac_fastapi.types.errors import ConflictError
 
 from ..conftest import create_collection, create_item
@@ -1722,3 +1723,62 @@ async def test_hide_private_data_from_item(app_client, txn_client, load_test_dat
     assert "private_data" not in item["properties"]
 
     del os.environ["EXCLUDED_FROM_ITEMS"]
+
+
+@pytest.mark.asyncio
+async def test_datetime_search_retry_succeeds_on_retry():
+    """Test retry works for datetime queries with NotFoundError."""
+    call_counter = {"count": 0}
+
+    @datetime_search_retry
+    async def mock_search(self, datetime_search=None):
+        call_counter["count"] += 1
+        if call_counter["count"] < 2:
+            raise Exception("not found: Index not found.")
+        return "success"
+
+    result = await mock_search(None, datetime_search={"datetime": "2025-01-01"})
+
+    assert result == "success"
+    assert call_counter["count"] == 2
+
+
+@pytest.mark.asyncio
+async def test_datetime_search_no_retry_for_none_datetime():
+    """Test that retry is invoked only for datetime queries."""
+    call_counter = {"count": 0}
+
+    @datetime_search_retry
+    async def mock_search(self, datetime_search=None):
+        call_counter["count"] += 1
+        if call_counter["count"] < 2:
+            raise Exception("not found: Index not found.")
+        return "success"
+
+    with pytest.raises(Exception):
+        await mock_search(None, datetime_search={})
+
+    assert call_counter["count"] == 1
+    call_counter["count"] = 0
+
+    with pytest.raises(Exception):
+        await mock_search(None, datetime_search=None)
+
+    assert call_counter["count"] == 1
+
+
+@pytest.mark.asyncio
+async def test_datetime_search_max_retries():
+    """Test that retry stops after max attempts."""
+    call_counter = {"count": 0}
+
+    @datetime_search_retry
+    async def mock_search(self, datetime_search=None):
+        call_counter["count"] += 1
+        raise Exception(f"Not found: Fails, attempt {call_counter['count']}")
+
+    with pytest.raises(Exception) as exc_info:
+        await mock_search(None, datetime_search={"start_datetime": "2025-01-01"})
+
+    assert call_counter["count"] == 3
+    assert "attempt 3" in str(exc_info.value).lower()