[AKS] az aks create/update: Add ephemeralDisk and elasticSan storag… (#32558)

Author: NedAnd1

src/azure-cli/azure/cli/command_modules/acs/_params.py

@@ -558,7 +558,7 @@ def load_arguments(self, _):
         # azure container storage
         c.argument(
             "enable_azure_container_storage",
-            arg_type=_get_enable_azure_container_storage_type(),
+            arg_type=_get_container_storage_enum_type(storage_pool_types),
             help="enable azure container storage. Can be used as a flag (defaults to True) or with a storage pool type value: (azureDisk, ephemeralDisk, elasticSan)",
         )
         c.argument(
@@ -764,12 +764,12 @@ def load_arguments(self, _):
         # azure container storage
         c.argument(
             "enable_azure_container_storage",
-            arg_type=_get_enable_azure_container_storage_type(),
+            arg_type=_get_container_storage_enum_type(storage_pool_types),
             help="enable azure container storage. Can be used as a flag (defaults to True) or with a storage pool type value: (azureDisk, ephemeralDisk, elasticSan)",
         )
         c.argument(
             "disable_azure_container_storage",
-            arg_type=_get_disable_azure_container_storage_type(),
+            arg_type=_get_container_storage_enum_type(disable_storage_pool_types),
             help="disable azure container storage or any one of the storage pool types. Can be used as a flag (defaults to True) or with a storagepool type value: azureDisk, ephemeralDisk, elasticSan, all (to disable all storage pools).",
         )
         c.argument(
@@ -1286,63 +1286,33 @@ def _get_default_install_location(exe_name):
     return install_location
 
 
-def _get_enable_azure_container_storage_type():
+def _get_container_storage_enum_type(choices):
     """Custom argument type that accepts both None and enum values"""
     import argparse
     from azure.cli.core.azclierror import InvalidArgumentValueError
 
     class AzureContainerStorageAction(argparse.Action):
         def __call__(self, parser, namespace, values, option_string=None):
-            if values is None:
+            if values in [[], None]:
                 # When used as a flag without value, set as True
                 setattr(namespace, self.dest, True)
                 return
 
-            if isinstance(values, str):
-                # Handle enum values (case insensitive)
-                for storage_type in storage_pool_types:
-                    if values.lower() == storage_type.lower():
-                        setattr(namespace, self.dest, storage_type)
-                        return
-
-            # Invalid value
-            valid_values = storage_pool_types
-            raise InvalidArgumentValueError(
-                f"Invalid value '{values}'. Valid values are: {', '.join(valid_values)}"
-            )
-
-    return CLIArgumentType(
-        nargs='?',  # Optional argument
-        action=AzureContainerStorageAction,
-    )
-
-
-def _get_disable_azure_container_storage_type():
-    """Custom argument type that accepts both None and enum values"""
-    import argparse
-    from azure.cli.core.azclierror import InvalidArgumentValueError
-
-    class AzureContainerStorageAction(argparse.Action):
-        def __call__(self, parser, namespace, values, option_string=None):
-            if values is None:
-                # When used as a flag without value, set as True
-                setattr(namespace, self.dest, True)
+            # Allow multiple enum values in a case insensitive manner
+            normalized_value_arr = values if isinstance(values, list) else str(values).split(',')
+            normalized_value_arr = [str(v).lower().strip() for v in normalized_value_arr]
+            valid_value_arr = [v for v in choices if v.lower() in normalized_value_arr]
+            if len(valid_value_arr) == len(normalized_value_arr):
+                normalized_values = valid_value_arr[0] if len(valid_value_arr) == 1 else valid_value_arr
+                setattr(namespace, self.dest, normalized_values)
                 return
 
-            if isinstance(values, str):
-                # Handle enum values (case insensitive)
-                for storage_type in disable_storage_pool_types:
-                    if values.lower() == storage_type.lower():
-                        setattr(namespace, self.dest, storage_type)
-                        return
-
             # Invalid value
-            valid_values = disable_storage_pool_types
             raise InvalidArgumentValueError(
-                f"Invalid value '{values}'. Valid values are: {', '.join(valid_values)}"
+                f"Invalid value '{values}'. Valid values are: {', '.join(choices)}"
             )
 
     return CLIArgumentType(
-        nargs='?',  # Optional argument
+        nargs='*',  # Allow multiple values
         action=AzureContainerStorageAction,
     )

src/azure-cli/azure/cli/command_modules/acs/azurecontainerstorage/_helpers.py

@@ -7,7 +7,9 @@
 
 from azure.cli.command_modules.acs.azurecontainerstorage._consts import (
     CONST_ACSTOR_ALL,
+    CONST_ACSTOR_EXT_INSTALLATION_NAME,
     CONST_ACSTOR_IO_ENGINE_LABEL_KEY,
+    CONST_ACSTOR_K8S_EXTENSION_NAME,
     CONST_ACSTOR_V1_K8S_EXTENSION_NAME,
     CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY,
     CONST_DISK_TYPE_PV_WITH_ANNOTATION,
@@ -146,7 +148,7 @@ def check_if_extension_is_installed(cmd, resource_group, cluster_name) -> bool:
     return return_val
 
 
-def get_extension_installed_and_cluster_configs(
+def get_extension_installed_and_cluster_configs_v1(
     cmd,
     resource_group,
     cluster_name,
@@ -245,6 +247,56 @@ def get_extension_installed_and_cluster_configs(
     )
 
 
+def get_extension_installed_and_cluster_configs(
+    cmd,
+    resource_group,
+    cluster_name
+) -> Tuple[bool, bool, bool]:
+    client_factory = get_k8s_extension_module(CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME)
+    client = client_factory.cf_k8s_extension_operation(cmd.cli_ctx)
+    k8s_extension_custom_mod = get_k8s_extension_module(CONST_K8S_EXTENSION_CUSTOM_MOD_NAME)
+    is_extension_installed = False
+    is_ephemeral_disk_enabled = False
+    is_elastic_san_enabled = False
+
+    try:
+        extension = k8s_extension_custom_mod.show_k8s_extension(
+            client,
+            resource_group,
+            cluster_name,
+            CONST_ACSTOR_EXT_INSTALLATION_NAME,
+            "managedClusters",
+        )
+
+        extension_type = extension.extension_type.lower()
+        is_extension_installed = extension_type == CONST_ACSTOR_K8S_EXTENSION_NAME
+        config_settings = extension.configuration_settings
+
+        if is_extension_installed and config_settings is not None:
+            is_ephemeral_disk_enabled = (
+                config_settings.get("csiDriverConfigs.local-csi-driver.enabled", "False") == "True"
+            )
+            is_elastic_san_enabled = (
+                config_settings.get("csiDriverConfigs.azuresan-csi-driver.enabled", "False") == "True"
+            )
+
+    except:  # pylint: disable=bare-except
+        is_extension_installed = False
+
+    return (
+        is_extension_installed,
+        is_ephemeral_disk_enabled,
+        is_elastic_san_enabled,
+    )
+
+
+def should_delete_extension(storage_options_to_remove) -> bool:
+    return (
+        storage_options_to_remove in [True, CONST_ACSTOR_ALL] or
+        (isinstance(storage_options_to_remove, list) and CONST_ACSTOR_ALL in storage_options_to_remove)
+    )
+
+
 def get_container_storage_extension_installed(
     cmd,
     resource_group,

src/azure-cli/azure/cli/command_modules/acs/azurecontainerstorage/_validators.py

@@ -432,10 +432,12 @@ def validate_enable_azure_container_storage_v1_params(  # pylint: disable=too-ma
 
 
 def validate_enable_azure_container_storage_params(
+    enablement_option,
     is_extension_installed,
+    is_ephemeral_disk_enabled,
+    is_elastic_san_enabled,
     is_v1_extension_installed,
     v1_extension_version,
-    storage_pool_type,
     storage_pool_name,
     storage_pool_sku,
     storage_pool_option,
@@ -450,18 +452,28 @@ def validate_enable_azure_container_storage_params(
             'that depend on Azure Container Storage.'
         )
 
-    if is_extension_installed:
-        raise InvalidArgumentValueError(
-            'Cannot enable Azure Container Storage as it is already enabled on the cluster.'
-        )
-
-    # Todo: Remove this for the 2.1.0 release
-    if storage_pool_type is not None and not isinstance(storage_pool_type, bool):
-        raise InvalidArgumentValueError(
-            'The latest version of Azure Container Storage only supports ephemeral nvme storage and does not '
-            'require or support a storage-pool-type value for --enable-azure-container-storage parameter. '
-            f'Please remove {storage_pool_type} from the command and try again.'
-        )
+    if not enablement_option or enablement_option is True:
+        if is_extension_installed:
+            raise InvalidArgumentValueError(
+                'Cannot enable Azure Container Storage as it is already enabled on the cluster.'
+            )
+    else:
+        enablement_option_arr = enablement_option if isinstance(enablement_option, list) else [enablement_option]
+        enable_ephemeral_disk = CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK in enablement_option
+        enable_elastic_san = CONST_STORAGE_POOL_TYPE_ELASTIC_SAN in enablement_option
+        for option in enablement_option_arr:
+            if option not in [CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK, CONST_STORAGE_POOL_TYPE_ELASTIC_SAN]:
+                raise InvalidArgumentValueError(
+                    f"Unsupported storage option '{option}'. "
+                    f"Supported values are '{CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK}' "
+                    f"and '{CONST_STORAGE_POOL_TYPE_ELASTIC_SAN}'."
+                )
+        if is_ephemeral_disk_enabled == enable_ephemeral_disk and is_elastic_san_enabled == enable_elastic_san:
+            options_display = "', '".join(enablement_option_arr)
+            raise InvalidArgumentValueError(
+                f"Cannot enable the requested storage options ('{options_display}') "
+                "as they are already enabled on the cluster."
+            )
 
     if storage_pool_name is not None:
         raise InvalidArgumentValueError(
@@ -493,25 +505,42 @@ def validate_enable_azure_container_storage_params(
 
 
 def validate_disable_azure_container_storage_params(
+    disablement_option,
     is_extension_installed,
-    storage_pool_type,
+    is_ephemeral_disk_enabled,
+    is_elastic_san_enabled,
     storage_pool_name,
     storage_pool_sku,
     storage_pool_option,
     storage_pool_size
 ):
     if not is_extension_installed:
         raise InvalidArgumentValueError(
-            'Cannot disable Azure Container Storage as it is not enabled on the cluster.'
+            'Cannot disable Azure Container Storage as it could not be found on the cluster.'
         )
 
-    # Todo: Remove this for the 2.1.0 release
-    if storage_pool_type is not None and not isinstance(storage_pool_type, bool):
-        raise InvalidArgumentValueError(
-            'The latest version of Azure Container Storage only supports ephemeral nvme storage and does not '
-            'require or support a storage-pool-type value for --disable-azure-container-storage parameter. '
-            f'Please remove {storage_pool_type} from the command and try again.'
-        )
+    if disablement_option and disablement_option not in [True, CONST_ACSTOR_ALL]:
+        actionable = False
+        disablement_option_arr = disablement_option if isinstance(disablement_option, list) else [disablement_option]
+        for disable_option in disablement_option_arr:
+            if disable_option == CONST_ACSTOR_ALL:
+                actionable = True
+            elif disable_option == CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK:
+                actionable = actionable or is_ephemeral_disk_enabled
+            elif disable_option == CONST_STORAGE_POOL_TYPE_ELASTIC_SAN:
+                actionable = actionable or is_elastic_san_enabled
+            else:
+                raise InvalidArgumentValueError(
+                    f"Cannot disable unsupported storage option '{disable_option}'. "
+                    f"Supported values are '{CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK}', "
+                    f"'{CONST_STORAGE_POOL_TYPE_ELASTIC_SAN}' and '{CONST_ACSTOR_ALL}'."
+                )
+        if not actionable:
+            options_display = "', '".join(disablement_option_arr)
+            raise InvalidArgumentValueError(
+                f"Cannot disable the requested storage options ('{options_display}') "
+                "as they could not be found on the cluster."
+            )
 
     if storage_pool_name is not None:
         raise InvalidArgumentValueError(