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

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",
+]

stac_fastapi/sfeos_helpers/search_engine/selection/base.py

@@ -0,0 +1,30 @@
+"""Base classes for index selection strategies."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional
+
+
+class BaseIndexSelector(ABC):
+    """Base class for async index selectors."""
+
+    @abstractmethod
+    async def select_indexes(
+        self,
+        collection_ids: Optional[List[str]],
+        datetime_search: Dict[str, Optional[str]],
+    ) -> str:
+        """Select appropriate indexes asynchronously.
+
+        Args:
+            collection_ids (Optional[List[str]]): List of collection IDs to filter by.
+            datetime_search (Dict[str, Optional[str]]): Datetime search criteria.
+
+        Returns:
+            str: Comma-separated string of selected index names.
+        """
+        pass
+
+    @abstractmethod
+    async def refresh_cache(self):
+        """Refresh cache (no-op for unfiltered selector)."""
+        pass

stac_fastapi/sfeos_helpers/search_engine/selection/cache_manager.py

@@ -0,0 +1,127 @@
+"""Cache management for index selection strategies."""
+
+import threading
+import time
+from collections import defaultdict
+from typing import Any, Dict, List, Optional
+
+from stac_fastapi.sfeos_helpers.database import index_alias_by_collection_id
+from stac_fastapi.sfeos_helpers.mappings import ITEMS_INDEX_PREFIX
+
+
+class IndexCacheManager:
+    """Manages caching of index aliases with expiration."""
+
+    def __init__(self, cache_ttl_seconds: int = 3600):
+        """Initialize the cache manager.
+
+        Args:
+            cache_ttl_seconds (int): Time-to-live for cache entries in seconds.
+        """
+        self._cache: Optional[Dict[str, List[str]]] = None
+        self._timestamp: float = 0
+        self._ttl = cache_ttl_seconds
+        self._lock = threading.Lock()
+
+    @property
+    def is_expired(self) -> bool:
+        """Check if the cache has expired.
+
+        Returns:
+            bool: True if cache is expired, False otherwise.
+        """
+        return time.time() - self._timestamp > self._ttl
+
+    def get_cache(self) -> Optional[Dict[str, List[str]]]:
+        """Get the current cache if not expired.
+
+        Returns:
+            Optional[Dict[str, List[str]]]: Cache data if valid, None if expired.
+        """
+        with self._lock:
+            if self.is_expired:
+                return None
+            return {k: v.copy() for k, v in self._cache.items()}
+
+    def set_cache(self, data: Dict[str, List[str]]) -> None:
+        """Set cache data and update timestamp.
+
+        Args:
+            data (Dict[str, List[str]]): Cache data to store.
+        """
+        self._cache = data
+        self._timestamp = time.time()
+
+    def clear_cache(self) -> None:
+        """Clear the cache and reset timestamp."""
+        self._cache = None
+        self._timestamp = 0
+
+
+class IndexAliasLoader:
+    """Asynchronous loader for index aliases."""
+
+    def __init__(self, client: Any, cache_manager: IndexCacheManager):
+        """Initialize the async alias loader.
+
+        Args:
+            client: Async search engine client instance.
+            cache_manager (IndexCacheManager): Cache manager instance.
+        """
+        self.client = client
+        self.cache_manager = cache_manager
+
+    async def load_aliases(self) -> Dict[str, List[str]]:
+        """Load index aliases from search engine.
+
+        Returns:
+            Dict[str, List[str]]: Mapping of base aliases to item aliases.
+        """
+        response = await self.client.indices.get_alias(index=f"{ITEMS_INDEX_PREFIX}*")
+        result = defaultdict(list)
+        for index_info in response.values():
+            aliases = index_info.get("aliases", {})
+            items_aliases = sorted(
+                [
+                    alias
+                    for alias in aliases.keys()
+                    if alias.startswith(ITEMS_INDEX_PREFIX)
+                ]
+            )
+
+            if items_aliases:
+                result[items_aliases[0]].extend(items_aliases[1:])
+
+        self.cache_manager.set_cache(result)
+        return result
+
+    async def get_aliases(self) -> Dict[str, List[str]]:
+        """Get aliases from cache or load if expired.
+
+        Returns:
+            Dict[str, List[str]]: Alias mapping data.
+        """
+        cached = self.cache_manager.get_cache()
+        if cached is not None:
+            return cached
+        return await self.load_aliases()
+
+    async def refresh_aliases(self) -> Dict[str, List[str]]:
+        """Force refresh aliases from search engine.
+
+        Returns:
+            Dict[str, List[str]]: Fresh alias mapping data.
+        """
+        return await self.load_aliases()
+
+    async def get_collection_indexes(self, collection_id: str) -> List[str]:
+        """Get all index aliases for a specific collection.
+
+        Args:
+            collection_id (str): Collection identifier.
+
+        Returns:
+            List[str]: List of index aliases for the collection.
+        """
+        aliases = await self.get_aliases()
+        return aliases.get(index_alias_by_collection_id(collection_id), [])