Stable entity registry id when a deleted entity is restored (#77710)

* Stable entity_id and registry id when a deleted entity is restored

* Don't restore area_id

* Don't restore entity_id

* Address review comments
This commit is contained in:
Erik Montnemery 2023-06-26 15:54:35 +02:00 committed by GitHub
parent f08f0fbb8b
commit 8fda56d2c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 283 additions and 9 deletions

View File

@ -11,7 +11,9 @@ from __future__ import annotations
from collections import UserDict
from collections.abc import Callable, Iterable, Mapping, ValuesView
from datetime import datetime, timedelta
import logging
import time
from typing import TYPE_CHECKING, Any, TypeVar, cast
import attr
@ -26,6 +28,7 @@ from homeassistant.const import (
ATTR_SUPPORTED_FEATURES,
ATTR_UNIT_OF_MEASUREMENT,
EVENT_HOMEASSISTANT_START,
EVENT_HOMEASSISTANT_STOP,
MAX_LENGTH_STATE_DOMAIN,
MAX_LENGTH_STATE_ENTITY_ID,
STATE_UNAVAILABLE,
@ -61,9 +64,12 @@ SAVE_DELAY = 10
_LOGGER = logging.getLogger(__name__)
STORAGE_VERSION_MAJOR = 1
STORAGE_VERSION_MINOR = 10
STORAGE_VERSION_MINOR = 11
STORAGE_KEY = "core.entity_registry"
CLEANUP_INTERVAL = 3600 * 24
ORPHANED_ENTITY_KEEP_SECONDS = 3600 * 24 * 30
ENTITY_CATEGORY_VALUE_TO_INDEX: dict[EntityCategory | None, int] = {
# mypy does not understand strenum
val: idx # type: ignore[misc]
@ -138,7 +144,10 @@ class RegistryEntry:
entity_category: EntityCategory | None = attr.ib(default=None)
hidden_by: RegistryEntryHider | None = attr.ib(default=None)
icon: str | None = attr.ib(default=None)
id: str = attr.ib(factory=uuid_util.random_uuid_hex)
id: str = attr.ib(
default=None,
converter=attr.converters.default_if_none(factory=uuid_util.random_uuid_hex), # type: ignore[misc]
)
has_entity_name: bool = attr.ib(default=False)
name: str | None = attr.ib(default=None)
options: ReadOnlyEntityOptionsType = attr.ib(
@ -297,6 +306,24 @@ class RegistryEntry:
hass.states.async_set(self.entity_id, STATE_UNAVAILABLE, attrs)
@attr.s(slots=True, frozen=True)
class DeletedRegistryEntry:
"""Deleted Entity Registry Entry."""
entity_id: str = attr.ib()
unique_id: str = attr.ib()
platform: str = attr.ib()
config_entry_id: str | None = attr.ib()
domain: str = attr.ib(init=False, repr=False)
id: str = attr.ib()
orphaned_timestamp: float | None = attr.ib()
@domain.default
def _domain_default(self) -> str:
"""Compute domain value."""
return split_entity_id(self.entity_id)[0]
class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
"""Store entity registry data."""
@ -372,6 +399,10 @@ class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
for entity in data["entities"]:
entity["aliases"] = []
if old_major_version == 1 and old_minor_version < 11:
# Version 1.11 adds deleted_entities
data["deleted_entities"] = data.get("deleted_entities", [])
if old_major_version > 1:
raise NotImplementedError
return data
@ -424,6 +455,7 @@ class EntityRegistryItems(UserDict[str, "RegistryEntry"]):
class EntityRegistry:
"""Class to hold a registry of entities."""
deleted_entities: dict[tuple[str, str, str], DeletedRegistryEntry]
entities: EntityRegistryItems
_entities_data: dict[str, RegistryEntry]
@ -496,6 +528,9 @@ class EntityRegistry:
- It's not registered
- It's not known by the entity component adding the entity
- It's not in the state machine
Note that an entity_id which belongs to a deleted entity is considered
available.
"""
if known_object_ids is None:
known_object_ids = {}
@ -591,8 +626,16 @@ class EntityRegistry:
unit_of_measurement=unit_of_measurement,
)
entity_registry_id: str | None = None
deleted_entity = self.deleted_entities.pop((domain, platform, unique_id), None)
if deleted_entity is not None:
# Restore id
entity_registry_id = deleted_entity.id
entity_id = self.async_generate_entity_id(
domain, suggested_object_id or f"{platform}_{unique_id}", known_object_ids
domain,
suggested_object_id or f"{platform}_{unique_id}",
known_object_ids,
)
if disabled_by and not isinstance(disabled_by, RegistryEntryDisabler):
@ -630,6 +673,7 @@ class EntityRegistry:
entity_id=entity_id,
hidden_by=hidden_by,
has_entity_name=none_if_undefined(has_entity_name) or False,
id=entity_registry_id,
options=initial_options,
original_device_class=none_if_undefined(original_device_class),
original_icon=none_if_undefined(original_icon),
@ -653,7 +697,19 @@ class EntityRegistry:
@callback
def async_remove(self, entity_id: str) -> None:
"""Remove an entity from registry."""
self.entities.pop(entity_id)
entity = self.entities.pop(entity_id)
config_entry_id = entity.config_entry_id
key = (entity.domain, entity.platform, entity.unique_id)
# If the entity does not belong to a config entry, mark it as orphaned
orphaned_timestamp = None if config_entry_id else time.time()
self.deleted_entities[key] = DeletedRegistryEntry(
config_entry_id=config_entry_id,
entity_id=entity_id,
id=entity.id,
orphaned_timestamp=orphaned_timestamp,
platform=entity.platform,
unique_id=entity.unique_id,
)
self.hass.bus.async_fire(
EVENT_ENTITY_REGISTRY_UPDATED, {"action": "remove", "entity_id": entity_id}
)
@ -954,10 +1010,12 @@ class EntityRegistry:
async def async_load(self) -> None:
"""Load the entity registry."""
async_setup_entity_restore(self.hass, self)
_async_setup_cleanup(self.hass, self)
_async_setup_entity_restore(self.hass, self)
data = await self._store.async_load()
entities = EntityRegistryItems()
deleted_entities: dict[tuple[str, str, str], DeletedRegistryEntry] = {}
if data is not None:
for entity in data["entities"]:
@ -996,7 +1054,22 @@ class EntityRegistry:
unique_id=entity["unique_id"],
unit_of_measurement=entity["unit_of_measurement"],
)
for entity in data["deleted_entities"]:
key = (
split_entity_id(entity["entity_id"])[0],
entity["platform"],
entity["unique_id"],
)
deleted_entities[key] = DeletedRegistryEntry(
config_entry_id=entity["config_entry_id"],
entity_id=entity["entity_id"],
id=entity["id"],
orphaned_timestamp=entity["orphaned_timestamp"],
platform=entity["platform"],
unique_id=entity["unique_id"],
)
self.deleted_entities = deleted_entities
self.entities = entities
self._entities_data = entities.data
@ -1038,18 +1111,54 @@ class EntityRegistry:
}
for entry in self.entities.values()
]
data["deleted_entities"] = [
{
"config_entry_id": entry.config_entry_id,
"entity_id": entry.entity_id,
"id": entry.id,
"orphaned_timestamp": entry.orphaned_timestamp,
"platform": entry.platform,
"unique_id": entry.unique_id,
}
for entry in self.deleted_entities.values()
]
return data
@callback
def async_clear_config_entry(self, config_entry: str) -> None:
def async_clear_config_entry(self, config_entry_id: str) -> None:
"""Clear config entry from registry entries."""
now_time = time.time()
for entity_id in [
entity_id
for entity_id, entry in self.entities.items()
if config_entry == entry.config_entry_id
if config_entry_id == entry.config_entry_id
]:
self.async_remove(entity_id)
for key, deleted_entity in list(self.deleted_entities.items()):
if config_entry_id != deleted_entity.config_entry_id:
continue
# Add a time stamp when the deleted entity became orphaned
self.deleted_entities[key] = attr.evolve(
deleted_entity, orphaned_timestamp=now_time, config_entry_id=None
)
self.async_schedule_save()
@callback
def async_purge_expired_orphaned_entities(self) -> None:
"""Purge expired orphaned entities from the registry.
We need to purge these periodically to avoid the database
growing without bound.
"""
now_time = time.time()
for key, deleted_entity in list(self.deleted_entities.items()):
if (orphaned_timestamp := deleted_entity.orphaned_timestamp) is None:
continue
if orphaned_timestamp + ORPHANED_ENTITY_KEEP_SECONDS < now_time:
self.deleted_entities.pop(key)
self.async_schedule_save()
@callback
def async_clear_area_id(self, area_id: str) -> None:
@ -1136,7 +1245,31 @@ def async_config_entry_disabled_by_changed(
@callback
def async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -> None:
def _async_setup_cleanup(hass: HomeAssistant, registry: EntityRegistry) -> None:
"""Clean up device registry when entities removed."""
from . import event # pylint: disable=import-outside-toplevel
@callback
def cleanup(_: datetime) -> None:
"""Clean up entity registry."""
# Periodic purge of orphaned entities to avoid the registry
# growing without bounds when there are lots of deleted entities
registry.async_purge_expired_orphaned_entities()
cancel = event.async_track_time_interval(
hass, cleanup, timedelta(seconds=CLEANUP_INTERVAL)
)
@callback
def _on_homeassistant_stop(event: Event) -> None:
"""Cancel cleanup."""
cancel()
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _on_homeassistant_stop)
@callback
def _async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -> None:
"""Set up the entity restore mechanism."""
@callback

View File

@ -511,6 +511,7 @@ def mock_registry(
registry = er.EntityRegistry(hass)
if mock_entries is None:
mock_entries = {}
registry.deleted_entities = {}
registry.entities = er.EntityRegistryItems()
registry._entities_data = registry.entities.data
for key, entry in mock_entries.items():

View File

@ -1,7 +1,9 @@
"""Tests for the Entity Registry."""
from datetime import timedelta
from typing import Any
from unittest.mock import patch
import attr
import pytest
import voluptuous as vol
@ -15,7 +17,7 @@ from homeassistant.core import CoreState, HomeAssistant, callback
from homeassistant.exceptions import MaxLengthExceeded
from homeassistant.helpers import device_registry as dr, entity_registry as er
from tests.common import MockConfigEntry, flush_store
from tests.common import MockConfigEntry, async_fire_time_changed, flush_store
YAML__OPEN_PATH = "homeassistant.util.yaml.loader.open"
@ -276,8 +278,13 @@ async def test_loading_saving_data(
orig_entry2.entity_id, "light", {"minimum_brightness": 20}
)
orig_entry2 = entity_registry.async_get(orig_entry2.entity_id)
orig_entry3 = entity_registry.async_get_or_create("light", "hue", "ABCD")
orig_entry4 = entity_registry.async_get_or_create("light", "hue", "EFGH")
entity_registry.async_remove(orig_entry3.entity_id)
entity_registry.async_remove(orig_entry4.entity_id)
assert len(entity_registry.entities) == 2
assert len(entity_registry.deleted_entities) == 2
# Now load written data in new registry
registry2 = er.EntityRegistry(hass)
@ -286,11 +293,16 @@ async def test_loading_saving_data(
# Ensure same order
assert list(entity_registry.entities) == list(registry2.entities)
assert list(entity_registry.deleted_entities) == list(registry2.deleted_entities)
new_entry1 = entity_registry.async_get_or_create("light", "hue", "1234")
new_entry2 = entity_registry.async_get_or_create("light", "hue", "5678")
new_entry3 = entity_registry.async_get_or_create("light", "hue", "ABCD")
new_entry4 = entity_registry.async_get_or_create("light", "hue", "EFGH")
assert orig_entry1 == new_entry1
assert orig_entry2 == new_entry2
assert orig_entry3 == new_entry3
assert orig_entry4 == new_entry4
assert new_entry2.area_id == "mock-area-id"
assert new_entry2.capabilities == {"max": 100}
@ -485,6 +497,42 @@ async def test_removing_config_entry_id(
assert update_events[1]["entity_id"] == entry.entity_id
async def test_deleted_entity_removing_config_entry_id(
hass, entity_registry: er.EntityRegistry
):
"""Test that we update config entry id in registry on deleted entity."""
mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1")
entry = entity_registry.async_get_or_create(
"light", "hue", "5678", config_entry=mock_config
)
assert entry.config_entry_id == "mock-id-1"
entity_registry.async_remove(entry.entity_id)
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 1
assert (
entity_registry.deleted_entities[("light", "hue", "5678")].config_entry_id
== "mock-id-1"
)
assert (
entity_registry.deleted_entities[("light", "hue", "5678")].orphaned_timestamp
is None
)
entity_registry.async_clear_config_entry("mock-id-1")
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 1
assert (
entity_registry.deleted_entities[("light", "hue", "5678")].config_entry_id
is None
)
assert (
entity_registry.deleted_entities[("light", "hue", "5678")].orphaned_timestamp
is not None
)
async def test_removing_area_id(entity_registry: er.EntityRegistry) -> None:
"""Make sure we can clear area id."""
entry = entity_registry.async_get_or_create("light", "hue", "5678")
@ -1537,3 +1585,95 @@ def test_migrate_entity_to_new_platform(
new_unique_id=new_unique_id,
new_config_entry_id=new_config_entry.entry_id,
)
async def test_restore_entity(hass, update_events, freezer):
"""Make sure entity registry id is stable and entity_id is reused if possible."""
registry = er.async_get(hass) # We need the real entity registry for this test
config_entry = MockConfigEntry(domain="light")
entry1 = registry.async_get_or_create(
"light", "hue", "1234", config_entry=config_entry
)
entry2 = registry.async_get_or_create(
"light", "hue", "5678", config_entry=config_entry
)
entry1 = registry.async_update_entity(
entry1.entity_id, new_entity_id="light.custom_1"
)
registry.async_remove(entry1.entity_id)
registry.async_remove(entry2.entity_id)
assert len(registry.entities) == 0
assert len(registry.deleted_entities) == 2
# Re-add entities
entry1_restored = registry.async_get_or_create(
"light", "hue", "1234", config_entry=config_entry
)
entry2_restored = registry.async_get_or_create("light", "hue", "5678")
assert len(registry.entities) == 2
assert len(registry.deleted_entities) == 0
assert entry1 != entry1_restored
# entity_id is not restored
assert attr.evolve(entry1, entity_id="light.hue_1234") == entry1_restored
assert entry2 != entry2_restored
# Config entry is not restored
assert attr.evolve(entry2, config_entry_id=None) == entry2_restored
# Remove two of the entities again, then bump time
registry.async_remove(entry1_restored.entity_id)
registry.async_remove(entry2.entity_id)
assert len(registry.entities) == 0
assert len(registry.deleted_entities) == 2
freezer.tick(timedelta(seconds=er.ORPHANED_ENTITY_KEEP_SECONDS + 1))
async_fire_time_changed(hass)
await hass.async_block_till_done()
# Re-add two entities, expect to get a new id after the purge for entity w/o config entry
entry1_restored = registry.async_get_or_create(
"light", "hue", "1234", config_entry=config_entry
)
entry2_restored = registry.async_get_or_create("light", "hue", "5678")
assert len(registry.entities) == 2
assert len(registry.deleted_entities) == 0
assert entry1.id == entry1_restored.id
assert entry2.id != entry2_restored.id
# Remove the first entity, then its config entry, finally bump time
registry.async_remove(entry1_restored.entity_id)
assert len(registry.entities) == 1
assert len(registry.deleted_entities) == 1
registry.async_clear_config_entry(config_entry.entry_id)
freezer.tick(timedelta(seconds=er.ORPHANED_ENTITY_KEEP_SECONDS + 1))
async_fire_time_changed(hass)
await hass.async_block_till_done()
# Re-add the entity, expect to get a new id after the purge
entry1_restored = registry.async_get_or_create(
"light", "hue", "1234", config_entry=config_entry
)
assert len(registry.entities) == 2
assert len(registry.deleted_entities) == 0
assert entry1.id != entry1_restored.id
# Check the events
await hass.async_block_till_done()
assert len(update_events) == 13
assert update_events[0] == {"action": "create", "entity_id": "light.hue_1234"}
assert update_events[1] == {"action": "create", "entity_id": "light.hue_5678"}
assert update_events[2]["action"] == "update"
assert update_events[3] == {"action": "remove", "entity_id": "light.custom_1"}
assert update_events[4] == {"action": "remove", "entity_id": "light.hue_5678"}
# Restore entities the 1st time
assert update_events[5] == {"action": "create", "entity_id": "light.hue_1234"}
assert update_events[6] == {"action": "create", "entity_id": "light.hue_5678"}
assert update_events[7] == {"action": "remove", "entity_id": "light.hue_1234"}
assert update_events[8] == {"action": "remove", "entity_id": "light.hue_5678"}
# Restore entities the 2nd time
assert update_events[9] == {"action": "create", "entity_id": "light.hue_1234"}
assert update_events[10] == {"action": "create", "entity_id": "light.hue_5678"}
assert update_events[11] == {"action": "remove", "entity_id": "light.hue_1234"}
# Restore entities the 3rd time
assert update_events[12] == {"action": "create", "entity_id": "light.hue_1234"}