[Storage] Get sas-token from container url and cleanup code (#6996)

Author: williexu

src/command_modules/azure-cli-storage/HISTORY.rst

@@ -9,6 +9,7 @@ Release History
 * Make '--resource-group' parameter optional for 'storage account' commands.
 * Remove 'Failed precondition' warnings for individual failures in batch commands for single aggregated message.
 * blob/file delete-batch commands no longer output array of nulls.
+* blob download/upload/delete-batch commands will read sas-token from container url
 
 2.1.1
 +++++

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

@@ -264,13 +264,12 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
 
     with self.argument_context('storage blob upload-batch') as c:
         from .sdkutil import get_blob_types
-        from ._validators import process_blob_upload_batch_parameters
 
         t_blob_content_settings = self.get_sdk('blob.models#ContentSettings')
         c.register_content_settings_argument(t_blob_content_settings, update=False, arg_group='Content Control')
         c.ignore('source_files', 'destination_container_name')
 
-        c.argument('source', options_list=('--source', '-s'), validator=process_blob_upload_batch_parameters)
+        c.argument('source', options_list=('--source', '-s'))
         c.argument('destination', options_list=('--destination', '-d'))
         c.argument('max_connections', type=int,
                    help='Maximum number of parallel connections to use when the blob size exceeds 64MB.')
@@ -290,11 +289,9 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.extra('socket_timeout', socket_timeout_type)
 
     with self.argument_context('storage blob download-batch') as c:
-        from azure.cli.command_modules.storage._validators import process_blob_download_batch_parameters
-
         c.ignore('source_container_name')
         c.argument('destination', options_list=('--destination', '-d'))
-        c.argument('source', options_list=('--source', '-s'), validator=process_blob_download_batch_parameters)
+        c.argument('source', options_list=('--source', '-s'))
         c.extra('no_progress', progress_type)
         c.extra('socket_timeout', socket_timeout_type)
         c.argument('max_connections', type=int,
@@ -305,10 +302,8 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.argument('delete_snapshots', arg_type=get_enum_type(get_delete_blob_snapshot_type_names()))
 
     with self.argument_context('storage blob delete-batch') as c:
-        from azure.cli.command_modules.storage._validators import process_blob_batch_source_parameters
-
         c.ignore('source_container_name')
-        c.argument('source', options_list=('--source', '-s'), validator=process_blob_batch_source_parameters)
+        c.argument('source', options_list=('--source', '-s'))
         c.argument('delete_snapshots', arg_type=get_enum_type(get_delete_blob_snapshot_type_names()),
                    help='Required if the blob has associated snapshots.')
         c.argument('lease_id', help='Required if the blob has an active lease.')
@@ -335,8 +330,7 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.register_source_uri_arguments(validator=validate_source_uri)
 
     with self.argument_context('storage blob copy start-batch', arg_group='Copy Source') as c:
-        from azure.cli.command_modules.storage._validators import (get_source_file_or_blob_service_client,
-                                                                   process_blob_copy_batch_namespace)
+        from azure.cli.command_modules.storage._validators import get_source_file_or_blob_service_client
 
         c.argument('source_client', ignore_type, validator=get_source_file_or_blob_service_client)
 
@@ -346,7 +340,6 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.argument('source_sas')
         c.argument('source_container')
         c.argument('source_share')
-        c.argument('prefix', validator=process_blob_copy_batch_namespace)
 
     with self.argument_context('storage blob incremental-copy start') as c:
         from azure.cli.command_modules.storage._validators import process_blob_source_uri

src/command_modules/azure-cli-storage/azure/cli/command_modules/storage/_validators.py

@@ -676,7 +676,7 @@ def _update_progress(current, total):
     del namespace.no_progress
 
 
-def process_blob_download_batch_parameters(namespace, cmd):
+def process_blob_download_batch_parameters(cmd, namespace):
     """Process the parameters for storage blob download command"""
     import os
 
@@ -685,65 +685,22 @@ def process_blob_download_batch_parameters(namespace, cmd):
         raise ValueError('incorrect usage: destination must be an existing directory')
 
     # 2. try to extract account name and container name from source string
-    process_blob_batch_source_parameters(cmd, namespace)
+    _process_blob_batch_container_parameters(cmd, namespace)
 
-
-def process_blob_batch_source_parameters(cmd, namespace):
-    """Process the parameters for storage blob download command"""
-
-    # try to extract account name and container name from source string
-    from .storage_url_helpers import StorageResourceIdentifier
-    identifier = StorageResourceIdentifier(cmd.cli_ctx.cloud, namespace.source)
-
-    if not identifier.is_url():
-        namespace.source_container_name = namespace.source
-    elif identifier.blob:
-        raise ValueError('incorrect usage: source should be either container URL or name')
-    else:
-        namespace.source_container_name = identifier.container
-        if namespace.account_name is None:
-            namespace.account_name = identifier.account_name
+    # 3. Call validators
+    add_progress_callback(cmd, namespace)
 
 
 def process_blob_upload_batch_parameters(cmd, namespace):
     """Process the source and destination of storage blob upload command"""
     import os
 
     # 1. quick check
-    if not os.path.exists(namespace.source):
-        raise ValueError('incorrect usage: source {} does not exist'.format(namespace.source))
-
-    if not os.path.isdir(namespace.source):
-        raise ValueError('incorrect usage: source must be a directory')
+    if not os.path.exists(namespace.source) or not os.path.isdir(namespace.source):
+        raise ValueError('incorrect usage: source must be an existing directory')
 
     # 2. try to extract account name and container name from destination string
-    from .storage_url_helpers import StorageResourceIdentifier
-    identifier = StorageResourceIdentifier(cmd.cli_ctx.cloud, namespace.destination)
-
-    if not identifier.is_url():
-        namespace.destination_container_name = namespace.destination
-    elif identifier.blob is not None:
-        raise ValueError('incorrect usage: destination cannot be a blob url')
-    else:
-        namespace.destination_container_name = identifier.container
-
-        if namespace.account_name:
-            if namespace.account_name != identifier.account_name:
-                raise ValueError(
-                    'The given storage account name is not consistent with the account name in the destination URI')
-        else:
-            namespace.account_name = identifier.account_name
-
-        if not (namespace.account_key or namespace.sas_token or namespace.connection_string):
-            validate_client_parameters(cmd, namespace)
-
-        # it is possible the account name be overwritten by the connection string
-        if namespace.account_name != identifier.account_name:
-            raise ValueError(
-                'The given storage account name is not consistent with the account name in the destination URI')
-
-        if not (namespace.account_key or namespace.sas_token or namespace.connection_string):
-            raise ValueError('Missing storage account credential information.')
+    _process_blob_batch_container_parameters(cmd, namespace, source=False)
 
     # 3. collect the files to be uploaded
     namespace.source = os.path.realpath(namespace.source)
@@ -765,10 +722,48 @@ def process_blob_upload_batch_parameters(cmd, namespace):
         else:
             namespace.blob_type = 'block'
 
+    # 5. call other validators
+    validate_metadata(namespace)
+    t_blob_content_settings = cmd.loader.get_sdk('blob.models#ContentSettings')
+    get_content_setting_validator(t_blob_content_settings, update=False)(cmd, namespace)
+    add_progress_callback(cmd, namespace)
+
+
+def process_blob_delete_batch_parameters(cmd, namespace):
+    _process_blob_batch_container_parameters(cmd, namespace)
+
+
+def _process_blob_batch_container_parameters(cmd, namespace, source=True):
+    """Process the container parameters for storage blob batch commands before populating args from environment."""
+    if source:
+        container_arg, container_name_arg = 'source', 'source_container_name'
+    else:
+        # destination
+        container_arg, container_name_arg = 'destination', 'destination_container_name'
+
+    # try to extract account name and container name from source string
+    from .storage_url_helpers import StorageResourceIdentifier
+    container_arg_val = getattr(namespace, container_arg)  # either a url or name
+    identifier = StorageResourceIdentifier(cmd.cli_ctx.cloud, container_arg_val)
+
+    if not identifier.is_url():
+        setattr(namespace, container_name_arg, container_arg_val)
+    elif identifier.blob:
+        raise ValueError('incorrect usage: {} should be either a container URL or name'.format(container_arg))
+    else:
+        setattr(namespace, container_name_arg, identifier.container)
+        if namespace.account_name is None:
+            namespace.account_name = identifier.account_name
+        elif namespace.account_name != identifier.account_name:
+            raise ValueError('The given storage account name is not consistent with the '
+                             'account name in the destination URL')
+
+        # if no sas-token is given and the container url contains one, use it
+        if not namespace.sas_token and identifier.sas_token:
+            namespace.sas_token = identifier.sas_token
 
-def process_blob_copy_batch_namespace(namespace):
-    if namespace.prefix is None and not namespace.recursive:
-        raise ValueError('incorrect usage: --recursive | --pattern PATTERN')
+    # Finally, grab missing storage connection parameters from environment variables
+    validate_client_parameters(cmd, namespace)
 
 
 def process_file_upload_batch_parameters(cmd, namespace):

src/command_modules/azure-cli-storage/azure/cli/command_modules/storage/commands.py

@@ -93,6 +93,8 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         from ._format import transform_boolean_for_table, transform_blob_output
         from ._transformers import (transform_storage_list_output, transform_url,
                                     create_boolean_result_output_transformer)
+        from ._validators import (process_blob_download_batch_parameters, process_blob_delete_batch_parameters,
+                                  process_blob_upload_batch_parameters)
 
         g.storage_command('list', 'list_blobs', transform=transform_storage_list_output,
                           table_transformer=transform_blob_output)
@@ -110,9 +112,12 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         g.storage_custom_command('set-tier', 'set_blob_tier')
         g.storage_custom_command('upload', 'upload_blob',
                                  doc_string_source='blob#BlockBlobService.create_blob_from_path')
-        g.storage_custom_command('upload-batch', 'storage_blob_upload_batch')
-        g.storage_custom_command('download-batch', 'storage_blob_download_batch')
-        g.storage_custom_command('delete-batch', 'storage_blob_delete_batch')
+        g.storage_custom_command('upload-batch', 'storage_blob_upload_batch',
+                                 validator=process_blob_upload_batch_parameters)
+        g.storage_custom_command('download-batch', 'storage_blob_download_batch',
+                                 validator=process_blob_download_batch_parameters)
+        g.storage_custom_command('delete-batch', 'storage_blob_delete_batch',
+                                 validator=process_blob_delete_batch_parameters)
         g.storage_custom_command('show', 'show_blob', table_transformer=transform_blob_output,
                                  client_factory=page_blob_service_factory,
                                  doc_string_source='blob#PageBlobService.get_blob_properties',

src/command_modules/azure-cli-storage/azure/cli/command_modules/storage/operations/blob.py

@@ -201,7 +201,8 @@ def _upload_blob(*args, **kwargs):
                 results.append(_create_return_result(dst, guessed_content_settings, result))
 
         num_failures = len(source_files) - len(results)
-        logger.warning('%s of %s files not uploaded due to "Failed Precondition"', num_failures, len(source_files))
+        if num_failures:
+            logger.warning('%s of %s files not uploaded due to "Failed Precondition"', num_failures, len(source_files))
     return results
 
 
@@ -322,10 +323,10 @@ def _delete_blob(blob_name):
         }
         return client.delete_blob(**delete_blob_args)
 
+    logger = get_logger(__name__)
     source_blobs = list(collect_blobs(client, source_container_name, pattern))
 
     if dryrun:
-        logger = get_logger(__name__)
         logger.warning('delete action: from %s', source)
         logger.warning('    pattern %s', pattern)
         logger.warning('  container %s', source_container_name)
@@ -337,7 +338,8 @@ def _delete_blob(blob_name):
 
     results = [result for include, result in (_delete_blob(blob) for blob in source_blobs) if include]
     num_failures = len(source_blobs) - len(results)
-    logger.warning('%s of %s blobs not deleted due to "Failed Precondition"', num_failures, len(source_blobs))
+    if num_failures:
+        logger.warning('%s of %s blobs not deleted due to "Failed Precondition"', num_failures, len(source_blobs))
 
 
 def _copy_blob_to_blob_container(blob_service, source_blob_service, destination_container, destination_path,

src/command_modules/azure-cli-storage/azure/cli/command_modules/storage/tests/latest/test_storage_batch_operations.py

@@ -421,6 +421,42 @@ def create_and_populate_share():
         self.storage_cmd('storage file delete-batch -s {} --pattern nonexists/*', storage_account_info,
                          src_share).assert_with_checks(JMESPathCheck('length(@)', 0))
 
+    @ResourceGroupPreparer()
+    @StorageAccountPreparer()
+    @StorageTestFilesPreparer()
+    def test_storage_blob_batch_sas_scenarios(self, test_dir, storage_account_info):
+        from datetime import datetime, timedelta
+
+        container_name = self.create_container(storage_account_info)
+        temp_dir = self.create_temp_dir()
+        expiry = (datetime.utcnow() + timedelta(hours=1)).strftime('%Y-%m-%dT%H:%MZ')
+
+        sas_token = self.storage_cmd('storage container generate-sas -n {} --permissions dwrl '
+                                     '--expiry {}', storage_account_info, container_name, expiry).get_output_in_json()
+
+        storage_account = self.cmd('storage account show -n {}'.format(storage_account_info[0])).get_output_in_json()
+
+        # create container url with sas token
+        container_url = storage_account.get('primaryEndpoints').get('blob') + container_name
+        container_url += '?' + sas_token
+
+        self.kwargs.update({
+            'test_dir': test_dir,
+            'container_url': container_url,
+            'temp_dir': temp_dir
+        })
+
+        self.cmd('storage blob upload-batch -s "{test_dir}" -d {container_url}')
+        self.cmd('storage blob download-batch -s {container_url} -d "{temp_dir}"')
+
+        self.storage_cmd('storage blob list -c {}', storage_account_info, container_name).assert_with_checks(
+            JMESPathCheck('length(@)', sum(len(files) for _, __, files in os.walk(test_dir))),
+            JMESPathCheck('length(@)', sum(len(files) for _, __, files in os.walk(temp_dir))))
+
+        self.cmd('storage blob delete-batch -s {container_url}')
+        self.storage_cmd('storage blob list -c {}', storage_account_info, container_name).assert_with_checks(
+            JMESPathCheck('length(@)', 0))
+
 
 if __name__ == '__main__':
     import unittest