From fe34ef976817f777ddfd49f1ad27925d91305e4a Mon Sep 17 00:00:00 2001
From: Pierre Riteau <pierre@stackhpc.com>
Date: Mon, 10 Jan 2022 15:20:29 +0100
Subject: [PATCH] Support customising Prometheus queries

This commit adds support for adding optional prefix and/or suffix to
Prometheus queries.

For example, this can be used to perform vector matches between the
collected metric and another one, to gather additional metadata.

Change-Id: I725f0f5ad00b67f55bcacaf8447e050af3815c73
---
 cloudkitty/collector/prometheus.py                | 15 ++++++++++++++-
 cloudkitty/tests/collectors/test_validation.py    |  8 ++++++++
 doc/source/admin/configuration/collector.rst      |  6 ++++++
 .../prometheus-custom-query-ab2dc00e97b14be2.yaml |  5 +++++
 4 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 releasenotes/notes/prometheus-custom-query-ab2dc00e97b14be2.yaml

diff --git a/cloudkitty/collector/prometheus.py b/cloudkitty/collector/prometheus.py
index 88cfce9e..8217fef3 100644
--- a/cloudkitty/collector/prometheus.py
+++ b/cloudkitty/collector/prometheus.py
@@ -18,6 +18,7 @@ from decimal import ROUND_HALF_UP
 
 from oslo_config import cfg
 from oslo_log import log
+from voluptuous import All
 from voluptuous import In
 from voluptuous import Optional
 from voluptuous import Required
@@ -83,7 +84,9 @@ PROMETHEUS_EXTRA_SCHEMA = {
                 'changes', 'delta', 'deriv',
                 'idelta', 'irange', 'irate',
                 'rate'
-            ])
+            ]),
+        Optional('query_prefix', default=''): All(str),
+        Optional('query_suffix', default=''): All(str),
     }
 }
 
@@ -159,6 +162,8 @@ class PrometheusCollector(collector.BaseCollector):
             'range_function')
         groupby = self.conf[metric_name].get('groupby', [])
         metadata = self.conf[metric_name].get('metadata', [])
+        query_prefix = self.conf[metric_name]['extra_args']['query_prefix']
+        query_suffix = self.conf[metric_name]['extra_args']['query_suffix']
         period = tzutils.diff_seconds(end, start)
         time = end
 
@@ -198,6 +203,14 @@ class PrometheusCollector(collector.BaseCollector):
             ', '.join(groupby + metadata)
         )
 
+        # Add custom query prefix
+        if query_prefix:
+            query = "{0} {1}".format(query_prefix, query)
+
+        # Add custom query suffix
+        if query_suffix:
+            query = "{0} {1}".format(query, query_suffix)
+
         try:
             res = self._conn.get_instant(
                 query,
diff --git a/cloudkitty/tests/collectors/test_validation.py b/cloudkitty/tests/collectors/test_validation.py
index bd08526d..20ebf0f7 100644
--- a/cloudkitty/tests/collectors/test_validation.py
+++ b/cloudkitty/tests/collectors/test_validation.py
@@ -135,6 +135,8 @@ class MetricConfigValidationTest(tests.TestCase):
         expected_output['metric_one']['groupby'].append('project_id')
         expected_output['metric_one']['extra_args'] = {
             'aggregation_method': 'max',
+            'query_prefix': '',
+            'query_suffix': '',
         }
         self.assertEqual(
             collector.prometheus.PrometheusCollector.check_configuration(data),
@@ -147,6 +149,8 @@ class MetricConfigValidationTest(tests.TestCase):
         expected_output['metric_one']['groupby'].append('project_id')
         expected_output['metric_one']['extra_args'] = {
             'aggregation_method': 'max',
+            'query_prefix': '',
+            'query_suffix': '',
         }
         self.assertEqual(
             collector.prometheus.PrometheusCollector.check_configuration(data),
@@ -158,6 +162,8 @@ class MetricConfigValidationTest(tests.TestCase):
         data['metrics']['metric_one']['extra_args'] = {
             'aggregation_method': 'max',
             'query_function': 'abs',
+            'query_prefix': 'custom_prefix',
+            'query_suffix': 'custom_suffix',
             'range_function': 'delta',
         }
         expected_output = copy.deepcopy(self.base_output)
@@ -165,6 +171,8 @@ class MetricConfigValidationTest(tests.TestCase):
         expected_output['metric_one']['extra_args'] = {
             'aggregation_method': 'max',
             'query_function': 'abs',
+            'query_prefix': 'custom_prefix',
+            'query_suffix': 'custom_suffix',
             'range_function': 'delta',
         }
 
diff --git a/doc/source/admin/configuration/collector.rst b/doc/source/admin/configuration/collector.rst
index 9291b631..8c334a33 100644
--- a/doc/source/admin/configuration/collector.rst
+++ b/doc/source/admin/configuration/collector.rst
@@ -359,6 +359,12 @@ Prometheus
   ``log10``, ``round``, ``sqrt``. For more information on these functions,
   you can check `this page`_
 
+* ``query_prefix``: Optional argument. An arbitrary prefix to add to the
+  Prometheus query generated by CloudKitty, separated by a space.
+
+* ``query_suffix``: Optional argument. An arbitrary suffix to add to the
+  Prometheus query generated by CloudKitty, separated by a space.
+
 * ``range_function``: Optional argument. The function to apply instead of the
   implicit ``{aggregation_method}_over_time``. Must be one of ``changes``,
   ``delta``, ``deriv``, ``idelta``, ``irange``, ``irate``, ``rate``. For more
diff --git a/releasenotes/notes/prometheus-custom-query-ab2dc00e97b14be2.yaml b/releasenotes/notes/prometheus-custom-query-ab2dc00e97b14be2.yaml
new file mode 100644
index 00000000..67477b9b
--- /dev/null
+++ b/releasenotes/notes/prometheus-custom-query-ab2dc00e97b14be2.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Adds support for specifying optional prefix and/or suffix to add to
+    Prometheus queries.