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_ZRS,
+    CONST_STORAGE_POOL_SKU_STANDARDSSD_ZRS,
+    CONST_STORAGE_POOL_OPTION_NVME,
+    CONST_STORAGE_POOL_OPTION_TEMP,
+    CONST_STORAGE_POOL_DEFAULT_SIZE,
+)
 from knack.arguments import CLIArgumentType
 
 # candidates for enumeration
@@ -256,6 +270,29 @@
     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_ZRS,
+    CONST_STORAGE_POOL_SKU_STANDARDSSD_ZRS,
+]
+
+storage_pool_options = [
+    CONST_STORAGE_POOL_OPTION_NVME,
+    CONST_STORAGE_POOL_OPTION_TEMP,
+]
+
 
 def load_arguments(self, _):
 
@@ -455,6 +492,17 @@ 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', action='store_true',
+                   help='Flag to enable azure container storage')
+        c.argument('storage_pool_type', arg_type=get_enum_type(storage_pool_types),
+                   help='set storage pool type for 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')
 
     with self.argument_context('aks update') as c:
         # managed cluster paramerters
@@ -583,6 +631,21 @@ 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', action='store_true',
+                   help='Flag to enable azure container storage')
+        c.argument('disable_azure_container_storage', action='store_true',
+                   help='Flag to disable azure container storage')
+        c.argument('storage_pool_type', arg_type=get_enum_type(storage_pool_types),
+                   help='set storage pool type for 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,27 @@
+# --------------------------------------------------------------------------------------------
+# 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 = "512Gi"
+CONST_STORAGE_POOL_NAME_PREFIX = "mypool-"
+CONST_STORAGE_POOL_OPTION_NVME = "LsV2"
+CONST_STORAGE_POOL_OPTION_TEMP = "temp"
+CONST_STORAGE_POOL_SKU_PREMIUM_LRS = "Premium_LRS"
+CONST_STORAGE_POOL_SKU_PREMIUM_ZRS = "Premium_ZRS"
+CONST_STORAGE_POOL_SKU_PREMIUMV2_ZRS = "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,124 @@
+# --------------------------------------------------------------------------------------------
+# 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_K8S_EXTENSION_NAME,
+    CONST_STORAGE_POOL_NAME_PREFIX,
+    CONST_STORAGE_POOL_RANDOM_LENGTH,
+    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):
+    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)
+    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:
+                raise UnknownError("Timed out while waiting for the {0} resource provider to be registered.".format(required_rp))
+
+    except Exception as e:
+        raise UnknownError(
+            "Installation of Azure Container Storage requires registering to the following resource providers: {0}. "
+            "We were unable to perform the registration on your behalf. "
+            "The following error was received: {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)
+        )
+
+
+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
+
+
+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"]
+    exception = False
+
+    for role in roles:
+        try:
+            if assign:
+                add_role_assignment(
+                    cmd,
+                    role,
+                    kubelet_identity_object_id,
+                    scope=managed_rg_role_scope,
+                )
+            else:
+                # NOTE: delete_role_assignments accepts cli_ctx
+                # instead of cmd unlike add_role_assignment.
+                delete_role_assignments(
+                    cmd.cli_ctx,
+                    role,
+                    kubelet_identity_object_id,
+                    scope=managed_rg_role_scope,
+                )
+        except Exception as ex:
+            exception = True
+
+    if assign and exception:
+        logger.warning(
+            "Unable to add Role Assignments needed for Elastic SAN storagepools to be functional. "
+            "Going ahead with the installation of Azure Container Storage..."
+        )
+
+
+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 --name k8s-extension`"
+        )