[dataprotection] Add Cosmos DB item-level recovery

src/dataprotection/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+1.14.0
+++++++
+* ``az dataprotection backup-instance initialize-restoreconfig`` and ``az dataprotection backup-instance restore initialize-for-item-recovery``: Added item-level restore support for AzureCosmosDB databases and containers.
+
 1.13.0
 ++++++
 * ``az dataprotection backup-vault create/update``: Added ``--cost-management-granularity`` (alias ``--cost-granularity``) to configure vault cost management settings. Supported values are ``VaultLevel``, ``ProtectedItemLevel``, and ``ProtectedItemWithParentTag``; vault creation defaults to ``VaultLevel``.

src/dataprotection/README.md

@@ -163,6 +163,47 @@ az dataprotection backup-instance restore trigger \
     --restore-request-object @esan_restore_request.json
 ```
 
+#### Azure Cosmos DB item-level restore scripts ####
+##### Restore selected databases and containers #####
+```
+# 1) Build restore configuration: select database or database/container names to restore
+az dataprotection backup-instance initialize-restoreconfig \
+    --datasource-type AzureCosmosDB \
+    --resource-identifiers db1/container1 db2 > cosmos_restore_config.json
+
+# 2) Initialize restore request (target-resource-id is the target Cosmos DB account ARM ID)
+az dataprotection backup-instance restore initialize-for-item-recovery \
+    --datasource-type AzureCosmosDB \
+    --source-datastore VaultStore \
+    --restore-location {location} \
+    --recovery-point-id {recovery_point_id} \
+    --target-resource-id {target_cosmosdb_account_id} \
+    --restore-configuration @cosmos_restore_config.json \
+    --mi-system-assigned true > cosmos_restore_request.json
+
+# 3) Grant missing restore permissions to vault MSI
+az dataprotection backup-instance update-msi-permissions \
+    --datasource-type AzureCosmosDB \
+    --operation Restore \
+    --permissions-scope Resource \
+    --resource-group {vault_rg} \
+    --vault-name {vault_name} \
+    --restore-request-object @cosmos_restore_request.json
+
+# 4) Validate and trigger restore
+az dataprotection backup-instance validate-for-restore \
+    --resource-group {vault_rg} \
+    --vault-name {vault_name} \
+    --backup-instance-name {backup_instance_name} \
+    --restore-request-object @cosmos_restore_request.json
+
+az dataprotection backup-instance restore trigger \
+    --resource-group {vault_rg} \
+    --vault-name {vault_name} \
+    --backup-instance-name {backup_instance_name} \
+    --restore-request-object @cosmos_restore_request.json
+```
+
 #### dataprotection backup-vault ####
 ##### Create #####
 ```

src/dataprotection/azext_dataprotection/manual/Manifests/AzureCosmosDB.py

@@ -12,7 +12,7 @@
   "datasourceType": "Microsoft.DocumentDB/databaseAccounts",
   "allowedRestoreModes": [ "RecoveryPointBased" ],
   "allowedRestoreTargetTypes": [ "AlternateLocation" ],
-  "itemLevelRecoveryEnabled": false,
+  "itemLevelRecoveryEnabled": true,
   "addBackupDatasourceParametersList": false,
   "backupConfigurationRequired": false,
   "addDataStoreParametersList": false,

src/dataprotection/azext_dataprotection/manual/_help.py

@@ -207,6 +207,8 @@
             --vaulted-blob-prefix-pattern '{"containers":[{"name":"container1","prefixmatch":["a","b"],"renameto":"container1renamed"},{"name":"container2","renameto":"container2renamed"}]}'
       - name: Initialize item-level restore request for an AzureElasticSAN volume group
         text: az dataprotection backup-instance restore initialize-for-item-recovery --datasource-type AzureElasticSAN --restore-location centraluseuap --source-datastore OperationalStore --recovery-point-id {recovery_point_id} --target-resource-id {target_volume_group_id} --restore-configuration {restore_configuration}
+      - name: Initialize item-level restore request for an AzureCosmosDB account
+        text: az dataprotection backup-instance restore initialize-for-item-recovery --datasource-type AzureCosmosDB --restore-location {restore_location} --source-datastore VaultStore --recovery-point-id {recovery_point_id} --target-resource-id {target_cosmosdb_account_id} --restore-configuration {restore_configuration} --mi-system-assigned true
 """
 
 helps['dataprotection resource-guard list-protected-operations'] = """
@@ -238,12 +240,14 @@
 
 helps['dataprotection backup-instance initialize-restoreconfig'] = """
     type: command
-    short-summary: Initialize JSON request body for initializing and configuring restore of an AzureKubernetesService or AzureElasticSAN resource. The generated JSON is meant for use with other CLI commands, and may not work as an input for non-CLI scenarios without modification.
+    short-summary: Initialize JSON request body for initializing and configuring restore of an AzureKubernetesService, AzureElasticSAN or AzureCosmosDB resource. The generated JSON is meant for use with other CLI commands, and may not work as an input for non-CLI scenarios without modification.
     examples:
       - name: Initialize restore configuration
         text: az dataprotection backup-instance initialize-restoreconfig --datasource-type AzureKubernetesService
       - name: Initialize restore configuration for AzureElasticSAN
         text: az dataprotection backup-instance initialize-restoreconfig --datasource-type AzureElasticSAN --resource-identifiers source-vol1 source-vol2 --resource-name-overrides '{"source-vol1":"restored-vol1","source-vol2":"restored-vol2"}'
+      - name: Initialize restore configuration for AzureCosmosDB
+        text: az dataprotection backup-instance initialize-restoreconfig --datasource-type AzureCosmosDB --resource-identifiers db1/container1 db2
 """
 
 helps['dataprotection backup-instance validate-for-backup'] = """

src/dataprotection/azext_dataprotection/manual/_params.py

@@ -319,7 +319,8 @@ def load_arguments(self, _):
                    help='Key value mapping for resource modifier reference')
         c.argument('resource_identifiers', type=str, nargs='+',
                    help='List of source volume identifiers (volume names or ARM IDs) to restore. '
-                        'Use this parameter for DatasourceType AzureElasticSAN.')
+                        'Use this parameter for DatasourceType AzureElasticSAN. '
+                        'For DatasourceType AzureCosmosDB, provide the database or database/container names to restore.')
         c.argument('resource_name_overrides', type=validate_file_or_dict,
                    options_list=['--resource-name-overrides', '--name-overrides'],
                    help='Map of source volume names to target volume names to restore into. Keys must match selected '
@@ -368,7 +369,7 @@ def load_arguments(self, _):
         c.argument('container_list', type=str, nargs='+', help="specify the list of containers to restore.")
         c.argument('from_prefix_pattern', type=str, nargs='+', help="specify the prefix pattern for start range.")
         c.argument('to_prefix_pattern', type=str, nargs='+', help="specify the prefix pattern for end range.")
-        c.argument('restore_configuration', type=validate_file_or_dict, help="Restore configuration for restore. Use this parameter to restore with AzureKubernetesService or AzureElasticSAN.")
+        c.argument('restore_configuration', type=validate_file_or_dict, help="Restore configuration for restore. Use this parameter to restore with AzureKubernetesService, AzureElasticSAN or AzureCosmosDB.")
         c.argument('vaulted_blob_prefix_pattern', options_list=['--vaulted-blob-prefix-pattern', '--vaulted-blob-prefix'],
                    type=validate_file_or_dict, help="Specify the prefix pattern for vaulted blobs.")
         c.argument('use_system_assigned_identity', options_list=['--mi-system-assigned', '--use-system-identity', '--use-system-assigned-identity'], arg_type=get_three_state_flag(), help="Use system assigned identity")

src/dataprotection/azext_dataprotection/manual/custom.py

@@ -1054,6 +1054,18 @@ def dataprotection_backup_instance_initialize_restoreconfig(datasource_type, exc
             "resource_selectors": resource_selectors
         }
 
+    if datasource_type == "AzureCosmosDB":
+        if not resource_identifiers:
+            raise RequiredArgumentMissingError('--resource-identifiers is required for AzureCosmosDB restore configuration. '
+                                               'Provide the list of database or database/container names to restore.')
+        return {
+            "object_type": "GenericRestoreDatasourceCriteria",
+            "resource_selectors": {
+                "object_type": "ResourceListSelectionCriteria",
+                "resource_identifiers": resource_identifiers
+            }
+        }
+
     if datasource_type != "AzureKubernetesService":
         raise InvalidArgumentValueError("This command is currently not supported for datasource types other than AzureKubernetesService")
 

src/dataprotection/azext_dataprotection/manual/helpers.py

@@ -7,6 +7,7 @@
 # pylint: disable=unnecessary-list-index-lookup
 # pylint: disable=unidiomatic-typecheck
 # pylint: disable=possibly-used-before-assignment
+# pylint: disable=chained-comparison
 
 import uuid
 import re
@@ -585,6 +586,16 @@ def get_resource_criteria_list(datasource_type, restore_configuration, container
         restore_criteria_list.append(restore_criteria)
         return restore_criteria_list
 
+    if datasource_type == "AzureCosmosDB":
+        if restore_configuration is not None:
+            restore_criteria = restore_configuration
+        else:
+            raise RequiredArgumentMissingError(
+                "Please input parameter restore_configuration for AzureCosmosDB restore.\n"
+                "Use command initialize-restoreconfig for creating the RestoreConfiguration")
+        restore_criteria_list.append(restore_criteria)
+        return restore_criteria_list
+
     # implicit else:
     # For non-AKS workloads (blobs (non-vaulted)), we need either a prefix-pattern or a container-list. Accordingly, the restore
     # criteria's min_matching_value and max_matching_value are set. We need to provide one, but can't provide both

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

@@ -13,11 +13,13 @@
   * default policy template generation
   * backup-instance initialize output
   * restore initialize-for-data-recovery output (AlternateLocation)
+  * restore initialize-for-item-recovery output (database/container selection)
 
 The end-to-end live scenario lives in
 ``test_dataprotection_backup_and_restore_workloads.py::BackupAndRestoreScenarioTest.test_dataprotection_backup_and_restore_cosmosdb``.
 """
 
+from azure.cli.core.azclierror import RequiredArgumentMissingError
 from azure.cli.testsdk import ScenarioTest
 from azure.cli.testsdk.scenario_tests import AllowLargeResponse
 
@@ -167,3 +169,64 @@ def test_dataprotection_cosmosdb_restore_initialize_alternate_location(test):
         # CosmosDB is not a proxy resource - restore target shouldn't carry a datasource_set_info.
         test.assertIsNone(rti.get('datasource_set_info'),
                           "datasource_set_info should not be present in restore target for non-proxy CosmosDB")
+
+
+class CosmosDBItemRestoreInitializeTest(ScenarioTest):
+    """Verify ``restore initialize-for-item-recovery`` payload for AzureCosmosDB (database/container selection)."""
+
+    @AllowLargeResponse()
+    def test_dataprotection_cosmosdb_restore_initialize_item_recovery(test):
+        test.kwargs.update({
+            'dataSourceType': 'AzureCosmosDB',
+            'sourceDataStore': 'VaultStore',
+            'restoreLocation': 'northcentralus',
+            'recoveryPointId': '33d445e6891444638ffca3d35b3c479f',
+            'targetResourceId': TARGET_COSMOS_ID,
+        })
+
+        restore_config = test.cmd(
+            'az dataprotection backup-instance initialize-restoreconfig '
+            '--datasource-type "{dataSourceType}" --resource-identifiers db1/container1 db1/container2 db2'
+        ).get_output_in_json()
+        test.assertEqual(restore_config, {
+            'object_type': 'GenericRestoreDatasourceCriteria',
+            'resource_selectors': {
+                'object_type': 'ResourceListSelectionCriteria',
+                'resource_identifiers': ['db1/container1', 'db1/container2', 'db2']
+            }
+        })
+        test.kwargs.update({'restoreConfig': restore_config})
+
+        restore_request = test.cmd(
+            'az dataprotection backup-instance restore initialize-for-item-recovery '
+            '--datasource-type "{dataSourceType}" --restore-location "{restoreLocation}" '
+            '--source-datastore "{sourceDataStore}" --recovery-point-id "{recoveryPointId}" '
+            '--target-resource-id "{targetResourceId}" --restore-configuration "{restoreConfig}" '
+            '--mi-system-assigned true'
+        ).get_output_in_json()
+
+        test.assertEqual(restore_request.get('object_type'), 'AzureBackupRecoveryPointBasedRestoreRequest')
+        test.assertEqual(restore_request.get('source_data_store_type'), 'VaultStore')
+        test.assertEqual(restore_request.get('recovery_point_id'), test.kwargs['recoveryPointId'])
+        test.assertEqual(restore_request.get('identity_details'), {'useSystemAssignedIdentity': True})
+
+        rti = restore_request.get('restore_target_info', {})
+        test.assertEqual(rti.get('object_type'), 'ItemLevelRestoreTargetInfo')
+        test.assertEqual(rti.get('restore_location'), 'northcentralus')
+        test.assertEqual(rti.get('datasource_info', {}).get('resource_id'), TARGET_COSMOS_ID)
+        test.assertIsNone(rti.get('datasource_set_info'),
+                          "datasource_set_info should not be present in restore target for non-proxy CosmosDB")
+        test.assertEqual(rti.get('restore_criteria'), [restore_config])
+
+    def test_dataprotection_cosmosdb_restore_initialize_item_recovery_requires_inputs(test):
+        test.kwargs.update({'targetResourceId': TARGET_COSMOS_ID})
+
+        with test.assertRaises(RequiredArgumentMissingError):
+            test.cmd('az dataprotection backup-instance initialize-restoreconfig --datasource-type AzureCosmosDB')
+
+        with test.assertRaises(RequiredArgumentMissingError):
+            test.cmd(
+                'az dataprotection backup-instance restore initialize-for-item-recovery '
+                '--datasource-type AzureCosmosDB --restore-location northcentralus --source-datastore VaultStore '
+                '--recovery-point-id 33d445e6891444638ffca3d35b3c479f --target-resource-id "{targetResourceId}"'
+            )

src/dataprotection/setup.py

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