From 7fc0e36b2f5ceedf5caeb4bd84cb806281b90dbb Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 4 Sep 2024 08:38:46 +0200 Subject: [PATCH] Move recorder EntityIDPostMigrationTask to migration (#125136) * Move recorder EntityIDPostMigrationTask to migration * Update test --- homeassistant/components/recorder/core.py | 4 ---- homeassistant/components/recorder/migration.py | 13 ++++++++++++- homeassistant/components/recorder/tasks.py | 13 ------------- .../recorder/test_migration_from_schema_32.py | 3 +-- tests/components/recorder/test_v32_migration.py | 10 +++++----- 5 files changed, 18 insertions(+), 25 deletions(-) diff --git a/homeassistant/components/recorder/core.py b/homeassistant/components/recorder/core.py index c0ac1fc1277..002d8937e3a 100644 --- a/homeassistant/components/recorder/core.py +++ b/homeassistant/components/recorder/core.py @@ -1283,10 +1283,6 @@ class Recorder(threading.Thread): self.event_session = self.get_session() self.event_session.expire_on_commit = False - def _post_migrate_entity_ids(self) -> bool: - """Post migrate entity_ids if needed.""" - return migration.post_migrate_entity_ids(self) - def _send_keep_alive(self) -> None: """Send a keep alive to keep the db connection open.""" assert self.event_session is not None diff --git a/homeassistant/components/recorder/migration.py b/homeassistant/components/recorder/migration.py index 890fc3045b2..d2d8fff136e 100644 --- a/homeassistant/components/recorder/migration.py +++ b/homeassistant/components/recorder/migration.py @@ -100,7 +100,7 @@ from .queries import ( migrate_single_statistics_row_to_timestamp, ) from .statistics import cleanup_statistics_timestamp_migration, get_start_time -from .tasks import EntityIDPostMigrationTask, RecorderTask +from .tasks import RecorderTask from .util import ( database_job_retry_wrapper, execute_stmt_lambda_element, @@ -2667,6 +2667,17 @@ class EventIDPostMigration(BaseRunTimeMigration): return NeedsMigrateResult(needs_migrate=False, migration_done=True) +@dataclass(slots=True) +class EntityIDPostMigrationTask(RecorderTask): + """An object to insert into the recorder queue to cleanup after entity_ids migration.""" + + def run(self, instance: Recorder) -> None: + """Run entity_id post migration task.""" + if not post_migrate_entity_ids(instance): + # Schedule a new migration task if this one didn't finish + instance.queue_task(EntityIDPostMigrationTask()) + + def _mark_migration_done( session: Session, migration: type[BaseRunTimeMigration] ) -> None: diff --git a/homeassistant/components/recorder/tasks.py b/homeassistant/components/recorder/tasks.py index c51ba2b16ca..2529e8012bf 100644 --- a/homeassistant/components/recorder/tasks.py +++ b/homeassistant/components/recorder/tasks.py @@ -333,19 +333,6 @@ class AdjustLRUSizeTask(RecorderTask): instance._adjust_lru_size() # noqa: SLF001 -@dataclass(slots=True) -class EntityIDPostMigrationTask(RecorderTask): - """An object to insert into the recorder queue to cleanup after entity_ids migration.""" - - def run(self, instance: Recorder) -> None: - """Run entity_id post migration task.""" - if ( - not instance._post_migrate_entity_ids() # noqa: SLF001 - ): - # Schedule a new migration task if this one didn't finish - instance.queue_task(EntityIDPostMigrationTask()) - - @dataclass(slots=True) class RefreshEventTypesTask(RecorderTask): """An object to insert into the recorder queue to refresh event types.""" diff --git a/tests/components/recorder/test_migration_from_schema_32.py b/tests/components/recorder/test_migration_from_schema_32.py index f1613909722..40d18ab51fd 100644 --- a/tests/components/recorder/test_migration_from_schema_32.py +++ b/tests/components/recorder/test_migration_from_schema_32.py @@ -32,7 +32,6 @@ from homeassistant.components.recorder.queries import ( get_migration_changes, select_event_type_ids, ) -from homeassistant.components.recorder.tasks import EntityIDPostMigrationTask from homeassistant.components.recorder.util import ( execute_stmt_lambda_element, get_index_by_name, @@ -746,7 +745,7 @@ async def test_post_migrate_entity_ids( await _async_wait_migration_done(hass) # This is a threadsafe way to add a task to the recorder - recorder_mock.queue_task(EntityIDPostMigrationTask()) + recorder_mock.queue_task(migration.EntityIDPostMigrationTask()) await _async_wait_migration_done(hass) def _fetch_migrated_states(): diff --git a/tests/components/recorder/test_v32_migration.py b/tests/components/recorder/test_v32_migration.py index 1932860d845..58bcabdff51 100644 --- a/tests/components/recorder/test_v32_migration.py +++ b/tests/components/recorder/test_v32_migration.py @@ -109,6 +109,7 @@ async def test_migrate_times( with ( patch.object(recorder, "db_schema", old_db_schema), patch.object(migration, "SCHEMA_VERSION", old_db_schema.SCHEMA_VERSION), + patch.object(migration, "post_migrate_entity_ids", return_value=False), patch.object(migration.EventsContextIDMigration, "migrate_data"), patch.object(migration.StatesContextIDMigration, "migrate_data"), patch.object(migration.EventTypeIDMigration, "migrate_data"), @@ -120,7 +121,6 @@ async def test_migrate_times( patch.object(core, "States", old_db_schema.States), patch.object(core, "Events", old_db_schema.Events), patch(CREATE_ENGINE_TARGET, new=_create_engine_test(SCHEMA_MODULE_30)), - patch("homeassistant.components.recorder.Recorder._post_migrate_entity_ids"), ): async with ( async_test_home_assistant() as hass, @@ -264,13 +264,13 @@ async def test_migrate_can_resume_entity_id_post_migration( patch.object(recorder, "db_schema", old_db_schema), patch.object(migration, "SCHEMA_VERSION", old_db_schema.SCHEMA_VERSION), patch.object(migration.EventIDPostMigration, "migrate_data"), + patch.object(migration, "post_migrate_entity_ids", return_value=False), patch.object(core, "StatesMeta", old_db_schema.StatesMeta), patch.object(core, "EventTypes", old_db_schema.EventTypes), patch.object(core, "EventData", old_db_schema.EventData), patch.object(core, "States", old_db_schema.States), patch.object(core, "Events", old_db_schema.Events), patch(CREATE_ENGINE_TARGET, new=_create_engine_test(SCHEMA_MODULE_32)), - patch("homeassistant.components.recorder.Recorder._post_migrate_entity_ids"), ): async with ( async_test_home_assistant() as hass, @@ -382,13 +382,13 @@ async def test_migrate_can_resume_ix_states_event_id_removed( patch.object(recorder, "db_schema", old_db_schema), patch.object(migration, "SCHEMA_VERSION", old_db_schema.SCHEMA_VERSION), patch.object(migration.EventIDPostMigration, "migrate_data"), + patch.object(migration, "post_migrate_entity_ids", return_value=False), patch.object(core, "StatesMeta", old_db_schema.StatesMeta), patch.object(core, "EventTypes", old_db_schema.EventTypes), patch.object(core, "EventData", old_db_schema.EventData), patch.object(core, "States", old_db_schema.States), patch.object(core, "Events", old_db_schema.Events), patch(CREATE_ENGINE_TARGET, new=_create_engine_test(SCHEMA_MODULE_32)), - patch("homeassistant.components.recorder.Recorder._post_migrate_entity_ids"), ): async with ( async_test_home_assistant() as hass, @@ -514,13 +514,13 @@ async def test_out_of_disk_space_while_rebuild_states_table( patch.object(recorder, "db_schema", old_db_schema), patch.object(migration, "SCHEMA_VERSION", old_db_schema.SCHEMA_VERSION), patch.object(migration.EventIDPostMigration, "migrate_data"), + patch.object(migration, "post_migrate_entity_ids", return_value=False), patch.object(core, "StatesMeta", old_db_schema.StatesMeta), patch.object(core, "EventTypes", old_db_schema.EventTypes), patch.object(core, "EventData", old_db_schema.EventData), patch.object(core, "States", old_db_schema.States), patch.object(core, "Events", old_db_schema.Events), patch(CREATE_ENGINE_TARGET, new=_create_engine_test(SCHEMA_MODULE_32)), - patch("homeassistant.components.recorder.Recorder._post_migrate_entity_ids"), ): async with ( async_test_home_assistant() as hass, @@ -690,13 +690,13 @@ async def test_out_of_disk_space_while_removing_foreign_key( patch.object(recorder, "db_schema", old_db_schema), patch.object(migration, "SCHEMA_VERSION", old_db_schema.SCHEMA_VERSION), patch.object(migration.EventIDPostMigration, "migrate_data"), + patch.object(migration, "post_migrate_entity_ids", return_value=False), patch.object(core, "StatesMeta", old_db_schema.StatesMeta), patch.object(core, "EventTypes", old_db_schema.EventTypes), patch.object(core, "EventData", old_db_schema.EventData), patch.object(core, "States", old_db_schema.States), patch.object(core, "Events", old_db_schema.Events), patch(CREATE_ENGINE_TARGET, new=_create_engine_test(SCHEMA_MODULE_32)), - patch("homeassistant.components.recorder.Recorder._post_migrate_entity_ids"), ): async with ( async_test_home_assistant() as hass,