diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 9575e6722ed..27a4eaa98de 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -36,7 +36,7 @@ from homeassistant.core import ( valid_entity_id, ) from homeassistant.exceptions import MaxLengthExceeded -from homeassistant.helpers import device_registry as dr +from homeassistant.helpers import device_registry as dr, storage from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED from homeassistant.loader import bind_hass from homeassistant.util import slugify @@ -58,7 +58,8 @@ DISABLED_HASS = "hass" DISABLED_INTEGRATION = "integration" DISABLED_USER = "user" -STORAGE_VERSION = 1 +STORAGE_VERSION_MAJOR = 1 +STORAGE_VERSION_MINOR = 2 STORAGE_KEY = "core.entity_registry" # Attributes relevant to describing entity @@ -147,6 +148,16 @@ class RegistryEntry: hass.states.async_set(self.entity_id, STATE_UNAVAILABLE, attrs) +class EntityRegistryStore(storage.Store): + """Store entity registry data.""" + + async def _async_migrate_func( + self, old_major_version: int, old_minor_version: int, old_data: dict + ) -> dict: + """Migrate to the new version.""" + return await _async_migrate(old_major_version, old_minor_version, old_data) + + class EntityRegistry: """Class to hold a registry of entities.""" @@ -155,8 +166,12 @@ class EntityRegistry: self.hass = hass self.entities: dict[str, RegistryEntry] self._index: dict[tuple[str, str, str], str] = {} - self._store = hass.helpers.storage.Store( - STORAGE_VERSION, STORAGE_KEY, atomic_writes=True + self._store = EntityRegistryStore( + hass, + STORAGE_VERSION_MAJOR, + STORAGE_KEY, + atomic_writes=True, + minor_version=STORAGE_VERSION_MINOR, ) self.hass.bus.async_listen( EVENT_DEVICE_REGISTRY_UPDATED, self.async_device_modified @@ -509,11 +524,12 @@ class EntityRegistry: """Load the entity registry.""" async_setup_entity_restore(self.hass, self) - data = await self.hass.helpers.storage.async_migrator( + data = await storage.async_migrator( + self.hass, self.hass.config.path(PATH_REGISTRY), self._store, old_conf_load_func=load_yaml, - old_conf_migrate_func=_async_migrate, + old_conf_migrate_func=_async_migrate_yaml_to_json, ) entities: dict[str, RegistryEntry] = OrderedDict() @@ -526,22 +542,22 @@ class EntityRegistry: continue entities[entity["entity_id"]] = RegistryEntry( - area_id=entity.get("area_id"), - capabilities=entity.get("capabilities") or {}, - config_entry_id=entity.get("config_entry_id"), - device_class=entity.get("device_class"), - device_id=entity.get("device_id"), - disabled_by=entity.get("disabled_by"), - entity_category=entity.get("entity_category"), + area_id=entity["area_id"], + capabilities=entity["capabilities"], + config_entry_id=entity["config_entry_id"], + device_class=entity["device_class"], + device_id=entity["device_id"], + disabled_by=entity["disabled_by"], + entity_category=entity["entity_category"], entity_id=entity["entity_id"], - icon=entity.get("icon"), - name=entity.get("name"), - original_icon=entity.get("original_icon"), - original_name=entity.get("original_name"), + icon=entity["icon"], + name=entity["name"], + original_icon=entity["original_icon"], + original_name=entity["original_name"], platform=entity["platform"], - supported_features=entity.get("supported_features", 0), + supported_features=entity["supported_features"], unique_id=entity["unique_id"], - unit_of_measurement=entity.get("unit_of_measurement"), + unit_of_measurement=entity["unit_of_measurement"], ) self.entities = entities @@ -703,13 +719,43 @@ def async_config_entry_disabled_by_changed( ) -async def _async_migrate(entities: dict[str, Any]) -> dict[str, list[dict[str, Any]]]: +async def _async_migrate( + old_major_version: int, old_minor_version: int, data: dict +) -> dict: + """Migrate to the new version.""" + if old_major_version < 2 and old_minor_version < 2: + # From version 1.1 + for entity in data["entities"]: + # Populate all keys + entity["area_id"] = entity.get("area_id") + entity["capabilities"] = entity.get("capabilities") or {} + entity["config_entry_id"] = entity.get("config_entry_id") + entity["device_class"] = entity.get("device_class") + entity["device_id"] = entity.get("device_id") + entity["disabled_by"] = entity.get("disabled_by") + entity["entity_category"] = entity.get("entity_category") + entity["icon"] = entity.get("icon") + entity["name"] = entity.get("name") + entity["original_icon"] = entity.get("original_icon") + entity["original_name"] = entity.get("original_name") + entity["platform"] = entity["platform"] + entity["supported_features"] = entity.get("supported_features", 0) + entity["unit_of_measurement"] = entity.get("unit_of_measurement") + if old_major_version > 1: + raise NotImplementedError + return data + + +async def _async_migrate_yaml_to_json( + entities: dict[str, Any] +) -> dict[str, list[dict[str, Any]]]: """Migrate the YAML config file to storage helper format.""" - return { + entities_1_1 = { "entities": [ {"entity_id": entity_id, **info} for entity_id, info in entities.items() ] } + return await _async_migrate(1, 1, entities_1_1) @callback diff --git a/homeassistant/helpers/instance_id.py b/homeassistant/helpers/instance_id.py index 5b6f645c55a..59a4cf39498 100644 --- a/homeassistant/helpers/instance_id.py +++ b/homeassistant/helpers/instance_id.py @@ -18,7 +18,7 @@ async def async_get(hass: HomeAssistant) -> str: """Get unique ID for the hass instance.""" store = storage.Store(hass, DATA_VERSION, DATA_KEY, True) - data: dict[str, str] | None = await storage.async_migrator( # type: ignore + data: dict[str, str] | None = await storage.async_migrator( hass, hass.config.path(LEGACY_UUID_FILE), store, diff --git a/homeassistant/helpers/storage.py b/homeassistant/helpers/storage.py index 87d8044c044..323737d5b98 100644 --- a/homeassistant/helpers/storage.py +++ b/homeassistant/helpers/storage.py @@ -13,7 +13,6 @@ from typing import Any from homeassistant.const import EVENT_HOMEASSISTANT_FINAL_WRITE from homeassistant.core import CALLBACK_TYPE, CoreState, Event, HomeAssistant, callback -from homeassistant.helpers.event import async_call_later from homeassistant.loader import MAX_LOAD_CONCURRENTLY, bind_hass from homeassistant.util import json as json_util @@ -28,13 +27,13 @@ STORAGE_SEMAPHORE = "storage_semaphore" @bind_hass async def async_migrator( - hass, - old_path, - store, + hass: HomeAssistant, + old_path: str, + store: Store, *, - old_conf_load_func=None, - old_conf_migrate_func=None, -): + old_conf_load_func: Callable | None = None, + old_conf_migrate_func: Callable | None = None, +) -> Any: """Migrate old data to a store and then load data. async def old_conf_migrate_func(old_data) @@ -157,10 +156,12 @@ class Store: stored = data["data"] else: _LOGGER.info( - "Migrating %s storage from %s to %s", + "Migrating %s storage from %s.%s to %s.%s", self.key, data["version"], + data["minor_version"], self.version, + self.minor_version, ) if len(inspect.signature(self._async_migrate_func).parameters) == 2: # pylint: disable-next=no-value-for-parameter @@ -195,6 +196,9 @@ class Store: @callback def async_delay_save(self, data_func: Callable[[], dict], delay: float = 0) -> None: """Save data with an optional delay.""" + # pylint: disable-next=import-outside-toplevel + from homeassistant.helpers.event import async_call_later + self._data = { "version": self.version, "minor_version": self.minor_version, diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index fe445e32c96..b07c4131795 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -223,7 +223,8 @@ def test_is_registered(registry): async def test_loading_extra_values(hass, hass_storage): """Test we load extra data from the registry.""" hass_storage[er.STORAGE_KEY] = { - "version": er.STORAGE_VERSION, + "version": er.STORAGE_VERSION_MAJOR, + "minor_version": 1, "data": { "entities": [ { @@ -387,7 +388,8 @@ async def test_migration(hass): async def test_loading_invalid_entity_id(hass, hass_storage): """Test we autofix invalid entity IDs.""" hass_storage[er.STORAGE_KEY] = { - "version": er.STORAGE_VERSION, + "version": er.STORAGE_VERSION_MAJOR, + "minor_version": er.STORAGE_VERSION_MINOR, "data": { "entities": [ {