From fabd89f11dd7e60135a964645d7e664703abb2d5 Mon Sep 17 00:00:00 2001
From: Jaromir Wysoglad <jwysogla@redhat.com>
Date: Mon, 30 Jun 2025 04:32:51 -0400
Subject: [PATCH] Allow for providing metric names to modify_query

This change allows for providing the metric name list to the
modify_query method, which can greatly increase performence
when calling the method in a loop. Old behavior with sending
a Prometheus API request is preserved when the new parameter
isn't used.

Change-Id: I151c150c47e39699889baa1564a974e163c6a944
Signed-off-by: Jaromir Wysoglad <jwysogla@redhat.com>
---
 observabilityclient/rbac.py                   | 19 ++++++--
 observabilityclient/tests/unit/test_rbac.py   | 46 +++++++++++++++----
 ...-names-list-for-rbac-e46a4fd0df3fad73.yaml | 10 ++++
 3 files changed, 63 insertions(+), 12 deletions(-)
 create mode 100644 releasenotes/notes/allow-for-providing-metric-names-list-for-rbac-e46a4fd0df3fad73.yaml

diff --git a/observabilityclient/rbac.py b/observabilityclient/rbac.py
index a1e9887..9a7f806 100644
--- a/observabilityclient/rbac.py
+++ b/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
diff --git a/observabilityclient/tests/unit/test_rbac.py b/observabilityclient/tests/unit/test_rbac.py
index e6fb21f..1fc2eb3 100644
--- a/observabilityclient/tests/unit/test_rbac.py
+++ b/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()
diff --git a/releasenotes/notes/allow-for-providing-metric-names-list-for-rbac-e46a4fd0df3fad73.yaml b/releasenotes/notes/allow-for-providing-metric-names-list-for-rbac-e46a4fd0df3fad73.yaml
new file mode 100644
index 0000000..f300822
--- /dev/null
+++ b/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.