AKS: Implement support to perform Azure container storage operations

src/aks-preview/HISTORY.rst

@@ -12,6 +12,11 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+0.5.164
++++++++
+* Add option `--enable-azure-container-storage` and supporting options `--storage-pool-name`, `--storage-pool-type`, `--storage-pool-sku`, `--storage-pool-size` for `az aks create` and `az aks update`. `az aks update` also supports `--azure-container-storage-nodepools` option.
+* Add option `--disable-azure-container-storage` to `az aks create` and `az aks update`.
+
 0.5.163
 +++++++
 * Add `get-upgrades` and `get-revisions` to the `az aks mesh` command.

src/aks-preview/azext_aks_preview/_client_factory.py

@@ -151,3 +151,7 @@ def get_resource_by_name(cli_ctx, resource_name, resource_type):
 def get_msi_client(cli_ctx, subscription_id=None):
     return get_mgmt_service_client(cli_ctx, ManagedServiceIdentityClient,
                                    subscription_id=subscription_id)
+
+
+def get_providers_client_factory(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES, subscription_id=subscription_id).providers

src/aks-preview/azext_aks_preview/_params.py

@@ -169,7 +169,21 @@
     tags_type,
     zones_type,
 )
-from azure.cli.core.profiles import ResourceType
+from azext_aks_preview.azurecontainerstorage._consts import (
+    CONST_STORAGE_POOL_TYPE_AZURE_DISK,
+    CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
+    CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_SKU_PREMIUM_LRS,
+    CONST_STORAGE_POOL_SKU_STANDARD_LRS,
+    CONST_STORAGE_POOL_SKU_STANDARDSSD_LRS,
+    CONST_STORAGE_POOL_SKU_ULTRASSD_LRS,
+    CONST_STORAGE_POOL_SKU_PREMIUM_ZRS,
+    CONST_STORAGE_POOL_SKU_PREMIUMV2_LRS,
+    CONST_STORAGE_POOL_SKU_STANDARDSSD_ZRS,
+    CONST_STORAGE_POOL_OPTION_NVME,
+    CONST_STORAGE_POOL_OPTION_SSD,
+    CONST_STORAGE_POOL_DEFAULT_SIZE,
+)
 from knack.arguments import CLIArgumentType
 
 # candidates for enumeration
@@ -256,6 +270,28 @@
     CONST_AZURE_SERVICE_MESH_INGRESS_MODE_INTERNAL,
 ]
 
+# azure container storage
+storage_pool_types = [
+    CONST_STORAGE_POOL_TYPE_AZURE_DISK,
+    CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
+    CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+]
+
+storage_pool_skus = [
+    CONST_STORAGE_POOL_SKU_PREMIUM_LRS,
+    CONST_STORAGE_POOL_SKU_STANDARD_LRS,
+    CONST_STORAGE_POOL_SKU_STANDARDSSD_LRS,
+    CONST_STORAGE_POOL_SKU_ULTRASSD_LRS,
+    CONST_STORAGE_POOL_SKU_PREMIUM_ZRS,
+    CONST_STORAGE_POOL_SKU_PREMIUMV2_LRS,
+    CONST_STORAGE_POOL_SKU_STANDARDSSD_ZRS,
+]
+
+storage_pool_options = [
+    CONST_STORAGE_POOL_OPTION_NVME,
+    CONST_STORAGE_POOL_OPTION_SSD,
+]
+
 
 def load_arguments(self, _):
 
@@ -455,6 +491,15 @@ def load_arguments(self, _):
         c.argument('grafana_resource_id', validator=validate_grafanaresourceid)
         c.argument('enable_windows_recording_rules', action='store_true')
         c.argument('enable_cost_analysis', is_preview=True, action='store_true')
+        # 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')
+        c.argument('storage_pool_name', help='set storage pool name for azure container storage')
+        c.argument('storage_pool_size', help='set storage pool size for azure container storage')
+        c.argument('storage_pool_sku', arg_type=get_enum_type(storage_pool_skus),
+                   help='set azure disk type storage pool sku for azure container storage')
+        c.argument('storage_pool_option', arg_type=get_enum_type(storage_pool_options),
+                   help='set ephemeral disk storage pool option for azure container storage')
 
     with self.argument_context('aks update') as c:
         # managed cluster paramerters
@@ -583,6 +628,19 @@ def load_arguments(self, _):
         c.argument('enable_network_observability', action='store_true', is_preview=True, help="enable network observability for cluster")
         c.argument('enable_cost_analysis', is_preview=True, action='store_true')
         c.argument('disable_cost_analysis', is_preview=True, action='store_true')
+        # 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')
+        c.argument('disable_azure_container_storage', action='store_true',
+                   help='Flag to disable azure container storage')
+        c.argument('storage_pool_name', help='set storage pool name for azure container storage')
+        c.argument('storage_pool_size', help='set storage pool size for azure container storage')
+        c.argument('storage_pool_sku', arg_type=get_enum_type(storage_pool_skus),
+                   help='set azure disk type storage pool sku for azure container storage')
+        c.argument('storage_pool_option', arg_type=get_enum_type(storage_pool_options),
+                   help='set ephemeral disk storage pool option for azure container storage')
+        c.argument('azure_container_storage_nodepools',
+                   help='define the comma separated nodepool list to install azure container storage')
 
     with self.argument_context('aks upgrade') as c:
         c.argument('kubernetes_version', completer=get_k8s_upgrades_completion_list)

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

@@ -0,0 +1,4 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------

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

@@ -0,0 +1,28 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+CONST_ACSTOR_K8S_EXTENSION_NAME = "microsoft.azurecontainerstorage"
+CONST_EXT_INSTALLATION_NAME = "azurecontainerstorage"
+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"
+CONST_STORAGE_POOL_DEFAULT_SIZE_ESAN = "1Ti"
+CONST_STORAGE_POOL_DEFAULT_SIZE = "512Gi"
+CONST_STORAGE_POOL_NAME_PREFIX = "storagepool-"
+CONST_STORAGE_POOL_OPTION_NVME = "NVMe"
+CONST_STORAGE_POOL_OPTION_SSD = "SSD"
+CONST_STORAGE_POOL_SKU_PREMIUM_LRS = "Premium_LRS"
+CONST_STORAGE_POOL_SKU_PREMIUM_ZRS = "Premium_ZRS"
+CONST_STORAGE_POOL_SKU_PREMIUMV2_LRS = "PremiumV2_LRS"
+CONST_STORAGE_POOL_SKU_STANDARD_LRS = "Standard_LRS"
+CONST_STORAGE_POOL_SKU_STANDARDSSD_LRS = "StandardSSD_LRS"
+CONST_STORAGE_POOL_SKU_STANDARDSSD_ZRS = "StandardSSD_ZRS"
+CONST_STORAGE_POOL_SKU_ULTRASSD_LRS = "UltraSSD_LRS"
+CONST_STORAGE_POOL_TYPE_AZURE_DISK = "azureDisk"
+CONST_STORAGE_POOL_TYPE_ELASTIC_SAN = "elasticSan"
+CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK = "ephemeralDisk"
+
+CONST_STORAGE_POOL_RANDOM_LENGTH = 7
+RP_REGISTRATION_POLLING_INTERVAL_IN_SEC = 5

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

@@ -0,0 +1,206 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.azclierror import UnknownError
+from azure.cli.command_modules.acs._roleassignments import (
+    add_role_assignment,
+    build_role_scope,
+    delete_role_assignments,
+)
+from azext_aks_preview._client_factory import get_providers_client_factory
+from azext_aks_preview.azurecontainerstorage._consts import (
+    CONST_ACSTOR_K8S_EXTENSION_NAME,
+    CONST_EXT_INSTALLATION_NAME,
+    CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME,
+    CONST_K8S_EXTENSION_CUSTOM_MOD_NAME,
+    CONST_K8S_EXTENSION_NAME,
+    CONST_STORAGE_POOL_NAME_PREFIX,
+    CONST_STORAGE_POOL_OPTION_NVME,
+    CONST_STORAGE_POOL_RANDOM_LENGTH,
+    CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
+    RP_REGISTRATION_POLLING_INTERVAL_IN_SEC,
+)
+
+from datetime import datetime
+from knack.log import get_logger
+import random
+import string
+import time
+
+logger = get_logger(__name__)
+
+
+def register_dependent_rps(cmd, subscription_id) -> bool:
+    required_rp = 'Microsoft.KubernetesConfiguration'
+    from azure.mgmt.resource.resources.models import ProviderRegistrationRequest, ProviderConsentDefinition
+
+    properties = ProviderRegistrationRequest(third_party_provider_consent=ProviderConsentDefinition(consent_to_authorization=False))
+    client = get_providers_client_factory(cmd.cli_ctx)
+    is_registered = False
+    try:
+        is_registered = _is_rp_registered(cmd, required_rp, subscription_id)
+        if is_registered:
+            return
+        client.register(required_rp, properties=properties)
+        # wait for registration to finish
+        timeout_secs = 120
+        start = datetime.utcnow()
+        is_registered = _is_rp_registered(cmd, required_rp, subscription_id)
+        while not is_registered:
+            is_registered = _is_rp_registered(cmd, required_rp, subscription_id)
+            time.sleep(RP_REGISTRATION_POLLING_INTERVAL_IN_SEC)
+            if (datetime.utcnow() - start).seconds >= timeout_secs:
+                logger.error("Timed out while waiting for the {0} resource provider to be registered.".format(required_rp))
+                break
+
+    except Exception as e:
+        logger.error(
+            "Installation of Azure Container Storage requires registering to the following resource provider: {0}. "
+            "We were unable to perform the registration on your behalf due to the following error: {1}\n"
+            "Please check with your admin on permissions, "
+            "or try running registration manually with: `az provider register --namespace {0}` command."
+            .format(required_rp, e.msg)
+        )
+
+    return is_registered
+
+
+def should_create_storagepool(
+    cmd,
+    subscription_id,
+    node_resource_group,
+    kubelet_identity_object_id,
+    storage_pool_type,
+    storage_pool_option,
+    agentpool_details,
+    nodepool_name,
+):
+    role_assignment_success = perform_role_operations_on_managed_rg(cmd, subscription_id, node_resource_group, kubelet_identity_object_id, True)
+    return_val = True
+
+    if not role_assignment_success:
+        msg = "\nUnable to add Role Assignments needed for Elastic SAN storagepools to be functional. " \
+            "Please check with your admin on permissions."
+        if storage_pool_type == CONST_STORAGE_POOL_TYPE_ELASTIC_SAN:
+            msg += "\nThis command will not create an Elastic SAN storagepool after installation."
+            return_val = False
+        msg += "\nGoing ahead with the installation of Azure Container Storage..."
+        logger.warning(msg)
+
+    if not return_val:
+        return return_val
+
+    if storage_pool_type == CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK and \
+       storage_pool_option == CONST_STORAGE_POOL_OPTION_NVME:
+        nodepool_list = nodepool_name.split(',')
+        for nodepool in nodepool_list:
+            agentpool_vm = agentpool_details.get(nodepool.lower())
+            if agentpool_vm is not None and agentpool_vm.lower().startswith('standard_l'):
+                break
+        else:
+            logger.warning(
+                "\nNo supporting nodepool found which can support ephemeral NVMe disk "
+                "so this command will not create an ephemeral NVMe disk storage pool after installation."
+                "\nGoing ahead with the installation of Azure Container Storage..."
+            )
+            return_val = False
+
+    return return_val
+
+
+def perform_role_operations_on_managed_rg(cmd, subscription_id, node_resource_group, kubelet_identity_object_id, assign):
+    managed_rg_role_scope = build_role_scope(node_resource_group, None, subscription_id)
+    roles = ["Reader", "Network Contributor", "Elastic SAN Owner", "Elastic SAN Volume Group Owner"]
+    result = True
+
+    for role in roles:
+        try:
+            if assign:
+                result = add_role_assignment(
+                    cmd,
+                    role,
+                    kubelet_identity_object_id,
+                    scope=managed_rg_role_scope,
+                    delay=0,
+                )
+            else:
+                # NOTE: delete_role_assignments accepts cli_ctx
+                # instead of cmd unlike add_role_assignment.
+                result = delete_role_assignments(
+                    cmd.cli_ctx,
+                    role,
+                    kubelet_identity_object_id,
+                    scope=managed_rg_role_scope,
+                    delay=0,
+                )
+
+            if not result:
+                break
+        except Exception as ex:
+            break
+    else:
+        return True
+
+    if not assign:
+        logger.error("\nUnable to revoke Role Assignments if any, added for Azure Container Storage.")
+
+    return False
+
+
+def generate_random_storage_pool_name():
+    random_name = CONST_STORAGE_POOL_NAME_PREFIX + ''.join(random.choices(string.ascii_lowercase, k=CONST_STORAGE_POOL_RANDOM_LENGTH))
+    return random_name
+
+
+def get_k8s_extension_module(module_name):
+    try:
+        # adding the installed extension in the path
+        from azure.cli.core.extension.operations import add_extension_to_path
+        add_extension_to_path(CONST_K8S_EXTENSION_NAME)
+        # import the extension module
+        from importlib import import_module
+        azext_custom = import_module(module_name)
+        return azext_custom
+    except ImportError as ie:
+        raise UnknownError(
+            "Please add CLI extension `k8s-extension` for performing Azure Container Storage operations.\n"
+            "Run command `az extension add --name k8s-extension`"
+        )
+
+
+def check_if_extension_is_installed(cmd, resource_group, cluster_name) -> 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)
+    return_val = True
+    try:
+        extension = k8s_extension_custom_mod.show_k8s_extension(
+            client,
+            resource_group,
+            cluster_name,
+            CONST_EXT_INSTALLATION_NAME,
+            "managedClusters",
+        )
+
+        extension_type = extension.extension_type.lower()
+        if extension_type != CONST_ACSTOR_K8S_EXTENSION_NAME:
+            return_val = False
+    except:
+        return_val = False
+
+    return return_val
+
+
+def _is_rp_registered(cmd, required_rp, subscription_id):
+    registered = False
+    try:
+        providers_client = get_providers_client_factory(cmd.cli_ctx, subscription_id)
+        registration_state = getattr(providers_client.get(required_rp), 'registration_state', "NotRegistered")
+
+        registered = (registration_state and registration_state.lower() == 'registered')
+    except Exception:  # pylint: disable=broad-except
+        pass
+    return registered