Add support for operator in labels

This allows us to match labels
with regular expressions.

Author: piptouque

prometheus_api_client/metric_query.py

@@ -0,0 +1,39 @@
+from collections.abc import Sequence
+from typing import Dict,Tuple,Union,Optional
+from enum import StrEnum
+from functools import reduce
+import logging
+
+class LabelQueryOp(StrEnum):
+    EQUAL='='
+    NOT_EQUAL='!='
+    REGEX_EQUAL='=~'
+    REGEX_NOT_EQUAL='!~'
+
+LabelQuery=Union[str,Tuple[LabelQueryOp, str]]
+MetricLabelQuery = Dict[str,LabelQuery]
+
+def query_to_str(metric_name: str, label_query: Optional[MetricLabelQuery]=None)->str:
+    """
+    Contruct query string from label query dictionary
+    
+    :param label_query: (MetricLabelQuery) The label query dictionary. Default is None
+    :return: (str) Query string inside brackets
+    :raises:
+        (ValueError) Raises an exception in case of an invalid label query operator
+    """
+    if not label_query:
+        return metric_name
+    def _format_label_query(label_key: str, label: LabelQuery)->str:
+        if isinstance(label, Sequence) and not isinstance(label, str):
+            if len(label) != 2:
+                raise ValueError(f"wrong number of elements in label query with operator: {len(label)} instead of 2")
+            label_op=label[0]
+            if label_op not in LabelQueryOp:
+                raise ValueError(f"unknown label operator: '{label_op}'")
+            label_value=label[1]
+            return f"{label_key}{label_op}'{label_value}'"
+        else:
+            return f"{label_key}{LabelQueryOp.EQUAL}'{label}'"
+    label_list=[_format_label_query(label_key, label) for label_key, label in label_query.items()]
+    return  metric_name + "{" + ",".join(label_list) + "}"

prometheus_api_client/prometheus_connect.py

@@ -10,6 +10,7 @@
 from requests.packages.urllib3.util.retry import Retry
 from requests import Session
 
+from .metric_query import MetricLabelQuery, query_to_str
 from .exceptions import PrometheusApiClientException
 
 # set up logging
@@ -229,14 +230,14 @@ def get_label_values(self, label_name: str, params: dict = None):
         return labels
 
     def get_current_metric_value(
-        self, metric_name: str, label_config: dict = None, params: dict = None
+        self, metric_name: str, label_config: MetricLabelQuery = None, params: dict = None
     ):
         r"""
         Get the current metric value for the specified metric and label configuration.
 
         :param metric_name: (str) The name of the metric
-        :param label_config: (dict) A dictionary that specifies metric labels and their
-            values
+        :param label_config: (MetricLabelQuery) A dictionary specifying metric labels and their
+            values, with optional operator (default is equality).
         :param params: (dict) Optional dictionary containing GET parameters to be sent
             along with the API request, such as "time"
         :returns: (list) A list of current metric values for the specified metric
@@ -249,17 +250,13 @@ def get_current_metric_value(
 
               prom = PrometheusConnect()
 
-              my_label_config = {'cluster': 'my_cluster_id', 'label_2': 'label_2_value'}
+              my_label_config = {'cluster': 'my_cluster_id', 'label_2': ('=~','label_2_.*')}
 
               prom.get_current_metric_value(metric_name='up', label_config=my_label_config)
         """
         params = params or {}
         data = []
-        if label_config:
-            label_list = [str(key + "=" + "'" + label_config[key] + "'") for key in label_config]
-            query = metric_name + "{" + ",".join(label_list) + "}"
-        else:
-            query = metric_name
+        query = query_to_str(metric_name, label_query=label_config)
 
         # using the query API to get raw data
         response = self._session.request(
@@ -284,7 +281,7 @@ def get_current_metric_value(
     def get_metric_range_data(
         self,
         metric_name: str,
-        label_config: dict = None,
+        label_config: MetricLabelQuery = None,
         start_time: datetime = (datetime.now() - timedelta(minutes=10)),
         end_time: datetime = datetime.now(),
         chunk_size: timedelta = None,
@@ -295,8 +292,8 @@ def get_metric_range_data(
         Get the current metric value for the specified metric and label configuration.
 
         :param metric_name: (str) The name of the metric.
-        :param label_config: (dict) A dictionary specifying metric labels and their
-            values.
+        :param label_config: (MetricLabelQuery) A dictionary specifying metric labels and their
+            values, with optional operator (default is equality).
         :param start_time:  (datetime) A datetime object that specifies the metric range start time.
         :param end_time: (datetime) A datetime object that specifies the metric range end time.
         :param chunk_size: (timedelta) Duration of metric data downloaded in one request. For
@@ -338,11 +335,7 @@ def get_metric_range_data(
             raise ValueError("specified chunk_size is too big")
         chunk_seconds = round(chunk_size.total_seconds())
 
-        if label_config:
-            label_list = [str(key + "=" + "'" + label_config[key] + "'") for key in label_config]
-            query = metric_name + "{" + ",".join(label_list) + "}"
-        else:
-            query = metric_name
+        query = query_to_str(metric_name, label_query=label_config)
         _LOGGER.debug("Prometheus Query: %s", query)
 
         while start < end:

tests/test_metric_query.py

@@ -0,0 +1,47 @@
+"""Test module for class PrometheusConnect."""
+import unittest
+
+from prometheus_api_client.metric_query import query_to_str
+
+class TestMetricQuery(unittest.TestCase):
+    """Test module for metric query."""
+
+    def test_query_to_str_with_wrong_label_query(self):  # noqa D102
+        # wrong op ('~=' instead of '=~')
+        with self.assertRaises(ValueError, msg=f"unknown label operator: '~='"):
+            _ = query_to_str(
+                metric_name="up",
+                label_query={"some_label": ("~=", "some-value-.*")}
+            )
+        # inverted label value and op
+        with self.assertRaises(ValueError, msg=f"unknown label operator: 'some-value-.*'"):
+            _ = query_to_str(
+                metric_name="up",
+                label_query={"some_label": ("some-value-.*", "=~")}
+            )
+        # Wrong number of label query arguments
+        with self.assertRaises(ValueError, msg=f"wrong number of elements in label query with operator: 3 instead of 2"):
+            _ = query_to_str(
+                metric_name="up",
+                label_query={"some_label": ("=~", "some-value-.*", "whatever")}
+            )
+    def test_query_to_str_with_correct_label_query(self):  # noqa D102
+        correct_label_queries = [
+            { "some_label": "some-value"}, # exact match
+            { "some_label": ("=", "some-value")}, # exact match, explicit op
+            { "some_label": ("!=", "some-value")}, # negative match
+            { "some_label": ("=~", "some-value-.*")}, # regex match
+            { "some_label": ("!~", "some-value-.*")}, # negative regex match
+        ]
+        for label_query in correct_label_queries:
+            try:
+                _ = query_to_str(
+                    metric_name="up",
+                    label_query=label_query
+                )
+            except Exception as e:
+                self.fail(f"query_to_str('up') with label_config raised an unexpected exception: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()