fix: Integrate energy_consumed_wh function using prometheus-api-client

- Replaced custom requests-based Prometheus query with prometheus-api-client library
- Implemented energy_consumed_wh function based on @jaimeib's sample code
- Updated configuration options:
  - Removed: prometheus_query, prometheus_timeout
  - Added: prometheus_metric_name, prometheus_label_type_instance,
    prometheus_step_seconds, prometheus_query_range, prometheus_verify_ssl
- Query now uses sum_over_time with configurable metric name and labels
- Calculates energy in Wh from microwatt samples using step_seconds factor
- Updated tests to mock PrometheusConnect instead of requests
- Added prometheus-api-client dependency to pyproject.toml
- Updated poetry.lock file
- All tests pass (4 prometheus tests + 2 energy record tests)

Co-authored-by: alvarolopez <[email protected]>

Author: Copilot

caso/extract/prometheus.py

@@ -18,7 +18,7 @@
 
 import uuid
 
-import requests
+from prometheus_api_client import PrometheusConnect
 from oslo_config import cfg
 from oslo_log import log
 
@@ -34,16 +34,29 @@
         help="Prometheus server endpoint URL.",
     ),
     cfg.StrOpt(
-        "prometheus_query",
-        default="sum(rate(libvirt_domain_info_energy_consumption_joules_total"
-        '{uuid=~"{{uuid}}"}[5m])) * 300 / 3600',
-        help="Prometheus query to retrieve energy consumption in Wh. "
-        "The query can use {{uuid}} as a template variable for the VM UUID.",
+        "prometheus_metric_name",
+        default="prometheus_value",
+        help="Name of the Prometheus metric to query for energy consumption.",
+    ),
+    cfg.StrOpt(
+        "prometheus_label_type_instance",
+        default="scaph_process_power_microwatts",
+        help="Value for the type_instance label in Prometheus queries.",
     ),
     cfg.IntOpt(
-        "prometheus_timeout",
+        "prometheus_step_seconds",
         default=30,
-        help="Timeout for Prometheus API requests in seconds.",
+        help="Frequency between samples in the time series (in seconds).",
+    ),
+    cfg.StrOpt(
+        "prometheus_query_range",
+        default="1h",
+        help="Query time range (e.g., '1h', '6h', '24h').",
+    ),
+    cfg.BoolOpt(
+        "prometheus_verify_ssl",
+        default=True,
+        help="Whether to verify SSL when connecting to Prometheus.",
     ),
 ]
 
@@ -72,39 +85,61 @@ def _get_flavors(self):
             LOG.warning(f"Could not get flavors: {e}")
         return flavors
 
-    def _query_prometheus(self, query, timestamp=None):
-        """Query Prometheus API and return results.
+    def _energy_consumed_wh(self, vm_uuid):
+        """Calculate the energy consumed (Wh) for a VM from Prometheus.
+
+        This function queries Prometheus for instantaneous power samples
+        (in microwatts) and calculates the energy consumed in Watt-hours.
 
-        :param query: PromQL query string
-        :param timestamp: Optional timestamp for query (datetime object)
-        :returns: Query results
+        :param vm_uuid: UUID of the VM to query energy for
+        :returns: Energy consumed in Watt-hours (Wh)
         """
-        endpoint = CONF.prometheus.prometheus_endpoint
-        url = f"{endpoint}/api/v1/query"
+        prom_url = CONF.prometheus.prometheus_endpoint
+        metric_name = CONF.prometheus.prometheus_metric_name
+        step_seconds = CONF.prometheus.prometheus_step_seconds
+        query_range = CONF.prometheus.prometheus_query_range
+        verify_ssl = CONF.prometheus.prometheus_verify_ssl
+
+        prom = PrometheusConnect(url=prom_url, disable_ssl=not verify_ssl)
+
+        # factor = step_seconds / 3600 converts µW·s to µWh
+        factor = step_seconds / 3600
+
+        # Build labels for this VM
+        labels = {
+            "type_instance": CONF.prometheus.prometheus_label_type_instance,
+            "uuid": vm_uuid,
+        }
+
+        # Build label string: {key="value", ...}
+        label_selector = ",".join(f'{k}="{v}"' for k, v in labels.items())
+
+        # Construct the PromQL query
+        query = (
+            f"sum_over_time({metric_name}{{{label_selector}}}[{query_range}]) "
+            f"* {factor} / 1000000"
+        )
 
-        params = {"query": query}
-        if timestamp:
-            params["time"] = int(timestamp.timestamp())
+        LOG.debug(f"Querying Prometheus for VM {vm_uuid} with query: {query}")
 
         try:
-            response = requests.get(
-                url, params=params, timeout=CONF.prometheus.prometheus_timeout
-            )
-            response.raise_for_status()
-            data = response.json()
-
-            if data.get("status") != "success":
-                error_msg = data.get("error", "Unknown error")
-                LOG.error(f"Prometheus query failed: {error_msg}")
-                return None
-
-            return data.get("data", {}).get("result", [])
-        except requests.exceptions.RequestException as e:
-            LOG.error(f"Failed to query Prometheus: {e}")
-            return None
+            # Run query
+            result = prom.custom_query(query=query)
+
+            if not result:
+                LOG.debug(f"No energy data returned for VM {vm_uuid}")
+                return 0.0
+
+            energy_wh = float(result[0]["value"][1])
+            LOG.debug(f"VM {vm_uuid} consumed {energy_wh:.4f} Wh")
+            return energy_wh
+
+        except (KeyError, IndexError, ValueError) as e:
+            LOG.warning(f"Error parsing Prometheus result for VM {vm_uuid}: {e}")
+            return 0.0
         except Exception as e:
-            LOG.error(f"Unexpected error querying Prometheus: {e}")
-            return None
+            LOG.error(f"Error querying Prometheus for VM {vm_uuid}: {e}")
+            return 0.0
 
     def _get_servers(self, extract_from):
         """Get all servers for a given date."""
@@ -239,48 +274,31 @@ def extract(self, extract_from, extract_to):
         )
 
         # Query Prometheus for each server
-        query_template = CONF.prometheus.prometheus_query
-
         for server in servers:
             vm_uuid = str(server.id)
             vm_name = server.name
 
-            # Replace template variables in the query
-            query = query_template.replace("{{uuid}}", vm_uuid)
+            LOG.debug(f"Querying energy consumption for VM {vm_name} ({vm_uuid})")
 
-            LOG.debug(
-                f"Querying Prometheus for VM {vm_name} ({vm_uuid}) "
-                f"with query: {query}"
-            )
-
-            results = self._query_prometheus(query, extract_to)
+            # Get energy consumption using the new method
+            energy_value = self._energy_consumed_wh(vm_uuid)
 
-            if results is None:
-                LOG.warning(
-                    f"No results returned from Prometheus for VM "
-                    f"{vm_name} ({vm_uuid})"
+            if energy_value <= 0:
+                LOG.debug(
+                    f"No energy consumption data for VM {vm_name} ({vm_uuid}), "
+                    "skipping record creation"
                 )
                 continue
 
-            # Process results and create records
-            for result in results:
-                value = result.get("value", [])
-
-                if len(value) < 2:
-                    continue
-
-                # value is [timestamp, value_string]
-                energy_value = float(value[1])
-
-                LOG.debug(
-                    f"Creating energy record: {energy_value} Wh "
-                    f"for VM {vm_name} ({vm_uuid})"
-                )
+            LOG.debug(
+                f"Creating energy record: {energy_value} Wh "
+                f"for VM {vm_name} ({vm_uuid})"
+            )
 
-                energy_record = self._build_energy_record(
-                    server, energy_value, extract_from, extract_to
-                )
-                records.append(energy_record)
+            energy_record = self._build_energy_record(
+                server, energy_value, extract_from, extract_to
+            )
+            records.append(energy_record)
 
         LOG.info(f"Extracted {len(records)} energy records for project {self.project}")
 

caso/tests/extract/test_prometheus.py

@@ -79,47 +79,28 @@ def configured_extractor(mock_flavors):
 
 
 @pytest.fixture
-def prometheus_success_response():
-    """Fixture for a successful Prometheus response."""
-    response = mock.Mock()
-    response.json.return_value = {
-        "status": "success",
-        "data": {
-            "result": [
-                {
-                    "metric": {"instance": "test"},
-                    "value": [1685051946, "5.0"],
-                }
-            ]
-        },
-    }
-    response.raise_for_status = mock.Mock()
-    return response
+def mock_prometheus_result_success():
+    """Fixture for a successful Prometheus API result."""
+    return [{"metric": {"uuid": "test-uuid"}, "value": [1685051946, "5.0"]}]
 
 
 @pytest.fixture
-def prometheus_error_response():
-    """Fixture for a failed Prometheus response."""
-    response = mock.Mock()
-    response.json.return_value = {
-        "status": "error",
-        "error": "query failed",
-    }
-    response.raise_for_status = mock.Mock()
-    return response
+def mock_prometheus_result_empty():
+    """Fixture for an empty Prometheus API result."""
+    return []
 
 
 class TestEnergyConsumptionExtractor:
     """Test the energy consumption extractor."""
 
-    @mock.patch("caso.extract.prometheus.requests.get")
+    @mock.patch("caso.extract.prometheus.PrometheusConnect")
     def test_extract_with_results(
         self,
-        mock_get,
+        mock_prom_connect,
         configured_extractor,
         mock_server,
         extract_dates,
-        prometheus_success_response,
+        mock_prometheus_result_success,
     ):
         """Test extraction with successful Prometheus query."""
         # Create a second server
@@ -137,8 +118,10 @@ def test_extract_with_results(
             mock_server2,
         ]
 
-        # Mock Prometheus response
-        mock_get.return_value = prometheus_success_response
+        # Mock Prometheus client
+        mock_prom = mock.Mock()
+        mock_prom.custom_query.return_value = mock_prometheus_result_success
+        mock_prom_connect.return_value = mock_prom
 
         # Extract records
         records = configured_extractor.extract(**extract_dates)
@@ -161,41 +144,50 @@ def test_extract_with_no_vms(self, configured_extractor, extract_dates):
         # Verify - no VMs, no records
         assert len(records) == 0
 
-    @mock.patch("caso.extract.prometheus.requests.get")
-    def test_extract_with_failed_query(
+    @mock.patch("caso.extract.prometheus.PrometheusConnect")
+    def test_extract_with_no_energy_data(
         self,
-        mock_get,
+        mock_prom_connect,
         configured_extractor,
         mock_server,
         extract_dates,
-        prometheus_error_response,
+        mock_prometheus_result_empty,
     ):
-        """Test extraction when Prometheus query fails."""
+        """Test extraction when Prometheus returns no energy data."""
         # Mock Nova client with servers
         configured_extractor.nova = mock.Mock()
         configured_extractor.nova.servers.list.return_value = [mock_server]
 
-        # Mock Prometheus error response
-        mock_get.return_value = prometheus_error_response
+        # Mock Prometheus client with empty result
+        mock_prom = mock.Mock()
+        mock_prom.custom_query.return_value = mock_prometheus_result_empty
+        mock_prom_connect.return_value = mock_prom
 
         # Extract records
         records = configured_extractor.extract(**extract_dates)
 
-        # Verify - query failed, no records
+        # Verify - no energy data, no records
         assert len(records) == 0
 
     @mock.patch("caso.extract.prometheus.LOG")
-    @mock.patch("caso.extract.prometheus.requests.get")
-    def test_extract_with_request_exception(
-        self, mock_get, mock_log, configured_extractor, mock_server, extract_dates
+    @mock.patch("caso.extract.prometheus.PrometheusConnect")
+    def test_extract_with_prometheus_exception(
+        self,
+        mock_prom_connect,
+        mock_log,
+        configured_extractor,
+        mock_server,
+        extract_dates,
     ):
-        """Test extraction when request to Prometheus fails."""
+        """Test extraction when Prometheus query raises an exception."""
         # Mock Nova client with servers
         configured_extractor.nova = mock.Mock()
         configured_extractor.nova.servers.list.return_value = [mock_server]
 
-        # Mock request exception
-        mock_get.side_effect = Exception("Connection error")
+        # Mock Prometheus client to raise exception
+        mock_prom = mock.Mock()
+        mock_prom.custom_query.side_effect = Exception("Connection error")
+        mock_prom_connect.return_value = mock_prom
 
         # Extract records
         records = configured_extractor.extract(**extract_dates)