sfeos-helpers 6.9.0__py3-none-any.whl → 6.10.1__py3-none-any.whl

stac_fastapi/sfeos_helpers/search_engine/inserters.py

@@ -1,10 +1,11 @@
 """Async index insertion strategies."""
+
 import logging
-from datetime import timedelta
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from fastapi import HTTPException, status
 
+from stac_fastapi.core.utilities import get_bool_env
 from stac_fastapi.sfeos_helpers.database import (
     extract_date,
     extract_first_date_from_index,
@@ -14,7 +15,7 @@ from stac_fastapi.sfeos_helpers.database import (
 
 from .base import BaseIndexInserter
 from .index_operations import IndexOperations
-from .managers import DatetimeIndexManager
+from .managers import DatetimeIndexManager, ProductDatetimes
 from .selection import DatetimeBasedIndexSelector
 
 logger = logging.getLogger(__name__)
@@ -33,6 +34,25 @@ class DatetimeIndexInserter(BaseIndexInserter):
         self.client = client
         self.index_operations = index_operations
         self.datetime_manager = DatetimeIndexManager(client, index_operations)
+        self.index_selector = DatetimeBasedIndexSelector(client)
+
+    @property
+    def use_datetime(self) -> bool:
+        """Get USE_DATETIME setting dynamically.
+
+        Returns:
+            bool: Current value of USE_DATETIME environment variable.
+        """
+        return get_bool_env("USE_DATETIME", default=True)
+
+    @property
+    def primary_datetime_name(self) -> str:
+        """Get primary datetime field name based on current USE_DATETIME setting.
+
+        Returns:
+            str: "datetime" if USE_DATETIME is True, else "start_datetime".
+        """
+        return "datetime" if self.use_datetime else "start_datetime"
 
     @staticmethod
     def should_create_collection_index() -> bool:
@@ -55,6 +75,48 @@ class DatetimeIndexInserter(BaseIndexInserter):
         """
         return await self.index_operations.create_simple_index(client, collection_id)
 
+    async def refresh_cache(self) -> None:
+        """Refresh the index selector cache.
+
+        This method refreshes the cached index information used for
+        datetime-based index selection.
+        """
+        await self.index_selector.refresh_cache()
+
+    def validate_datetime_field_update(self, field_path: str) -> None:
+        """Validate if a datetime field can be updated.
+
+        For datetime-based indexing, the primary datetime field cannot be modified
+        because it determines the index where the item is stored.
+
+        When USE_DATETIME=True, 'properties.datetime' is protected.
+        When USE_DATETIME=False, 'properties.start_datetime' and 'properties.end_datetime' are protected.
+
+        Args:
+            field_path (str): The path of the field being updated.
+        """
+        # TODO: In the future, updating these fields will be able to move an item between indices by changing the time-based aliases
+        if self.use_datetime:
+            if field_path == "properties/datetime":
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail=(
+                        "Updating 'properties.datetime' is not yet supported for datetime-based indexing. "
+                        "This feature will be available in a future release, enabling automatic "
+                        "index and time-based alias updates when datetime values change."
+                    ),
+                )
+        else:
+            if field_path in ("properties/start_datetime", "properties/end_datetime"):
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail=(
+                        f"Updating '{field_path}' is not yet supported for datetime-based indexing. "
+                        "This feature will be available in a future release, enabling automatic "
+                        "index and time-based alias updates when datetime values change."
+                    ),
+                )
+
     async def get_target_index(
         self, collection_id: str, product: Dict[str, Any]
     ) -> str:
@@ -67,9 +129,8 @@ class DatetimeIndexInserter(BaseIndexInserter):
         Returns:
             str: Target index name for the product.
         """
-        index_selector = DatetimeBasedIndexSelector(self.client)
         return await self._get_target_index_internal(
-            index_selector, collection_id, product, check_size=True
+            collection_id, product, check_size=True
         )
 
     async def prepare_bulk_actions(
@@ -89,18 +150,12 @@ class DatetimeIndexInserter(BaseIndexInserter):
             logger.error(msg)
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=msg)
 
-        items.sort(key=lambda item: item["properties"]["datetime"])
-        index_selector = DatetimeBasedIndexSelector(self.client)
-
-        await self._ensure_indexes_exist(index_selector, collection_id, items)
-        await self._check_and_handle_oversized_index(
-            index_selector, collection_id, items
-        )
+        items.sort(key=lambda item: item["properties"][self.primary_datetime_name])
 
         actions = []
         for item in items:
             target_index = await self._get_target_index_internal(
-                index_selector, collection_id, item, check_size=False
+                collection_id, item, check_size=True
             )
             actions.append(
                 {
@@ -114,15 +169,13 @@ class DatetimeIndexInserter(BaseIndexInserter):
 
     async def _get_target_index_internal(
         self,
-        index_selector,
         collection_id: str,
         product: Dict[str, Any],
         check_size: bool = True,
-    ) -> str:
+    ) -> Optional[str]:
         """Get target index with size checking internally.
 
         Args:
-            index_selector: Index selector instance.
             collection_id (str): Collection identifier.
             product (Dict[str, Any]): Product data.
             check_size (bool): Whetheru to check index size limits.
@@ -130,112 +183,139 @@ class DatetimeIndexInserter(BaseIndexInserter):
         Returns:
             str: Target index name.
         """
-        product_datetime = self.datetime_manager.validate_product_datetime(product)
-        datetime_range = {"gte": product_datetime, "lte": product_datetime}
-        target_index = await index_selector.select_indexes(
-            [collection_id], datetime_range
+        product_datetimes = self.datetime_manager.validate_product_datetimes(
+            product, self.use_datetime
         )
-        all_indexes = await index_selector.get_collection_indexes(collection_id)
+        primary_datetime_value = (
+            product_datetimes.datetime
+            if self.use_datetime
+            else product_datetimes.start_datetime
+        )
+
+        all_indexes = await self.index_selector.get_collection_indexes(collection_id)
 
         if not all_indexes:
             target_index = await self.datetime_manager.handle_new_collection(
-                collection_id, product_datetime
+                collection_id, self.primary_datetime_name, product_datetimes
             )
-            await index_selector.refresh_cache()
+            await self.refresh_cache()
             return target_index
 
-        all_indexes.sort()
-        start_date = extract_date(product_datetime)
-        end_date = extract_first_date_from_index(all_indexes[0])
+        all_indexes = sorted(
+            all_indexes, key=lambda x: x[0][self.primary_datetime_name]
+        )
+
+        target_index = await self.index_selector.select_indexes(
+            [collection_id], primary_datetime_value, for_insertion=True
+        )
 
-        if start_date < end_date:
-            alias = await self.datetime_manager.handle_early_date(
-                collection_id, start_date, end_date
+        start_date = extract_date(primary_datetime_value)
+        earliest_index_date = extract_first_date_from_index(
+            all_indexes[0][0][self.primary_datetime_name]
+        )
+
+        if start_date < earliest_index_date:
+            target_index = await self.datetime_manager.handle_early_date(
+                collection_id,
+                self.primary_datetime_name,
+                product_datetimes,
+                all_indexes[0][0],
+                True,
             )
-            await index_selector.refresh_cache()
+            await self.refresh_cache()
+            return target_index
 
-            return alias
+        if not target_index:
+            target_index = all_indexes[-1][0][self.primary_datetime_name]
 
-        if target_index != all_indexes[-1]:
+        aliases_dict, is_first_index = self._find_aliases_for_index(
+            all_indexes, target_index
+        )
+
+        if target_index != all_indexes[-1][0][self.primary_datetime_name]:
+            await self.datetime_manager.handle_early_date(
+                collection_id,
+                self.primary_datetime_name,
+                product_datetimes,
+                aliases_dict,
+                is_first_index,
+            )
+            await self.refresh_cache()
             return target_index
 
         if check_size and await self.datetime_manager.size_manager.is_index_oversized(
             target_index
         ):
-            target_index = await self.datetime_manager.handle_oversized_index(
-                collection_id, target_index, product_datetime
+            latest_item = await self.index_operations.find_latest_item_in_index(
+                self.client, target_index
+            )
+            latest_index_datetimes = ProductDatetimes(
+                start_datetime=str(
+                    extract_date(latest_item["_source"]["properties"]["start_datetime"])
+                ),
+                datetime=str(
+                    extract_date(latest_item["_source"]["properties"]["datetime"])
+                ),
+                end_datetime=str(
+                    extract_first_date_from_index(aliases_dict["end_datetime"])
+                )
+                if aliases_dict.get("end_datetime")
+                else None,
             )
-            await index_selector.refresh_cache()
-
-        return target_index
-
-    async def _ensure_indexes_exist(
-        self, index_selector, collection_id: str, items: List[Dict[str, Any]]
-    ):
-        """Ensure necessary indexes exist for the items.
-
-        Args:
-            index_selector: Index selector instance.
-            collection_id (str): Collection identifier.
-            items (List[Dict[str, Any]]): List of items to process.
-        """
-        all_indexes = await index_selector.get_collection_indexes(collection_id)
 
-        if not all_indexes:
-            first_item = items[0]
-            await self.index_operations.create_datetime_index(
-                self.client,
+            await self.datetime_manager.handle_oversized_index(
                 collection_id,
-                extract_date(first_item["properties"]["datetime"]),
+                self.primary_datetime_name,
+                product_datetimes,
+                latest_index_datetimes,
+                aliases_dict,
+            )
+            await self.refresh_cache()
+            all_indexes = await self.index_selector.get_collection_indexes(
+                collection_id
+            )
+            all_indexes = sorted(
+                all_indexes, key=lambda x: x[0][self.primary_datetime_name]
+            )
+            return (
+                await self.index_selector.select_indexes(
+                    [collection_id], primary_datetime_value, for_insertion=True
+                )
+                or all_indexes[-1][0][self.primary_datetime_name]
             )
-            await index_selector.refresh_cache()
 
-    async def _check_and_handle_oversized_index(
-        self, index_selector, collection_id: str, items: List[Dict[str, Any]]
-    ) -> None:
-        """Check if index is oversized and create new index if needed.
+        await self.datetime_manager.handle_early_date(
+            collection_id,
+            self.primary_datetime_name,
+            product_datetimes,
+            aliases_dict,
+            is_first_index,
+        )
+        await self.refresh_cache()
+        all_indexes = await self.index_selector.get_collection_indexes(collection_id)
+        all_indexes = sorted(
+            all_indexes, key=lambda x: x[0][self.primary_datetime_name]
+        )
+        return all_indexes[-1][0][self.primary_datetime_name]
 
-        Checks if the index where the first item would be inserted is oversized.
-        If so, creates a new index starting from the next day.
+    @staticmethod
+    def _find_aliases_for_index(
+        all_indexes: List, target_index: str
+    ) -> tuple[Optional[Dict[str, Any]], bool]:
+        """Find aliases for a given index.
 
         Args:
-            index_selector: Index selector instance.
-            collection_id (str): Collection identifier.
-            items (List[Dict[str, Any]]): List of items to process.
+            all_indexes: List of index alias dictionaries.
+            target_index: Target index name to find.
 
         Returns:
-            None
+            Tuple of (aliases_dict or None, is_first_element).
         """
-        first_item = items[0]
-        first_item_index = await self._get_target_index_internal(
-            index_selector, collection_id, first_item, check_size=False
-        )
-
-        all_indexes = await index_selector.get_collection_indexes(collection_id)
-        all_indexes.sort()
-        latest_index = all_indexes[-1]
-
-        if first_item_index != latest_index:
-            return None
-
-        if not await self.datetime_manager.size_manager.is_index_oversized(
-            first_item_index
-        ):
-            return None
-
-        latest_item = await self.index_operations.find_latest_item_in_index(
-            self.client, latest_index
-        )
-        product_datetime = latest_item["_source"]["properties"]["datetime"]
-        end_date = extract_date(product_datetime)
-        await self.index_operations.update_index_alias(
-            self.client, str(end_date), latest_index
-        )
-        next_day_start = end_date + timedelta(days=1)
-        await self.index_operations.create_datetime_index(
-            self.client, collection_id, str(next_day_start)
-        )
-        await index_selector.refresh_cache()
+        for idx, item in enumerate(all_indexes):
+            aliases_dict = item[0]
+            if target_index in aliases_dict.values():
+                return aliases_dict, idx == 0
+        return None, False
 
 
 class SimpleIndexInserter(BaseIndexInserter):