Improve error handling when recorder schema migration fails (#122397)

homeassistant/components/recorder/core.py

@@ -725,6 +725,10 @@ class Recorder(threading.Thread):
             "recorder_database_migration",
         )
 
+    def _dismiss_migration_in_progress(self) -> None:
+        """Dismiss notification about migration in progress."""
+        persistent_notification.dismiss(self.hass, "recorder_database_migration")
+
     def _run(self) -> None:
         """Start processing events to save."""
         thread_id = threading.get_ident()
@@ -787,9 +791,16 @@ class Recorder(threading.Thread):
                     # was True, we need to reinitialize the listener.
                     self.hass.add_job(self.async_initialize)
             else:
+                self.migration_in_progress = False
+                self._dismiss_migration_in_progress()
                 self._notify_migration_failed()
                 return
 
+        # Schema migration and repair is now completed
+        if self.migration_in_progress:
+            self.migration_in_progress = False
+            self._dismiss_migration_in_progress()
+
         # Catch up with missed statistics
         self._schedule_compile_missing_statistics()
         _LOGGER.debug("Recorder processing the queue")
@@ -984,14 +995,10 @@ class Recorder(threading.Thread):
             "recorder_database_migration",
         )
         self.hass.add_job(self._async_migration_started)
-        try:
-            migration_result, schema_status = self._migrate_schema(schema_status, True)
-            if migration_result:
-                self._setup_run()
-            return migration_result, schema_status
-        finally:
-            self.migration_in_progress = False
-            persistent_notification.dismiss(self.hass, "recorder_database_migration")
+        migration_result, schema_status = self._migrate_schema(schema_status, True)
+        if migration_result:
+            self._setup_run()
+        return migration_result, schema_status
 
     def _migrate_schema(
         self,
@@ -1010,7 +1017,15 @@ class Recorder(threading.Thread):
             )
         except exc.DatabaseError as err:
             if self._handle_database_error(err):
-                return (True, schema_status)
+                # If _handle_database_error returns True, we have a new database
+                # which does not need migration or repair.
+                new_schema_status = migration.SchemaValidationStatus(
+                    current_version=SCHEMA_VERSION,
+                    migration_needed=False,
+                    schema_errors=set(),
+                    start_version=SCHEMA_VERSION,
+                )
+                return (True, new_schema_status)
             _LOGGER.exception("Database error during schema migration")
             return (False, schema_status)
         except Exception:

tests/components/recorder/test_migrate.py

@@ -160,7 +160,7 @@ async def test_database_migration_failed(
 
 @pytest.mark.skip_on_db_engine(["mysql", "postgresql"])
 @pytest.mark.usefixtures("skip_by_db_engine")
-async def test_database_migration_encounters_corruption(
+async def test_live_database_migration_encounters_corruption(
     hass: HomeAssistant,
     recorder_db_url: str,
     async_setup_recorder_instance: RecorderInstanceGenerator,
@@ -183,6 +183,51 @@ async def test_database_migration_encounters_corruption(
             "homeassistant.components.recorder.migration._schema_is_current",
             side_effect=[False],
         ),
+        patch(
+            "homeassistant.components.recorder.migration.migrate_schema_live",
+            side_effect=sqlite3_exception,
+        ),
+        patch(
+            "homeassistant.components.recorder.core.move_away_broken_database"
+        ) as move_away,
+    ):
+        await async_setup_recorder_instance(hass)
+        hass.states.async_set("my.entity", "on", {})
+        hass.states.async_set("my.entity", "off", {})
+        await async_wait_recording_done(hass)
+
+    assert recorder.util.async_migration_in_progress(hass) is False
+    move_away.assert_called_once()
+
+
+@pytest.mark.skip_on_db_engine(["mysql", "postgresql"])
+@pytest.mark.usefixtures("skip_by_db_engine")
+async def test_non_live_database_migration_encounters_corruption(
+    hass: HomeAssistant,
+    recorder_db_url: str,
+    async_setup_recorder_instance: RecorderInstanceGenerator,
+) -> None:
+    """Test we move away the database if its corrupt.
+
+    This test is specific for SQLite, wiping the database on error only happens
+    with SQLite.
+    """
+
+    assert recorder.util.async_migration_in_progress(hass) is False
+
+    sqlite3_exception = DatabaseError("statement", {}, [])
+    sqlite3_exception.__cause__ = sqlite3.DatabaseError(
+        "database disk image is malformed"
+    )
+
+    with (
+        patch(
+            "homeassistant.components.recorder.migration._schema_is_current",
+            side_effect=[False],
+        ),
+        patch(
+            "homeassistant.components.recorder.migration.migrate_schema_live",
+        ) as migrate_schema_live,
         patch(
             "homeassistant.components.recorder.migration.migrate_schema_non_live",
             side_effect=sqlite3_exception,
@@ -197,7 +242,8 @@ async def test_database_migration_encounters_corruption(
         await async_wait_recording_done(hass)
 
     assert recorder.util.async_migration_in_progress(hass) is False
-    assert move_away.called
+    move_away.assert_called_once()
+    migrate_schema_live.assert_not_called()
 
 
 @pytest.mark.parametrize(