feat: Add Prometheus extractor and EnergyRecord for energy consumption metrics

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

Author: Copilot

caso/extract/__init__.py

@@ -19,9 +19,11 @@
 from caso.extract.openstack import CinderExtractor
 from caso.extract.openstack import NeutronExtractor
 from caso.extract.openstack import NovaExtractor
+from caso.extract.prometheus import PrometheusExtractor
 
 __all__ = [
     "NovaExtractor",
     "CinderExtractor",
     "NeutronExtractor",
+    "PrometheusExtractor",
 ]

caso/extract/prometheus.py

@@ -0,0 +1,168 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2014 Spanish National Research Council (CSIC)
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+"""Module containing the Prometheus extractor for energy consumption metrics."""
+
+import uuid
+
+import requests
+from oslo_config import cfg
+from oslo_log import log
+
+from caso.extract import base
+from caso import record
+
+CONF = cfg.CONF
+
+opts = [
+    cfg.StrOpt(
+        "prometheus_endpoint",
+        default="http://localhost:9090",
+        help="Prometheus server endpoint URL.",
+    ),
+    cfg.StrOpt(
+        "prometheus_query",
+        default="sum(rate(node_energy_joules_total[5m])) * 300 / 3600000",
+        help="Prometheus query to retrieve energy consumption in kWh. "
+        "The query should return energy consumption metrics.",
+    ),
+    cfg.IntOpt(
+        "prometheus_timeout",
+        default=30,
+        help="Timeout for Prometheus API requests in seconds.",
+    ),
+]
+
+CONF.import_opt("site_name", "caso.extract.base")
+CONF.register_opts(opts, group="prometheus")
+
+LOG = log.getLogger(__name__)
+
+
+class PrometheusExtractor(base.BaseProjectExtractor):
+    """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
+
+    def _query_prometheus(self, query, timestamp=None):
+        """Query Prometheus API and return results.
+
+        :param query: PromQL query string
+        :param timestamp: Optional timestamp for query (datetime object)
+        :returns: Query results
+        """
+        endpoint = CONF.prometheus.prometheus_endpoint
+        url = f"{endpoint}/api/v1/query"
+
+        params = {"query": query}
+        if timestamp:
+            params["time"] = int(timestamp.timestamp())
+
+        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
+        except Exception as e:
+            LOG.error(f"Unexpected error querying Prometheus: {e}")
+            return None
+
+    def _build_energy_record(self, energy_value, measurement_time):
+        """Build an energy consumption record.
+
+        :param energy_value: Energy consumption value in kWh
+        :param measurement_time: Time of measurement
+        :returns: EnergyRecord object
+        """
+        r = record.EnergyRecord(
+            uuid=uuid.uuid4(),
+            measurement_time=measurement_time,
+            site_name=CONF.site_name,
+            user_id=None,
+            group_id=self.project_id,
+            user_dn=None,
+            fqan=self.vo,
+            energy_consumption=energy_value,
+            energy_unit="kWh",
+            compute_service=CONF.service_name,
+        )
+
+        return r
+
+    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.
+
+        :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
+        """
+        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}"
+        )
+
+        results = self._query_prometheus(query, extract_to)
+
+        if results is None:
+            LOG.warning(
+                f"No results returned from Prometheus for project {self.project}"
+            )
+            return records
+
+        # 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} kWh "
+                f"for project {self.project}"
+            )
+
+            energy_record = self._build_energy_record(energy_value, extract_to)
+            records.append(energy_record)
+
+        LOG.info(f"Extracted {len(records)} energy records for project {self.project}")
+
+        return records

caso/record.py

@@ -555,3 +555,81 @@ def ssm_message(self):
         populate_by_name=True,
         extra="forbid",
     )
+
+
+def map_energy_fields(field: str) -> str:
+    """Map object fields to accounting Energy Usage Record fields."""
+    d = {
+        "measurement_time_epoch": "MeasurementTime",
+        "site_name": "SiteName",
+        "cloud_type": "CloudType",
+        "user_id": "LocalUser",
+        "group_id": "LocalGroup",
+        "fqan": "FQAN",
+        "user_dn": "GlobalUserName",
+        "energy_consumption": "EnergyConsumption",
+        "energy_unit": "EnergyUnit",
+        "compute_service": "CloudComputeService",
+    }
+    return d.get(field, field)
+
+
+class EnergyRecord(_BaseRecord):
+    """The EnergyRecord class holds energy consumption information.
+
+    This class is used to report energy consumption metrics gathered from
+    external monitoring systems like Prometheus.
+    """
+
+    version: str = pydantic.Field("0.1", exclude=True)
+
+    uuid: m_uuid.UUID
+
+    user_id: typing.Optional[str]
+    user_dn: typing.Optional[str]
+    group_id: str
+    fqan: str
+
+    # Make these fields private, and deal with them as properties. This is done as  all
+    # the accounting infrastructure needs start and end times as integers, but it is
+    # easier for us to maintain them as datetime objects internally.
+    _measurement_time: datetime.datetime
+
+    energy_consumption: float
+    energy_unit: str = "kWh"
+
+    def __init__(self, measurement_time: datetime.datetime, *args, **kwargs):
+        """Initialize the record."""
+        super(EnergyRecord, self).__init__(*args, **kwargs)
+
+        self._measurement_time = measurement_time
+
+    @property
+    def measurement_time(self) -> datetime.datetime:
+        """Get measurement time."""
+        return self._measurement_time
+
+    @measurement_time.setter
+    def measurement_time(self, measurement_time: datetime.datetime) -> None:
+        """Set measurement time."""
+        self._measurement_time = measurement_time
+
+    @pydantic.computed_field()  # type: ignore[misc]
+    @property
+    def measurement_time_epoch(self) -> int:
+        """Get measurement time as epoch."""
+        return int(self._measurement_time.timestamp())
+
+    def ssm_message(self):
+        """Render record as the expected SSM message."""
+        opts = {
+            "by_alias": True,
+            "exclude_none": True,
+        }
+        return self.model_dump_json(**opts)
+
+    model_config = dict(
+        alias_generator=map_energy_fields,
+        populate_by_name=True,
+        extra="forbid",
+    )

caso/tests/conftest.py

@@ -757,3 +757,51 @@ def expected_message_storage() -> str:
         "</sr:StorageUsageRecords>"
     )
     return message
+
+
+# Energy record fixtures
+
+valid_energy_records_fields = [
+    dict(
+        uuid="e3c5aeef-37b8-4332-ad9f-9d068f156dc2",
+        measurement_time=now,
+        site_name="TEST-Site",
+        user_id="a4519d7d-f60a-4908-9d63-7d9e17422188",
+        group_id="03b6a6c4-cf2b-48b9-82f1-69c52b9f30af",
+        user_dn="User 1 DN",
+        fqan="VO 1 FQAN",
+        energy_consumption=125.5,
+        energy_unit="kWh",
+        compute_service="Fake Cloud Service",
+        cloud_type=cloud_type,
+    ),
+]
+
+valid_energy_records_dict = [
+    {
+        "CloudComputeService": "Fake Cloud Service",
+        "FQAN": "VO 1 FQAN",
+        "GlobalUserName": "User 1 DN",
+        "EnergyConsumption": 125.5,
+        "EnergyUnit": "kWh",
+        "LocalGroup": "03b6a6c4-cf2b-48b9-82f1-69c52b9f30af",
+        "LocalUser": "a4519d7d-f60a-4908-9d63-7d9e17422188",
+        "MeasurementTime": 1685051946,
+        "SiteName": "TEST-Site",
+        "uuid": "e3c5aeef-37b8-4332-ad9f-9d068f156dc2",
+        "CloudType": cloud_type,
+    },
+]
+
+
+@pytest.fixture()
+def energy_record() -> caso.record.EnergyRecord:
+    """Get a fixture for an Energy record."""
+    record = caso.record.EnergyRecord(**valid_energy_records_fields[0])
+    return record
+
+
+@pytest.fixture()
+def valid_energy_record() -> dict:
+    """Get a fixture for a valid Energy record as a dict."""
+    return valid_energy_records_dict[0]