move to_es filter update to sfeos_helpers

Author: jonhealy1

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

@@ -147,111 +147,3 @@ def to_es_field(queryables_mapping: Dict[str, Any], field: str) -> str:
         str: The mapped field name suitable for Elasticsearch queries.
     """
     return queryables_mapping.get(field, field)
-
-
-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.
-
-    Args:
-        query (Dict[str, Any]): The query dictionary containing 'op' and 'args'.
-
-    Returns:
-        Dict[str, Any]: The corresponding Elasticsearch query in the form of a dictionary.
-    """
-    if query["op"] in [LogicalOp.AND, LogicalOp.OR, LogicalOp.NOT]:
-        bool_type = {
-            LogicalOp.AND: "must",
-            LogicalOp.OR: "should",
-            LogicalOp.NOT: "must_not",
-        }[query["op"]]
-        return {
-            "bool": {
-                bool_type: [
-                    to_es(queryables_mapping, sub_query) for sub_query in query["args"]
-                ]
-            }
-        }
-
-    elif query["op"] in [
-        ComparisonOp.EQ,
-        ComparisonOp.NEQ,
-        ComparisonOp.LT,
-        ComparisonOp.LTE,
-        ComparisonOp.GT,
-        ComparisonOp.GTE,
-    ]:
-        range_op = {
-            ComparisonOp.LT: "lt",
-            ComparisonOp.LTE: "lte",
-            ComparisonOp.GT: "gt",
-            ComparisonOp.GTE: "gte",
-        }
-
-        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"]
-            if query["op"] == ComparisonOp.EQ:
-                return {"range": {field: {"gte": value, "lte": value}}}
-            elif query["op"] == ComparisonOp.NEQ:
-                return {
-                    "bool": {
-                        "must_not": [{"range": {field: {"gte": value, "lte": value}}}]
-                    }
-                }
-            else:
-                return {"range": {field: {range_op[query["op"]]: value}}}
-        else:
-            if query["op"] == ComparisonOp.EQ:
-                return {"term": {field: value}}
-            elif query["op"] == ComparisonOp.NEQ:
-                return {"bool": {"must_not": [{"term": {field: value}}]}}
-            else:
-                return {"range": {field: {range_op[query["op"]]: value}}}
-
-    elif query["op"] == ComparisonOp.IS_NULL:
-        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(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"]
-        if isinstance(lte, dict) and "timestamp" in lte:
-            lte = lte["timestamp"]
-        return {"range": {field: {"gte": gte, "lte": lte}}}
-
-    elif query["op"] == AdvancedComparisonOp.IN:
-        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(queryables_mapping, query["args"][0]["property"])
-        pattern = cql2_like_to_es(query["args"][1])
-        return {"wildcard": {field: {"value": pattern, "case_insensitive": True}}}
-
-    elif query["op"] in [
-        SpatialOp.S_INTERSECTS,
-        SpatialOp.S_CONTAINS,
-        SpatialOp.S_WITHIN,
-        SpatialOp.S_DISJOINT,
-    ]:
-        field = to_es_field(queryables_mapping, query["args"][0]["property"])
-        geometry = query["args"][1]
-
-        relation_mapping = {
-            SpatialOp.S_INTERSECTS: "intersects",
-            SpatialOp.S_CONTAINS: "contains",
-            SpatialOp.S_WITHIN: "within",
-            SpatialOp.S_DISJOINT: "disjoint",
-        }
-
-        relation = relation_mapping[query["op"]]
-        return {"geo_shape": {field: {"shape": geometry, "relation": relation}}}
-
-    return {}

stac_fastapi/sfeos_helpers/stac_fastapi/sfeos_helpers/filter.py

@@ -20,7 +20,7 @@
 from .mappings import ES_MAPPING_TYPE_TO_JSON
 
 
-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.
 
@@ -36,7 +36,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,
@@ -53,7 +59,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"]
@@ -76,11 +82,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"]
@@ -89,14 +95,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}}}
 
@@ -106,7 +112,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 = {