84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
"""UniFi Protect data migrations."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from pyunifiprotect import ProtectApiClient
|
|
|
|
from homeassistant.config_entries import ConfigEntry
|
|
from homeassistant.const import Platform
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.helpers import entity_registry as er
|
|
|
|
from .const import DEVICES_THAT_ADOPT
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
async def async_migrate_data(
|
|
hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient
|
|
) -> None:
|
|
"""Run all valid UniFi Protect data migrations."""
|
|
|
|
_LOGGER.debug("Start Migrate: async_migrate_buttons")
|
|
await async_migrate_buttons(hass, entry, protect)
|
|
_LOGGER.debug("Completed Migrate: async_migrate_buttons")
|
|
|
|
|
|
async def async_migrate_buttons(
|
|
hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient
|
|
) -> None:
|
|
"""
|
|
Migrate existing Reboot button unique IDs from {device_id} to {deivce_id}_reboot.
|
|
|
|
This allows for additional types of buttons that are outside of just a reboot button.
|
|
|
|
Added in 2022.6.0.
|
|
"""
|
|
|
|
registry = er.async_get(hass)
|
|
to_migrate = []
|
|
for entity in er.async_entries_for_config_entry(registry, entry.entry_id):
|
|
if entity.domain == Platform.BUTTON and "_" not in entity.unique_id:
|
|
_LOGGER.debug("Button %s needs migration", entity.entity_id)
|
|
to_migrate.append(entity)
|
|
|
|
if len(to_migrate) == 0:
|
|
_LOGGER.debug("No button entities need migration")
|
|
return
|
|
|
|
bootstrap = await protect.get_bootstrap()
|
|
count = 0
|
|
for button in to_migrate:
|
|
device = None
|
|
for model in DEVICES_THAT_ADOPT:
|
|
attr = f"{model.value}s"
|
|
device = getattr(bootstrap, attr).get(button.unique_id)
|
|
if device is not None:
|
|
break
|
|
|
|
if device is None:
|
|
continue
|
|
|
|
new_unique_id = f"{device.id}_reboot"
|
|
_LOGGER.debug(
|
|
"Migrating entity %s (old unique_id: %s, new unique_id: %s)",
|
|
button.entity_id,
|
|
button.unique_id,
|
|
new_unique_id,
|
|
)
|
|
try:
|
|
registry.async_update_entity(button.entity_id, new_unique_id=new_unique_id)
|
|
except ValueError:
|
|
_LOGGER.warning(
|
|
"Could not migrate entity %s (old unique_id: %s, new unique_id: %s)",
|
|
button.entity_id,
|
|
button.unique_id,
|
|
new_unique_id,
|
|
)
|
|
else:
|
|
count += 1
|
|
|
|
if count < len(to_migrate):
|
|
_LOGGER.warning("Failed to migate %s reboot buttons", len(to_migrate) - count)
|