[dataprotection] Added UAMI support while setting permissions for AKS…

… backup/restore (#6928) * UAMI support for AKS permissions * Updated version * Static analysis fix --------- Co-authored-by: Zubair Abid <zubairabid@microsoft.com>
Azure · Nov 1, 2023 · f16ed00 · f16ed00
1 parent 3d70a49
commit f16ed00
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/src/dataprotection/HISTORY.rst b/src/dataprotection/HISTORY.rst
@@ -2,6 +2,10 @@
 
 Release History
 ===============
+0.11.2
+++++++
+* `az dataprotection backup-instance update-msi-permissions`: Added UAMI support for AKS backup/restore.
+
 0.11.1
 ++++++
 * Added '-v' option for all --vault-name parameters

diff --git a/src/dataprotection/azext_dataprotection/manual/custom.py b/src/dataprotection/azext_dataprotection/manual/custom.py
@@ -397,7 +397,14 @@ def dataprotection_backup_instance_update_msi_permissions(cmd, resource_group_na
                     aks_rg_id = helper.get_rg_id_from_arm_id(datasource_arm_id)
                     aks_rg = aks_rg_id.split('/')[-1]
                     aks_cluster = aks_client.get(aks_rg, aks_name)
-                    datasource_principal_id = aks_cluster.identity.principal_id
+
+                    if "UserAssigned" in aks_cluster.identity.type:
+                        uami_key = list(aks_cluster.identity.user_assigned_identities.keys())[0]
+                        if uami_key == "" or uami_key is None:
+                            raise CLIInternalError("User assigned identity not found for AKS Cluster")
+                        datasource_principal_id = aks_cluster.identity.user_assigned_identities[uami_key].principal_id
+                    else:
+                        datasource_principal_id = aks_cluster.identity.principal_id
                 else:
                     raise InvalidArgumentValueError("Datasource-over-X permissions can currently only be set for Datasource type AzureKubernetesService")
 
@@ -490,7 +497,14 @@ def dataprotection_backup_instance_update_msi_permissions(cmd, resource_group_na
                     aks_rg_id = helper.get_rg_id_from_arm_id(datasource_arm_id)
                     aks_rg = aks_rg_id.split('/')[-1]
                     aks_cluster = aks_client.get(aks_rg, aks_name)
-                    datasource_principal_id = aks_cluster.identity.principal_id
+
+                    if "UserAssigned" in aks_cluster.identity.type:
+                        uami_key = list(aks_cluster.identity.user_assigned_identities.keys())[0]
+                        if uami_key == "" or uami_key is None:
+                            raise CLIInternalError("User assigned identity not found for AKS Cluster")
+                        datasource_principal_id = aks_cluster.identity.user_assigned_identities[uami_key].principal_id
+                    else:
+                        datasource_principal_id = aks_cluster.identity.principal_id
                 else:
                     raise InvalidArgumentValueError("Datasource-over-X permissions can currently only be set for Datasource type AzureKubernetesService")
 

diff --git a/src/dataprotection/setup.py b/src/dataprotection/setup.py
@@ -10,7 +10,7 @@
 from setuptools import setup, find_packages
 
 # HISTORY.rst entry.
-VERSION = '0.11.1'
+VERSION = '0.11.2'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers