From 0073227785dd456089bc0bdc6a345f5b102b07d2 Mon Sep 17 00:00:00 2001 From: Mike Degatano Date: Thu, 16 Jan 2025 12:15:06 -0500 Subject: [PATCH] Add env on core restart due to restore (#5548) * Add env on core restart due to restore * Move is_restore to backup manager --- supervisor/backups/backup.py | 2 +- supervisor/backups/const.py | 1 - supervisor/backups/manager.py | 24 +++++++++---- supervisor/core.py | 6 ++-- supervisor/docker/homeassistant.py | 20 ++++++----- supervisor/homeassistant/core.py | 6 ++-- supervisor/jobs/__init__.py | 12 +++---- tests/api/test_backups.py | 57 ++++++++++++++++++++++++++++++ tests/backups/test_manager.py | 9 ++--- 9 files changed, 102 insertions(+), 35 deletions(-) diff --git a/supervisor/backups/backup.py b/supervisor/backups/backup.py index 70736c0c9..ef0b038ea 100644 --- a/supervisor/backups/backup.py +++ b/supervisor/backups/backup.py @@ -772,7 +772,7 @@ class Backup(JobGroup): @Job(name="backup_restore_homeassistant", cleanup=False) async def restore_homeassistant(self) -> Awaitable[None]: """Restore Home Assistant Core configuration folder.""" - await self.sys_homeassistant.core.stop() + await self.sys_homeassistant.core.stop(remove_container=True) # Restore Home Assistant Core config directory tar_name = Path( diff --git a/supervisor/backups/const.py b/supervisor/backups/const.py index fe37c1752..ea40c1f03 100644 --- a/supervisor/backups/const.py +++ b/supervisor/backups/const.py @@ -39,7 +39,6 @@ class RestoreJobStage(StrEnum): ADDONS = "addons" AWAIT_ADDON_RESTARTS = "await_addon_restarts" AWAIT_HOME_ASSISTANT_RESTART = "await_home_assistant_restart" - CHECK_HOME_ASSISTANT = "check_home_assistant" DOCKER_CONFIG = "docker_config" FOLDERS = "folders" HOME_ASSISTANT = "home_assistant" diff --git a/supervisor/backups/manager.py b/supervisor/backups/manager.py index 7a655af53..1ebd47700 100644 --- a/supervisor/backups/manager.py +++ b/supervisor/backups/manager.py @@ -47,6 +47,9 @@ from .validate import ALL_FOLDERS, SCHEMA_BACKUPS_CONFIG _LOGGER: logging.Logger = logging.getLogger(__name__) +JOB_FULL_RESTORE = "backup_manager_full_restore" +JOB_PARTIAL_RESTORE = "backup_manager_partial_restore" + class BackupManager(FileConfiguration, JobGroup): """Manage backups.""" @@ -86,6 +89,16 @@ class BackupManager(FileConfiguration, JobGroup): if mount.state == UnitActiveState.ACTIVE } + @property + def current_restore(self) -> str | None: + """Return id of current restore job if a restore job is in progress.""" + job = self.sys_jobs.current + while job.parent_id: + job = self.sys_jobs.get_job(job.parent_id) + if job.name in {JOB_FULL_RESTORE, JOB_PARTIAL_RESTORE}: + return job.uuid + return None + def get(self, slug: str) -> Backup: """Return backup object.""" return self._backups.get(slug) @@ -619,9 +632,6 @@ class BackupManager(FileConfiguration, JobGroup): # Wait for Home Assistant Core update/downgrade if task_hass: - self._change_stage( - RestoreJobStage.AWAIT_HOME_ASSISTANT_RESTART, backup - ) await task_hass except BackupError: raise @@ -644,7 +654,7 @@ class BackupManager(FileConfiguration, JobGroup): finally: # Leave Home Assistant alone if it wasn't part of the restore if homeassistant: - self._change_stage(RestoreJobStage.CHECK_HOME_ASSISTANT, backup) + self._change_stage(RestoreJobStage.AWAIT_HOME_ASSISTANT_RESTART, backup) # Do we need start Home Assistant Core? if not await self.sys_homeassistant.core.is_running(): @@ -660,7 +670,7 @@ class BackupManager(FileConfiguration, JobGroup): ) @Job( - name="backup_manager_full_restore", + name=JOB_FULL_RESTORE, conditions=[ JobCondition.FREE_SPACE, JobCondition.HEALTHY, @@ -706,7 +716,7 @@ class BackupManager(FileConfiguration, JobGroup): try: # Stop Home-Assistant / Add-ons - await self.sys_core.shutdown() + await self.sys_core.shutdown(remove_homeassistant_container=True) success = await self._do_restore( backup, @@ -724,7 +734,7 @@ class BackupManager(FileConfiguration, JobGroup): return success @Job( - name="backup_manager_partial_restore", + name=JOB_PARTIAL_RESTORE, conditions=[ JobCondition.FREE_SPACE, JobCondition.HEALTHY, diff --git a/supervisor/core.py b/supervisor/core.py index 203cf71df..8a07ffc5e 100644 --- a/supervisor/core.py +++ b/supervisor/core.py @@ -333,7 +333,7 @@ class Core(CoreSysAttributes): _LOGGER.info("Supervisor is down - %d", self.exit_code) self.sys_loop.stop() - async def shutdown(self): + async def shutdown(self, *, remove_homeassistant_container: bool = False): """Shutdown all running containers in correct order.""" # don't process scheduler anymore if self.state == CoreState.RUNNING: @@ -344,7 +344,9 @@ class Core(CoreSysAttributes): # Close Home Assistant with suppress(HassioError): - await self.sys_homeassistant.core.stop() + await self.sys_homeassistant.core.stop( + remove_container=remove_homeassistant_container + ) # Shutdown System Add-ons await self.sys_addons.shutdown(AddonStartup.SERVICES) diff --git a/supervisor/docker/homeassistant.py b/supervisor/docker/homeassistant.py index 2d6199dd9..a33579268 100644 --- a/supervisor/docker/homeassistant.py +++ b/supervisor/docker/homeassistant.py @@ -35,6 +35,7 @@ _LOGGER: logging.Logger = logging.getLogger(__name__) _VERIFY_TRUST: AwesomeVersion = AwesomeVersion("2021.5.0") _HASS_DOCKER_NAME: str = "homeassistant" ENV_S6_GRACETIME = re.compile(r"^S6_SERVICES_GRACETIME=([0-9]+)$") +ENV_RESTORE_JOB_ID = "SUPERVISOR_RESTORE_JOB_ID" class DockerHomeAssistant(DockerInterface): @@ -163,8 +164,17 @@ class DockerHomeAssistant(DockerInterface): limit=JobExecutionLimit.GROUP_ONCE, on_condition=DockerJobError, ) - async def run(self) -> None: + async def run(self, *, restore_job_id: str | None = None) -> None: """Run Docker image.""" + environment = { + "SUPERVISOR": self.sys_docker.network.supervisor, + "HASSIO": self.sys_docker.network.supervisor, + ENV_TIME: self.sys_timezone, + ENV_TOKEN: self.sys_homeassistant.supervisor_token, + ENV_TOKEN_OLD: self.sys_homeassistant.supervisor_token, + } + if restore_job_id: + environment[ENV_RESTORE_JOB_ID] = restore_job_id await self._run( tag=(self.sys_homeassistant.version), name=self.name, @@ -180,13 +190,7 @@ class DockerHomeAssistant(DockerInterface): "supervisor": self.sys_docker.network.supervisor, "observer": self.sys_docker.network.observer, }, - environment={ - "SUPERVISOR": self.sys_docker.network.supervisor, - "HASSIO": self.sys_docker.network.supervisor, - ENV_TIME: self.sys_timezone, - ENV_TOKEN: self.sys_homeassistant.supervisor_token, - ENV_TOKEN_OLD: self.sys_homeassistant.supervisor_token, - }, + environment=environment, tmpfs={"/tmp": ""}, # noqa: S108 oom_score_adj=-300, ) diff --git a/supervisor/homeassistant/core.py b/supervisor/homeassistant/core.py index 911a1f844..29e829085 100644 --- a/supervisor/homeassistant/core.py +++ b/supervisor/homeassistant/core.py @@ -345,7 +345,7 @@ class HomeAssistantCore(JobGroup): self.sys_homeassistant.write_pulse() try: - await self.instance.run() + await self.instance.run(restore_job_id=self.sys_backups.current_restore) except DockerError as err: raise HomeAssistantError() from err @@ -356,10 +356,10 @@ class HomeAssistantCore(JobGroup): limit=JobExecutionLimit.GROUP_ONCE, on_condition=HomeAssistantJobError, ) - async def stop(self) -> None: + async def stop(self, *, remove_container: bool = False) -> None: """Stop Home Assistant Docker.""" try: - return await self.instance.stop(remove_container=False) + return await self.instance.stop(remove_container=remove_container) except DockerError as err: raise HomeAssistantError() from err diff --git a/supervisor/jobs/__init__.py b/supervisor/jobs/__init__.py index 01bad3615..f513b3f30 100644 --- a/supervisor/jobs/__init__.py +++ b/supervisor/jobs/__init__.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from datetime import datetime import logging from typing import Any -from uuid import UUID, uuid4 +from uuid import uuid4 from attrs import Attribute, define, field from attrs.setters import convert as attr_convert, frozen, validate as attr_validate @@ -28,7 +28,7 @@ from .validate import SCHEMA_JOBS_CONFIG # When a new asyncio task is started the current context is copied over. # Modifications to it in one task are not visible to others though. # This allows us to track what job is currently in progress in each task. -_CURRENT_JOB: ContextVar[UUID] = ContextVar("current_job") +_CURRENT_JOB: ContextVar[str] = ContextVar("current_job") _LOGGER: logging.Logger = logging.getLogger(__name__) @@ -85,7 +85,7 @@ class SupervisorJob: """Representation of a job running in supervisor.""" created: datetime = field(init=False, factory=utcnow, on_setattr=frozen) - uuid: UUID = field(init=False, factory=lambda: uuid4().hex, on_setattr=frozen) + uuid: str = field(init=False, factory=lambda: uuid4().hex, on_setattr=frozen) name: str | None = field(default=None, validator=[_invalid_if_started]) reference: str | None = field(default=None, on_setattr=_on_change) progress: float = field( @@ -97,7 +97,7 @@ class SupervisorJob: stage: str | None = field( default=None, validator=[_invalid_if_done], on_setattr=_on_change ) - parent_id: UUID | None = field( + parent_id: str | None = field( factory=lambda: _CURRENT_JOB.get(None), on_setattr=frozen ) done: bool | None = field(init=False, default=None, on_setattr=_on_change) @@ -146,7 +146,7 @@ class SupervisorJob: raise JobStartException("Job has a different parent from current job") self.done = False - token: Token[UUID] | None = None + token: Token[str] | None = None try: token = _CURRENT_JOB.set(self.uuid) yield self @@ -237,7 +237,7 @@ class JobManager(FileConfiguration, CoreSysAttributes): self._jobs[job.uuid] = job return job - def get_job(self, uuid: UUID) -> SupervisorJob: + def get_job(self, uuid: str) -> SupervisorJob: """Return a job by uuid. Raises if it does not exist.""" if uuid not in self._jobs: raise JobNotFound(f"No job found with id {uuid}") diff --git a/tests/api/test_backups.py b/tests/api/test_backups.py index 7b8710a90..6a4ba9bb6 100644 --- a/tests/api/test_backups.py +++ b/tests/api/test_backups.py @@ -15,9 +15,11 @@ from supervisor.addons.addon import Addon from supervisor.backups.backup import Backup from supervisor.const import CoreState from supervisor.coresys import CoreSys +from supervisor.docker.manager import DockerAPI from supervisor.exceptions import AddonsError, HomeAssistantBackupError from supervisor.homeassistant.core import HomeAssistantCore from supervisor.homeassistant.module import HomeAssistant +from supervisor.homeassistant.websocket import HomeAssistantWebSocket from supervisor.mounts.mount import Mount from supervisor.supervisor import Supervisor @@ -857,3 +859,58 @@ async def test_restore_backup_from_location( ) assert resp.status == 200 assert test_file.is_file() + + +@pytest.mark.parametrize( + ("backup_type", "postbody"), [("partial", {"homeassistant": True}), ("full", {})] +) +@pytest.mark.usefixtures("tmp_supervisor_data", "path_extern") +async def test_restore_homeassistant_adds_env( + api_client: TestClient, + coresys: CoreSys, + docker: DockerAPI, + backup_type: str, + postbody: dict[str, Any], +): + """Test restoring home assistant from backup adds env to container.""" + event = asyncio.Event() + coresys.core.state = CoreState.RUNNING + coresys.hardware.disk.get_disk_free_space = lambda x: 5000 + coresys.homeassistant.version = AwesomeVersion("2025.1.0") + backup = await coresys.backups.do_backup_full() + + async def mock_async_send_message(_, message: dict[str, Any]): + """Mock of async send message in ws client.""" + if ( + message["data"]["event"] == "job" + and message["data"]["data"]["name"] + == f"backup_manager_{backup_type}_restore" + and message["data"]["data"]["reference"] == backup.slug + and message["data"]["data"]["done"] + ): + event.set() + + with ( + patch.object(HomeAssistantCore, "_block_till_run"), + patch.object( + HomeAssistantWebSocket, "async_send_message", new=mock_async_send_message + ), + ): + resp = await api_client.post( + f"/backups/{backup.slug}/restore/{backup_type}", + json={"background": True} | postbody, + ) + assert resp.status == 200 + body = await resp.json() + job = coresys.jobs.get_job(body["data"]["job_id"]) + + if not job.done: + await asyncio.wait_for(event.wait(), 5) + + assert docker.containers.create.call_args.kwargs["name"] == "homeassistant" + assert ( + docker.containers.create.call_args.kwargs["environment"][ + "SUPERVISOR_RESTORE_JOB_ID" + ] + == job.uuid + ) diff --git a/tests/backups/test_manager.py b/tests/backups/test_manager.py index dacdb96db..90e3ebb73 100644 --- a/tests/backups/test_manager.py +++ b/tests/backups/test_manager.py @@ -1249,11 +1249,6 @@ async def test_restore_progress( _make_backup_message_for_assert( action="full_restore", reference=full_backup.slug, stage="addons" ), - _make_backup_message_for_assert( - action="full_restore", - reference=full_backup.slug, - stage="await_home_assistant_restart", - ), _make_backup_message_for_assert( action="full_restore", reference=full_backup.slug, @@ -1262,12 +1257,12 @@ async def test_restore_progress( _make_backup_message_for_assert( action="full_restore", reference=full_backup.slug, - stage="check_home_assistant", + stage="await_home_assistant_restart", ), _make_backup_message_for_assert( action="full_restore", reference=full_backup.slug, - stage="check_home_assistant", + stage="await_home_assistant_restart", done=True, ), ]