Merge "Allow for providing metric names to modify_query"

observabilityclient/rbac.py

@@ -97,14 +97,27 @@ class PromQLRbac(object):
                 f"{labels_str}"
                 f"{query[location:]}")
 
-    def modify_query(self, query):
+    def modify_query(self, query, metric_names=None):
         """Add rbac labels to a query.
 
         :param query: The query to modify
         :type query: str
+
+        :param metric_names: List of metric names currently stored in
+                             prometheus. For correct function of the query
+                             modification, it's important for the list to be
+                             accurate and to include all metrics across all
+                             tenants. This parameter can be unspecified or
+                             None, in which case the list will be retrieved
+                             from Prometheus by sending an API request to it.
+                             It's advised to provide the metric list when
+                             doing a query modification in a loop for
+                             performance purposes.
+        :type metric_names: list
         """
-        # We need to get all metric names, no matter the rbac
-        metric_names = self.client.label_values("__name__")
+
+        if metric_names is None:
+            metric_names = self.client.label_values("__name__")
 
         # We need to detect the locations of metric names
         # inside the query

observabilityclient/tests/unit/test_rbac.py

@@ -27,15 +27,17 @@ class PromQLRbacTest(testtools.TestCase):
         self.rbac.labels = {
             "project": "project123"
         }
-        self.rbac.client.label_values = lambda label: [
-            'test_query',
-            'cpu_temp_celsius',
-            'http_requests',
-            'test:query:with:colon:',
-            'test_query_with_digit1',
-            'method_code:http_errors:rate5m',
-            'method:http_requests:rate5m'
-        ]
+        self.rbac.client.label_values = mock.MagicMock(
+            return_value=[
+                'test_query',
+                'cpu_temp_celsius',
+                'http_requests',
+                'test:query:with:colon:',
+                'test_query_with_digit1',
+                'method_code:http_errors:rate5m',
+                'method:http_requests:rate5m'
+            ]
+        )
         self.test_cases = [
             (
                 "test_query",
@@ -227,3 +229,29 @@ class PromQLRbacTest(testtools.TestCase):
         kwargs = {'project_label': project_label}
 
         self.assertRaises(ValueError, rbac.PromQLRbac, *args, **kwargs)
+
+    def test_specifying_metric_names(self):
+        query = "test_query"
+        metric_names = [
+            "metric_name1",
+            "foo",
+            "bar",
+            "test_query"
+        ]
+
+        expected = "test_query{project='project123'}"
+
+        ret = self.rbac.modify_query(query, metric_names)
+
+        self.assertEqual(expected, ret)
+        self.rbac.client.label_values.assert_not_called()
+
+    def test_not_specifying_metric_names(self):
+        query = "test_query"
+
+        expected = "test_query{project='project123'}"
+
+        ret = self.rbac.modify_query(query)
+
+        self.assertEqual(expected, ret)
+        self.rbac.client.label_values.assert_called_once()

releasenotes/notes/allow-for-providing-metric-names-list-for-rbac-e46a4fd0df3fad73.yaml

@@ -0,0 +1,10 @@
+---
+features:
+  - |
+    The rbac.PromQLRbac.modify_query method now optionally allows
+    for providing a list of metric names thanks to a new parameter.
+    When the new parameter isn't used, the old behavior with sending
+    a Prometheus API request to retrieve the metric names is used.
+    When the metric names are provided, no Prometheus API request is
+    done, which greatly increases performance when using the
+    modify_query repeatedly, for example in loops.