Add migration support for Unraid integration and update versioning

custom_components/unraid/__init__.py

@@ -16,6 +16,7 @@
 from cryptography.utils import CryptographyDeprecationWarning  # type: ignore
 from homeassistant.helpers.importlib import async_import_module  # type: ignore
 
+from .migrations import async_migrate_with_rollback
 from .const import (
     CONF_HOSTNAME,
     DOMAIN,
@@ -25,6 +26,7 @@
     DEFAULT_DISK_INTERVAL,
     DEFAULT_PORT,
     CONF_HAS_UPS,
+    MIGRATION_VERSION
 )
 
 # Suppress deprecation warnings for paramiko
@@ -46,6 +48,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     _LOGGER.debug("Setting up Unraid integration")
 
     try:
+        # Migrate data if necessary
+        if entry.version < MIGRATION_VERSION:
+            if not await async_migrate_with_rollback(hass, entry):
+                raise ConfigEntryNotReady("Migration failed")
+
         # Import required modules asynchronously
         modules = {}
         for module_name in ["unraid", "coordinator", "services", "migrations"]:
@@ -105,9 +112,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             except Exception as hostname_err:
                 _LOGGER.warning("Could not get hostname: %s", hostname_err)
 
-        # Run entity migrations before setting up new entities
-        await modules["migrations"].async_migrate_entities(hass, entry)
-
         # Create coordinator using imported module
         coordinator = modules["coordinator"].UnraidDataUpdateCoordinator(hass, api, entry)
 

custom_components/unraid/config_flow.py

@@ -7,10 +7,13 @@
 
 import voluptuous as vol # type: ignore
 from homeassistant import config_entries # type: ignore
+from homeassistant.config_entries import ConfigEntry # type: ignore
 from homeassistant.const import CONF_HOST, CONF_USERNAME, CONF_PASSWORD, CONF_PORT # type: ignore
 from homeassistant.data_entry_flow import FlowResult # type: ignore
 from homeassistant.exceptions import HomeAssistantError # type: ignore
-from homeassistant.core import callback # type: ignore
+from homeassistant.core import HomeAssistant, callback # type: ignore
+
+from .migrations import async_migrate_with_rollback
 
 from .const import (
     DOMAIN,
@@ -24,6 +27,7 @@
     MIN_DISK_INTERVAL,
     MAX_DISK_INTERVAL,
     CONF_HAS_UPS,
+    MIGRATION_VERSION,
 )
 from .unraid import UnraidAPI
 
@@ -148,7 +152,7 @@ async def validate_input(data: dict[str, Any]) -> dict[str, Any]:
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Unraid."""
 
-    VERSION = 1
+    VERSION = MIGRATION_VERSION
 
     def __init__(self) -> None:
         """Initialize the config flow."""
@@ -247,6 +251,25 @@ async def async_step_user(
     async def async_step_import(self, import_data: dict[str, Any]) -> FlowResult:
         """Handle import from configuration.yaml."""
         return await self.async_step_user(import_data)
+
+    async def async_migrate_entry(self, hass: HomeAssistant, entry: ConfigEntry) -> bool:
+        """Handle migration of config entries."""
+        _LOGGER.debug("Migrating from version %s", entry.version)
+
+        if entry.version == 1:
+            try:
+                if await async_migrate_with_rollback(hass, entry):
+                    # Update entry version after successful migration
+                    self.hass.config_entries.async_update_entry(
+                        entry,
+                        version=MIGRATION_VERSION
+                    )
+                    return True
+            except Exception as err:
+                _LOGGER.error("Migration failed: %s", err)
+                return False
+
+        return True
 
     @staticmethod
     def async_get_options_flow(

custom_components/unraid/const.py

@@ -14,6 +14,9 @@
 DOMAIN = "unraid"
 DEFAULT_PORT = 22
 
+# Migration version
+MIGRATION_VERSION = 2
+
 # Unraid Server Hostname
 CONF_HOSTNAME = "hostname"
 MAX_HOSTNAME_LENGTH = 32

custom_components/unraid/migrations.py

@@ -3,6 +3,7 @@
 
 import logging
 import re
+from typing import Dict, Any
 
 from homeassistant.core import HomeAssistant # type: ignore
 from homeassistant.helpers import entity_registry as er # type: ignore
@@ -12,96 +13,124 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-def clean_entity_id(entity_id: str, hostname: str) -> str:
-    """Clean entity ID by removing duplicate hostname occurrences.
+def clean_and_validate_unique_id(unique_id: str, hostname: str) -> str:
+    """Clean and validate unique ID format.
     
     Args:
-        entity_id: Original entity ID
-        hostname: Current hostname
+        unique_id: Original unique ID
+        hostname: Server hostname
         
     Returns:
-        Cleaned entity ID without hostname duplicates
+        Cleaned and validated unique ID
     """
-    # Convert everything to lowercase for comparison
     hostname_lower = hostname.lower()
+    parts = unique_id.split('_')
 
-    # Split into parts
-    parts = entity_id.split('_')
+    # Remove hostname duplicates while preserving structure
+    cleaned_parts = [part for part in parts if part.lower() != hostname_lower]
 
-    # Remove any parts that are exactly the hostname
-    cleaned_parts = [part for part in parts if part != hostname_lower]
+    # Ensure proper structure
+    if not cleaned_parts[0:2] == ['unraid', 'server']:
+        cleaned_parts = ['unraid', 'server'] + cleaned_parts
+
+    # Insert hostname after unraid_server prefix
+    if len(cleaned_parts) >= 2:
+        cleaned_parts.insert(2, hostname_lower)
 
-    # Reconstruct the ID
     return '_'.join(cleaned_parts)
 
-async def async_migrate_entities(hass: HomeAssistant, entry: ConfigEntry) -> None:
-    """Migrate entities to use hostname-based naming."""
-    ent_reg = er.async_get(hass)
-    hostname = entry.data.get(CONF_HOSTNAME, DEFAULT_NAME).capitalize()
+async def perform_rollback(
+    hass: HomeAssistant,
+    ent_reg: er.EntityRegistry,
+    original_states: Dict[str, Any]
+) -> None:
+    """Rollback entities to original state."""
+    for entity_id, original in original_states.items():
+        try:
+            ent_reg.async_update_entity(
+                entity_id,
+                new_unique_id=original["unique_id"],
+                new_name=original["name"]
+            )
+        except Exception as err:
+            _LOGGER.error("Rollback failed for %s: %s", entity_id, err)
 
-    # Find all entities for this config entry
-    entity_entries = er.async_entries_for_config_entry(ent_reg, entry.entry_id)
+async def update_entity_safely(
+    ent_reg: er.EntityRegistry,
+    entity_entry: er.RegistryEntry,
+    new_unique_id: str
+) -> None:
+    """Safely update entity with new unique ID."""
+    # Check for existing entity with new ID
+    existing = ent_reg.async_get_entity_id(
+        entity_entry.domain,
+        DOMAIN,
+        new_unique_id
+    )
 
-    migrated_count = 0
+    if existing and existing != entity_entry.entity_id:
+        _LOGGER.debug("Removing duplicate entity %s", existing)
+        ent_reg.async_remove(existing)
 
-    for entity_entry in entity_entries:
-        try:
-            # Check if entity needs migration (has IP-based naming or duplicate hostname)
-            ip_pattern = r'unraid_server_\d+_\d+_\d+_\d+_'
-            needs_migration = bool(re.search(ip_pattern, entity_entry.unique_id))
-
-            # Also check for duplicate hostname
-            hostname_count = entity_entry.unique_id.lower().count(hostname.lower())
-            needs_migration = needs_migration or hostname_count > 1
-
-            if needs_migration:
-                # Clean the entity type by removing any hostname duplicates
-                cleaned_id = clean_entity_id(entity_entry.unique_id, hostname)
-                # Ensure we have the correct prefix
-                if not cleaned_id.startswith("unraid_server_"):
-                    cleaned_id = f"unraid_server_{cleaned_id}"
-                # Create new unique_id with single hostname instance
-                new_unique_id = f"unraid_server_{hostname}_{cleaned_id.split('_')[-1]}"
+    ent_reg.async_update_entity(
+        entity_entry.entity_id,
+        new_unique_id=new_unique_id
+    )
+
+async def async_migrate_with_rollback(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate entities with rollback support."""
+    ent_reg = er.async_get(hass)
+
+    # Store original states for rollback
+    original_states = {
+        entity.entity_id: {
+            "unique_id": entity.unique_id,
+            "name": entity.name
+        }
+        for entity in er.async_entries_for_config_entry(ent_reg, entry.entry_id)
+    }
+
+    try:
+        hostname = entry.data.get(CONF_HOSTNAME, DEFAULT_NAME)
+        migrated_count = 0
+
+        for entity_entry in er.async_entries_for_config_entry(ent_reg, entry.entry_id):
+            try:
+                # Check if migration is needed (has IP-based naming or duplicate hostname)
+                ip_pattern = r'unraid_server_\d+_\d+_\d+_\d+_'
+                needs_migration = bool(re.search(ip_pattern, entity_entry.unique_id))
 
-                if new_unique_id != entity_entry.unique_id:
-                    _LOGGER.debug(
-                        "Migrating entity %s to new unique_id %s",
-                        entity_entry.entity_id,
-                        new_unique_id
-                    )
-
-                    # Check if target ID already exists
-                    existing = ent_reg.async_get_entity_id(
-                        entity_entry.domain,
-                        DOMAIN,
-                        new_unique_id
+                hostname_count = entity_entry.unique_id.lower().count(hostname.lower())
+                needs_migration = needs_migration or hostname_count > 1
+
+                if needs_migration:
+                    new_unique_id = clean_and_validate_unique_id(
+                        entity_entry.unique_id,
+                        hostname
                     )
 
-                    if existing and existing != entity_entry.entity_id:
-                        _LOGGER.debug(
-                            "Removing existing entity %s before migration",
-                            existing
-                        )
-                        ent_reg.async_remove(existing)
-
-                    # Update entity with new unique_id
-                    ent_reg.async_update_entity(
-                        entity_entry.entity_id,
-                        new_unique_id=new_unique_id
-                    )
-                    migrated_count += 1
+                    if new_unique_id != entity_entry.unique_id:
+                        await update_entity_safely(ent_reg, entity_entry, new_unique_id)
+                        migrated_count += 1
 
-        except Exception as err:
-            _LOGGER.error(
-                "Error migrating entity %s: %s",
-                entity_entry.entity_id,
-                err
+            except Exception as entity_err:
+                _LOGGER.error(
+                    "Failed to migrate entity %s: %s",
+                    entity_entry.entity_id,
+                    entity_err
+                )
+                continue
+
+        if migrated_count:
+            _LOGGER.info(
+                "Successfully migrated %s entities for %s",
+                migrated_count,
+                hostname
             )
-            continue
-
-    if migrated_count > 0:
-        _LOGGER.info(
-            "Successfully migrated %s entities to use hostname %s",
-            migrated_count,
-            hostname
-        )
+
+        return True
+
+    except Exception as err:
+        _LOGGER.error("Migration failed, rolling back: %s", err)
+        await perform_rollback(hass, ent_reg, original_states)
+        return False

custom_components/unraid/naming.py

@@ -6,9 +6,11 @@
 from typing import Callable, Dict, Union
 from dataclasses import dataclass
 
+from .migrations import clean_and_validate_unique_id
+
 _LOGGER = logging.getLogger(__name__)
 
-# Patterns for entity name formatting
+# Patterns for entity name formatting - keep unchanged
 ENTITY_NAME_PATTERNS: Dict[str, Union[str, Callable[[str], str]]] = {
     "disk": lambda num: f"Array {num}",
     "cache": lambda _: "Pool Cache",
@@ -40,25 +42,39 @@ class EntityNaming:
 
     def get_entity_id(self, name: str) -> str:
         """Get normalized entity ID."""
+        # Use normalized name
         clean_name = normalize_name(name)
-        entity_id = f"{self.domain}_server_{normalize_name(self.hostname)}_{self.component}_{clean_name}"
-        _LOGGER.debug("Generated entity_id: %s | hostname: %s | component: %s | clean_name: %s",
-                    entity_id, self.hostname, self.component, clean_name)
+
+        # Create base unique ID
+        base_id = f"unraid_server_{self.component}_{clean_name}"
+
+        # Use migration's cleaning function for consistency
+        entity_id = clean_and_validate_unique_id(base_id, self.hostname)
+
+        _LOGGER.debug(
+            "Generated entity_id: %s | hostname: %s | component: %s | clean_name: %s",
+            entity_id, self.hostname, self.component, clean_name
+        )
         return entity_id
 
     def get_entity_name(self, name: str, component_type: str = None) -> str:
         """Get formatted entity name."""
         if component_type and component_type in ENTITY_NAME_PATTERNS:
             pattern = ENTITY_NAME_PATTERNS[component_type]
             entity_name = pattern(name) if callable(pattern) else pattern
-            _LOGGER.debug("Generated entity_name: %s | hostname: %s | name: %s | component_type: %s",
-                        entity_name, self.hostname, name, component_type)
+            _LOGGER.debug(
+                "Generated entity_name: %s | hostname: %s | name: %s | component_type: %s",
+                entity_name, self.hostname, name, component_type
+            )
             return entity_name
+
         entity_name = name.title()
-        _LOGGER.debug("Generated entity_name: %s | name: %s | component_type: %s",
-                    entity_name, name, component_type)
+        _LOGGER.debug(
+            "Generated entity_name: %s | name: %s | component_type: %s",
+            entity_name, name, component_type
+        )
         return entity_name
 
     def clean_hostname(self) -> str:
         """Get cleaned hostname."""
-        return self.hostname.capitalize()
+        return normalize_name(self.hostname).capitalize()