From 4bca9596ee5d0e5af7a4faa1d896fe7e88ecc849 Mon Sep 17 00:00:00 2001 From: Corbeno Date: Thu, 14 Jan 2021 04:31:37 -0600 Subject: [PATCH] Rework Proxmoxve to use a DataUpdateCoordinator (#45068) Co-authored-by: Paulus Schoutsen --- .../components/proxmoxve/__init__.py | 235 ++++++++++++++---- .../components/proxmoxve/binary_sensor.py | 159 ++++++------ 2 files changed, 264 insertions(+), 130 deletions(-) diff --git a/homeassistant/components/proxmoxve/__init__.py b/homeassistant/components/proxmoxve/__init__.py index 0919feb15e3..2f42ca8fe9e 100644 --- a/homeassistant/components/proxmoxve/__init__.py +++ b/homeassistant/components/proxmoxve/__init__.py @@ -1,9 +1,10 @@ """Support for Proxmox VE.""" -from enum import Enum +from datetime import timedelta import logging from proxmoxer import ProxmoxAPI from proxmoxer.backends.https import AuthenticationError +from proxmoxer.core import ResourceException from requests.exceptions import SSLError import voluptuous as vol @@ -14,11 +15,14 @@ from homeassistant.const import ( CONF_USERNAME, CONF_VERIFY_SSL, ) +from homeassistant.core import HomeAssistant import homeassistant.helpers.config_validation as cv +from homeassistant.helpers.update_coordinator import ( + CoordinatorEntity, + DataUpdateCoordinator, +) -_LOGGER = logging.getLogger(__name__) - - +PLATFORMS = ["binary_sensor"] DOMAIN = "proxmoxve" PROXMOX_CLIENTS = "proxmox_clients" CONF_REALM = "realm" @@ -27,9 +31,17 @@ CONF_NODES = "nodes" CONF_VMS = "vms" CONF_CONTAINERS = "containers" +COORDINATOR = "coordinator" +API_DATA = "api_data" + DEFAULT_PORT = 8006 DEFAULT_REALM = "pam" DEFAULT_VERIFY_SSL = True +TYPE_VM = 0 +TYPE_CONTAINER = 1 +UPDATE_INTERVAL = 60 + +_LOGGER = logging.getLogger(__name__) CONFIG_SCHEMA = vol.Schema( { @@ -71,52 +83,191 @@ CONFIG_SCHEMA = vol.Schema( ) -def setup(hass, config): - """Set up the component.""" +async def async_setup(hass: HomeAssistant, config: dict): + """Set up the platform.""" + hass.data.setdefault(DOMAIN, {}) - # Create API Clients for later use - hass.data[PROXMOX_CLIENTS] = {} - for entry in config[DOMAIN]: - host = entry[CONF_HOST] - port = entry[CONF_PORT] - user = entry[CONF_USERNAME] - realm = entry[CONF_REALM] - password = entry[CONF_PASSWORD] - verify_ssl = entry[CONF_VERIFY_SSL] + def build_client() -> ProxmoxAPI: + """Build the Proxmox client connection.""" + hass.data[PROXMOX_CLIENTS] = {} + for entry in config[DOMAIN]: + host = entry[CONF_HOST] + port = entry[CONF_PORT] + user = entry[CONF_USERNAME] + realm = entry[CONF_REALM] + password = entry[CONF_PASSWORD] + verify_ssl = entry[CONF_VERIFY_SSL] - try: - # Construct an API client with the given data for the given host - proxmox_client = ProxmoxClient( - host, port, user, realm, password, verify_ssl + try: + # Construct an API client with the given data for the given host + proxmox_client = ProxmoxClient( + host, port, user, realm, password, verify_ssl + ) + proxmox_client.build_client() + except AuthenticationError: + _LOGGER.warning( + "Invalid credentials for proxmox instance %s:%d", host, port + ) + continue + except SSLError: + _LOGGER.error( + 'Unable to verify proxmox server SSL. Try using "verify_ssl: false"' + ) + continue + + return proxmox_client + + proxmox_client = await hass.async_add_executor_job(build_client) + + async def async_update_data() -> dict: + """Fetch data from API endpoint.""" + + proxmox = proxmox_client.get_api_client() + + def poll_api() -> dict: + data = {} + + for host_config in config[DOMAIN]: + host_name = host_config["host"] + + data[host_name] = {} + + for node_config in host_config["nodes"]: + node_name = node_config["node"] + data[host_name][node_name] = {} + + for vm_id in node_config["vms"]: + data[host_name][node_name][vm_id] = {} + + vm_status = call_api_container_vm( + proxmox, node_name, vm_id, TYPE_VM + ) + + if vm_status is None: + _LOGGER.warning("Vm/Container %s unable to be found", vm_id) + data[host_name][node_name][vm_id] = None + continue + + data[host_name][node_name][vm_id] = parse_api_container_vm( + vm_status + ) + + for container_id in node_config["containers"]: + data[host_name][node_name][container_id] = {} + + container_status = call_api_container_vm( + proxmox, node_name, container_id, TYPE_CONTAINER + ) + + if container_status is None: + _LOGGER.error( + "Vm/Container %s unable to be found", container_id + ) + data[host_name][node_name][container_id] = None + continue + + data[host_name][node_name][ + container_id + ] = parse_api_container_vm(container_status) + + return data + + return await hass.async_add_executor_job(poll_api) + + coordinator = DataUpdateCoordinator( + hass, + _LOGGER, + name="proxmox_coordinator", + update_method=async_update_data, + update_interval=timedelta(seconds=UPDATE_INTERVAL), + ) + + hass.data[DOMAIN][COORDINATOR] = coordinator + + # Fetch initial data + await coordinator.async_refresh() + + for component in PLATFORMS: + await hass.async_create_task( + hass.helpers.discovery.async_load_platform( + component, DOMAIN, {"config": config}, config ) - proxmox_client.build_client() - except AuthenticationError: - _LOGGER.warning( - "Invalid credentials for proxmox instance %s:%d", host, port - ) - continue - except SSLError: - _LOGGER.error( - 'Unable to verify proxmox server SSL. Try using "verify_ssl: false"' - ) - continue - - hass.data[PROXMOX_CLIENTS][f"{host}:{port}"] = proxmox_client - - if hass.data[PROXMOX_CLIENTS]: - hass.helpers.discovery.load_platform( - "binary_sensor", DOMAIN, {"entries": config[DOMAIN]}, config ) - return True - return False + return True -class ProxmoxItemType(Enum): - """Represents the different types of machines in Proxmox.""" +def parse_api_container_vm(status): + """Get the container or vm api data and return it formatted in a dictionary. - qemu = 0 - lxc = 1 + It is implemented in this way to allow for more data to be added for sensors + in the future. + """ + + return {"status": status["status"], "name": status["name"]} + + +def call_api_container_vm(proxmox, node_name, vm_id, machine_type): + """Make proper api calls.""" + status = None + + try: + if machine_type == TYPE_VM: + status = proxmox.nodes(node_name).qemu(vm_id).status.current.get() + elif machine_type == TYPE_CONTAINER: + status = proxmox.nodes(node_name).lxc(vm_id).status.current.get() + except ResourceException: + return None + + return status + + +class ProxmoxEntity(CoordinatorEntity): + """Represents any entity created for the Proxmox VE platform.""" + + def __init__( + self, + coordinator: DataUpdateCoordinator, + unique_id, + name, + icon, + host_name, + node_name, + vm_id=None, + ): + """Initialize the Proxmox entity.""" + super().__init__(coordinator) + + self.coordinator = coordinator + self._unique_id = unique_id + self._name = name + self._host_name = host_name + self._icon = icon + self._available = True + self._node_name = node_name + self._vm_id = vm_id + + self._state = None + + @property + def unique_id(self) -> str: + """Return the unique ID for this sensor.""" + return self._unique_id + + @property + def name(self) -> str: + """Return the name of the entity.""" + return self._name + + @property + def icon(self) -> str: + """Return the mdi icon of the entity.""" + return self._icon + + @property + def available(self) -> bool: + """Return True if entity is available.""" + return self.coordinator.last_update_success and self._available class ProxmoxClient: diff --git a/homeassistant/components/proxmoxve/binary_sensor.py b/homeassistant/components/proxmoxve/binary_sensor.py index 698a2c35ae1..014766b532e 100644 --- a/homeassistant/components/proxmoxve/binary_sensor.py +++ b/homeassistant/components/proxmoxve/binary_sensor.py @@ -1,112 +1,95 @@ """Binary sensor to read Proxmox VE data.""" import logging -from homeassistant.components.binary_sensor import BinarySensorEntity -from homeassistant.const import ATTR_ATTRIBUTION, CONF_HOST, CONF_PORT +from homeassistant.const import STATE_OFF, STATE_ON +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import CONF_CONTAINERS, CONF_NODES, CONF_VMS, PROXMOX_CLIENTS, ProxmoxItemType +from . import COORDINATOR, DOMAIN, ProxmoxEntity -ATTRIBUTION = "Data provided by Proxmox VE" _LOGGER = logging.getLogger(__name__) -def setup_platform(hass, config, add_entities, discovery_info=None): - """Set up the sensor platform.""" +async def async_setup_platform(hass, config, add_entities, discovery_info=None): + """Set up binary sensors.""" + if discovery_info is None: + return + + coordinator = hass.data[DOMAIN][COORDINATOR] sensors = [] - for entry in discovery_info["entries"]: - port = entry[CONF_PORT] + for host_config in discovery_info["config"][DOMAIN]: + host_name = host_config["host"] - for node in entry[CONF_NODES]: - for virtual_machine in node[CONF_VMS]: - sensors.append( - ProxmoxBinarySensor( - hass.data[PROXMOX_CLIENTS][f"{entry[CONF_HOST]}:{port}"], - node["node"], - ProxmoxItemType.qemu, - virtual_machine, - ) + for node_config in host_config["nodes"]: + node_name = node_config["node"] + + for vm_id in node_config["vms"]: + coordinator_data = coordinator.data[host_name][node_name][vm_id] + + # unfound vm case + if coordinator_data is None: + continue + + vm_name = coordinator_data["name"] + vm_status = create_binary_sensor( + coordinator, host_name, node_name, vm_id, vm_name ) + sensors.append(vm_status) - for container in node[CONF_CONTAINERS]: - sensors.append( - ProxmoxBinarySensor( - hass.data[PROXMOX_CLIENTS][f"{entry[CONF_HOST]}:{port}"], - node["node"], - ProxmoxItemType.lxc, - container, - ) + for container_id in node_config["containers"]: + coordinator_data = coordinator.data[host_name][node_name][container_id] + + # unfound container case + if coordinator_data is None: + continue + + container_name = coordinator_data["name"] + container_status = create_binary_sensor( + coordinator, host_name, node_name, container_id, container_name ) + sensors.append(container_status) - add_entities(sensors, True) + add_entities(sensors) -class ProxmoxBinarySensor(BinarySensorEntity): +def create_binary_sensor(coordinator, host_name, node_name, vm_id, name): + """Create a binary sensor based on the given data.""" + return ProxmoxBinarySensor( + coordinator=coordinator, + unique_id=f"proxmox_{node_name}_{vm_id}_running", + name=f"{node_name}_{name}_running", + icon="", + host_name=host_name, + node_name=node_name, + vm_id=vm_id, + ) + + +class ProxmoxBinarySensor(ProxmoxEntity): """A binary sensor for reading Proxmox VE data.""" - def __init__(self, proxmox_client, item_node, item_type, item_id): - """Initialize the binary sensor.""" - self._proxmox_client = proxmox_client - self._item_node = item_node - self._item_type = item_type - self._item_id = item_id - - self._vmname = None - self._name = None + def __init__( + self, + coordinator: DataUpdateCoordinator, + unique_id, + name, + icon, + host_name, + node_name, + vm_id, + ): + """Create the binary sensor for vms or containers.""" + super().__init__( + coordinator, unique_id, name, icon, host_name, node_name, vm_id + ) self._state = None @property - def name(self): - """Return the name of the entity.""" - return self._name - - @property - def is_on(self): - """Return true if VM/container is running.""" - return self._state - - @property - def device_state_attributes(self): - """Return device attributes of the entity.""" - return { - "node": self._item_node, - "vmid": self._item_id, - "vmname": self._vmname, - "type": self._item_type.name, - ATTR_ATTRIBUTION: ATTRIBUTION, - } - - def update(self): - """Check if the VM/Container is running.""" - item = self.poll_item() - - if item is None: - _LOGGER.warning("Failed to poll VM/container %s", self._item_id) - return - - self._state = item["status"] == "running" - - def poll_item(self): - """Find the VM/Container with the set item_id.""" - items = ( - self._proxmox_client.get_api_client() - .nodes(self._item_node) - .get(self._item_type.name) - ) - item = next( - (item for item in items if item["vmid"] == str(self._item_id)), None - ) - - if item is None: - _LOGGER.warning("Couldn't find VM/Container with the ID %s", self._item_id) - return None - - if self._vmname is None: - self._vmname = item["name"] - - if self._name is None: - self._name = f"{self._item_node} {self._vmname} running" - - return item + def state(self): + """Return the state of the binary sensor.""" + data = self.coordinator.data[self._host_name][self._node_name][self._vm_id] + if data["status"] == "running": + return STATE_ON + return STATE_OFF