[storage-blob-preview] Support blob immutability policy and legal hold

src/storage-blob-preview/HISTORY.rst

@@ -2,6 +2,11 @@
 
 Release History
 ===============
+0.6.1
+++++++
+* `az storage blob immutability-policy set/delete`: Extend/Lock/Unlock/Delete blob's immutability policy
+* `az storage blob set-legal-hold`: Configure/Clear blob legal hold
+
 0.6.0
 ++++++
 * Remove `az storage account blob-service-properties` since all the preview arguments are supported in main azure cli

src/storage-blob-preview/azext_storage_blob_preview/_params.py

@@ -8,7 +8,7 @@
 from azure.cli.core.commands.parameters import (file_type, get_enum_type, get_three_state_flag)
 
 from ._validators import (validate_metadata, get_permission_validator, get_permission_help_string,
-                          validate_blob_type, validate_included_datasets_v2,
+                          validate_blob_type, validate_included_datasets_v2, get_datetime_type,
                           add_download_progress_callback, add_upload_progress_callback,
                           validate_storage_data_plane_list, as_user_validator, blob_tier_validator)
 
@@ -248,6 +248,20 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
     with self.argument_context('storage blob exists') as c:
         c.register_blob_arguments()
 
+    with self.argument_context('storage blob set-legal-hold') as c:
+        c.register_blob_arguments()
+        c.argument('legal_hold', arg_type=get_three_state_flag(),
+                   help='Specified if a legal hold should be set on the blob.')
+
+    with self.argument_context('storage blob immutability-policy delete') as c:
+        c.register_blob_arguments()
+
+    with self.argument_context('storage blob immutability-policy set') as c:
+        c.register_blob_arguments()
+        c.argument('expiry_time', type=get_datetime_type(False),
+                   help='expiration UTC datetime in (Y-m-d\'T\'H:M:S\'Z\')')
+        c.argument('policy_mode', arg_type=get_enum_type(['Locked', 'Unlocked']), help='Lock or Unlock the policy')
+
     with self.argument_context('storage blob filter') as c:
         c.argument('filter_expression', options_list=['--tag-filter'])
 

src/storage-blob-preview/azext_storage_blob_preview/commands.py

@@ -76,6 +76,11 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         g.storage_custom_command_oauth('copy start-batch', 'storage_blob_copy_batch', client_factory=cf_blob_service)
         g.storage_custom_command_oauth('query', 'query_blob',
                                        is_preview=True, min_api='2020-10-02')
+        g.storage_command_oauth('set-legal-hold', 'set_legal_hold', min_api='2020-10-02', is_preview=True)
+        g.storage_custom_command_oauth('immutability-policy set', 'set_blob_immutability_policy',
+                                       min_api='2020-10-02', is_preview=True)
+        g.storage_command_oauth('immutability-policy delete', 'delete_immutability_policy',
+                                min_api='2020-10-02', is_preview=True)
 
     with self.command_group('storage blob', blob_service_sdk, resource_type=CUSTOM_DATA_STORAGE_BLOB,
                             min_api='2019-12-12',

src/storage-blob-preview/azext_storage_blob_preview/operations/blob.py

@@ -96,6 +96,15 @@ def set_service_properties(client, delete_retention=None, delete_retention_perio
     return client.get_service_properties()
 
 
+def set_blob_immutability_policy(cmd, client, expiry_time=None, policy_mode=None, **kwargs):
+    ImmutabilityPolicy = cmd.get_models("_models#ImmutabilityPolicy", resource_type=CUSTOM_DATA_STORAGE_BLOB)
+    if not expiry_time and not policy_mode:
+        from azure.cli.core.azclierror import InvalidArgumentValueError
+        raise InvalidArgumentValueError('Please specify --expiry-time | --policy-mode')
+    immutability_policy = ImmutabilityPolicy(expiry_time=expiry_time, policy_mode=policy_mode)
+    return client.set_immutability_policy(immutability_policy=immutability_policy, **kwargs)
+
+
 def storage_blob_copy_batch(cmd, client, source_client, container_name=None,
                             destination_path=None, source_container=None, source_share=None,
                             source_sas=None, pattern=None, dryrun=False, source_account_name=None,

src/storage-blob-preview/azext_storage_blob_preview/tests/latest/test_storage_blob_preview_scenario.py

@@ -149,6 +149,39 @@ def test_storage_blob_versioning(self, resource_group, storage_account):
         self.storage_cmd('storage blob list -c {} --include v', account_info, container)\
             .assert_with_checks(JMESPathCheck('length(@)', count - 1))
 
+    @ResourceGroupPreparer(name_prefix='clitest')
+    @StorageAccountPreparer(name_prefix='version', kind='StorageV2', location='centraluseuap')
+    def test_storage_blob_vlm(self, resource_group, storage_account_info):
+        container = self.create_random_name(prefix='container', length=18)
+        blob = self.create_random_name(prefix='blob', length=18)
+        self.kwargs.update({
+            'container': container,
+            'blob': blob
+        })
+        # Enable blob versioning
+        self.cmd('storage account blob-service-properties update -n {sa} -g {rg} --enable-versioning')
+        # Enable vlm on container creation
+        self.cmd('storage container-rm create -n {container} --storage-account {sa} -g {rg} --enable-vlw')
+        # Prepare blob resource
+        file = self.create_temp_file(10)
+        self.storage_cmd('storage blob upload -c {} -f "{}" -n {} ', storage_account_info, container, file, blob)
+
+        # Test set immutability policy
+        from datetime import datetime, timedelta
+        expiry = (datetime.utcnow() + timedelta(hours=1)).strftime('%Y-%m-%dT%H:%MZ')
+        result = self.storage_cmd('storage blob immutability-policy set -n {} -c {} '
+                                  '--expiry-time {} --policy-mode Unlocked',
+                                  storage_account_info, blob, container, expiry).get_output_in_json()
+        self.assertEqual(result.get('immutability_policy_mode'), 'unlocked')
+        self.assertIsNotNone(result.get('immutability_policy_until_date'))
+        # Test delete immutability policy
+        self.storage_cmd('storage blob immutability-policy delete -n {} -c {}', storage_account_info, blob, container)
+        # Test set legal hold
+        self.storage_cmd('storage blob set-legal-hold --legal-hold -n {} -c {}', storage_account_info, blob, container)\
+            .assert_with_checks(JMESPathCheck('legal_hold', True))
+        self.storage_cmd('storage blob set-legal-hold --legal-hold false -n {} -c {}', storage_account_info, blob, container) \
+            .assert_with_checks(JMESPathCheck('legal_hold', False))
+
     @ResourceGroupPreparer(name_prefix='clitest')
     @StorageAccountPreparer(name_prefix='blobtag', kind='StorageV2', location='eastus2euap')
     def test_storage_blob_tags_scenario(self, resource_group, storage_account):

src/storage-blob-preview/setup.py

@@ -16,7 +16,7 @@
 
 # TODO: Confirm this is the right version number you want and it matches your
 # HISTORY.rst entry.
-VERSION = '0.6.0'
+VERSION = '0.6.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers