Initialize docker stack repo

homeassistant/config/custom_components/prometheus_sensor/sensor.py

@@ -0,0 +1,112 @@
+"""Prometheus Sensor component."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+import voluptuous as vol
+
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
+    SensorDeviceClass,
+    SensorEntity,
+    SensorStateClass,
+)
+from homeassistant.const import (
+    CONF_DEVICE_CLASS,
+    CONF_NAME,
+    CONF_UNIQUE_ID,
+    CONF_UNIT_OF_MEASUREMENT,
+    CONF_URL,
+)
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+import homeassistant.helpers.config_validation as cv
+
+from . import Prometheus
+
+if TYPE_CHECKING:
+    from homeassistant.core import HomeAssistant
+    from homeassistant.helpers.entity_platform import AddEntitiesCallback
+    from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+
+    from . import QueryResult
+
+from .const import CONF_EXPR, CONF_QUERIES, DEFAULT_URL, SCAN_INTERVAL as SCAN_INTERVAL
+
+_QUERY_SCHEMA: Final = vol.Schema(
+    {
+        vol.Required(CONF_NAME): cv.string,
+        vol.Optional(CONF_UNIQUE_ID): cv.string,
+        vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
+        vol.Required(CONF_EXPR): cv.string,
+        vol.Optional(CONF_DEVICE_CLASS): vol.Coerce(SensorDeviceClass),
+        vol.Optional(CONF_STATE_CLASS): vol.Coerce(SensorStateClass),
+    }
+)
+
+PLATFORM_SCHEMA: Final = SENSOR_PLATFORM_SCHEMA.extend(
+    {
+        vol.Optional(CONF_URL, default=DEFAULT_URL): cv.string,
+        vol.Required(CONF_QUERIES): [_QUERY_SCHEMA],
+    }
+)
+
+
+async def async_setup_platform(
+    hass: HomeAssistant,
+    config: ConfigType,
+    async_add_entities: AddEntitiesCallback,
+    discovery_info: DiscoveryInfoType | None = None,
+):
+    """Set up the sensor platform."""
+    session = async_get_clientsession(hass)
+    url = config[CONF_URL]
+    prometheus = Prometheus(url, session)
+
+    async_add_entities(
+        new_entities=[
+            PrometheusSensor(
+                prometheus=prometheus,
+                expression=query[CONF_EXPR],
+                unique_id=query.get(CONF_UNIQUE_ID),
+                device_name=query[CONF_NAME],
+                device_class=query.get(CONF_DEVICE_CLASS),
+                state_class=query.get(CONF_STATE_CLASS),
+                unit_of_measurement=query.get(CONF_UNIT_OF_MEASUREMENT),
+            )
+            for query in config[CONF_QUERIES]
+        ],
+        update_before_add=True,
+    )
+
+
+class PrometheusSensor(SensorEntity):
+    """Sensor entity representing the result of a PromQL expression."""
+
+    def __init__(
+        self,
+        *,
+        prometheus: Prometheus,
+        expression: str,
+        unique_id: str | None,
+        device_name: str,
+        device_class: SensorDeviceClass | None,
+        state_class: SensorStateClass | None,
+        unit_of_measurement: str | None,
+    ) -> None:
+        """Initialize the sensor."""
+        self._prometheus: Prometheus = prometheus
+        self._expression = expression
+
+        self._attr_device_class = device_class
+        self._attr_name = device_name
+        self._attr_native_unit_of_measurement = unit_of_measurement
+        self._attr_state_class = state_class
+        self._attr_unique_id = unique_id
+
+    async def async_update(self) -> None:
+        """Update state by executing query."""
+        result: QueryResult = await self._prometheus.query(self._expression)
+        self._attr_available = result.error is None
+        self._attr_native_value = result.value