Add support for offline DB migration (#5202)

* Add support for offline DB migration

* Format code

supervisor/homeassistant/api.py

@@ -1,6 +1,7 @@
 """Home Assistant control object."""
 import asyncio
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress
+from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
 import logging
 from typing import Any
@@ -21,6 +22,14 @@ _LOGGER: logging.Logger = logging.getLogger(__name__)
 GET_CORE_STATE_MIN_VERSION: AwesomeVersion = AwesomeVersion("2023.8.0.dev20230720")
 
 
+@dataclass(frozen=True)
+class APIState:
+    """Container for API state response."""
+
+    core_state: str
+    offline_db_migration: bool
+
+
 class HomeAssistantAPI(CoreSysAttributes):
     """Home Assistant core object for handle it."""
 
@@ -132,7 +141,7 @@ class HomeAssistantAPI(CoreSysAttributes):
         """Return Home Assistant core state."""
         return await self._get_json("api/core/state")
 
-    async def get_api_state(self) -> str | None:
+    async def get_api_state(self) -> APIState | None:
         """Return state of Home Assistant Core or None."""
         # Skip check on landingpage
         if (
@@ -161,12 +170,17 @@ class HomeAssistantAPI(CoreSysAttributes):
                 data = await self.get_config()
             # Older versions of home assistant does not expose the state
             if data:
-                return data.get("state", "RUNNING")
+                state = data.get("state", "RUNNING")
+                # Recorder state was added in HA Core 2024.8
+                recorder_state = data.get("recorder_state", {})
+                migrating = recorder_state.get("migration_in_progress", False)
+                live_migration = recorder_state.get("migration_is_live", False)
+                return APIState(state, migrating and not live_migration)
 
         return None
 
     async def check_api_state(self) -> bool:
         """Return Home Assistant Core state if up."""
         if state := await self.get_api_state():
-            return state == "RUNNING"
+            return state.core_state == "RUNNING" or state.offline_db_migration
         return False

supervisor/homeassistant/core.py

@@ -49,6 +49,10 @@ SECONDS_BETWEEN_API_CHECKS: Final[int] = 5
 STARTUP_API_RESPONSE_TIMEOUT: Final[timedelta] = timedelta(minutes=3)
 # All stages plus event start timeout and some wiggle rooom
 STARTUP_API_CHECK_RUNNING_TIMEOUT: Final[timedelta] = timedelta(minutes=15)
+# While database migration is running, the timeout will be extended
+DATABASE_MIGRATION_TIMEOUT: Final[timedelta] = timedelta(
+    seconds=SECONDS_BETWEEN_API_CHECKS * 10
+)
 RE_YAML_ERROR = re.compile(r"homeassistant\.util\.yaml")
 
 
@@ -490,11 +494,15 @@ class HomeAssistantCore(JobGroup):
                     _LOGGER.info("Home Assistant Core state changed to %s", state)
                     last_state = state
 
-                if state == "RUNNING":
+                if state.core_state == "RUNNING":
                     _LOGGER.info("Detect a running Home Assistant instance")
                     self._error_state = False
                     return
 
+                if state.offline_db_migration:
+                    # Keep extended the deadline while database migration is active
+                    deadline = datetime.now() + DATABASE_MIGRATION_TIMEOUT
+
         self._error_state = True
         if timeout:
             raise HomeAssistantStartupTimeout(

tests/conftest.py

@@ -42,6 +42,7 @@ from supervisor.coresys import CoreSys
 from supervisor.dbus.network import NetworkManager
 from supervisor.docker.manager import DockerAPI
 from supervisor.docker.monitor import DockerMonitor
+from supervisor.homeassistant.api import APIState
 from supervisor.host.logs import LogsControl
 from supervisor.os.manager import OSManager
 from supervisor.store.addon import AddonStore
@@ -360,7 +361,9 @@ async def coresys(
     )
 
     # WebSocket
-    coresys_obj.homeassistant.api.get_api_state = AsyncMock(return_value="RUNNING")
+    coresys_obj.homeassistant.api.get_api_state = AsyncMock(
+        return_value=APIState("RUNNING", False)
+    )
     coresys_obj.homeassistant._websocket._client = AsyncMock(
         ha_version=AwesomeVersion("2021.2.4")
     )

tests/homeassistant/test_core.py

@@ -21,6 +21,7 @@ from supervisor.exceptions import (
     HomeAssistantError,
     HomeAssistantJobError,
 )
+from supervisor.homeassistant.api import APIState
 from supervisor.homeassistant.core import HomeAssistantCore
 from supervisor.homeassistant.module import HomeAssistant
 from supervisor.updater import Updater
@@ -316,6 +317,42 @@ async def test_api_check_success(
     assert "Detect a running Home Assistant instance" in caplog.text
 
 
+async def test_api_check_database_migration(
+    coresys: CoreSys, container: MagicMock, caplog: pytest.LogCaptureFixture
+):
+    """Test attempts to contact the API timeout."""
+    calls = []
+
+    def mock_api_state(*args):
+        calls.append(None)
+        if len(calls) > 50:
+            return APIState("RUNNING", False)
+        else:
+            return APIState("NOT_RUNNING", True)
+
+    container.status = "stopped"
+    coresys.homeassistant.version = AwesomeVersion("2023.9.0")
+    coresys.homeassistant.api.get_api_state.side_effect = mock_api_state
+
+    async def mock_instance_start(*_):
+        container.status = "running"
+
+    with (
+        patch.object(DockerHomeAssistant, "start", new=mock_instance_start),
+        patch.object(DockerAPI, "container_is_initialized", return_value=True),
+        travel(datetime(2023, 10, 2, 0, 0, 0), tick=False) as traveller,
+    ):
+
+        async def mock_sleep(*args):
+            traveller.shift(timedelta(minutes=1))
+
+        with patch("supervisor.homeassistant.core.asyncio.sleep", new=mock_sleep):
+            await coresys.homeassistant.core.start()
+
+    assert coresys.homeassistant.api.get_api_state.call_count == 51
+    assert "Detect a running Home Assistant instance" in caplog.text
+
+
 async def test_core_loads_wrong_image_for_machine(
     coresys: CoreSys, container: MagicMock
 ):