|
| 1 | +"""Prometheus Binary Sensor component.""" |
| 2 | + |
| 3 | +from typing import TYPE_CHECKING, Final |
| 4 | + |
| 5 | +import voluptuous as vol |
| 6 | + |
| 7 | +from homeassistant.components.binary_sensor import ( |
| 8 | + PLATFORM_SCHEMA as BINARY_SENSOR_PLATFORM_SCHEMA, |
| 9 | + BinarySensorDeviceClass, |
| 10 | + BinarySensorEntity, |
| 11 | +) |
| 12 | +from homeassistant.const import ( |
| 13 | + CONF_DEVICE_CLASS, |
| 14 | + CONF_NAME, |
| 15 | + CONF_UNIQUE_ID, |
| 16 | + CONF_URL, |
| 17 | + CONF_VALUE_TEMPLATE, |
| 18 | +) |
| 19 | +from homeassistant.helpers.aiohttp_client import async_get_clientsession |
| 20 | +import homeassistant.helpers.config_validation as cv |
| 21 | + |
| 22 | +if TYPE_CHECKING: |
| 23 | + from homeassistant.core import HomeAssistant |
| 24 | + from homeassistant.helpers.entity_platform import AddEntitiesCallback |
| 25 | + from homeassistant.helpers.template import Template |
| 26 | + from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType |
| 27 | + |
| 28 | + from . import Prometheus, QueryResult |
| 29 | + |
| 30 | +from .const import CONF_EXPR, CONF_QUERIES, DEFAULT_URL, SCAN_INTERVAL as SCAN_INTERVAL |
| 31 | + |
| 32 | +_QUERY_SCHEMA: Final = vol.Schema( |
| 33 | + { |
| 34 | + vol.Required(CONF_NAME): cv.string, |
| 35 | + vol.Optional(CONF_UNIQUE_ID): cv.string, |
| 36 | + vol.Required(CONF_EXPR): cv.string, |
| 37 | + vol.Required(CONF_VALUE_TEMPLATE): cv.template, |
| 38 | + vol.Optional(CONF_DEVICE_CLASS): vol.Coerce(BinarySensorDeviceClass), |
| 39 | + } |
| 40 | +) |
| 41 | + |
| 42 | +PLATFORM_SCHEMA: Final = BINARY_SENSOR_PLATFORM_SCHEMA.extend( |
| 43 | + { |
| 44 | + vol.Optional(CONF_URL, default=DEFAULT_URL): cv.string, # type: ignore |
| 45 | + vol.Required(CONF_QUERIES): [_QUERY_SCHEMA], |
| 46 | + } |
| 47 | +) |
| 48 | + |
| 49 | + |
| 50 | +async def async_setup_platform( |
| 51 | + hass: HomeAssistant, |
| 52 | + config: ConfigType, |
| 53 | + async_add_entities: AddEntitiesCallback, |
| 54 | + discovery_info: DiscoveryInfoType | None = None, |
| 55 | +): |
| 56 | + """Set up the sensor platform.""" |
| 57 | + session = async_get_clientsession(hass) |
| 58 | + url = config[CONF_URL] |
| 59 | + prometheus = Prometheus(url, session) |
| 60 | + |
| 61 | + async_add_entities( |
| 62 | + new_entities=[ |
| 63 | + PrometheusBinarySensor( |
| 64 | + prometheus=prometheus, |
| 65 | + unique_id=query.get(CONF_UNIQUE_ID), |
| 66 | + device_name=query[CONF_NAME], |
| 67 | + expression=query[CONF_EXPR], |
| 68 | + value_template=query.get(CONF_VALUE_TEMPLATE), |
| 69 | + device_class=query.get(CONF_DEVICE_CLASS), |
| 70 | + ) |
| 71 | + for query in config[CONF_QUERIES] |
| 72 | + ], |
| 73 | + update_before_add=True, |
| 74 | + ) |
| 75 | + |
| 76 | + |
| 77 | +class PrometheusBinarySensor(BinarySensorEntity): |
| 78 | + """Sensor entity representing the result of a PromQL expression.""" |
| 79 | + |
| 80 | + def __init__( |
| 81 | + self, |
| 82 | + *, |
| 83 | + prometheus: Prometheus, |
| 84 | + unique_id: str | None, |
| 85 | + device_name: str, |
| 86 | + expression: str, |
| 87 | + value_template: Template, |
| 88 | + device_class: BinarySensorDeviceClass | None, |
| 89 | + ): |
| 90 | + """Initialize the sensor.""" |
| 91 | + self._prometheus: Prometheus = prometheus |
| 92 | + self._expression = expression |
| 93 | + self._value_template = value_template |
| 94 | + |
| 95 | + self._attr_device_class = device_class |
| 96 | + self._attr_name = device_name |
| 97 | + self._attr_unique_id = unique_id |
| 98 | + |
| 99 | + async def async_update(self) -> None: |
| 100 | + """Update state by executing query.""" |
| 101 | + result: QueryResult = await self._prometheus.query(self._expression) |
| 102 | + self._attr_available = result.error is None |
| 103 | + |
| 104 | + render_result = ( |
| 105 | + await self._value_template.async_render_with_possible_json_value( |
| 106 | + value=result.value, error_value=None |
| 107 | + ) |
| 108 | + ) |
| 109 | + |
| 110 | + if render_result is not None: |
| 111 | + self._attr_is_on = bool(render_result) |
| 112 | + else: |
| 113 | + self._attr_is_on = None |
0 commit comments