Restore the device id after deleting and re-adding an integration (#45348)

homeassistant/helpers/device_registry.py

@@ -1,6 +1,7 @@
 """Provide a way to connect entities belonging to one device."""
 from collections import OrderedDict
 import logging
+import time
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
 
 import attr
@@ -39,6 +40,8 @@ DELETED_DEVICE = "deleted"
 DISABLED_INTEGRATION = "integration"
 DISABLED_USER = "user"
 
+ORPHANED_DEVICE_KEEP_SECONDS = 86400 * 30
+
 
 @attr.s(slots=True, frozen=True)
 class DeviceEntry:
@@ -83,6 +86,7 @@ class DeletedDeviceEntry:
     connections: Set[Tuple[str, str]] = attr.ib()
     identifiers: Set[Tuple[str, str]] = attr.ib()
     id: str = attr.ib()
+    orphaned_timestamp: Optional[float] = attr.ib()
 
     def to_device_entry(
         self,
@@ -440,6 +444,7 @@ class DeviceRegistry:
                 connections=device.connections,
                 identifiers=device.identifiers,
                 id=device.id,
+                orphaned_timestamp=None,
             )
         )
         self.hass.bus.async_fire(
@@ -489,6 +494,8 @@ class DeviceRegistry:
                     connections={tuple(conn) for conn in device["connections"]},  # type: ignore[misc]
                     identifiers={tuple(iden) for iden in device["identifiers"]},  # type: ignore[misc]
                     id=device["id"],
+                    # Introduced in 2021.2
+                    orphaned_timestamp=device.get("orphaned_timestamp"),
                 )
 
         self.devices = devices
@@ -529,6 +536,7 @@ class DeviceRegistry:
                 "connections": list(entry.connections),
                 "identifiers": list(entry.identifiers),
                 "id": entry.id,
+                "orphaned_timestamp": entry.orphaned_timestamp,
             }
             for entry in self.deleted_devices.values()
         ]
@@ -538,6 +546,7 @@ class DeviceRegistry:
     @callback
     def async_clear_config_entry(self, config_entry_id: str) -> None:
         """Clear config entry from registry entries."""
+        now_time = time.time()
         for device in list(self.devices.values()):
             self._async_update_device(device.id, remove_config_entry_id=config_entry_id)
         for deleted_device in list(self.deleted_devices.values()):
@@ -545,8 +554,10 @@ class DeviceRegistry:
             if config_entry_id not in config_entries:
                 continue
             if config_entries == {config_entry_id}:
-                # Permanently remove the device from the device registry.
-                self._remove_device(deleted_device)
+                # Add a time stamp when the deleted device became orphaned
+                self.deleted_devices[deleted_device.id] = attr.evolve(
+                    deleted_device, orphaned_timestamp=now_time, config_entries=set()
+                )
             else:
                 config_entries = config_entries - {config_entry_id}
                 # No need to reindex here since we currently
@@ -556,6 +567,24 @@ class DeviceRegistry:
                 )
             self.async_schedule_save()
 
+    @callback
+    def async_purge_expired_orphaned_devices(self) -> None:
+        """Purge expired orphaned devices from the registry.
+
+        We need to purge these periodically to avoid the database
+        growing without bound.
+        """
+        now_time = time.time()
+        for deleted_device in list(self.deleted_devices.values()):
+            if deleted_device.orphaned_timestamp is None:
+                continue
+
+            if (
+                deleted_device.orphaned_timestamp + ORPHANED_DEVICE_KEEP_SECONDS
+                < now_time
+            ):
+                self._remove_device(deleted_device)
+
     @callback
     def async_clear_area_id(self, area_id: str) -> None:
         """Clear area id from registry entries."""
@@ -623,6 +652,10 @@ def async_cleanup(
                     device.id, remove_config_entry_id=config_entry_id
                 )
 
+    # Periodic purge of orphaned devices to avoid the registry
+    # growing without bounds when there are lots of deleted devices
+    dev_reg.async_purge_expired_orphaned_devices()
+
 
 @callback
 def async_setup_cleanup(hass: HomeAssistantType, dev_reg: DeviceRegistry) -> None:

tests/helpers/test_device_registry.py

@@ -1,5 +1,6 @@
 """Tests for the Device Registry."""
 import asyncio
+import time
 from unittest.mock import patch
 
 import pytest
@@ -301,26 +302,41 @@ async def test_deleted_device_removing_config_entries(hass, registry, update_eve
 
     registry.async_clear_config_entry("123")
     assert len(registry.devices) == 0
-    assert len(registry.deleted_devices) == 1
+    assert len(registry.deleted_devices) == 2
 
     registry.async_clear_config_entry("456")
     assert len(registry.devices) == 0
-    assert len(registry.deleted_devices) == 0
+    assert len(registry.deleted_devices) == 2
 
     # No event when a deleted device is purged
     await hass.async_block_till_done()
     assert len(update_events) == 5
 
-    # Re-add, expect new device id
+    # Re-add, expect to keep the device id
     entry2 = registry.async_get_or_create(
-        config_entry_id="123",
+        config_entry_id="456",
         connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
         identifiers={("bridgeid", "0123")},
         manufacturer="manufacturer",
         model="model",
     )
 
-    assert entry.id != entry2.id
+    assert entry.id == entry2.id
+
+    future_time = time.time() + device_registry.ORPHANED_DEVICE_KEEP_SECONDS + 1
+
+    with patch("time.time", return_value=future_time):
+        registry.async_purge_expired_orphaned_devices()
+
+    # Re-add, expect to get a new device id after the purge
+    entry4 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+    assert entry3.id != entry4.id
 
 
 async def test_removing_area_id(registry):
@@ -734,6 +750,40 @@ async def test_cleanup_device_registry(hass, registry):
     assert registry.async_get_device({("something", "d4")}) is None
 
 
+async def test_cleanup_device_registry_removes_expired_orphaned_devices(hass, registry):
+    """Test cleanup removes expired orphaned devices."""
+    config_entry = MockConfigEntry(domain="hue")
+    config_entry.add_to_hass(hass)
+
+    registry.async_get_or_create(
+        identifiers={("hue", "d1")}, config_entry_id=config_entry.entry_id
+    )
+    registry.async_get_or_create(
+        identifiers={("hue", "d2")}, config_entry_id=config_entry.entry_id
+    )
+    registry.async_get_or_create(
+        identifiers={("hue", "d3")}, config_entry_id=config_entry.entry_id
+    )
+
+    registry.async_clear_config_entry(config_entry.entry_id)
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 3
+
+    ent_reg = await entity_registry.async_get_registry(hass)
+    device_registry.async_cleanup(hass, registry, ent_reg)
+
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 3
+
+    future_time = time.time() + device_registry.ORPHANED_DEVICE_KEEP_SECONDS + 1
+
+    with patch("time.time", return_value=future_time):
+        device_registry.async_cleanup(hass, registry, ent_reg)
+
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 0
+
+
 async def test_cleanup_startup(hass):
     """Test we run a cleanup on startup."""
     hass.state = CoreState.not_running