Improve migration to entity registry version 1.18 (home-assistant#151570)

Author: emontnemery

homeassistant/helpers/entity_registry.py

@@ -79,7 +79,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 STORAGE_VERSION_MAJOR = 1
-STORAGE_VERSION_MINOR = 18
+STORAGE_VERSION_MINOR = 19
 STORAGE_KEY = "core.entity_registry"
 
 CLEANUP_INTERVAL = 3600 * 24
@@ -164,6 +164,17 @@ def _protect_entity_options(
     return ReadOnlyDict({key: ReadOnlyDict(val) for key, val in data.items()})
 
 
+def _protect_optional_entity_options(
+    data: EntityOptionsType | UndefinedType | None,
+) -> ReadOnlyEntityOptionsType | UndefinedType:
+    """Protect entity options from being modified."""
+    if data is UNDEFINED:
+        return UNDEFINED
+    if data is None:
+        return ReadOnlyDict({})
+    return ReadOnlyDict({key: ReadOnlyDict(val) for key, val in data.items()})
+
+
 @attr.s(frozen=True, kw_only=True, slots=True)
 class RegistryEntry:
     """Entity Registry Entry."""
@@ -414,15 +425,17 @@ class DeletedRegistryEntry:
     config_subentry_id: str | None = attr.ib()
     created_at: datetime = attr.ib()
     device_class: str | None = attr.ib()
-    disabled_by: RegistryEntryDisabler | None = attr.ib()
+    disabled_by: RegistryEntryDisabler | UndefinedType | None = attr.ib()
     domain: str = attr.ib(init=False, repr=False)
-    hidden_by: RegistryEntryHider | None = attr.ib()
+    hidden_by: RegistryEntryHider | UndefinedType | None = attr.ib()
     icon: str | None = attr.ib()
     id: str = attr.ib()
     labels: set[str] = attr.ib()
     modified_at: datetime = attr.ib()
     name: str | None = attr.ib()
-    options: ReadOnlyEntityOptionsType = attr.ib(converter=_protect_entity_options)
+    options: ReadOnlyEntityOptionsType | UndefinedType = attr.ib(
+        converter=_protect_optional_entity_options
+    )
     orphaned_timestamp: float | None = attr.ib()
 
     _cache: dict[str, Any] = attr.ib(factory=dict, eq=False, init=False)
@@ -445,15 +458,22 @@ def as_storage_fragment(self) -> json_fragment:
                     "config_subentry_id": self.config_subentry_id,
                     "created_at": self.created_at,
                     "device_class": self.device_class,
-                    "disabled_by": self.disabled_by,
+                    "disabled_by": self.disabled_by
+                    if self.disabled_by is not UNDEFINED
+                    else None,
+                    "disabled_by_undefined": self.disabled_by is UNDEFINED,
                     "entity_id": self.entity_id,
-                    "hidden_by": self.hidden_by,
+                    "hidden_by": self.hidden_by
+                    if self.hidden_by is not UNDEFINED
+                    else None,
+                    "hidden_by_undefined": self.hidden_by is UNDEFINED,
                     "icon": self.icon,
                     "id": self.id,
                     "labels": list(self.labels),
                     "modified_at": self.modified_at,
                     "name": self.name,
-                    "options": self.options,
+                    "options": self.options if self.options is not UNDEFINED else {},
+                    "options_undefined": self.options is UNDEFINED,
                     "orphaned_timestamp": self.orphaned_timestamp,
                     "platform": self.platform,
                     "unique_id": self.unique_id,
@@ -590,6 +610,14 @@ async def _async_migrate_func(  # noqa: C901
                     entity["labels"] = []
                     entity["name"] = None
                     entity["options"] = {}
+            if old_minor_version < 19:
+                # Version 1.19 adds undefined flags to deleted entities, this is a bugfix
+                # of version 1.18
+                set_to_undefined = old_minor_version < 18
+                for entity in data["deleted_entities"]:
+                    entity["disabled_by_undefined"] = set_to_undefined
+                    entity["hidden_by_undefined"] = set_to_undefined
+                    entity["options_undefined"] = set_to_undefined
 
         if old_major_version > 1:
             raise NotImplementedError
@@ -958,25 +986,30 @@ def async_get_or_create(
             categories = deleted_entity.categories
             created_at = deleted_entity.created_at
             device_class = deleted_entity.device_class
-            disabled_by = deleted_entity.disabled_by
-            # Adjust disabled_by based on config entry state
-            if config_entry and config_entry is not UNDEFINED:
-                if config_entry.disabled_by:
-                    if disabled_by is None:
-                        disabled_by = RegistryEntryDisabler.CONFIG_ENTRY
+            if deleted_entity.disabled_by is not UNDEFINED:
+                disabled_by = deleted_entity.disabled_by
+                # Adjust disabled_by based on config entry state
+                if config_entry and config_entry is not UNDEFINED:
+                    if config_entry.disabled_by:
+                        if disabled_by is None:
+                            disabled_by = RegistryEntryDisabler.CONFIG_ENTRY
+                    elif disabled_by == RegistryEntryDisabler.CONFIG_ENTRY:
+                        disabled_by = None
                 elif disabled_by == RegistryEntryDisabler.CONFIG_ENTRY:
                     disabled_by = None
-            elif disabled_by == RegistryEntryDisabler.CONFIG_ENTRY:
-                disabled_by = None
             # Restore entity_id if it's available
             if self._entity_id_available(deleted_entity.entity_id):
                 entity_id = deleted_entity.entity_id
             entity_registry_id = deleted_entity.id
-            hidden_by = deleted_entity.hidden_by
+            if deleted_entity.hidden_by is not UNDEFINED:
+                hidden_by = deleted_entity.hidden_by
             icon = deleted_entity.icon
             labels = deleted_entity.labels
             name = deleted_entity.name
-            options = deleted_entity.options
+            if deleted_entity.options is not UNDEFINED:
+                options = deleted_entity.options
+            else:
+                options = get_initial_options() if get_initial_options else None
         else:
             aliases = set()
             area_id = None
@@ -1529,6 +1562,20 @@ async def async_load(self) -> None:
                     previous_unique_id=entity["previous_unique_id"],
                     unit_of_measurement=entity["unit_of_measurement"],
                 )
+
+            def get_optional_enum[_EnumT: StrEnum](
+                cls: type[_EnumT], value: str | None, undefined: bool
+            ) -> _EnumT | UndefinedType | None:
+                """Convert string to the passed enum, UNDEFINED or None."""
+                if undefined:
+                    return UNDEFINED
+                if value is None:
+                    return None
+                try:
+                    return cls(value)
+                except ValueError:
+                    return None
+
             for entity in data["deleted_entities"]:
                 try:
                     domain = split_entity_id(entity["entity_id"])[0]
@@ -1554,23 +1601,25 @@ async def async_load(self) -> None:
                     config_subentry_id=entity["config_subentry_id"],
                     created_at=datetime.fromisoformat(entity["created_at"]),
                     device_class=entity["device_class"],
-                    disabled_by=(
-                        RegistryEntryDisabler(entity["disabled_by"])
-                        if entity["disabled_by"]
-                        else None
+                    disabled_by=get_optional_enum(
+                        RegistryEntryDisabler,
+                        entity["disabled_by"],
+                        entity["disabled_by_undefined"],
                     ),
                     entity_id=entity["entity_id"],
-                    hidden_by=(
-                        RegistryEntryHider(entity["hidden_by"])
-                        if entity["hidden_by"]
-                        else None
+                    hidden_by=get_optional_enum(
+                        RegistryEntryHider,
+                        entity["hidden_by"],
+                        entity["hidden_by_undefined"],
                     ),
                     icon=entity["icon"],
                     id=entity["id"],
                     labels=set(entity["labels"]),
                     modified_at=datetime.fromisoformat(entity["modified_at"]),
                     name=entity["name"],
-                    options=entity["options"],
+                    options=entity["options"]
+                    if not entity["options_undefined"]
+                    else UNDEFINED,
                     orphaned_timestamp=entity["orphaned_timestamp"],
                     platform=entity["platform"],
                     unique_id=entity["unique_id"],