sfeos-helpers 6.0.0__py3-none-any.whl → 6.2.0__py3-none-any.whl

stac_fastapi/sfeos_helpers/search_engine/index_operations.py

@@ -0,0 +1,167 @@
+"""Search engine adapters for different implementations."""
+
+import uuid
+from typing import Any, Dict
+
+from stac_fastapi.sfeos_helpers.database import (
+    index_alias_by_collection_id,
+    index_by_collection_id,
+)
+from stac_fastapi.sfeos_helpers.mappings import (
+    _ES_INDEX_NAME_UNSUPPORTED_CHARS_TABLE,
+    ES_ITEMS_MAPPINGS,
+    ES_ITEMS_SETTINGS,
+    ITEMS_INDEX_PREFIX,
+)
+
+
+class IndexOperations:
+    """Base class for search engine adapters with common implementations."""
+
+    async def create_simple_index(self, client: Any, collection_id: str) -> str:
+        """Create a simple index for the given collection.
+
+        Args:
+            client: Search engine client instance.
+            collection_id (str): Collection identifier.
+
+        Returns:
+            str: Created index name.
+        """
+        index_name = f"{index_by_collection_id(collection_id)}-000001"
+        alias_name = index_alias_by_collection_id(collection_id)
+
+        await client.indices.create(
+            index=index_name,
+            body=self._create_index_body({alias_name: {}}),
+            params={"ignore": [400]},
+        )
+        return index_name
+
+    async def create_datetime_index(
+        self, client: Any, collection_id: str, start_date: str
+    ) -> str:
+        """Create a datetime-based index for the given collection.
+
+        Args:
+            client: Search engine client instance.
+            collection_id (str): Collection identifier.
+            start_date (str): Start date for the alias.
+
+        Returns:
+            str: Created index alias name.
+        """
+        index_name = self.create_index_name(collection_id)
+        alias_name = self.create_alias_name(collection_id, start_date)
+        collection_alias = index_alias_by_collection_id(collection_id)
+        await client.indices.create(
+            index=index_name,
+            body=self._create_index_body({collection_alias: {}, alias_name: {}}),
+        )
+        return alias_name
+
+    @staticmethod
+    async def update_index_alias(client: Any, end_date: str, old_alias: str) -> str:
+        """Update index alias with new end date.
+
+        Args:
+            client: Search engine client instance.
+            end_date (str): End date for the alias.
+            old_alias (str): Current alias name.
+
+        Returns:
+            str: New alias name.
+        """
+        new_alias = f"{old_alias}-{end_date}"
+        aliases_info = await client.indices.get_alias(name=old_alias)
+        actions = []
+
+        for index_name in aliases_info.keys():
+            actions.append({"remove": {"index": index_name, "alias": old_alias}})
+            actions.append({"add": {"index": index_name, "alias": new_alias}})
+
+        await client.indices.update_aliases(body={"actions": actions})
+        return new_alias
+
+    @staticmethod
+    async def change_alias_name(client: Any, old_alias: str, new_alias: str) -> None:
+        """Change alias name from old to new.
+
+        Args:
+            client: Search engine client instance.
+            old_alias (str): Current alias name.
+            new_alias (str): New alias name.
+
+        Returns:
+            None
+        """
+        aliases_info = await client.indices.get_alias(name=old_alias)
+        actions = []
+
+        for index_name in aliases_info.keys():
+            actions.append({"remove": {"index": index_name, "alias": old_alias}})
+            actions.append({"add": {"index": index_name, "alias": new_alias}})
+        await client.indices.update_aliases(body={"actions": actions})
+
+    @staticmethod
+    def create_index_name(collection_id: str) -> str:
+        """Create index name from collection ID and uuid4.
+
+        Args:
+            collection_id (str): Collection identifier.
+
+        Returns:
+            str: Formatted index name.
+        """
+        cleaned = collection_id.translate(_ES_INDEX_NAME_UNSUPPORTED_CHARS_TABLE)
+        return f"{ITEMS_INDEX_PREFIX}{cleaned.lower()}_{uuid.uuid4()}"
+
+    @staticmethod
+    def create_alias_name(collection_id: str, start_date: str) -> str:
+        """Create index name from collection ID and uuid4.
+
+        Args:
+            collection_id (str): Collection identifier.
+            start_date (str): Start date for the alias.
+
+        Returns:
+            str: Alias name with initial date.
+        """
+        cleaned = collection_id.translate(_ES_INDEX_NAME_UNSUPPORTED_CHARS_TABLE)
+        return f"{ITEMS_INDEX_PREFIX}{cleaned.lower()}_{start_date}"
+
+    @staticmethod
+    def _create_index_body(aliases: Dict[str, Dict]) -> Dict[str, Any]:
+        """Create index body with common settings.
+
+        Args:
+            aliases (Dict[str, Dict]): Aliases configuration.
+
+        Returns:
+            Dict[str, Any]: Index body configuration.
+        """
+        return {
+            "aliases": aliases,
+            "mappings": ES_ITEMS_MAPPINGS,
+            "settings": ES_ITEMS_SETTINGS,
+        }
+
+    @staticmethod
+    async def find_latest_item_in_index(client: Any, index_name: str) -> dict[str, Any]:
+        """Find the latest item date in the specified index.
+
+        Args:
+            client: Search engine client instance.
+            index_name (str): Name of the index to query.
+
+        Returns:
+            datetime: Date of the latest item in the index.
+        """
+        query = {
+            "size": 1,
+            "sort": [{"properties.datetime": {"order": "desc"}}],
+            "_source": ["properties.datetime"],
+        }
+
+        response = await client.search(index=index_name, body=query)
+        return response["hits"]["hits"][0]

stac_fastapi/sfeos_helpers/search_engine/inserters.py

@@ -0,0 +1,309 @@
+"""Async index insertion strategies."""
+import logging
+from datetime import timedelta
+from typing import Any, Dict, List
+
+from fastapi import HTTPException, status
+
+from stac_fastapi.sfeos_helpers.database import (
+    extract_date,
+    extract_first_date_from_index,
+    index_alias_by_collection_id,
+    mk_item_id,
+)
+
+from .base import BaseIndexInserter
+from .index_operations import IndexOperations
+from .managers import DatetimeIndexManager
+from .selection import DatetimeBasedIndexSelector
+
+logger = logging.getLogger(__name__)
+
+
+class DatetimeIndexInserter(BaseIndexInserter):
+    """Async datetime-based index insertion strategy."""
+
+    def __init__(self, client: Any, index_operations: IndexOperations):
+        """Initialize the async datetime index inserter.
+
+        Args:
+            client: Async search engine client instance.
+            index_operations (IndexOperations): Search engine adapter instance.
+        """
+        self.client = client
+        self.index_operations = index_operations
+        self.datetime_manager = DatetimeIndexManager(client, index_operations)
+
+    @staticmethod
+    def should_create_collection_index() -> bool:
+        """Whether this strategy requires collection index creation.
+
+        Returns:
+            bool: False, as datetime strategy doesn't create collection indexes.
+        """
+        return False
+
+    async def create_simple_index(self, client: Any, collection_id: str) -> str:
+        """Create a simple index asynchronously.
+
+        Args:
+            client: Search engine client instance.
+            collection_id (str): Collection identifier.
+
+        Returns:
+            str: Created index name.
+        """
+        return await self.index_operations.create_simple_index(client, collection_id)
+
+    async def get_target_index(
+        self, collection_id: str, product: Dict[str, Any]
+    ) -> str:
+        """Get target index for a single product.
+
+        Args:
+            collection_id (str): Collection identifier.
+            product (Dict[str, Any]): Product data containing datetime information.
+
+        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
+        )
+
+    async def prepare_bulk_actions(
+        self, collection_id: str, items: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Prepare bulk actions for multiple items.
+
+        Args:
+            collection_id (str): Collection identifier.
+            items (List[Dict[str, Any]]): List of items to process.
+
+        Returns:
+            List[Dict[str, Any]]: List of bulk actions ready for execution.
+        """
+        if not items:
+            msg = "The product list cannot be empty."
+            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
+        )
+
+        actions = []
+        for item in items:
+            target_index = await self._get_target_index_internal(
+                index_selector, collection_id, item, check_size=False
+            )
+            actions.append(
+                {
+                    "_index": target_index,
+                    "_id": mk_item_id(item["id"], item["collection"]),
+                    "_source": item,
+                }
+            )
+
+        return actions
+
+    async def _get_target_index_internal(
+        self,
+        index_selector,
+        collection_id: str,
+        product: Dict[str, Any],
+        check_size: bool = True,
+    ) -> 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.
+
+        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
+        )
+        all_indexes = await index_selector.get_collection_indexes(collection_id)
+
+        if not all_indexes:
+            target_index = await self.datetime_manager.handle_new_collection(
+                collection_id, product_datetime
+            )
+            await index_selector.refresh_cache()
+            return target_index
+
+        all_indexes.sort()
+        start_date = extract_date(product_datetime)
+        end_date = extract_first_date_from_index(all_indexes[0])
+
+        if start_date < end_date:
+            alias = await self.datetime_manager.handle_early_date(
+                collection_id, start_date, end_date
+            )
+            await index_selector.refresh_cache()
+
+            return alias
+
+        if target_index != all_indexes[-1]:
+            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
+            )
+            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,
+                collection_id,
+                extract_date(first_item["properties"]["datetime"]),
+            )
+            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.
+
+        Checks if the index where the first item would be inserted is oversized.
+        If so, creates a new index starting from the next day.
+
+        Args:
+            index_selector: Index selector instance.
+            collection_id (str): Collection identifier.
+            items (List[Dict[str, Any]]): List of items to process.
+
+        Returns:
+            None
+        """
+        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()
+
+
+class SimpleIndexInserter(BaseIndexInserter):
+    """Simple async index insertion strategy."""
+
+    def __init__(self, index_operations: IndexOperations, client: Any):
+        """Initialize the async simple index inserter.
+
+        Args:
+            index_operations (IndexOperations): Search engine adapter instance.
+            client: Async search engine client instance.
+        """
+        self.search_adapter = index_operations
+        self.client = client
+
+    @staticmethod
+    def should_create_collection_index() -> bool:
+        """Whether this strategy requires collection index creation.
+
+        Returns:
+            bool: True, as simple strategy creates collection indexes.
+        """
+        return True
+
+    async def create_simple_index(self, client: Any, collection_id: str) -> str:
+        """Create a simple index asynchronously.
+
+        Args:
+            client: Search engine client instance.
+            collection_id (str): Collection identifier.
+
+        Returns:
+            str: Created index name.
+        """
+        return await self.search_adapter.create_simple_index(client, collection_id)
+
+    async def get_target_index(
+        self, collection_id: str, product: Dict[str, Any]
+    ) -> str:
+        """Get target index (always the collection alias).
+
+        Args:
+            collection_id (str): Collection identifier.
+            product (Dict[str, Any]): Product data (not used in simple strategy).
+
+        Returns:
+            str: Collection alias name.
+        """
+        return index_alias_by_collection_id(collection_id)
+
+    async def prepare_bulk_actions(
+        self, collection_id: str, items: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Prepare bulk actions for simple indexing.
+
+        Args:
+            collection_id (str): Collection identifier.
+            items (List[Dict[str, Any]]): List of items to process.
+
+        Returns:
+            List[Dict[str, Any]]: List of bulk actions with collection alias as target.
+        """
+        target_index = index_alias_by_collection_id(collection_id)
+        return [
+            {
+                "_index": target_index,
+                "_id": mk_item_id(item["id"], item["collection"]),
+                "_source": item,
+            }
+            for item in items
+        ]

stac_fastapi/sfeos_helpers/search_engine/managers.py

@@ -0,0 +1,198 @@
+"""Index management utilities."""
+
+import logging
+import os
+from datetime import datetime, timedelta
+from typing import Any, Dict
+
+from fastapi import HTTPException, status
+
+from stac_fastapi.sfeos_helpers.database import (
+    extract_date,
+    extract_first_date_from_index,
+)
+
+from .index_operations import IndexOperations
+
+logger = logging.getLogger(__name__)
+
+
+class IndexSizeManager:
+    """Manages index size limits and operations."""
+
+    def __init__(self, client: Any):
+        """Initialize the index size manager.
+
+        Args:
+            client: Search engine client instance.
+        """
+        self.client = client
+        self.max_size_gb = self._get_max_size_from_env()
+
+    async def get_index_size_in_gb(self, index_name: str) -> float:
+        """Get index size in gigabytes asynchronously.
+
+        Args:
+            index_name (str): Name of the index to check.
+
+        Returns:
+            float: Size of the index in gigabytes.
+        """
+        data = await self.client.indices.stats(index=index_name)
+        return data["_all"]["primaries"]["store"]["size_in_bytes"] / 1e9
+
+    async def is_index_oversized(self, index_name: str) -> bool:
+        """Check if index exceeds size limit asynchronously.
+
+        Args:
+            index_name (str): Name of the index to check.
+
+        Returns:
+            bool: True if index exceeds size limit, False otherwise.
+        """
+        size_gb = await self.get_index_size_in_gb(index_name)
+        is_oversized = size_gb > self.max_size_gb
+        gb_milestone = int(size_gb)
+        if gb_milestone > 0:
+            logger.info(f"Index '{index_name}' size: {gb_milestone}GB")
+
+        if is_oversized:
+            logger.warning(
+                f"Index '{index_name}' is oversized: {size_gb:.2f} GB "
+                f"(limit: {self.max_size_gb} GB)"
+            )
+
+        return is_oversized
+
+    @staticmethod
+    def _get_max_size_from_env() -> float:
+        """Get max size from environment variable with error handling.
+
+        Returns:
+            float: Maximum index size in GB.
+
+        Raises:
+            ValueError: If environment variable contains invalid value.
+        """
+        env_value = os.getenv("DATETIME_INDEX_MAX_SIZE_GB", "25")
+
+        try:
+            max_size = float(env_value)
+            if max_size <= 0:
+                raise ValueError(
+                    f"DATETIME_INDEX_MAX_SIZE_GB must be positive, got: {max_size}"
+                )
+            return max_size
+        except (ValueError, TypeError):
+            error_msg = (
+                f"Invalid value for DATETIME_INDEX_MAX_SIZE_GB environment variable: "
+                f"'{env_value}'. Must be a positive number. Using default value 25.0 GB."
+            )
+            logger.warning(error_msg)
+
+        return 25.0
+
+
+class DatetimeIndexManager:
+    """Manages datetime-based index operations."""
+
+    def __init__(self, client: Any, index_operations: IndexOperations):
+        """Initialize the datetime index manager.
+
+        Args:
+            client: Search engine client instance.
+            index_operations (IndexOperations): Search engine adapter instance.
+        """
+        self.client = client
+        self.index_operations = index_operations
+        self.size_manager = IndexSizeManager(client)
+
+    @staticmethod
+    def validate_product_datetime(product: Dict[str, Any]) -> str:
+        """Validate and extract datetime from product.
+
+        Args:
+            product (Dict[str, Any]): Product data containing datetime information.
+
+        Returns:
+            str: Validated product datetime.
+
+        Raises:
+            HTTPException: If product datetime is missing or invalid.
+        """
+        product_datetime = product["properties"]["datetime"]
+        if not product_datetime:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Product datetime is required for indexing",
+            )
+        return product_datetime
+
+    async def handle_new_collection(
+        self, collection_id: str, product_datetime: str
+    ) -> str:
+        """Handle index creation for new collection asynchronously.
+
+        Args:
+            collection_id (str): Collection identifier.
+            product_datetime (str): Product datetime for index naming.
+
+
+        Returns:
+            str: Created index name.
+        """
+        target_index = await self.index_operations.create_datetime_index(
+            self.client, collection_id, extract_date(product_datetime)
+        )
+        logger.info(
+            f"Successfully created index '{target_index}' for collection '{collection_id}'"
+        )
+        return target_index
+
+    async def handle_early_date(
+        self, collection_id: str, start_date: datetime, end_date: datetime
+    ) -> str:
+        """Handle product with date earlier than existing indexes asynchronously.
+
+        Args:
+            collection_id (str): Collection identifier.
+            start_date (datetime): Start date for the new index.
+            end_date (datetime): End date for alias update.
+
+        Returns:
+            str: Updated alias name.
+        """
+        old_alias = self.index_operations.create_alias_name(
+            collection_id, str(end_date)
+        )
+        new_alias = self.index_operations.create_alias_name(
+            collection_id, str(start_date)
+        )
+        await self.index_operations.change_alias_name(self.client, old_alias, new_alias)
+        return new_alias
+
+    async def handle_oversized_index(
+        self, collection_id: str, target_index: str, product_datetime: str
+    ) -> str:
+        """Handle index that exceeds size limit asynchronously.
+
+        Args:
+            collection_id (str): Collection identifier.
+            target_index (str): Current target index name.
+            product_datetime (str): Product datetime for new index.
+
+        Returns:
+            str: New or updated index name.
+        """
+        end_date = extract_date(product_datetime)
+        latest_index_start = extract_first_date_from_index(target_index)
+
+        if end_date != latest_index_start:
+            await self.index_operations.update_index_alias(
+                self.client, str(end_date), target_index
+            )
+            target_index = await self.index_operations.create_datetime_index(
+                self.client, collection_id, str(end_date + timedelta(days=1))
+            )
+
+        return target_index

stac_fastapi/sfeos_helpers/search_engine/selection/__init__.py

@@ -0,0 +1,15 @@
+"""Index selection strategies package."""
+
+from .base import BaseIndexSelector
+from .cache_manager import IndexAliasLoader, IndexCacheManager
+from .factory import IndexSelectorFactory
+from .selectors import DatetimeBasedIndexSelector, UnfilteredIndexSelector
+
+__all__ = [
+    "IndexCacheManager",
+    "IndexAliasLoader",
+    "DatetimeBasedIndexSelector",
+    "UnfilteredIndexSelector",
+    "IndexSelectorFactory",
+    "BaseIndexSelector",
+]