add wait_for, reorg es db logic

Author: jonhealy1

stac_fastapi/core/stac_fastapi/core/core.py

@@ -673,27 +673,6 @@ class TransactionsClient(AsyncBaseTransactionsClient):
     settings: ApiBaseSettings = attr.ib()
     session: Session = attr.ib(default=attr.Factory(Session.create_from_env))
 
-    def _resolve_refresh(self, **kwargs) -> bool:
-        """
-        Resolve the `refresh` parameter from kwargs or the environment variable.
-
-        Args:
-            **kwargs: Additional keyword arguments, including `refresh`.
-
-        Returns:
-            bool: The resolved value of the `refresh` parameter.
-        """
-        refresh = kwargs.get(
-            "refresh", self.database.async_settings.database_refresh == "true"
-        )
-        if "refresh" in kwargs:
-            logger.info(f"`refresh` parameter explicitly passed in kwargs: {refresh}")
-        else:
-            logger.info(
-                f"`refresh` parameter derived from environment variable: {refresh}"
-            )
-        return refresh
-
     @overrides
     async def create_item(
         self, collection_id: str, item: Union[Item, ItemCollection], **kwargs
@@ -717,9 +696,6 @@ async def create_item(
         request = kwargs.get("request")
         base_url = str(request.base_url)
 
-        # Resolve the `refresh` parameter
-        refresh = self._resolve_refresh(**kwargs)
-
         # Convert Pydantic model to dict for uniform processing
         item_dict = item.model_dump(mode="json")
 
@@ -738,9 +714,9 @@ async def create_item(
             attempted = len(processed_items)
 
             success, errors = await self.database.bulk_async(
-                collection_id,
-                processed_items,
-                refresh=refresh,
+                collection_id=collection_id,
+                processed_items=processed_items,
+                **kwargs,
             )
             if errors:
                 logger.error(
@@ -754,10 +730,7 @@ async def create_item(
 
         # Handle single item
         await self.database.create_item(
-            item_dict,
-            refresh=refresh,
-            base_url=base_url,
-            exist_ok=False,
+            item_dict, base_url=base_url, exist_ok=False, **kwargs
         )
         return ItemSerializer.db_to_stac(item_dict, base_url)
 
@@ -783,14 +756,11 @@ async def update_item(
         item = item.model_dump(mode="json")
         base_url = str(kwargs["request"].base_url)
 
-        # Resolve the `refresh` parameter
-        refresh = self._resolve_refresh(**kwargs)
-
         now = datetime_type.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
         item["properties"]["updated"] = now
 
         await self.database.create_item(
-            item, refresh=refresh, base_url=base_url, exist_ok=True
+            item, base_url=base_url, exist_ok=True, **kwargs
         )
 
         return ItemSerializer.db_to_stac(item, base_url)
@@ -806,11 +776,8 @@ async def delete_item(self, item_id: str, collection_id: str, **kwargs) -> None:
         Returns:
             None: Returns 204 No Content on successful deletion
         """
-        # Resolve the `refresh` parameter
-        refresh = self._resolve_refresh(**kwargs)
-
         await self.database.delete_item(
-            item_id=item_id, collection_id=collection_id, refresh=refresh
+            item_id=item_id, collection_id=collection_id, **kwargs
         )
         return None
 
@@ -833,11 +800,8 @@ async def create_collection(
         collection = collection.model_dump(mode="json")
         request = kwargs["request"]
 
-        # Resolve the `refresh` parameter
-        refresh = self._resolve_refresh(**kwargs)
-
         collection = self.database.collection_serializer.stac_to_db(collection, request)
-        await self.database.create_collection(collection=collection, refresh=refresh)
+        await self.database.create_collection(collection=collection, **kwargs)
         return CollectionSerializer.db_to_stac(
             collection,
             request,
@@ -871,12 +835,9 @@ async def update_collection(
 
         request = kwargs["request"]
 
-        # Resolve the `refresh` parameter
-        refresh = self._resolve_refresh(**kwargs)
-
         collection = self.database.collection_serializer.stac_to_db(collection, request)
         await self.database.update_collection(
-            collection_id=collection_id, collection=collection, refresh=refresh
+            collection_id=collection_id, collection=collection, **kwargs
         )
 
         return CollectionSerializer.db_to_stac(
@@ -901,12 +862,7 @@ async def delete_collection(self, collection_id: str, **kwargs) -> None:
         Raises:
             NotFoundError: If the collection doesn't exist
         """
-        # Resolve the `refresh` parameter
-        refresh = self._resolve_refresh(**kwargs)
-
-        await self.database.delete_collection(
-            collection_id=collection_id, refresh=refresh
-        )
+        await self.database.delete_collection(collection_id=collection_id, **kwargs)
         return None
 
 
@@ -965,19 +921,6 @@ def bulk_item_insert(
         else:
             base_url = ""
 
-        # Use `refresh` from kwargs if provided, otherwise fall back to the environment variable
-        refresh = kwargs.get(
-            "refresh", self.database.sync_settings.database_refresh == "true"
-        )
-
-        # Log the value of `refresh` and its source
-        if "refresh" in kwargs:
-            logger.info(f"`refresh` parameter explicitly passed in kwargs: {refresh}")
-        else:
-            logger.info(
-                f"`refresh` parameter derived from environment variable: {refresh}"
-            )
-
         processed_items = []
         for item in items.items.values():
             try:
@@ -996,7 +939,7 @@ def bulk_item_insert(
         success, errors = self.database.bulk_sync(
             collection_id,
             processed_items,
-            refresh=refresh,
+            **kwargs,
         )
         if errors:
             logger.error(f"Bulk sync operation encountered errors: {errors}")

stac_fastapi/core/stac_fastapi/core/utilities.py

@@ -39,6 +39,30 @@ def get_bool_env(name: str, default: bool = False) -> bool:
         return default
 
 
+def resolve_refresh(refresh: str) -> str:
+    """
+    Resolve the `refresh` parameter from kwargs or the environment variable.
+
+    Args:
+        refresh (str): The `refresh` parameter value.
+
+    Returns:
+        str: The resolved value of the `refresh` parameter, which can be "true", "false", or "wait_for".
+    """
+    logger = logging.getLogger(__name__)
+
+    # Normalize and validate the `refresh` value
+    refresh = refresh.lower()
+    if refresh not in {"true", "false", "wait_for"}:
+        raise ValueError(
+            "Invalid value for `refresh`. Must be 'true', 'false', or 'wait_for'."
+        )
+
+    # Log the resolved value
+    logger.info(f"`refresh` parameter resolved to: {refresh}")
+    return refresh
+
+
 def bbox2polygon(b0: float, b1: float, b2: float, b3: float) -> List[List[List[float]]]:
     """Transform a bounding box represented by its four coordinates `b0`, `b1`, `b2`, and `b3` into a polygon.
 

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/config.py

@@ -3,7 +3,7 @@
 import logging
 import os
 import ssl
-from typing import Any, Dict, Set
+from typing import Any, Dict, Set, Union
 
 import certifi
 from elasticsearch._async.client import AsyncElasticsearch
@@ -87,7 +87,24 @@ class ElasticsearchSettings(ApiSettings, ApiBaseSettings):
     enable_response_models: bool = False
     enable_direct_response: bool = get_bool_env("ENABLE_DIRECT_RESPONSE", default=False)
     raise_on_bulk_error: bool = get_bool_env("RAISE_ON_BULK_ERROR", default=False)
-    database_refresh: bool = get_bool_env("DATABASE_REFRESH", default=False)
+
+    @property
+    def database_refresh(self) -> Union[bool, str]:
+        """
+        Get the value of the DATABASE_REFRESH environment variable.
+
+        Returns:
+            Union[bool, str]: The value of DATABASE_REFRESH, which can be True, False, or "wait_for".
+        """
+        value = os.getenv("DATABASE_REFRESH", "false").lower()
+        if value in {"true", "false"}:
+            return value == "true"
+        elif value == "wait_for":
+            return "wait_for"
+        else:
+            raise ValueError(
+                "Invalid value for DATABASE_REFRESH. Must be 'true', 'false', or 'wait_for'."
+            )
 
     @property
     def create_client(self):
@@ -109,7 +126,24 @@ class AsyncElasticsearchSettings(ApiSettings, ApiBaseSettings):
     enable_response_models: bool = False
     enable_direct_response: bool = get_bool_env("ENABLE_DIRECT_RESPONSE", default=False)
     raise_on_bulk_error: bool = get_bool_env("RAISE_ON_BULK_ERROR", default=False)
-    database_refresh: bool = get_bool_env("DATABASE_REFRESH", default=False)
+
+    @property
+    def database_refresh(self) -> Union[bool, str]:
+        """
+        Get the value of the DATABASE_REFRESH environment variable.
+
+        Returns:
+            Union[bool, str]: The value of DATABASE_REFRESH, which can be True, False, or "wait_for".
+        """
+        value = os.getenv("DATABASE_REFRESH", "false").lower()
+        if value in {"true", "false"}:
+            return value == "true"
+        elif value == "wait_for":
+            return "wait_for"
+        else:
+            raise ValueError(
+                "Invalid value for DATABASE_REFRESH. Must be 'true', 'false', or 'wait_for'."
+            )
 
     @property
     def create_client(self):