AKS: Change --enable-azure-container-storage --disable-azure-container -storage behavior and add --container-storage-version

src/aks-preview/HISTORY.rst

@@ -11,7 +11,15 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+
+18.0.0b40
++++++++
 * Add option `Windows2025` to `--os-sku` for `az aks nodepool add`.
+* `az aks create`: Add new parameter `--container-storage-version` to enable the given version of Azure Container Storage.
+* `az aks update`: Add new parameter `--container-storage-version` to enable the given version of Azure Container Storage.
+* `az aks create`: Change behavior of `--enable-azure-container-storage` to enable latest Azure Container Storage by default.
+* `az aks update`: Change behavior of `--enable-azure-container-storage` to enable latest Azure Container Storage by default.
+* `az aks update`: Change behavior of `--disable-azure-container-storage` to handle disable based on the installed version.
 
 18.0.0b39
 +++++++

src/aks-preview/azext_aks_preview/_params.py

@@ -441,6 +441,12 @@
 ]
 
 # azure container storage
+# azure container storage
+container_storage_versions = [
+    "1",
+    "2"
+]
+
 storage_pool_types = [
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
@@ -1028,8 +1034,14 @@ def load_arguments(self, _):
         # azure container storage
         c.argument(
             "enable_azure_container_storage",
-            arg_type=get_enum_type(storage_pool_types),
-            help="enable azure container storage and define storage pool type",
+            arg_type=_get_enable_azure_container_storage_type(),
+            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(
+            "container_storage_version",
+            arg_type=get_enum_type(container_storage_versions),
+            help="set azure container storage version, the latest version will be installed by default",
         )
         c.argument(
             "storage_pool_name",
@@ -1488,13 +1500,21 @@ def load_arguments(self, _):
         # azure container storage
         c.argument(
             "enable_azure_container_storage",
-            arg_type=get_enum_type(storage_pool_types),
-            help="enable azure container storage and define storage pool type",
+            arg_type=_get_enable_azure_container_storage_type(),
+            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_enum_type(disable_storage_pool_types),
-            help="disable azure container storage or any one of the storage pool types",
+            arg_type=_get_disable_azure_container_storage_type(),
+            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(
+            "container_storage_version",
+            arg_type=get_enum_type(container_storage_versions),
+            help="set azure container storage version, the latest version will be installed by default",
         )
         c.argument(
             "storage_pool_name",
@@ -2922,3 +2942,65 @@ def _get_default_install_location(exe_name):
     else:
         install_location = None
     return install_location
+
+
+def _get_enable_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)
+                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)
+                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)}"
+            )
+
+    return CLIArgumentType(
+        nargs='?',  # Optional argument
+        action=AzureContainerStorageAction,
+    )

src/aks-preview/azext_aks_preview/azurecontainerstorage/_consts.py

@@ -6,13 +6,17 @@
 CONST_ACSTOR_ALL = "all"
 CONST_ACSTOR_IO_ENGINE_LABEL_KEY = "acstor.azure.com/io-engine"
 CONST_ACSTOR_IO_ENGINE_LABEL_VAL = "acstor"
-CONST_ACSTOR_K8S_EXTENSION_NAME = "microsoft.azurecontainerstorage"
+CONST_ACSTOR_V1_K8S_EXTENSION_NAME = "microsoft.azurecontainerstorage"
+CONST_ACSTOR_V1_EXT_INSTALLATION_NAME = "azurecontainerstorage"
+CONST_ACSTOR_VERSION_V1 = "1"
 CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY = "EphemeralVolumeOnly"
 CONST_DISK_TYPE_PV_WITH_ANNOTATION = "PersistentVolumeWithAnnotation"
 CONST_EPHEMERAL_NVME_PERF_TIER_BASIC = "Basic"
 CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM = "Premium"
 CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD = "Standard"
-CONST_EXT_INSTALLATION_NAME = "azurecontainerstorage"
+CONST_ACSTOR_EXT_INSTALLATION_NAME = "acstor"
+CONST_ACSTOR_EXT_INSTALLATION_NAMESPACE = "kube-system"
+CONST_ACSTOR_K8S_EXTENSION_NAME = "microsoft.azurecontainerstoragev2"
 CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME = "azext_k8s_extension._client_factory"
 CONST_K8S_EXTENSION_CUSTOM_MOD_NAME = "azext_k8s_extension.custom"
 CONST_K8S_EXTENSION_NAME = "k8s-extension"

src/aks-preview/azext_aks_preview/azurecontainerstorage/_helpers.py

@@ -8,13 +8,11 @@
 from azext_aks_preview.azurecontainerstorage._consts import (
     CONST_ACSTOR_ALL,
     CONST_ACSTOR_IO_ENGINE_LABEL_KEY,
-    CONST_ACSTOR_K8S_EXTENSION_NAME,
     CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY,
     CONST_DISK_TYPE_PV_WITH_ANNOTATION,
     CONST_EPHEMERAL_NVME_PERF_TIER_BASIC,
     CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM,
     CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD,
-    CONST_EXT_INSTALLATION_NAME,
     CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME,
     CONST_K8S_EXTENSION_CUSTOM_MOD_NAME,
     CONST_K8S_EXTENSION_NAME,
@@ -23,13 +21,15 @@
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
+    CONST_ACSTOR_V1_K8S_EXTENSION_NAME,
+    CONST_ACSTOR_V1_EXT_INSTALLATION_NAME,
 )
 from azure.cli.command_modules.acs._roleassignments import (
     add_role_assignment,
     build_role_scope,
     delete_role_assignments,
 )
-from azure.cli.core.azclierror import UnknownError
+from azure.cli.core.azclierror import ResourceNotFoundError, UnknownError
 from knack.log import get_logger
 
 logger = get_logger(__name__)
@@ -133,12 +133,12 @@ def check_if_extension_is_installed(cmd, resource_group, cluster_name) -> bool:
             client,
             resource_group,
             cluster_name,
-            CONST_EXT_INSTALLATION_NAME,
+            CONST_ACSTOR_V1_EXT_INSTALLATION_NAME,
             "managedClusters",
         )
 
         extension_type = extension.extension_type.lower()
-        if extension_type != CONST_ACSTOR_K8S_EXTENSION_NAME:
+        if extension_type != CONST_ACSTOR_V1_K8S_EXTENSION_NAME:
             return_val = False
     except:  # pylint: disable=bare-except
         return_val = False
@@ -177,11 +177,11 @@ def get_extension_installed_and_cluster_configs(
             client,
             resource_group,
             cluster_name,
-            CONST_EXT_INSTALLATION_NAME,
+            CONST_ACSTOR_V1_EXT_INSTALLATION_NAME,
             "managedClusters",
         )
 
-        is_extension_installed = extension.extension_type.lower() == CONST_ACSTOR_K8S_EXTENSION_NAME
+        is_extension_installed = extension.extension_type.lower() == CONST_ACSTOR_V1_K8S_EXTENSION_NAME
         config_settings = extension.configuration_settings
 
         if is_extension_installed and config_settings is not None:
@@ -252,6 +252,35 @@ def get_extension_installed_and_cluster_configs(
     )
 
 
+def get_container_storage_extension_installed(
+    cmd,
+    resource_group,
+    cluster_name,
+    extension_name,
+) -> Tuple[bool, str]:
+
+    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
+    extension_version = ""
+
+    try:
+        extension = k8s_extension_custom_mod.show_k8s_extension(
+            client,
+            resource_group,
+            cluster_name,
+            extension_name,
+            "managedClusters",
+        )
+        is_extension_installed = True
+        extension_version = extension.current_version
+    except ResourceNotFoundError:
+        # Extension not found, which is expected if not installed.
+        is_extension_installed = False
+    return is_extension_installed, extension_version
+
+
 def get_initial_resource_value_args(
     storage_pool_type,
     storage_pool_option,