Create Google Generative AI sub entries for an enabled entry (home-assistant#148161)

Co-authored-by: Erik Montnemery <[email protected]>

Author: joostlek

homeassistant/components/google_generative_ai_conversation/__init__.py

@@ -195,11 +195,15 @@ async def async_update_options(
 async def async_migrate_integration(hass: HomeAssistant) -> None:
     """Migrate integration entry structure."""
 
-    entries = hass.config_entries.async_entries(DOMAIN)
+    # Make sure we get enabled config entries first
+    entries = sorted(
+        hass.config_entries.async_entries(DOMAIN),
+        key=lambda e: e.disabled_by is not None,
+    )
     if not any(entry.version == 1 for entry in entries):
         return
 
-    api_keys_entries: dict[str, ConfigEntry] = {}
+    api_keys_entries: dict[str, tuple[ConfigEntry, bool]] = {}
     entity_registry = er.async_get(hass)
     device_registry = dr.async_get(hass)
 
@@ -213,9 +217,14 @@ async def async_migrate_integration(hass: HomeAssistant) -> None:
         )
         if entry.data[CONF_API_KEY] not in api_keys_entries:
             use_existing = True
-            api_keys_entries[entry.data[CONF_API_KEY]] = entry
+            all_disabled = all(
+                e.disabled_by is not None
+                for e in entries
+                if e.data[CONF_API_KEY] == entry.data[CONF_API_KEY]
+            )
+            api_keys_entries[entry.data[CONF_API_KEY]] = (entry, all_disabled)
 
-        parent_entry = api_keys_entries[entry.data[CONF_API_KEY]]
+        parent_entry, all_disabled = api_keys_entries[entry.data[CONF_API_KEY]]
 
         hass.config_entries.async_add_subentry(parent_entry, subentry)
         if use_existing:
@@ -228,25 +237,51 @@ async def async_migrate_integration(hass: HomeAssistant) -> None:
                     unique_id=None,
                 ),
             )
-        conversation_entity = entity_registry.async_get_entity_id(
+        conversation_entity_id = entity_registry.async_get_entity_id(
             "conversation",
             DOMAIN,
             entry.entry_id,
         )
-        if conversation_entity is not None:
+        device = device_registry.async_get_device(
+            identifiers={(DOMAIN, entry.entry_id)}
+        )
+
+        if conversation_entity_id is not None:
+            conversation_entity_entry = entity_registry.entities[conversation_entity_id]
+            entity_disabled_by = conversation_entity_entry.disabled_by
+            if (
+                entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY
+                and not all_disabled
+            ):
+                # Device and entity registries don't update the disabled_by flag
+                # when moving a device or entity from one config entry to another,
+                # so we need to do it manually.
+                entity_disabled_by = (
+                    er.RegistryEntryDisabler.DEVICE
+                    if device
+                    else er.RegistryEntryDisabler.USER
+                )
             entity_registry.async_update_entity(
-                conversation_entity,
+                conversation_entity_id,
                 config_entry_id=parent_entry.entry_id,
                 config_subentry_id=subentry.subentry_id,
+                disabled_by=entity_disabled_by,
                 new_unique_id=subentry.subentry_id,
             )
 
-        device = device_registry.async_get_device(
-            identifiers={(DOMAIN, entry.entry_id)}
-        )
         if device is not None:
+            # Device and entity registries don't update the disabled_by flag when
+            # moving a device or entity from one config entry to another, so we
+            # need to do it manually.
+            device_disabled_by = device.disabled_by
+            if (
+                device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY
+                and not all_disabled
+            ):
+                device_disabled_by = dr.DeviceEntryDisabler.USER
             device_registry.async_update_device(
                 device.id,
+                disabled_by=device_disabled_by,
                 new_identifiers={(DOMAIN, subentry.subentry_id)},
                 add_config_subentry_id=subentry.subentry_id,
                 add_config_entry_id=parent_entry.entry_id,
@@ -266,12 +301,13 @@ async def async_migrate_integration(hass: HomeAssistant) -> None:
         if not use_existing:
             await hass.config_entries.async_remove(entry.entry_id)
         else:
+            _add_ai_task_subentry(hass, entry)
             hass.config_entries.async_update_entry(
                 entry,
                 title=DEFAULT_TITLE,
                 options={},
                 version=2,
-                minor_version=2,
+                minor_version=4,
             )
 
 
@@ -315,19 +351,58 @@ async def async_migrate_entry(
 
     if entry.version == 2 and entry.minor_version == 2:
         # Add AI Task subentry with default options
-        hass.config_entries.async_add_subentry(
-            entry,
-            ConfigSubentry(
-                data=MappingProxyType(RECOMMENDED_AI_TASK_OPTIONS),
-                subentry_type="ai_task_data",
-                title=DEFAULT_AI_TASK_NAME,
-                unique_id=None,
-            ),
-        )
+        _add_ai_task_subentry(hass, entry)
         hass.config_entries.async_update_entry(entry, minor_version=3)
 
+    if entry.version == 2 and entry.minor_version == 3:
+        # Fix migration where the disabled_by flag was not set correctly.
+        # We can currently only correct this for enabled config entries,
+        # because migration does not run for disabled config entries. This
+        # is asserted in tests, and if that behavior is changed, we should
+        # correct also disabled config entries.
+        device_registry = dr.async_get(hass)
+        entity_registry = er.async_get(hass)
+        devices = dr.async_entries_for_config_entry(device_registry, entry.entry_id)
+        entity_entries = er.async_entries_for_config_entry(
+            entity_registry, entry.entry_id
+        )
+        if entry.disabled_by is None:
+            # If the config entry is not disabled, we need to set the disabled_by
+            # flag on devices to USER, and on entities to DEVICE, if they are set
+            # to CONFIG_ENTRY.
+            for device in devices:
+                if device.disabled_by is not dr.DeviceEntryDisabler.CONFIG_ENTRY:
+                    continue
+                device_registry.async_update_device(
+                    device.id,
+                    disabled_by=dr.DeviceEntryDisabler.USER,
+                )
+            for entity in entity_entries:
+                if entity.disabled_by is not er.RegistryEntryDisabler.CONFIG_ENTRY:
+                    continue
+                entity_registry.async_update_entity(
+                    entity.entity_id,
+                    disabled_by=er.RegistryEntryDisabler.DEVICE,
+                )
+        hass.config_entries.async_update_entry(entry, minor_version=4)
+
     LOGGER.debug(
         "Migration to version %s:%s successful", entry.version, entry.minor_version
     )
 
     return True
+
+
+def _add_ai_task_subentry(
+    hass: HomeAssistant, entry: GoogleGenerativeAIConfigEntry
+) -> None:
+    """Add AI Task subentry to the config entry."""
+    hass.config_entries.async_add_subentry(
+        entry,
+        ConfigSubentry(
+            data=MappingProxyType(RECOMMENDED_AI_TASK_OPTIONS),
+            subentry_type="ai_task_data",
+            title=DEFAULT_AI_TASK_NAME,
+            unique_id=None,
+        ),
+    )

homeassistant/components/google_generative_ai_conversation/config_flow.py

@@ -97,7 +97,7 @@ class GoogleGenerativeAIConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Google Generative AI Conversation."""
 
     VERSION = 2
-    MINOR_VERSION = 3
+    MINOR_VERSION = 4
 
     async def async_step_api(
         self, user_input: dict[str, Any] | None = None