From 532e277894c3feeaf4736a6208b745e858b749d8 Mon Sep 17 00:00:00 2001
From: ylamgarchal <yassine.lamgarchal@redhat.com>
Date: Wed, 30 Oct 2024 11:59:22 +0100
Subject: [PATCH] analytics: add comparison operators: <, <=, >, >=

Change-Id: Ib970576d7635f113a15d8ce9dfdad815b0a4ca22
---
 dci/analytics/query_es_dsl.py        | 20 ++++++++++++++-
 tests/analytics/test_query_es_dsl.py | 38 ++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/dci/analytics/query_es_dsl.py b/dci/analytics/query_es_dsl.py
index b370e6d2d..92e8f4092 100644
--- a/dci/analytics/query_es_dsl.py
+++ b/dci/analytics/query_es_dsl.py
@@ -29,7 +29,7 @@
 _comma_string = _comma + _word
 _list = _lb + _word + pp.ZeroOrMore(_comma_string) + _rb
 
-_comparison_operators = {"=", "!=", "<=" "<", ">=", ">", "=~"}
+_comparison_operators = {"=", "!=", "<=", "<", ">=", ">", "=~"}
 _comparison_operators = pp.oneOf(" ".join(_comparison_operators))
 _comparison = _field + _comparison_operators + _value
 
@@ -58,6 +58,20 @@ def parse(q):
     return query.parseString(q).asList()
 
 
+_op_to_es_op = {"<": "lt", "<=": "lte", ">": "gt", ">=": "gte"}
+
+
+def _handle_comparison_operator(handle_nested, operator, operand_1, operand_2):
+    if handle_nested and "." in operand_1:
+        return {
+            "nested": {
+                "path": operand_1.split(".")[0],
+                "query": {"range": {operand_1: {_op_to_es_op[operator]: operand_2}}},
+            }
+        }
+    return {"range": {operand_1: {_op_to_es_op[operator]: operand_2}}}
+
+
 def _generate_from_operators(parsed_query, handle_nested=False):
     operand_1 = parsed_query[0]
     operator = parsed_query[1]
@@ -72,6 +86,10 @@ def _generate_from_operators(parsed_query, handle_nested=False):
                 }
             }
         return {"term": {operand_1: operand_2}}
+    if operator in _op_to_es_op.keys():
+        return _handle_comparison_operator(
+            handle_nested, operator, operand_1, operand_2
+        )
     elif operator == "=~":
         return {
             "regexp": {
diff --git a/tests/analytics/test_query_es_dsl.py b/tests/analytics/test_query_es_dsl.py
index 069f593cc..8da95478a 100644
--- a/tests/analytics/test_query_es_dsl.py
+++ b/tests/analytics/test_query_es_dsl.py
@@ -345,3 +345,41 @@ def test_query_build_regex():
             ]
         }
     }
+
+
+def test_query_build_comparison_operator():
+    ret = qed.build(
+        "(((keys_values.a>0) and (keys_values.a<10)) or ((keys_values.b>0) and (keys_values.b<=10)))"
+    )
+    assert ret == {
+        "bool": {
+            "should": [
+                {
+                    "nested": {
+                        "path": "keys_values",
+                        "query": {
+                            "bool": {
+                                "filter": [
+                                    {"range": {"keys_values.a": {"gt": "0"}}},
+                                    {"range": {"keys_values.a": {"lt": "10"}}},
+                                ]
+                            }
+                        },
+                    }
+                },
+                {
+                    "nested": {
+                        "path": "keys_values",
+                        "query": {
+                            "bool": {
+                                "filter": [
+                                    {"range": {"keys_values.b": {"gt": "0"}}},
+                                    {"range": {"keys_values.b": {"lte": "10"}}},
+                                ]
+                            }
+                        },
+                    }
+                },
+            ]
+        }
+    }