fix: Update Prometheus extractor to scan VMs and support query templating

- Changed PrometheusExtractor to inherit from BaseOpenStackExtractor instead of BaseProjectExtractor
- Added _get_servers() method to retrieve VMs from Nova for each project
- Updated extract() to iterate over VMs and query Prometheus per VM
- Added template variable support: {{uuid}} can be used in queries to reference VM UUID
- Updated default query to use libvirt domain metrics with UUID templating
- Updated all tests to mock the OpenStack base class
- Updated documentation with new query examples and VM scanning behavior
- Updated sample configuration file with new default query

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

Author: Copilot

caso/extract/prometheus.py

@@ -22,7 +22,7 @@
 from oslo_config import cfg
 from oslo_log import log
 
-from caso.extract import base
+from caso.extract.openstack import base
 from caso import record
 
 CONF = cfg.CONF
@@ -35,9 +35,10 @@
     ),
     cfg.StrOpt(
         "prometheus_query",
-        default="sum(rate(node_energy_joules_total[5m])) * 300 / 3600000",
+        default="sum(rate(libvirt_domain_info_energy_consumption_joules_total"
+        '{uuid=~"{{uuid}}"}[5m])) * 300 / 3600000',
         help="Prometheus query to retrieve energy consumption in kWh. "
-        "The query should return energy consumption metrics.",
+        "The query can use {{uuid}} as a template variable for the VM UUID.",
     ),
     cfg.IntOpt(
         "prometheus_timeout",
@@ -52,14 +53,13 @@
 LOG = log.getLogger(__name__)
 
 
-class PrometheusExtractor(base.BaseProjectExtractor):
+class PrometheusExtractor(base.BaseOpenStackExtractor):
     """A Prometheus extractor for energy consumption metrics in cASO."""
 
     def __init__(self, project, vo):
         """Initialize a Prometheus extractor for a given project."""
-        super(PrometheusExtractor, self).__init__(project)
-        self.vo = vo
-        self.project_id = project
+        super(PrometheusExtractor, self).__init__(project, vo)
+        self.nova = self._get_nova_client()
 
     def _query_prometheus(self, query, timestamp=None):
         """Query Prometheus API and return results.
@@ -95,9 +95,35 @@ def _query_prometheus(self, query, timestamp=None):
             LOG.error(f"Unexpected error querying Prometheus: {e}")
             return None
 
-    def _build_energy_record(self, energy_value, measurement_time):
-        """Build an energy consumption record.
+    def _get_servers(self, extract_from):
+        """Get all servers for a given date."""
+        servers = []
+        limit = 200
+        marker = None
+        # Use a marker and iter over results until we do not have more to get
+        while True:
+            aux = self.nova.servers.list(
+                search_opts={
+                    "changes-since": extract_from,
+                    "project_id": self.project_id,
+                    "all_tenants": True,
+                },
+                limit=limit,
+                marker=marker,
+            )
+            servers.extend(aux)
+
+            if len(aux) < limit:
+                break
+            marker = aux[-1].id
+
+        return servers
+
+    def _build_energy_record(self, vm_uuid, vm_name, energy_value, measurement_time):
+        """Build an energy consumption record for a VM.
 
+        :param vm_uuid: VM UUID
+        :param vm_name: VM name
         :param energy_value: Energy consumption value in kWh
         :param measurement_time: Time of measurement
         :returns: EnergyRecord object
@@ -121,47 +147,72 @@ def extract(self, extract_from, extract_to):
         """Extract energy consumption records from Prometheus.
 
         This method queries Prometheus for energy consumption metrics
-        in the specified time range.
+        for each VM in the project.
 
         :param extract_from: datetime.datetime object indicating the date to
                              extract records from
         :param extract_to: datetime.datetime object indicating the date to
                            extract records to
         :returns: A list of energy records
         """
+        # Remove timezone as Nova doesn't expect it
+        extract_from = extract_from.replace(tzinfo=None)
+        extract_to = extract_to.replace(tzinfo=None)
+
         records = []
 
-        # Query Prometheus at the extract_to timestamp
-        query = CONF.prometheus.prometheus_query
-        LOG.debug(
-            f"Querying Prometheus for project {self.project} " f"with query: {query}"
+        # Get all servers for the project
+        LOG.debug(f"Getting servers for project {self.project}")
+        servers = self._get_servers(extract_from)
+
+        LOG.info(
+            f"Found {len(servers)} VMs for project {self.project}, "
+            f"querying Prometheus for energy metrics"
         )
 
-        results = self._query_prometheus(query, 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
 
-        if results is None:
-            LOG.warning(
-                f"No results returned from Prometheus for project {self.project}"
+            # Replace template variables in the query
+            query = query_template.replace("{{uuid}}", vm_uuid)
+
+            LOG.debug(
+                f"Querying Prometheus for VM {vm_name} ({vm_uuid}) "
+                f"with query: {query}"
             )
-            return records
 
-        # Process results and create records
-        for result in results:
-            value = result.get("value", [])
+            results = self._query_prometheus(query, extract_to)
 
-            if len(value) < 2:
+            if results is None:
+                LOG.warning(
+                    f"No results returned from Prometheus for VM "
+                    f"{vm_name} ({vm_uuid})"
+                )
                 continue
 
-            # value is [timestamp, value_string]
-            energy_value = float(value[1])
+            # Process results and create records
+            for result in results:
+                value = result.get("value", [])
 
-            LOG.debug(
-                f"Creating energy record: {energy_value} kWh "
-                f"for project {self.project}"
-            )
+                if len(value) < 2:
+                    continue
+
+                # value is [timestamp, value_string]
+                energy_value = float(value[1])
+
+                LOG.debug(
+                    f"Creating energy record: {energy_value} kWh "
+                    f"for VM {vm_name} ({vm_uuid})"
+                )
 
-            energy_record = self._build_energy_record(energy_value, extract_to)
-            records.append(energy_record)
+                energy_record = self._build_energy_record(
+                    vm_uuid, vm_name, energy_value, extract_to
+                )
+                records.append(energy_record)
 
         LOG.info(f"Extracted {len(records)} energy records for project {self.project}")
 

caso/tests/extract/test_prometheus.py

@@ -30,13 +30,37 @@
 class TestPrometheusExtractor:
     """Test the Prometheus extractor."""
 
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor._get_nova_client")
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor.__init__")
     @mock.patch("caso.extract.prometheus.requests.get")
-    def test_extract_with_results(self, mock_get):
+    def test_extract_with_results(self, mock_get, mock_base_init, mock_get_nova):
         """Test extraction with successful Prometheus query."""
         # Configure CONF
         CONF.set_override("site_name", "TEST-Site")
         CONF.set_override("service_name", "TEST-Service")
 
+        # Mock the base class __init__ to do nothing
+        mock_base_init.return_value = None
+
+        # Mock Nova client and servers
+        mock_server1 = mock.Mock()
+        mock_server1.id = "vm-uuid-1"
+        mock_server1.name = "test-vm-1"
+
+        mock_server2 = mock.Mock()
+        mock_server2.id = "vm-uuid-2"
+        mock_server2.name = "test-vm-2"
+
+        mock_nova = mock.Mock()
+        mock_nova.servers.list.return_value = [mock_server1, mock_server2]
+        mock_get_nova.return_value = mock_nova
+
+        # Create extractor and manually set required attributes
+        extractor = PrometheusExtractor("test-project", "test-vo")
+        extractor.project = "test-project"
+        extractor.vo = "test-vo"
+        extractor.project_id = "test-project-id"
+
         # Mock Prometheus response
         mock_response = mock.Mock()
         mock_response.json.return_value = {
@@ -53,54 +77,74 @@ def test_extract_with_results(self, mock_get):
         mock_response.raise_for_status = mock.Mock()
         mock_get.return_value = mock_response
 
-        # Create extractor
-        extractor = PrometheusExtractor("test-project", "test-vo")
-
         # Extract records
         extract_from = datetime.datetime(2023, 5, 25, 0, 0, 0)
         extract_to = datetime.datetime(2023, 5, 25, 23, 59, 59)
         records = extractor.extract(extract_from, extract_to)
 
-        # Verify
-        assert len(records) == 1
+        # Verify - should create 2 records (one per VM)
+        assert len(records) == 2
         assert records[0].energy_consumption == 125.5
         assert records[0].energy_unit == "kWh"
         assert records[0].fqan == "test-vo"
 
-    @mock.patch("caso.extract.prometheus.requests.get")
-    def test_extract_with_no_results(self, mock_get):
-        """Test extraction when Prometheus returns no results."""
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor._get_nova_client")
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor.__init__")
+    def test_extract_with_no_vms(self, mock_base_init, mock_get_nova):
+        """Test extraction when there are no VMs."""
         # Configure CONF
         CONF.set_override("site_name", "TEST-Site")
         CONF.set_override("service_name", "TEST-Service")
 
-        # Mock Prometheus response with no results
-        mock_response = mock.Mock()
-        mock_response.json.return_value = {
-            "status": "success",
-            "data": {"result": []},
-        }
-        mock_response.raise_for_status = mock.Mock()
-        mock_get.return_value = mock_response
+        # Mock the base class __init__ to do nothing
+        mock_base_init.return_value = None
 
-        # Create extractor
+        # Mock Nova client with no servers
+        mock_nova = mock.Mock()
+        mock_nova.servers.list.return_value = []
+        mock_get_nova.return_value = mock_nova
+
+        # Create extractor and manually set required attributes
         extractor = PrometheusExtractor("test-project", "test-vo")
+        extractor.project = "test-project"
+        extractor.vo = "test-vo"
+        extractor.project_id = "test-project-id"
 
         # Extract records
         extract_from = datetime.datetime(2023, 5, 25, 0, 0, 0)
         extract_to = datetime.datetime(2023, 5, 25, 23, 59, 59)
         records = extractor.extract(extract_from, extract_to)
 
-        # Verify
+        # Verify - no VMs, no records
         assert len(records) == 0
 
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor._get_nova_client")
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor.__init__")
     @mock.patch("caso.extract.prometheus.requests.get")
-    def test_extract_with_failed_query(self, mock_get):
+    def test_extract_with_failed_query(self, mock_get, mock_base_init, mock_get_nova):
         """Test extraction when Prometheus query fails."""
         # Configure CONF
         CONF.set_override("site_name", "TEST-Site")
         CONF.set_override("service_name", "TEST-Service")
 
+        # Mock the base class __init__ to do nothing
+        mock_base_init.return_value = None
+
+        # Mock Nova client and servers
+        mock_server = mock.Mock()
+        mock_server.id = "vm-uuid-1"
+        mock_server.name = "test-vm-1"
+
+        mock_nova = mock.Mock()
+        mock_nova.servers.list.return_value = [mock_server]
+        mock_get_nova.return_value = mock_nova
+
+        # Create extractor and manually set required attributes
+        extractor = PrometheusExtractor("test-project", "test-vo")
+        extractor.project = "test-project"
+        extractor.vo = "test-vo"
+        extractor.project_id = "test-project-id"
+
         # Mock Prometheus error response
         mock_response = mock.Mock()
         mock_response.json.return_value = {
@@ -110,36 +154,52 @@ def test_extract_with_failed_query(self, mock_get):
         mock_response.raise_for_status = mock.Mock()
         mock_get.return_value = mock_response
 
-        # Create extractor
-        extractor = PrometheusExtractor("test-project", "test-vo")
-
         # Extract records
         extract_from = datetime.datetime(2023, 5, 25, 0, 0, 0)
         extract_to = datetime.datetime(2023, 5, 25, 23, 59, 59)
         records = extractor.extract(extract_from, extract_to)
 
-        # Verify
+        # Verify - query failed, no records
         assert len(records) == 0
 
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor._get_nova_client")
+    @mock.patch("caso.extract.openstack.base.BaseOpenStackExtractor.__init__")
     @mock.patch("caso.extract.prometheus.requests.get")
     @mock.patch("caso.extract.prometheus.LOG")
-    def test_extract_with_request_exception(self, mock_log, mock_get):
+    def test_extract_with_request_exception(
+        self, mock_log, mock_get, mock_base_init, mock_get_nova
+    ):
         """Test extraction when request to Prometheus fails."""
         # Configure CONF
         CONF.set_override("site_name", "TEST-Site")
         CONF.set_override("service_name", "TEST-Service")
 
-        # Mock request exception
-        mock_get.side_effect = Exception("Connection error")
+        # Mock the base class __init__ to do nothing
+        mock_base_init.return_value = None
+
+        # Mock Nova client and servers
+        mock_server = mock.Mock()
+        mock_server.id = "vm-uuid-1"
+        mock_server.name = "test-vm-1"
+
+        mock_nova = mock.Mock()
+        mock_nova.servers.list.return_value = [mock_server]
+        mock_get_nova.return_value = mock_nova
 
-        # Create extractor
+        # Create extractor and manually set required attributes
         extractor = PrometheusExtractor("test-project", "test-vo")
+        extractor.project = "test-project"
+        extractor.vo = "test-vo"
+        extractor.project_id = "test-project-id"
+
+        # Mock request exception
+        mock_get.side_effect = Exception("Connection error")
 
         # Extract records
         extract_from = datetime.datetime(2023, 5, 25, 0, 0, 0)
         extract_to = datetime.datetime(2023, 5, 25, 23, 59, 59)
         records = extractor.extract(extract_from, extract_to)
 
-        # Verify
+        # Verify - exception caught, no records
         assert len(records) == 0
         mock_log.error.assert_called()