Add env on core restart due to restore (#5548)

* Add env on core restart due to restore

* Move is_restore to backup manager
This commit is contained in:
Mike Degatano 2025-01-16 12:15:06 -05:00 committed by GitHub
parent 89a215cc1f
commit 0073227785
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 102 additions and 35 deletions

View File

@ -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(

View File

@ -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"

View File

@ -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,

View File

@ -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)

View File

@ -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,
)

View File

@ -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

View File

@ -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}")

View File

@ -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
)

View File

@ -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,
),
]