core/homeassistant/components/portainer/button.py

"""Support for Portainer buttons."""

from abc import abstractmethod
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, override

from pyportainer import DockerContainerState, Portainer
from pyportainer.exceptions import (
    PortainerAuthenticationError,
    PortainerConnectionError,
    PortainerTimeoutError,
)
from pyportainer.models.docker import DockerContainer

from homeassistant.components.button import (
    ButtonDeviceClass,
    ButtonEntity,
    ButtonEntityDescription,
)
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback

from . import PortainerConfigEntry
from .const import DOMAIN
from .coordinator import (
    PortainerContainerData,
    PortainerCoordinator,
    PortainerCoordinatorData,
)
from .entity import PortainerContainerEntity, PortainerEndpointEntity

PARALLEL_UPDATES = 1


@dataclass(frozen=True, kw_only=True)
class PortainerEndpointButtonDescription(ButtonEntityDescription):
    """Class to describe a Portainer endpoint button entity."""

    press_action: Callable[
        [Portainer, int],
        Coroutine[Any, Any, DockerContainer | None],
    ]


@dataclass(frozen=True, kw_only=True)
class PortainerContainerButtonDescription(ButtonEntityDescription):
    """Class to describe a Portainer container button entity."""

    press_action: Callable[
        [Portainer, int, str],
        Coroutine[Any, Any, DockerContainer | None],
    ]
    available_fn: Callable[[PortainerContainerData], bool]


ENDPOINT_BUTTONS: tuple[PortainerEndpointButtonDescription, ...] = (
    PortainerEndpointButtonDescription(
        key="images_prune",
        translation_key="images_prune",
        device_class=ButtonDeviceClass.RESTART,
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id: portainer.images_prune(
                endpoint_id=endpoint_id, dangling=False, until=timedelta(days=0)
            )
        ),
    ),
    PortainerEndpointButtonDescription(
        key="volumes_prune",
        translation_key="volumes_prune",
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id: portainer.prune_volumes(endpoint_id)
        ),
    ),
)

CONTAINER_BUTTONS: tuple[PortainerContainerButtonDescription, ...] = (
    PortainerContainerButtonDescription(
        key="restart",
        translation_key="restart_container",
        device_class=ButtonDeviceClass.RESTART,
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id, container_id: portainer.restart_container(
                endpoint_id, container_id
            )
        ),
        available_fn=lambda container: (
            container.container.state != DockerContainerState.PAUSED
        ),
    ),
    PortainerContainerButtonDescription(
        key="pause",
        translation_key="pause_container",
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id, container_id: portainer.pause_container(
                endpoint_id, container_id
            )
        ),
        available_fn=lambda container: (
            container.container.state == DockerContainerState.RUNNING
        ),
    ),
    PortainerContainerButtonDescription(
        key="resume",
        translation_key="resume_container",
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id, container_id: portainer.unpause_container(
                endpoint_id, container_id
            )
        ),
        available_fn=lambda container: (
            container.container.state == DockerContainerState.PAUSED
        ),
    ),
    PortainerContainerButtonDescription(
        key="recreate",
        translation_key="recreate_container",
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id, container_id: portainer.container_recreate(
                endpoint_id=endpoint_id,
                container_id=container_id,
                timeout=timedelta(minutes=10),
                pull_image=True,
            )
        ),
        available_fn=lambda container: (
            container.container.state
            not in (DockerContainerState.REMOVING, DockerContainerState.DEAD)
        ),
    ),
    PortainerContainerButtonDescription(
        key="kill",
        translation_key="kill_container",
        entity_category=EntityCategory.CONFIG,
        press_action=(
            lambda portainer, endpoint_id, container_id: portainer.kill_container(
                endpoint_id, container_id
            )
        ),
        available_fn=lambda container: (
            container.container.state
            in (DockerContainerState.RUNNING, DockerContainerState.PAUSED)
        ),
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: PortainerConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up Portainer buttons."""
    coordinator = entry.runtime_data

    def _async_add_new_endpoints(endpoints: list[PortainerCoordinatorData]) -> None:
        """Add new endpoint binary sensors."""
        async_add_entities(
            PortainerEndpointButton(
                coordinator,
                entity_description,
                endpoint,
            )
            for entity_description in ENDPOINT_BUTTONS
            for endpoint in endpoints
        )

    def _async_add_new_containers(
        containers: list[tuple[PortainerCoordinatorData, PortainerContainerData]],
    ) -> None:
        """Add new container button sensors."""
        async_add_entities(
            PortainerContainerButton(
                coordinator,
                entity_description,
                container,
                endpoint,
            )
            for (endpoint, container) in containers
            for entity_description in CONTAINER_BUTTONS
        )

    coordinator.new_endpoints_callbacks.append(_async_add_new_endpoints)
    coordinator.new_containers_callbacks.append(_async_add_new_containers)

    _async_add_new_endpoints(
        [
            endpoint
            for endpoint in coordinator.data.values()
            if endpoint.id in coordinator.known_endpoints
        ]
    )
    _async_add_new_containers(
        [
            (endpoint, container)
            for endpoint in coordinator.data.values()
            for container in endpoint.containers.values()
        ]
    )


class PortainerBaseButton(ButtonEntity):
    """Common base for Portainer buttons.

    Ensures the async_press logic isn't duplicated.
    """

    coordinator: PortainerCoordinator

    @abstractmethod
    async def _async_press_call(self) -> None:
        """Abstract method used per Portainer button class."""

    @override
    async def async_press(self) -> None:
        """Trigger the Portainer button press service."""
        try:
            await self._async_press_call()
        except PortainerConnectionError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="cannot_connect",
            ) from err
        except PortainerAuthenticationError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="invalid_auth",
            ) from err
        except PortainerTimeoutError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="timeout_connect",
            ) from err

        await self.coordinator.async_request_refresh()


class PortainerEndpointButton(PortainerEndpointEntity, PortainerBaseButton):
    """Defines a Portainer endpoint button."""

    entity_description: PortainerEndpointButtonDescription

    @override
    async def _async_press_call(self) -> None:
        """Call the endpoint button press action."""
        await self.entity_description.press_action(
            self.coordinator.portainer, self.device_id
        )


class PortainerContainerButton(PortainerContainerEntity, PortainerBaseButton):
    """Defines a Portainer button."""

    entity_description: PortainerContainerButtonDescription

    @property
    @override
    def available(self) -> bool:
        """Return if the button is available."""
        return super().available and self.entity_description.available_fn(
            self.container_data
        )

    @override
    async def _async_press_call(self) -> None:
        """Call the container button press action."""
        await self.entity_description.press_action(
            self.coordinator.portainer,
            self.endpoint_id,
            self.container_data.container.id,
        )