[AKS] support csi drivers extensibility

Signed-off-by: Ji An Liu <jiliu8@microsoft.com>

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

@@ -544,6 +544,15 @@
   - name: --azure-keyvault-kms-key-vault-resource-id
     type: string
     short-summary: Resource ID of Azure Key Vault.
+  - name: --disable-disk-driver
+    type: bool
+    short-summary: Disable AzureDisk CSI Driver.
+  - name: --disable-file-driver
+    type: bool
+    short-summary: Disable AzureFile CSI Driver.
+  - name: --disable-snapshot-controller
+    type: bool
+    short-summary: Disable CSI Snapshot Controller.
   - name: --http-proxy-config
     type: string
     short-summary: HTTP Proxy configuration for this cluster.
@@ -796,6 +805,24 @@
   - name: --azure-keyvault-kms-key-vault-resource-id
     type: string
     short-summary: Resource ID of Azure Key Vault.
+  - name: --enable-disk-driver
+    type: bool
+    short-summary: Enable AzureDisk CSI Driver.
+  - name: --disable-disk-driver
+    type: bool
+    short-summary: Disable AzureDisk CSI Driver.
+  - name: --enable-file-driver
+    type: bool
+    short-summary: Enable AzureFile CSI Driver.
+  - name: --disable-file-driver
+    type: bool
+    short-summary: Disable AzureFile CSI Driver.
+  - name: --enable-snapshot-controller
+    type: bool
+    short-summary: Enable Snapshot Controller.
+  - name: --disable-snapshot-controller
+    type: bool
+    short-summary: Disable CSI Snapshot Controller.
   - name: --http-proxy-config
     type: string
     short-summary: HTTP Proxy configuration for this cluster.

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

@@ -290,6 +290,9 @@ def load_arguments(self, _):
         c.argument('node_resource_group')
         c.argument('enable_defender', action='store_true')
         c.argument('defender_config', validator=validate_defender_config_parameter)
+        c.argument('disable_disk_driver', action='store_true')
+        c.argument('disable_file_driver', action='store_true')
+        c.argument('disable_snapshot_controller', action='store_true')
         c.argument('enable_azure_keyvault_kms', action='store_true')
         c.argument('azure_keyvault_kms_key_id', validator=validate_azure_keyvault_kms_key_id)
         c.argument('azure_keyvault_kms_key_vault_network_access', arg_type=get_enum_type(keyvault_network_access_types))
@@ -375,6 +378,12 @@ def load_arguments(self, _):
         c.argument('gmsa_root_domain_name')
         c.argument('attach_acr', acr_arg_type, validator=validate_acr)
         c.argument('detach_acr', acr_arg_type, validator=validate_acr)
+        c.argument('enable_disk_driver', action='store_true')
+        c.argument('disable_disk_driver', action='store_true')
+        c.argument('enable_file_driver', action='store_true')
+        c.argument('disable_file_driver', action='store_true')
+        c.argument('enable_snapshot_controller', action='store_true')
+        c.argument('disable_snapshot_controller', action='store_true')
         c.argument('disable_defender', action='store_true', validator=validate_defender_disable_and_enable_parameters)
         c.argument('enable_defender', action='store_true')
         c.argument('defender_config', validator=validate_defender_config_parameter)

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

@@ -1532,6 +1532,9 @@ def aks_create(
     node_resource_group=None,
     enable_defender=False,
     defender_config=None,
+    disable_disk_driver=False,
+    disable_file_driver=False,
+    disable_snapshot_controller=False,
     enable_azure_keyvault_kms=False,
     azure_keyvault_kms_key_id=None,
     azure_keyvault_kms_key_vault_network_access=None,
@@ -1644,6 +1647,12 @@ def aks_update(
     enable_defender=False,
     disable_defender=False,
     defender_config=None,
+    enable_disk_driver=False,
+    disable_disk_driver=False,
+    enable_file_driver=False,
+    disable_file_driver=False,
+    enable_snapshot_controller=False,
+    disable_snapshot_controller=False,
     enable_azure_keyvault_kms=False,
     disable_azure_keyvault_kms=False,
     azure_keyvault_kms_key_id=None,

src/azure-cli/azure/cli/command_modules/acs/linter_exclusions.yml

@@ -13,6 +13,10 @@ aks create:
       nat_gateway_managed_outbound_ip_count:
           rule_exclusions:
           - option_length_too_long
+      disable_snapshot_controller:
+          rule_exclusions:
+          - option_length_too_long
+
 aks enable-addons:
   parameters:
       appgw_watch_namespace:
@@ -36,6 +40,12 @@ aks update:
           rule_exclusions:
           - option_length_too_long
       assign_kubelet_identity:
+          rule_exclusions:
+          - option_length_too_long    
+      enable_snapshot_controller:
+          rule_exclusions:
+          - option_length_too_long
+      disable_snapshot_controller:
           rule_exclusions:
           - option_length_too_long
 ...

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

@@ -7,7 +7,7 @@
 import re
 import time
 from types import SimpleNamespace
-from typing import Dict, List, Tuple, TypeVar, Union
+from typing import Dict, List, Optional, Tuple, TypeVar, Union
 
 from azure.cli.command_modules.acs._consts import (
     CONST_LOAD_BALANCER_SKU_BASIC,
@@ -100,6 +100,10 @@
 KubeletConfig = TypeVar("KubeletConfig")
 LinuxOSConfig = TypeVar("LinuxOSConfig")
 ManagedClusterSecurityProfileDefender = TypeVar("ManagedClusterSecurityProfileDefender")
+ManagedClusterStorageProfile = TypeVar('ManagedClusterStorageProfile')
+ManagedClusterStorageProfileDiskCSIDriver = TypeVar('ManagedClusterStorageProfileDiskCSIDriver')
+ManagedClusterStorageProfileFileCSIDriver = TypeVar('ManagedClusterStorageProfileFileCSIDriver')
+ManagedClusterStorageProfileSnapshotController = TypeVar('ManagedClusterStorageProfileSnapshotController')
 
 # TODO
 # 1. remove enable_rbac related implementation
@@ -515,6 +519,132 @@ def get_kubernetes_version(self) -> str:
         """
         return self.agentpool_context.get_kubernetes_version()
 
+    def get_disk_driver(self) -> Optional[ManagedClusterStorageProfileDiskCSIDriver]:
+        """Obtain the value of storage_profile.disk_csi_driver
+
+        :return: Optional[ManagedClusterStorageProfileDiskCSIDriver]
+        """
+        enable_disk_driver = self.raw_param.get("enable_disk_driver")
+        disable_disk_driver = self.raw_param.get("disable_disk_driver")
+
+        if not enable_disk_driver and not disable_disk_driver:
+            return None
+        profile = self.models.ManagedClusterStorageProfileDiskCSIDriver()
+
+        if enable_disk_driver and disable_disk_driver:
+            raise MutuallyExclusiveArgumentError(
+                "Cannot specify --enable-disk-driver and "
+                "--disable-disk-driver at the same time."
+            )
+
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if disable_disk_driver:
+                profile.enabled = False
+            else:
+                profile.enabled = True
+
+        if self.decorator_mode == DecoratorMode.UPDATE:
+            if enable_disk_driver:
+                profile.enabled = True
+            elif disable_disk_driver:
+                msg = (
+                    "Please make sure there are no existing PVs and PVCs "
+                    "that are used by AzureDisk CSI driver before disabling."
+                )
+                if not self.get_yes() and not prompt_y_n(msg, default="n"):
+                    raise DecoratorEarlyExitException()
+                profile.enabled = False
+
+        return profile
+
+    def get_file_driver(self) -> Optional[ManagedClusterStorageProfileFileCSIDriver]:
+        """Obtain the value of storage_profile.file_csi_driver
+
+        :return: Optional[ManagedClusterStorageProfileFileCSIDriver]
+        """
+        enable_file_driver = self.raw_param.get("enable_file_driver")
+        disable_file_driver = self.raw_param.get("disable_file_driver")
+
+        if not enable_file_driver and not disable_file_driver:
+            return None
+        profile = self.models.ManagedClusterStorageProfileFileCSIDriver()
+
+        if enable_file_driver and disable_file_driver:
+            raise MutuallyExclusiveArgumentError(
+                "Cannot specify --enable-file-driver and "
+                "--disable-file-driver at the same time."
+            )
+
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if disable_file_driver:
+                profile.enabled = False
+
+        if self.decorator_mode == DecoratorMode.UPDATE:
+            if enable_file_driver:
+                profile.enabled = True
+            elif disable_file_driver:
+                msg = (
+                    "Please make sure there are no existing PVs and PVCs "
+                    "that are used by AzureFile CSI driver before disabling."
+                )
+                if not self.get_yes() and not prompt_y_n(msg, default="n"):
+                    raise DecoratorEarlyExitException()
+                profile.enabled = False
+
+        return profile
+
+    def get_snapshot_controller(self) -> Optional[ManagedClusterStorageProfileSnapshotController]:
+        """Obtain the value of storage_profile.snapshot_controller
+
+        :return: Optional[ManagedClusterStorageProfileSnapshotController]
+        """
+        enable_snapshot_controller = self.raw_param.get("enable_snapshot_controller")
+        disable_snapshot_controller = self.raw_param.get("disable_snapshot_controller")
+
+        if not enable_snapshot_controller and not disable_snapshot_controller:
+            return None
+
+        profile = self.models.ManagedClusterStorageProfileSnapshotController()
+
+        if enable_snapshot_controller and disable_snapshot_controller:
+            raise MutuallyExclusiveArgumentError(
+                "Cannot specify --enable-snapshot_controller and "
+                "--disable-snapshot_controller at the same time."
+            )
+
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if disable_snapshot_controller:
+                profile.enabled = False
+
+        if self.decorator_mode == DecoratorMode.UPDATE:
+            if enable_snapshot_controller:
+                profile.enabled = True
+            elif disable_snapshot_controller:
+                msg = (
+                    "Please make sure there are no existing "
+                    "VolumeSnapshots, VolumeSnapshotClasses and VolumeSnapshotContents "
+                    "that are used by the snapshot controller before disabling."
+                )
+                if not self.get_yes() and not prompt_y_n(msg, default="n"):
+                    raise DecoratorEarlyExitException()
+                profile.enabled = False
+
+        return profile
+
+    def get_storage_profile(self) -> Optional[ManagedClusterStorageProfile]:
+        """Obtain the value of storage_profile.
+
+        :return: Optional[ManagedClusterStorageProfile]
+        """
+        profile = self.models.ManagedClusterStorageProfile()
+        if self.mc.storage_profile is not None:
+            profile = self.mc.storage_profile
+        profile.disk_csi_driver = self.get_disk_driver()
+        profile.file_csi_driver = self.get_file_driver()
+        profile.snapshot_controller = self.get_snapshot_controller()
+
+        return profile
+
     def get_vnet_subnet_id(self) -> Union[str, None]:
         """Obtain the value of vnet_subnet_id.
 
@@ -4388,6 +4518,18 @@ def set_up_windows_profile(self, mc: ManagedCluster) -> ManagedCluster:
             mc.windows_profile = windows_profile
         return mc
 
+    def set_up_storage_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Set up storage profile for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        if hasattr(self.models, "ManagedClusterStorageProfile"):
+            mc.storage_profile = self.context.get_storage_profile()
+
+        return mc
+
     def set_up_service_principal_profile(self, mc: ManagedCluster) -> ManagedCluster:
         """Set up service principal profile for the ManagedCluster object.
 
@@ -5118,6 +5260,8 @@ def construct_mc_profile_default(self, bypass_restore_defaults: bool = False) ->
         mc = self.set_up_node_resource_group(mc)
         # set up defender
         mc = self.set_up_defender(mc)
+        # set up storage profile
+        mc = self.set_up_storage_profile(mc)
         # set up azure keyvalut kms
         mc = self.set_up_azure_keyvault_kms(mc)
         mc = self.set_up_http_proxy_config(mc)
@@ -5931,6 +6075,17 @@ def update_addon_profiles(self, mc: ManagedCluster) -> ManagedCluster:
             ] = azure_keyvault_secrets_provider_addon_profile
         return mc
 
+    def update_storage_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Update storage profile for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        mc.storage_profile = self.context.get_storage_profile()
+
+        return mc
+
     def update_defender(self, mc: ManagedCluster) -> ManagedCluster:
         """Update defender for the ManagedCluster object.
         :return: the ManagedCluster object
@@ -6062,6 +6217,8 @@ def update_mc_profile_default(self) -> ManagedCluster:
         mc = self.update_identity(mc)
         # update addon profiles
         mc = self.update_addon_profiles(mc)
+        # update stroage profile
+        mc = self.update_storage_profile(mc)
         # update defender
         mc = self.update_defender(mc)
         # update azure keyvalut kms