Add support for optional enum queryables

Author: Zaczero

stac_fastapi/core/stac_fastapi/core/base_database_logic.py

@@ -1,7 +1,7 @@
 """Base database logic."""
 
 import abc
-from typing import Any, Dict, Iterable, Optional
+from typing import Any, Dict, Iterable, List, Optional
 
 
 class BaseDatabaseLogic(abc.ABC):
@@ -36,6 +36,18 @@ async def delete_item(
         """Delete an item from the database."""
         pass
 
+    @abc.abstractmethod
+    async def get_items_mapping(self, collection_id: str) -> Dict[str, Dict[str, Any]]:
+        """Get the mapping for the items in the collection."""
+        pass
+
+    @abc.abstractmethod
+    async def get_items_unique_values(
+        self, collection_id: str, field_names: Iterable[str], *, limit: int = ...
+    ) -> Dict[str, List[str]]:
+        """Get the unique values for the given fields in the collection."""
+        pass
+
     @abc.abstractmethod
     async def create_collection(self, collection: Dict, refresh: bool = False) -> None:
         """Create a collection in the database."""

stac_fastapi/core/stac_fastapi/core/extensions/filter.py

@@ -60,6 +60,17 @@
         "maximum": 100,
     },
 }
+"""Queryables that are present in all collections."""
+
+OPTIONAL_QUERYABLES: Dict[str, Dict[str, Any]] = {
+    "platform": {
+        "$enum": True,
+        "description": "Satellite platform identifier",
+    },
+}
+"""Queryables that are present in some collections."""
+
+ALL_QUERYABLES: Dict[str, Dict[str, Any]] = DEFAULT_QUERYABLES | OPTIONAL_QUERYABLES
 
 
 class LogicalOp(str, Enum):

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -895,6 +895,31 @@ async def get_items_mapping(self, collection_id: str) -> Dict[str, Any]:
         except ESNotFoundError:
             raise NotFoundError(f"Mapping for index {index_name} not found")
 
+    async def get_items_unique_values(
+        self, collection_id: str, field_names: Iterable[str], *, limit: int = 100
+    ) -> Dict[str, List[str]]:
+        """Get the unique values for the given fields in the collection."""
+        limit_plus_one = limit + 1
+        index_name = index_alias_by_collection_id(collection_id)
+
+        query = await self.client.search(
+            index=index_name,
+            body={
+                "size": 0,
+                "aggs": {
+                    field: {"terms": {"field": field, "size": limit_plus_one}}
+                    for field in field_names
+                },
+            },
+        )
+
+        result: Dict[str, List[str]] = {}
+        for field, agg in query["aggregations"].items():
+            if len(agg["buckets"]) > limit:
+                raise ValueError(f"Field {field} has more than {limit} unique values.")
+            result[field] = [bucket["key"] for bucket in agg["buckets"]]
+        return result
+
     async def create_collection(self, collection: Collection, **kwargs: Any):
         """Create a single collection in the database.
 

stac_fastapi/opensearch/stac_fastapi/opensearch/database_logic.py

@@ -904,6 +904,31 @@ async def get_items_mapping(self, collection_id: str) -> Dict[str, Any]:
         except exceptions.NotFoundError:
             raise NotFoundError(f"Mapping for index {index_name} not found")
 
+    async def get_items_unique_values(
+        self, collection_id: str, field_names: Iterable[str], *, limit: int = 100
+    ) -> Dict[str, List[str]]:
+        """Get the unique values for the given fields in the collection."""
+        limit_plus_one = limit + 1
+        index_name = index_alias_by_collection_id(collection_id)
+
+        query = await self.client.search(
+            index=index_name,
+            body={
+                "size": 0,
+                "aggs": {
+                    field: {"terms": {"field": field, "size": limit_plus_one}}
+                    for field in field_names
+                },
+            },
+        )
+
+        result: Dict[str, List[str]] = {}
+        for field, agg in query["aggregations"].items():
+            if len(agg["buckets"]) > limit:
+                raise ValueError(f"Field {field} has more than {limit} unique values.")
+            result[field] = [bucket["key"] for bucket in agg["buckets"]]
+        return result
+
     async def create_collection(self, collection: Collection, **kwargs: Any):
         """Create a single collection in the database.
 

stac_fastapi/sfeos_helpers/stac_fastapi/sfeos_helpers/filter/client.py

@@ -1,12 +1,12 @@
 """Filter client implementation for Elasticsearch/OpenSearch."""
 
 from collections import deque
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 import attr
 
 from stac_fastapi.core.base_database_logic import BaseDatabaseLogic
-from stac_fastapi.core.extensions.filter import DEFAULT_QUERYABLES
+from stac_fastapi.core.extensions.filter import ALL_QUERYABLES, DEFAULT_QUERYABLES
 from stac_fastapi.extensions.core.filter.client import AsyncBaseFiltersClient
 from stac_fastapi.sfeos_helpers.mappings import ES_MAPPING_TYPE_TO_JSON
 
@@ -59,31 +59,31 @@ async def get_queryables(
 
         mapping_data = await self.database.get_items_mapping(collection_id)
         mapping_properties = next(iter(mapping_data.values()))["mappings"]["properties"]
-        stack = deque(mapping_properties.items())
+        stack: deque[Tuple[str, Dict[str, Any]]] = deque(mapping_properties.items())
+        enum_fields: Dict[str, Dict[str, Any]] = {}
 
         while stack:
-            field_name, field_def = stack.popleft()
+            field_fqn, field_def = stack.popleft()
 
             # Iterate over nested fields
             field_properties = field_def.get("properties")
             if field_properties:
-                # Fields in Item Properties should be exposed with their un-prefixed names,
-                # and not require expressions to prefix them with properties,
-                # e.g., eo:cloud_cover instead of properties.eo:cloud_cover.
-                if field_name == "properties":
-                    stack.extend(field_properties.items())
-                else:
-                    stack.extend(
-                        (f"{field_name}.{k}", v) for k, v in field_properties.items()
-                    )
+                stack.extend(
+                    (f"{field_fqn}.{k}", v) for k, v in field_properties.items()
+                )
 
             # Skip non-indexed or disabled fields
             field_type = field_def.get("type")
             if not field_type or not field_def.get("enabled", True):
                 continue
 
+            # Fields in Item Properties should be exposed with their un-prefixed names,
+            # and not require expressions to prefix them with properties,
+            # e.g., eo:cloud_cover instead of properties.eo:cloud_cover.
+            field_name = field_fqn.removeprefix("properties.")
+
             # Generate field properties
-            field_result = DEFAULT_QUERYABLES.get(field_name, {})
+            field_result = ALL_QUERYABLES.get(field_name, {})
             properties[field_name] = field_result
 
             field_name_human = field_name.replace("_", " ").title()
@@ -95,4 +95,13 @@ async def get_queryables(
             if field_type in {"date", "date_nanos"}:
                 field_result.setdefault("format", "date-time")
 
+            if field_result.pop("$enum", False):
+                enum_fields[field_fqn] = field_result
+
+        if enum_fields:
+            for field_fqn, unique_values in (
+                await self.database.get_items_unique_values(collection_id, enum_fields)
+            ).items():
+                enum_fields[field_fqn]["enum"] = unique_values
+
         return queryables