Prevent calling stop or restart services during db upgrade

homeassistant/components/homeassistant/__init__.py

@@ -6,6 +6,7 @@
 import voluptuous as vol
 
 from homeassistant.auth.permissions.const import CAT_ENTITIES, POLICY_CONTROL
+from homeassistant.components import recorder
 import homeassistant.config as conf_util
 from homeassistant.const import (
     ATTR_ENTITY_ID,
@@ -45,6 +46,7 @@
     ),
     cv.has_at_least_one_key(ATTR_ENTRY_ID, *cv.ENTITY_SERVICE_FIELDS),
 )
+SHUTDOWN_SERVICES = (SERVICE_HOMEASSISTANT_STOP, SERVICE_HOMEASSISTANT_RESTART)
 
 
 async def async_setup(hass: ha.HomeAssistant, config: dict) -> bool:
@@ -125,6 +127,14 @@ async def async_handle_turn_service(service):
 
     async def async_handle_core_service(call):
         """Service handler for handling core services."""
+        if (
+            call.service in SHUTDOWN_SERVICES
+            and await recorder.async_migration_in_progress(hass)
+        ):
+            raise HomeAssistantError(
+                f"The system cannot {call.service} while a database upgrade in progress."
+            )
+
         if call.service == SERVICE_HOMEASSISTANT_STOP:
             hass.async_create_task(hass.async_stop())
             return

homeassistant/components/recorder/__init__.py

@@ -36,6 +36,7 @@
 )
 from homeassistant.helpers.event import async_track_time_interval, track_time_change
 from homeassistant.helpers.typing import ConfigType
+from homeassistant.loader import bind_hass
 import homeassistant.util.dt as dt_util
 
 from . import migration, purge
@@ -132,6 +133,18 @@
 )
 
 
+@bind_hass
+async def async_migration_in_progress(hass):
+    """Determine is a migration is in progress.
+
+    This is a thin wrapper that allows us to change
+    out the implementation later.
+    """
+    if DATA_INSTANCE not in hass.data:
+        return False
+    return hass.data[DATA_INSTANCE].migration_in_progress
+
+
 def run_information(hass, point_in_time: datetime | None = None):
     """Return information about current run.
 
@@ -291,7 +304,8 @@ def __init__(
         self.get_session = None
         self._completed_database_setup = None
         self._event_listener = None
-
+        self.async_migration_event = asyncio.Event()
+        self.migration_in_progress = False
         self._queue_watcher = None
 
         self.enabled = True
@@ -510,6 +524,11 @@ def _setup_recorder(self) -> None | int:
 
         return None
 
+    @callback
+    def _async_migration_started(self):
+        """Set the migration started event."""
+        self.async_migration_event.set()
+
     def _migrate_schema_and_setup_run(self, current_version) -> bool:
         """Migrate schema to the latest version."""
         persistent_notification.create(
@@ -518,6 +537,8 @@ def _migrate_schema_and_setup_run(self, current_version) -> bool:
             "Database upgrade in progress",
             "recorder_database_migration",
         )
+        self.migration_in_progress = True
+        self.hass.add_job(self._async_migration_started)
 
         try:
             migration.migrate_schema(self, current_version)
@@ -533,6 +554,7 @@ def _migrate_schema_and_setup_run(self, current_version) -> bool:
             self._setup_run()
             return True
         finally:
+            self.migration_in_progress = False
             persistent_notification.dismiss(self.hass, "recorder_database_migration")
 
     def _run_purge(self, keep_days, repack, apply_filter):

tests/components/homeassistant/test_init.py

@@ -447,3 +447,41 @@ async def test_reload_config_entry_by_entry_id(hass):
 
     assert len(mock_reload.mock_calls) == 1
     assert mock_reload.mock_calls[0][1][0] == "8955375327824e14ba89e4b29cc3ec9a"
+
+
+@pytest.mark.parametrize(
+    "service", [SERVICE_HOMEASSISTANT_RESTART, SERVICE_HOMEASSISTANT_STOP]
+)
+async def test_raises_when_db_upgrade_in_progress(hass, service):
+    """Test services raise an exception when a db upgrade is in progress."""
+    await async_setup_component(hass, "persistent_notification", {})
+    await async_setup_component(hass, "homeassistant", {})
+
+    with pytest.raises(HomeAssistantError), patch(
+        "homeassistant.components.recorder.async_migration_in_progress",
+        return_value=True,
+    ) as mock_async_migration_in_progress:
+        await hass.services.async_call(
+            "homeassistant",
+            service,
+            blocking=True,
+        )
+
+    assert mock_async_migration_in_progress.called
+
+    with patch(
+        "homeassistant.components.recorder.async_migration_in_progress",
+        return_value=False,
+    ) as mock_async_migration_in_progress, patch(
+        "homeassistant.core.HomeAssistant.async_stop"
+    ) as stop_mock, patch(
+        "homeassistant.config.async_check_ha_config_file", return_value=False
+    ):
+        await hass.services.async_call(
+            "homeassistant",
+            service,
+            blocking=True,
+        )
+
+    assert stop_mock.called
+    assert mock_async_migration_in_progress.called

tests/components/recorder/test_migrate.py

@@ -48,6 +48,7 @@ def create_engine_test(*args, **kwargs):
 
 async def test_schema_update_calls(hass):
     """Test that schema migrations occur in correct order."""
+    assert await recorder.async_migration_in_progress(hass) is False
     await async_setup_component(hass, "persistent_notification", {})
     with patch(
         "homeassistant.components.recorder.create_engine", new=create_engine_test
@@ -60,6 +61,7 @@ async def test_schema_update_calls(hass):
         )
         await async_wait_recording_done_without_instance(hass)
 
+    assert await recorder.async_migration_in_progress(hass) is False
     update.assert_has_calls(
         [
             call(hass.data[DATA_INSTANCE].engine, version + 1, 0)
@@ -68,11 +70,30 @@ async def test_schema_update_calls(hass):
     )
 
 
+async def test_migration_in_progress(hass):
+    """Test that we can check for migration in progress."""
+    assert await recorder.async_migration_in_progress(hass) is False
+    await async_setup_component(hass, "persistent_notification", {})
+
+    with patch(
+        "homeassistant.components.recorder.create_engine", new=create_engine_test
+    ):
+        await async_setup_component(
+            hass, "recorder", {"recorder": {"db_url": "sqlite://"}}
+        )
+        await hass.data[DATA_INSTANCE].async_migration_event.wait()
+        assert await recorder.async_migration_in_progress(hass) is True
+        await async_wait_recording_done_without_instance(hass)
+
+    assert await recorder.async_migration_in_progress(hass) is False
+
+
 async def test_database_migration_failed(hass):
     """Test we notify if the migration fails."""
     await async_setup_component(hass, "persistent_notification", {})
     create_calls = async_mock_service(hass, "persistent_notification", "create")
     dismiss_calls = async_mock_service(hass, "persistent_notification", "dismiss")
+    assert await recorder.async_migration_in_progress(hass) is False
 
     with patch(
         "homeassistant.components.recorder.create_engine", new=create_engine_test
@@ -89,13 +110,15 @@ async def test_database_migration_failed(hass):
         await hass.async_add_executor_job(hass.data[DATA_INSTANCE].join)
         await hass.async_block_till_done()
 
+    assert await recorder.async_migration_in_progress(hass) is False
     assert len(create_calls) == 2
     assert len(dismiss_calls) == 1
 
 
 async def test_database_migration_encounters_corruption(hass):
     """Test we move away the database if its corrupt."""
     await async_setup_component(hass, "persistent_notification", {})
+    assert await recorder.async_migration_in_progress(hass) is False
 
     sqlite3_exception = DatabaseError("statement", {}, [])
     sqlite3_exception.__cause__ = sqlite3.DatabaseError()
@@ -116,6 +139,7 @@ async def test_database_migration_encounters_corruption(hass):
         hass.states.async_set("my.entity", "off", {})
         await async_wait_recording_done_without_instance(hass)
 
+    assert await recorder.async_migration_in_progress(hass) is False
     assert move_away.called
 
 
@@ -124,6 +148,7 @@ async def test_database_migration_encounters_corruption_not_sqlite(hass):
     await async_setup_component(hass, "persistent_notification", {})
     create_calls = async_mock_service(hass, "persistent_notification", "create")
     dismiss_calls = async_mock_service(hass, "persistent_notification", "dismiss")
+    assert await recorder.async_migration_in_progress(hass) is False
 
     with patch(
         "homeassistant.components.recorder.migration.schema_is_current",
@@ -143,6 +168,7 @@ async def test_database_migration_encounters_corruption_not_sqlite(hass):
         await hass.async_add_executor_job(hass.data[DATA_INSTANCE].join)
         await hass.async_block_till_done()
 
+    assert await recorder.async_migration_in_progress(hass) is False
     assert not move_away.called
     assert len(create_calls) == 2
     assert len(dismiss_calls) == 1
@@ -151,6 +177,7 @@ async def test_database_migration_encounters_corruption_not_sqlite(hass):
 async def test_events_during_migration_are_queued(hass):
     """Test that events during migration are queued."""
 
+    assert await recorder.async_migration_in_progress(hass) is False
     await async_setup_component(hass, "persistent_notification", {})
     with patch(
         "homeassistant.components.recorder.create_engine", new=create_engine_test
@@ -167,13 +194,15 @@ async def test_events_during_migration_are_queued(hass):
         await hass.data[DATA_INSTANCE].async_recorder_ready.wait()
         await async_wait_recording_done_without_instance(hass)
 
+    assert await recorder.async_migration_in_progress(hass) is False
     db_states = await hass.async_add_executor_job(_get_native_states, hass, "my.entity")
     assert len(db_states) == 2
 
 
 async def test_events_during_migration_queue_exhausted(hass):
     """Test that events during migration takes so long the queue is exhausted."""
     await async_setup_component(hass, "persistent_notification", {})
+    assert await recorder.async_migration_in_progress(hass) is False
 
     with patch(
         "homeassistant.components.recorder.create_engine", new=create_engine_test
@@ -191,6 +220,7 @@ async def test_events_during_migration_queue_exhausted(hass):
         await hass.data[DATA_INSTANCE].async_recorder_ready.wait()
         await async_wait_recording_done_without_instance(hass)
 
+    assert await recorder.async_migration_in_progress(hass) is False
     db_states = await hass.async_add_executor_job(_get_native_states, hass, "my.entity")
     assert len(db_states) == 1
     hass.states.async_set("my.entity", "on", {})