addon_manager.py

Go to the documentation of this file.

"""Provide add-on management."""
 
from __future__ import annotations
 
import asyncio
from collections.abc import Awaitable, Callable, Coroutine
from dataclasses import dataclass
from enum import Enum
from functools import partial, wraps
import logging
from typing import Any, Concatenate
 
from aiohasupervisor import SupervisorError
from aiohasupervisor.models import (
    AddonsOptions,
    AddonState as SupervisorAddonState,
    InstalledAddonComplete,
    StoreAddonUpdate,
)
 
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
 
from .handler import HassioAPIError, async_create_backup, get_supervisor_client
 
type _FuncType[_T, **_P, _R] = Callable[Concatenate[_T, _P], Awaitable[_R]]
type _ReturnFuncType[_T, **_P, _R] = Callable[
    Concatenate[_T, _P], Coroutine[Any, Any, _R]
]
 
 
def api_error[_AddonManagerT: AddonManager, **_P, _R](
    error_message: str,
    *,
    expected_error_type: type[HassioAPIError | SupervisorError] | None = None,
) -> Callable[
    [_FuncType[_AddonManagerT, _P, _R]], _ReturnFuncType[_AddonManagerT, _P, _R]
]:
    """Handle HassioAPIError and raise a specific AddonError."""
    error_type = expected_error_type or (HassioAPIError, SupervisorError)
 
    def handle_hassio_api_error(
        func: _FuncType[_AddonManagerT, _P, _R],
    ) -> _ReturnFuncType[_AddonManagerT, _P, _R]:
        """Handle a HassioAPIError."""
 
        @wraps(func)
        async def wrapper(
            self: _AddonManagerT, *args: _P.args, **kwargs: _P.kwargs
        ) -> _R:
            """Wrap an add-on manager method."""
            try:
                return_value = await func(self, *args, **kwargs)
            except error_type as err:
                raise AddonError(
                    f"{error_message.format(addon_name=self.addon_name)}: {err}"
                ) from err
 
            return return_value
 
        return wrapper
 
    return handle_hassio_api_error
 
 
@dataclass
class AddonInfo:
    """Represent the current add-on info state."""
 
    available: bool
    hostname: str | None
    options: dict[str, Any]
    state: AddonState
    update_available: bool
    version: str | None
 
 
class AddonState(Enum):
    """Represent the current state of the add-on."""
 
    NOT_INSTALLED = "not_installed"
    INSTALLING = "installing"
    UPDATING = "updating"
    NOT_RUNNING = "not_running"
    RUNNING = "running"
 
 
class AddonManager:
    """Manage the add-on.
 
    Methods may raise AddonError.
    Only one instance of this class may exist per add-on
    to keep track of running add-on tasks.
    """
 
    def __init__(
        self,
        hass: HomeAssistant,
        logger: logging.Logger,
        addon_name: str,
        addon_slug: str,
    ) -> None:
        """Set up the add-on manager."""
        self.addon_nameaddon_name = addon_name
        self.addon_slugaddon_slug = addon_slug
        self._hass_hass = hass
        self._logger_logger = logger
        self._install_task_install_task: asyncio.Task | None = None
        self._restart_task_restart_task: asyncio.Task | None = None
        self._start_task_start_task: asyncio.Task | None = None
        self._update_task_update_task: asyncio.Task | None = None
        self._supervisor_client_supervisor_client = get_supervisor_client(hass)
 
    def task_in_progress(self) -> bool:
        """Return True if any of the add-on tasks are in progress."""
        return any(
            task and not task.done()
            for task in (
                self._restart_task_restart_task,
                self._install_task_install_task,
                self._start_task_start_task,
                self._update_task_update_task,
            )
        )
 
    @api_error(
        "Failed to get the {addon_name} add-on discovery info",
        expected_error_type=SupervisorError,
    )
    async def async_get_addon_discovery_info(self) -> dict:
        """Return add-on discovery info."""
        discovery_info = next(
            (
                msg
                for msg in await self._supervisor_client_supervisor_client.discovery.list()
                if msg.addon == self.addon_slugaddon_slug
            ),
            None,
        )
 
        if not discovery_info:
            raise AddonError(f"Failed to get {self.addon_name} add-on discovery info")
 
        return discovery_info.config
 
    @api_error(
        "Failed to get the {addon_name} add-on info",
        expected_error_type=SupervisorError,
    )
    async def async_get_addon_info(self) -> AddonInfo:
        """Return and cache manager add-on info."""
        addon_store_info = await self._supervisor_client_supervisor_client.store.addon_info(
            self.addon_slugaddon_slug
        )
        self._logger_logger.debug("Add-on store info: %s", addon_store_info.to_dict())
        if not addon_store_info.installed:
            return AddonInfo(
                available=addon_store_info.available,
                hostname=None,
                options={},
                state=AddonState.NOT_INSTALLED,
                update_available=False,
                version=None,
            )
 
        addon_info = await self._supervisor_client_supervisor_client.addons.addon_info(self.addon_slugaddon_slug)
        addon_state = self.async_get_addon_stateasync_get_addon_state(addon_info)
        return AddonInfo(
            available=addon_info.available,
            hostname=addon_info.hostname,
            options=addon_info.options,
            state=addon_state,
            update_available=addon_info.update_available,
            version=addon_info.version,
        )
 
    @callback
    def async_get_addon_state(self, addon_info: InstalledAddonComplete) -> AddonState:
        """Return the current state of the managed add-on."""
        addon_state = AddonState.NOT_RUNNING
 
        if addon_info.state == SupervisorAddonState.STARTED:
            addon_state = AddonState.RUNNING
        if self._install_task_install_task and not self._install_task_install_task.done():
            addon_state = AddonState.INSTALLING
        if self._update_task_update_task and not self._update_task_update_task.done():
            addon_state = AddonState.UPDATING
 
        return addon_state
 
    @api_error(
        "Failed to set the {addon_name} add-on options",
        expected_error_type=SupervisorError,
    )
    async def async_set_addon_options(self, config: dict) -> None:
        """Set manager add-on options."""
        await self._supervisor_client_supervisor_client.addons.set_addon_options(
            self.addon_slugaddon_slug, AddonsOptions(config=config)
        )
 
    def _check_addon_available(self, addon_info: AddonInfo) -> None:
        """Check if the managed add-on is available."""
        if not addon_info.available:
            raise AddonError(f"{self.addon_name} add-on is not available")
 
    @api_error(
        "Failed to install the {addon_name} add-on", expected_error_type=SupervisorError
    )
    async def async_install_addon(self) -> None:
        """Install the managed add-on."""
        addon_info = await self.async_get_addon_infoasync_get_addon_info()
 
        self._check_addon_available_check_addon_available(addon_info)
 
        await self._supervisor_client_supervisor_client.store.install_addon(self.addon_slugaddon_slug)
 
    @api_error(
        "Failed to uninstall the {addon_name} add-on",
        expected_error_type=SupervisorError,
    )
    async def async_uninstall_addon(self) -> None:
        """Uninstall the managed add-on."""
        await self._supervisor_client_supervisor_client.addons.uninstall_addon(self.addon_slugaddon_slug)
 
    @api_error("Failed to update the {addon_name} add-on")
    async def async_update_addon(self) -> None:
        """Update the managed add-on if needed."""
        addon_info = await self.async_get_addon_infoasync_get_addon_info()
 
        self._check_addon_available_check_addon_available(addon_info)
 
        if addon_info.state is AddonState.NOT_INSTALLED:
            raise AddonError(f"{self.addon_name} add-on is not installed")
 
        if not addon_info.update_available:
            return
 
        await self.async_create_backupasync_create_backup()
        await self._supervisor_client_supervisor_client.store.update_addon(
            self.addon_slugaddon_slug, StoreAddonUpdate(backup=False)
        )
 
    @api_error(
        "Failed to start the {addon_name} add-on", expected_error_type=SupervisorError
    )
    async def async_start_addon(self) -> None:
        """Start the managed add-on."""
        await self._supervisor_client_supervisor_client.addons.start_addon(self.addon_slugaddon_slug)
 
    @api_error(
        "Failed to restart the {addon_name} add-on", expected_error_type=SupervisorError
    )
    async def async_restart_addon(self) -> None:
        """Restart the managed add-on."""
        await self._supervisor_client_supervisor_client.addons.restart_addon(self.addon_slugaddon_slug)
 
    @api_error(
        "Failed to stop the {addon_name} add-on", expected_error_type=SupervisorError
    )
    async def async_stop_addon(self) -> None:
        """Stop the managed add-on."""
        await self._supervisor_client_supervisor_client.addons.stop_addon(self.addon_slugaddon_slug)
 
    @api_error("Failed to create a backup of the {addon_name} add-on")
    async def async_create_backup(self) -> None:
        """Create a partial backup of the managed add-on."""
        addon_info = await self.async_get_addon_infoasync_get_addon_info()
        name = f"addon_{self.addon_slug}_{addon_info.version}"
 
        self._logger_logger.debug("Creating backup: %s", name)
        await async_create_backup(
            self._hass_hass,
            {"name": name, "addons": [self.addon_slugaddon_slug]},
            partial=True,
        )
 
    async def async_configure_addon(
        self,
        addon_config: dict[str, Any],
    ) -> None:
        """Configure the manager add-on, if needed."""
        addon_info = await self.async_get_addon_infoasync_get_addon_info()
 
        if addon_info.state is AddonState.NOT_INSTALLED:
            raise AddonError(f"{self.addon_name} add-on is not installed")
 
        if addon_config != addon_info.options:
            await self.async_set_addon_optionsasync_set_addon_options(addon_config)
 
    @callback
    def async_schedule_install_addon(self, catch_error: bool = False) -> asyncio.Task:
        """Schedule a task that installs the managed add-on.
 
        Only schedule a new install task if the there's no running task.
        """
        if not self._install_task_install_task or self._install_task_install_task.done():
            self._logger_logger.info(
                "%s add-on is not installed. Installing add-on", self.addon_nameaddon_name
            )
            self._install_task_install_task = self._async_schedule_addon_operation_async_schedule_addon_operation(
                self.async_install_addonasync_install_addon, catch_error=catch_error
            )
        return self._install_task_install_task
 
    @callback
    def async_schedule_install_setup_addon(
        self,
        addon_config: dict[str, Any],
        catch_error: bool = False,
    ) -> asyncio.Task:
        """Schedule a task that installs and sets up the managed add-on.
 
        Only schedule a new install task if the there's no running task.
        """
        if not self._install_task_install_task or self._install_task_install_task.done():
            self._logger_logger.info(
                "%s add-on is not installed. Installing add-on", self.addon_nameaddon_name
            )
            self._install_task_install_task = self._async_schedule_addon_operation_async_schedule_addon_operation(
                self.async_install_addonasync_install_addon,
                partial(
                    self.async_configure_addonasync_configure_addon,
                    addon_config,
                ),
                self.async_start_addonasync_start_addon,
                catch_error=catch_error,
            )
        return self._install_task_install_task
 
    @callback
    def async_schedule_update_addon(self, catch_error: bool = False) -> asyncio.Task:
        """Schedule a task that updates and sets up the managed add-on.
 
        Only schedule a new update task if the there's no running task.
        """
        if not self._update_task_update_task or self._update_task_update_task.done():
            self._logger_logger.info("Trying to update the %s add-on", self.addon_nameaddon_name)
            self._update_task_update_task = self._async_schedule_addon_operation_async_schedule_addon_operation(
                self.async_update_addonasync_update_addon,
                catch_error=catch_error,
            )
        return self._update_task_update_task
 
    @callback
    def async_schedule_start_addon(self, catch_error: bool = False) -> asyncio.Task:
        """Schedule a task that starts the managed add-on.
 
        Only schedule a new start task if the there's no running task.
        """
        if not self._start_task_start_task or self._start_task_start_task.done():
            self._logger_logger.info(
                "%s add-on is not running. Starting add-on", self.addon_nameaddon_name
            )
            self._start_task_start_task = self._async_schedule_addon_operation_async_schedule_addon_operation(
                self.async_start_addonasync_start_addon, catch_error=catch_error
            )
        return self._start_task_start_task
 
    @callback
    def async_schedule_restart_addon(self, catch_error: bool = False) -> asyncio.Task:
        """Schedule a task that restarts the managed add-on.
 
        Only schedule a new restart task if the there's no running task.
        """
        if not self._restart_task_restart_task or self._restart_task_restart_task.done():
            self._logger_logger.info("Restarting %s add-on", self.addon_nameaddon_name)
            self._restart_task_restart_task = self._async_schedule_addon_operation_async_schedule_addon_operation(
                self.async_restart_addonasync_restart_addon, catch_error=catch_error
            )
        return self._restart_task_restart_task
 
    @callback
    def async_schedule_setup_addon(
        self,
        addon_config: dict[str, Any],
        catch_error: bool = False,
    ) -> asyncio.Task:
        """Schedule a task that configures and starts the managed add-on.
 
        Only schedule a new setup task if there's no running task.
        """
        if not self._start_task_start_task or self._start_task_start_task.done():
            self._logger_logger.info(
                "%s add-on is not running. Starting add-on", self.addon_nameaddon_name
            )
            self._start_task_start_task = self._async_schedule_addon_operation_async_schedule_addon_operation(
                partial(
                    self.async_configure_addonasync_configure_addon,
                    addon_config,
                ),
                self.async_start_addonasync_start_addon,
                catch_error=catch_error,
            )
        return self._start_task_start_task
 
    @callback
    def _async_schedule_addon_operation(
        self, *funcs: Callable, catch_error: bool = False
    ) -> asyncio.Task:
        """Schedule an add-on task."""
 
        async def addon_operation() -> None:
            """Do the add-on operation and catch AddonError."""
            for func in funcs:
                try:
                    await func()
                except AddonError as err:
                    if not catch_error:
                        raise
                    self._logger_logger.error(err)
                    break
 
        return self._hass_hass.async_create_task(addon_operation(), eager_start=False)
 
 
class AddonError(HomeAssistantError):
    """Represent an error with the managed add-on."""