add metricrangedf and tests (4n4nd#137)

* add feature to convert range vectors into pandas dataframe; add metricrangedf class and tests; closes 4n4nd#130 * Update __init__.py remove extra empty line * Update metric_range_df.py remove extra empty line

Author: allenmqcymp

prometheus_api_client/__init__.py

@@ -7,3 +7,4 @@
 from .metric import Metric
 from .metrics_list import MetricsList
 from .metric_snapshot_df import MetricSnapshotDataFrame
+from .metric_range_df import MetricRangeDataFrame

prometheus_api_client/metric_range_df.py

@@ -0,0 +1,77 @@
+"""A pandas.DataFrame subclass for Prometheus range vector responses."""
+from pandas import DataFrame
+from pandas._typing import Axes, Dtype
+from typing import Optional, Sequence
+
+
+class MetricRangeDataFrame(DataFrame):
+ """Subclass to format and represent Prometheus query response as pandas.DataFrame.
+
+ Assumes response is either a json or sequence of jsons.
+
+ This class should be used specifically to instantiate a query response,
+ where the query response has several timestamp values per series.
+ That is, a range vector is expected.
+ If the data is an instant vector, use MetricSnapshotDataFrame instead.
+
+ Some argument descriptions in this docstring were copied from pandas.core.frame.DataFrame.
+
+ :param data: (list|json) A single metric (json with keys "metric" and "values"/"value")
+ or list of such metrics received from Prometheus as a response to query
+ :param index: (pandas.Index|array-like) Index to use for resulting dataframe. Will default to
+ pandas.RangeIndex if no indexing information part of input data and no index provided.
+ :param columns: (pandas.Index|array-like) Column labels to use for resulting dataframe. Will
+ default to list of labels + "timestamp" + "value" if not provided.
+ :param dtype: (dtype) default None. Data type to force. Only a single dtype is allowed. If None, infer.
+ :param copy: (bool) default False. Copy data from inputs. Only affects DataFrame / 2d ndarray input.
+
+ Example Usage:
+ .. code-block:: python
+
+ prom = PrometheusConnect()
+ metric_data = prom.get_current_metric_value(metric_name='up', label_config=my_label_config)
+ metric_df = MetricRangeDataFrame(metric_data)
+ metric_df.head()
+ '''
+ +------------+------------+-----------------+--------------------+-------+
+ | | __name__ | cluster | label_2 | value |
+ +-------------------------+-----------------+--------------------+-------+
+ | timestamp | | | | |
+ +============+============+=================+====================+=======+
+ | 1577836800 | __up__ | cluster_id_0 | label_2_value_2 | 0 |
+ +-------------------------+-----------------+--------------------+-------+
+ | 1577836801 | __up__ | cluster_id_1 | label_2_value_3 | 1 |
+ +-------------------------+-----------------+------------=-------+-------+
+ '''
+ """
+
+ def __init__(
+ self,
+ data=None,
+ index: Optional[Axes] = None,
+ columns: Optional[Axes] = None,
+ dtype: Optional[Dtype] = None,
+ copy: bool = False,
+ ):
+ """Constructor for MetricRangeDataFrame class."""
+ if data is not None:
+ # if just a single json instead of list/set/other sequence of jsons,
+ # treat as list with single entry
+ if not isinstance(data, Sequence):
+ data = [data]
+
+ row_data = []
+ for v in data:
+ if "value" in v:
+ raise TypeError(
+ "data must be a range vector. Expected range vector, got instant vector"
+ )
+ for t in v["values"]:
+ row_data.append({**v["metric"], "timestamp": t[0], "value": t[1]})
+
+ # init df normally now
+ super(MetricRangeDataFrame, self).__init__(
+ data=row_data, index=index, columns=columns, dtype=dtype, copy=copy
+ )
+
+ self.set_index(["timestamp"], inplace=True)

tests/test_metric_range_df.py

@@ -0,0 +1,97 @@
+import unittest
+import json
+import os
+from prometheus_api_client import MetricRangeDataFrame
+
+
+class TestMetricRangeDataFrame(unittest.TestCase):
+ def setUp(self):
+ """
+ read metrics stored as jsons in './tests/metrics'
+ """
+ self.raw_metrics_list = list()
+ self.raw_metrics_labels = list()
+ for (dir_path, _, file_names) in os.walk("./tests/metrics"):
+ for fname in file_names:
+ with open(os.path.join(dir_path, fname), "rb") as f:
+ metric_jsons = json.load(f)
+
+ # save json list
+ self.raw_metrics_list.extend([metric_jsons])
+
+ # save label configs
+ labels = set()
+ for i in metric_jsons:
+ labels.update(set(i["metric"].keys()))
+ self.raw_metrics_labels.append(labels)
+
+ def test_setup(self):
+ """
+ Check if setup was done correctly
+ """
+ self.assertEqual(
+ 8, len(self.raw_metrics_list), "incorrect number json files read (incorrect test setup)"
+ )
+
+ def test_init_shape(self):
+ """
+ Test if dataframe initialized is of correct shape
+ """
+ # check shape
+ # each metric json contains number of timestamps equal to number entries * number of timestamps in each series
+ # we're assuming each series has the same number of timestamps
+ # 3 labels
+ for current_metric_list in self.raw_metrics_list:
+ df = MetricRangeDataFrame(current_metric_list)
+ num_values = sum([len(v["values"]) for v in current_metric_list])
+ self.assertEqual(
+ (len(df.index.values), df.shape[1]), # shape[1] = 4xlabels + value
+ (num_values, 5),
+ "incorrect dataframe shape",
+ )
+
+ def test_init_timestamps(self):
+ """
+ Test if dataframe contains the correct timestamp indices
+ """
+ # check that the timestamp indices in each series are the same
+ for curr_metric_list in self.raw_metrics_list:
+ self.assertEqual(
+ set(MetricRangeDataFrame(curr_metric_list).index.values),
+ set([v[0] for s in curr_metric_list for v in s["values"] ])
+ )
+
+ def test_init_columns(self):
+ """
+ Test if dataframe initialized has correct columns
+ """
+ for curr_metric_labels, curr_metric_list in zip(self.raw_metrics_labels, self.raw_metrics_list):
+ self.assertEqual(
+ curr_metric_labels.union({"value"}),
+ set(MetricRangeDataFrame(curr_metric_list).columns),
+ "incorrect dataframe columns",
+ )
+
+ def test_init_single_metric(self):
+ """
+ Test if dataframe initialized is of correct shape when
+ 1. json object is passed as data
+ 2. list with single json object is passed as data
+ """
+ # check shape when single json passed
+ num_vals = len(self.raw_metrics_list[0][0]["values"])
+ self.assertEqual(
+ (num_vals, 5),
+ MetricRangeDataFrame(self.raw_metrics_list[0][0]).shape,
+ "incorrect dataframe shape when initialized with json",
+ )
+ # check shape when list with single json passed
+ self.assertEqual(
+ (num_vals, 5),
+ MetricRangeDataFrame([self.raw_metrics_list[0][0]]).shape,
+ "incorrect dataframe shape when initialized with single json list",
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()