[Dataprotection] az dataprotection backup-instance update-msi-permissions: Fix --operation Restore for --datasource-type AzureCosmosDB

src/dataprotection/HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+1.11.1
+++++++
+* Fixed `az dataprotection backup-instance update-msi-permissions --datasource-type AzureCosmosDB --operation Restore` erroring with "Set permissions for restore is currently not supported for given DataSourceType". The command now correctly assigns `Cosmos DB Operator` on the target Cosmos DB account to the backup vault's managed identity. Added live regression test `test_dataprotection_update_msi_permissions_cosmosdb_restore`.
+
 1.11.0
 ++++++
 * Added dataprotection support for AzureCosmosDB workload: new manifest (Microsoft.DocumentDB/databaseAccounts), registration in supported datasource types, datasource map, and permission help text. Added end-to-end backup/restore scenario test, unit tests for default policy template, backup-instance initialize, and restore initialize, and an update-msi-permissions live test that grants Reader/Cosmos DB Operator on the data source RG / account.

src/dataprotection/azext_dataprotection/manual/custom.py

@@ -550,7 +550,7 @@ def dataprotection_backup_instance_update_msi_permissions(cmd, resource_group_na
                 role_assignments_arr.append(helper.get_permission_object_from_server_firewall_rule(rule.result()))
     elif operation == "Restore":
         if datasource_type not in ("AzureKubernetesService", "AzureDatabaseForMySQL",
-                                   "AzureDatabaseForPostgreSQLFlexibleServer"):
+                                   "AzureDatabaseForPostgreSQLFlexibleServer", "AzureCosmosDB"):
             raise InvalidArgumentValueError("Set permissions for restore is currently not supported for given DataSourceType")
 
         for role_object in manifest['backupVaultRestorePermissions']:

src/dataprotection/azext_dataprotection/tests/latest/test_dataprotection_update_msi_permissions.py

@@ -186,6 +186,48 @@ def test_dataprotection_update_msi_permissions_cosmosdb(test):
                  '--backup-instance "{backupInstance}" --yes')
         # time.sleep(10)
 
+    # Uses persistent Cosmos DB accounts pre-provisioned in cosmos-bugbash-CLIrg-6 (subscription
+    # 97cda027-4279-4cde-b4ff-19afa0021d87). The test provisions a fresh backup vault in a
+    # ResourceGroupPreparer-managed RG and exercises update-msi-permissions --operation Restore
+    # for the AzureCosmosDB workload, which assigns Cosmos DB Operator on the target Cosmos DB
+    # account (see Manifests/AzureCosmosDB.py::backupVaultRestorePermissions). This is the
+    # regression test for the fix that adds AzureCosmosDB to the Restore allow-list in
+    # custom.py::dataprotection_backup_instance_update_msi_permissions.
+    @live_only()
+    @AllowLargeResponse()
+    @ResourceGroupPreparer(name_prefix='clitest-dpp-updatemsipermissions-', location='eastus2euap')
+    def test_dataprotection_update_msi_permissions_cosmosdb_restore(test):
+        test.kwargs.update({
+            'location': 'eastus2euap',
+            'vaultName': 'clitest-bkp-vault',
+            'policyName': 'cosmospolicy',
+            'dataSourceType': 'AzureCosmosDB',
+            'operation': 'Restore',
+            'permissionsScope': 'Resource',
+            'sourceDataStore': 'VaultStore',
+            'recoveryPointId': 'dummy-recovery-point-id',
+            'targetCosmosDbId': '/subscriptions/97cda027-4279-4cde-b4ff-19afa0021d87/resourceGroups/cosmos-bugbash-CLIrg-6/providers/Microsoft.DocumentDB/databaseAccounts/cosmosbugbash-cli6-tgt',
+        })
+        create_vault_and_policy(test)
+
+        restore_request_json = test.cmd('az dataprotection backup-instance restore initialize-for-data-recovery '
+                                        '--datasource-type "{dataSourceType}" '
+                                        '--restore-location "{location}" '
+                                        '--source-datastore "{sourceDataStore}" '
+                                        '--recovery-point-id "{recoveryPointId}" '
+                                        '--target-resource-id "{targetCosmosDbId}"').get_output_in_json()
+        test.kwargs.update({
+            "restoreRequest": restore_request_json,
+        })
+
+        test.cmd('az dataprotection backup-instance update-msi-permissions '
+                 '-g "{rg}" --vault-name "{vaultName}" '
+                 '--datasource-type "{dataSourceType}" '
+                 '--operation "{operation}" '
+                 '--permissions-scope "{permissionsScope}" '
+                 '--restore-request-object "{restoreRequest}" --yes')
+        # time.sleep(10)
+
     @live_only()
     @AllowLargeResponse()
     @ResourceGroupPreparer(name_prefix='clitest-dpp-updatemsipermissions-', location='eastus2euap')

src/dataprotection/setup.py

@@ -10,7 +10,7 @@
 from setuptools import setup, find_packages
 
 # HISTORY.rst entry.
-VERSION = '1.11.0'
+VERSION = '1.11.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers