core/homeassistant/components/hue/migration.py

"""Various helpers to handle config entry and api schema migrations."""

import logging

from aiohue import HueBridgeV2
from aiohue.discovery import is_v2_bridge
from aiohue.v2.models.resource import ResourceTypes

from homeassistant import core
from homeassistant.components.binary_sensor import DEVICE_CLASS_MOTION
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
    CONF_API_KEY,
    CONF_HOST,
    CONF_USERNAME,
    DEVICE_CLASS_BATTERY,
    DEVICE_CLASS_ILLUMINANCE,
    DEVICE_CLASS_TEMPERATURE,
)
from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.device_registry import async_get as async_get_device_registry
from homeassistant.helpers.entity_registry import (
    async_entries_for_config_entry as entities_for_config_entry,
    async_entries_for_device,
    async_get as async_get_entity_registry,
)

from .const import CONF_API_VERSION, DOMAIN

LOGGER = logging.getLogger(__name__)


async def check_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> None:
    """Check if config entry needs any migration actions."""
    host = entry.data[CONF_HOST]

    # migrate CONF_USERNAME --> CONF_API_KEY
    if CONF_USERNAME in entry.data:
        LOGGER.info("Migrate %s to %s in schema", CONF_USERNAME, CONF_API_KEY)
        data = dict(entry.data)
        data[CONF_API_KEY] = data.pop(CONF_USERNAME)
        hass.config_entries.async_update_entry(entry, data=data)

    conf_api_version = entry.data.get(CONF_API_VERSION, 1)
    if conf_api_version == 1:
        # a bridge might have upgraded firmware since last run so
        # we discover its capabilities at every startup
        websession = aiohttp_client.async_get_clientsession(hass)
        if await is_v2_bridge(host, websession):
            supported_api_version = 2
        else:
            supported_api_version = 1
        LOGGER.debug(
            "Configured api version is %s and supported api version %s for bridge %s",
            conf_api_version,
            supported_api_version,
            host,
        )

        # the call to `is_v2_bridge` returns (silently) False even on connection error
        # so if a migration is needed it will be done on next startup

        if conf_api_version == 1 and supported_api_version == 2:
            # run entity/device schema migration for v2
            await handle_v2_migration(hass, entry)

        # store api version in entry data
        if (
            CONF_API_VERSION not in entry.data
            or conf_api_version != supported_api_version
        ):
            data = dict(entry.data)
            data[CONF_API_VERSION] = supported_api_version
            hass.config_entries.async_update_entry(entry, data=data)


async def handle_v2_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> None:
    """Perform migration of devices and entities to V2 Id's."""
    host = entry.data[CONF_HOST]
    api_key = entry.data[CONF_API_KEY]
    websession = aiohttp_client.async_get_clientsession(hass)
    dev_reg = async_get_device_registry(hass)
    ent_reg = async_get_entity_registry(hass)
    LOGGER.info("Start of migration of devices and entities to support API schema 2")
    # initialize bridge connection just for the migration
    async with HueBridgeV2(host, api_key, websession) as api:

        sensor_class_mapping = {
            DEVICE_CLASS_BATTERY: ResourceTypes.DEVICE_POWER,
            DEVICE_CLASS_MOTION: ResourceTypes.MOTION,
            DEVICE_CLASS_ILLUMINANCE: ResourceTypes.LIGHT_LEVEL,
            DEVICE_CLASS_TEMPERATURE: ResourceTypes.TEMPERATURE,
        }

        # handle entities attached to device
        for hue_dev in api.devices:
            zigbee = api.devices.get_zigbee_connectivity(hue_dev.id)
            if not zigbee or not zigbee.mac_address:
                # not a zigbee device or invalid mac
                continue
            # get/update existing device by V1 identifier (mac address)
            # the device will now have both the old and the new identifier
            identifiers = {(DOMAIN, hue_dev.id), (DOMAIN, zigbee.mac_address)}
            hass_dev = dev_reg.async_get_or_create(
                config_entry_id=entry.entry_id, identifiers=identifiers
            )
            LOGGER.info("Migrated device %s (%s)", hass_dev.name, hass_dev.id)
            # loop through al entities for device and find match
            for ent in async_entries_for_device(ent_reg, hass_dev.id, True):
                # migrate light
                if ent.entity_id.startswith("light"):
                    # should always return one lightid here
                    new_unique_id = next(iter(hue_dev.lights))
                    if ent.unique_id == new_unique_id:
                        continue  # just in case
                    LOGGER.info(
                        "Migrating %s from unique id %s to %s",
                        ent.entity_id,
                        ent.unique_id,
                        new_unique_id,
                    )
                    ent_reg.async_update_entity(
                        ent.entity_id, new_unique_id=new_unique_id
                    )
                    continue
                # migrate sensors
                matched_dev_class = sensor_class_mapping.get(
                    ent.original_device_class or "unknown"
                )
                if matched_dev_class is None:
                    # this may happen if we're looking at orphaned or unsupported entity
                    LOGGER.warning(
                        "Skip migration of %s because it no longer exists on the bridge",
                        ent.entity_id,
                    )
                    continue
                for sensor in api.devices.get_sensors(hue_dev.id):
                    if sensor.type != matched_dev_class:
                        continue
                    new_unique_id = sensor.id
                    if ent.unique_id == new_unique_id:
                        break  # just in case
                    LOGGER.info(
                        "Migrating %s from unique id %s to %s",
                        ent.entity_id,
                        ent.unique_id,
                        new_unique_id,
                    )
                    try:
                        ent_reg.async_update_entity(
                            ent.entity_id, new_unique_id=sensor.id
                        )
                    except ValueError:
                        # assume edge case where the entity was already migrated in a previous run
                        # which got aborted somehow and we do not want
                        # to crash the entire integration init
                        LOGGER.warning(
                            "Skip migration of %s because it already exists",
                            ent.entity_id,
                        )
                    break

        # migrate entities that are not connected to a device (groups)
        for ent in entities_for_config_entry(ent_reg, entry.entry_id):
            if ent.device_id is not None:
                continue
            v1_id = f"/groups/{ent.unique_id}"
            hue_group = api.groups.room.get_by_v1_id(v1_id)
            if hue_group is None or hue_group.grouped_light is None:
                # try again with zone
                hue_group = api.groups.zone.get_by_v1_id(v1_id)
            if hue_group is None or hue_group.grouped_light is None:
                # this may happen if we're looking at some orphaned entity
                LOGGER.warning(
                    "Skip migration of %s because it no longer exist on the bridge",
                    ent.entity_id,
                )
                continue
            new_unique_id = hue_group.grouped_light
            LOGGER.info(
                "Migrating %s from unique id %s to %s ",
                ent.entity_id,
                ent.unique_id,
                new_unique_id,
            )
            try:
                ent_reg.async_update_entity(ent.entity_id, new_unique_id=new_unique_id)
            except ValueError:
                # assume edge case where the entity was already migrated in a previous run
                # which got aborted somehow and we do not want
                # to crash the entire integration init
                LOGGER.warning(
                    "Skip migration of %s because it already exists",
                    ent.entity_id,
                )
    LOGGER.info("Migration of devices and entities to support API schema 2 finished")