Dynamic queryables mapping for properties.

Author: rhysrevans3

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

@@ -91,20 +91,20 @@ class SpatialOp(str, Enum):
     S_DISJOINT = "s_disjoint"
 
 
-queryables_mapping = {
-    "id": "id",
-    "collection": "collection",
-    "geometry": "geometry",
-    "datetime": "properties.datetime",
-    "created": "properties.created",
-    "updated": "properties.updated",
-    "cloud_cover": "properties.eo:cloud_cover",
-    "cloud_shadow_percentage": "properties.s2:cloud_shadow_percentage",
-    "nodata_pixel_percentage": "properties.s2:nodata_pixel_percentage",
-}
-
-
-def to_es_field(field: str) -> str:
+# queryables_mapping = {
+#     "id": "id",
+#     "collection": "collection",
+#     "geometry": "geometry",
+#     "datetime": "properties.datetime",
+#     "created": "properties.created",
+#     "updated": "properties.updated",
+#     "cloud_cover": "properties.eo:cloud_cover",
+#     "cloud_shadow_percentage": "properties.s2:cloud_shadow_percentage",
+#     "nodata_pixel_percentage": "properties.s2:nodata_pixel_percentage",
+# }
+
+
+def to_es_field(queryables_mapping: Dict[str, Any], field: str) -> str:
     """
     Map a given field to its corresponding Elasticsearch field according to a predefined mapping.
 
@@ -117,7 +117,7 @@ def to_es_field(field: str) -> str:
     return queryables_mapping.get(field, field)
 
 
-def to_es(query: Dict[str, Any]) -> Dict[str, Any]:
+def to_es(queryables_mapping: Dict[str, Any], query: Dict[str, Any]) -> Dict[str, Any]:
     """
     Transform a simplified CQL2 query structure to an Elasticsearch compatible query DSL.
 
@@ -133,7 +133,13 @@ def to_es(query: Dict[str, Any]) -> Dict[str, Any]:
             LogicalOp.OR: "should",
             LogicalOp.NOT: "must_not",
         }[query["op"]]
-        return {"bool": {bool_type: [to_es(sub_query) for sub_query in query["args"]]}}
+        return {
+            "bool": {
+                bool_type: [
+                    to_es(queryables_mapping, sub_query) for sub_query in query["args"]
+                ]
+            }
+        }
 
     elif query["op"] in [
         ComparisonOp.EQ,
@@ -150,7 +156,7 @@ def to_es(query: Dict[str, Any]) -> Dict[str, Any]:
             ComparisonOp.GTE: "gte",
         }
 
-        field = to_es_field(query["args"][0]["property"])
+        field = to_es_field(queryables_mapping, query["args"][0]["property"])
         value = query["args"][1]
         if isinstance(value, dict) and "timestamp" in value:
             value = value["timestamp"]
@@ -173,11 +179,11 @@ def to_es(query: Dict[str, Any]) -> Dict[str, Any]:
                 return {"range": {field: {range_op[query["op"]]: value}}}
 
     elif query["op"] == ComparisonOp.IS_NULL:
-        field = to_es_field(query["args"][0]["property"])
+        field = to_es_field(queryables_mapping, query["args"][0]["property"])
         return {"bool": {"must_not": {"exists": {"field": field}}}}
 
     elif query["op"] == AdvancedComparisonOp.BETWEEN:
-        field = to_es_field(query["args"][0]["property"])
+        field = to_es_field(queryables_mapping, query["args"][0]["property"])
         gte, lte = query["args"][1], query["args"][2]
         if isinstance(gte, dict) and "timestamp" in gte:
             gte = gte["timestamp"]
@@ -186,14 +192,14 @@ def to_es(query: Dict[str, Any]) -> Dict[str, Any]:
         return {"range": {field: {"gte": gte, "lte": lte}}}
 
     elif query["op"] == AdvancedComparisonOp.IN:
-        field = to_es_field(query["args"][0]["property"])
+        field = to_es_field(queryables_mapping, query["args"][0]["property"])
         values = query["args"][1]
         if not isinstance(values, list):
             raise ValueError(f"Arg {values} is not a list")
         return {"terms": {field: values}}
 
     elif query["op"] == AdvancedComparisonOp.LIKE:
-        field = to_es_field(query["args"][0]["property"])
+        field = to_es_field(queryables_mapping, query["args"][0]["property"])
         pattern = cql2_like_to_es(query["args"][1])
         return {"wildcard": {field: {"value": pattern, "case_insensitive": True}}}
 
@@ -203,7 +209,7 @@ def to_es(query: Dict[str, Any]) -> Dict[str, Any]:
         SpatialOp.S_WITHIN,
         SpatialOp.S_DISJOINT,
     ]:
-        field = to_es_field(query["args"][0]["property"])
+        field = to_es_field(queryables_mapping, query["args"][0]["property"])
         geometry = query["args"][1]
 
         relation_mapping = {

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/database_logic.py

@@ -290,6 +290,34 @@ async def get_one_item(self, collection_id: str, item_id: str) -> Dict:
             )
         return item["_source"]
 
+    async def get_queryables_mapping(self, collection_id: str = "*") -> dict:
+        """Retrieve mapping of Queryables for search.
+
+        Args:
+            collection_id (str, optional): The id of the Collection the Queryables
+            belongs to. Defaults to "*".
+
+        Returns:
+            dict: A dictionary containing the Queryables mappings.
+        """
+        queryables_mapping = {}
+
+        mappings = await self.client.indices.get_mapping(
+            index=f"{ITEMS_INDEX_PREFIX}{collection_id}",
+        )
+
+        for mapping in mappings.values():
+            fields = mapping["mappings"]["properties"]
+            properties = fields.pop("properties")
+
+            for field_key in fields:
+                queryables_mapping[field_key] = field_key
+
+            for property_key in properties["properties"]:
+                queryables_mapping[property_key] = f"properties.{property_key}"
+
+        return queryables_mapping
+
     @staticmethod
     def make_search():
         """Database logic to create a Search instance."""
@@ -518,8 +546,7 @@ def apply_free_text_filter(search: Search, free_text_queries: Optional[List[str]
 
         return search
 
-    @staticmethod
-    def apply_cql2_filter(search: Search, _filter: Optional[Dict[str, Any]]):
+    def apply_cql2_filter(self, search: Search, _filter: Optional[Dict[str, Any]]):
         """
         Apply a CQL2 filter to an Elasticsearch Search object.
 
@@ -539,7 +566,7 @@ def apply_cql2_filter(search: Search, _filter: Optional[Dict[str, Any]]):
                     otherwise the original Search object.
         """
         if _filter is not None:
-            es_query = filter.to_es(_filter)
+            es_query = filter.to_es(self.get_queryables_mapping(), _filter)
             search = search.query(es_query)
 
         return search