[Storage] az storage fs directory upload/download: Support adls gen2 file system directory upload&download

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

@@ -1928,6 +1928,30 @@
       crafted: true
 """
 
+helps['storage fs directory upload'] = """
+    type: command
+    short-summary: Upload files or subdirectories to a directory in ADLS Gen2 file system.
+    examples:
+        - name: Upload a single file to a storage blob directory.
+          text: az storage fs directory upload -f myfilesystem --account-name mystorageaccount -s "path/to/file" -d directory
+        - name: Upload a local directory to a directory in ADLS Gen2 file system.
+          text: az storage fs directory upload -f myfilesystem --account-name mystorageaccount -s "path/to/directory" -d directory --recursive
+        - name: Upload a set of files in a local directory to a directory in ADLS Gen2 file system.
+          text: az storage fs directory upload -f myfilesystem --account-name mystorageaccount -s "path/to/file*" -d directory --recursive
+"""
+
+helps['storage fs directory download'] = """
+    type: command
+    short-summary: Download files from the directory in ADLS Gen2 file system to a local file path.
+    examples:
+        - name: Download a single file in a directory in ADLS Gen2 file system.
+          text: az storage fs directory download -f myfilesystem --account-name mystorageaccount -s "path/to/file" -d "<local-path>"
+        - name: Download the entire directory in a directory in ADLS Gen2 file system.
+          text: az storage fs directory download -f myfilesystem --account-name mystorageaccount -s SourceDirectoryPath -d "<local-path>" --recursive
+        - name: Download an entire subdirectory of a directory in ADLS Gen2 file system.
+          text: az storage fs directory download -f myfilesystem --account-name mystorageaccount -s "path/to/subdirectory" -d "<local-path>" --recursive
+"""
+
 helps['storage fs file'] = """
 type: group
 short-summary: Manage files in Azure Data Lake Storage Gen2 account.

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

@@ -1561,6 +1561,35 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
                    help='The new directory name the users want to move to. The value must have the following format: '
                         '"{filesystem}/{directory}/{subdirectory}".')
 
+    with self.argument_context('storage fs directory upload') as c:
+        from ._validators import validate_fs_directory_upload_destination_url
+        c.extra('destination_fs', options_list=['--file-system', '-f'], required=True,
+                help='The upload destination file system.')
+        c.extra('destination_path', options_list=['--destination-path', '-d'], required=True,
+                validator=validate_fs_directory_upload_destination_url,
+                help='The upload destination directory path. It should be an absolute path to file system. If the '
+                     'specified destination path does not exist, a new directory path will be created.')
+        c.argument('source', options_list=['--source', '-s'],
+                   help='The source file path to upload from.')
+        c.argument('recursive', options_list=['--recursive', '-r'], arg_type=get_three_state_flag(),
+                   help='Recursively upload files. If enabled, all the files including the files in subdirectories will'
+                        ' be uploaded.')
+        c.ignore('destination')
+
+    with self.argument_context('storage fs directory download') as c:
+        from ._validators import validate_fs_directory_download_source_url
+        c.extra('source_fs', options_list=['--file-system', '-f'], required=True,
+                help='The download source file system.')
+        c.extra('source_path', options_list=['--source-path', '-s'], required=True,
+                validator=validate_fs_directory_download_source_url,
+                help='The download source directory path. It should be an absolute path to file system.')
+        c.argument('destination', options_list=['--destination-path', '-d'],
+                   help='The destination local directory path to download.')
+        c.argument('recursive', options_list=['--recursive', '-r'], arg_type=get_three_state_flag(),
+                   help='Recursively download files. If enabled, all the files including the files in subdirectories '
+                        'will be downloaded.')
+        c.ignore('source')
+
     with self.argument_context('storage fs file list') as c:
         c.extra('file_system_name', options_list=['-f', '--file-system'], help="File system name.", required=True)
         c.argument('recursive', arg_type=get_three_state_flag(), default=True,

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

@@ -1577,6 +1577,55 @@ def validate_azcopy_credential(cmd, namespace):
                                             service=service, resource_types='sco', permissions='rl')
 
 
+def is_directory(props):
+    return 'hdi_isfolder' in props.metadata.keys() and props.metadata['hdi_isfolder'] == 'true'
+
+
+def validate_fs_directory_upload_destination_url(cmd, namespace):
+    destination = namespace.destination_path
+    container = namespace.destination_fs
+    kwargs = {'account_name': namespace.account_name,
+              'account_key': namespace.account_key,
+              'connection_string': namespace.connection_string,
+              'sas_token': namespace.sas_token}
+    client = blob_data_service_factory(cmd.cli_ctx, kwargs)
+    destination = destination[1:] if destination.startswith('/') else destination
+    from azure.common import AzureException
+    from azure.cli.core.azclierror import InvalidArgumentValueError
+    try:
+        props = client.get_blob_properties(container, destination)
+        if not is_directory(props):
+            raise InvalidArgumentValueError('usage error: You are specifying --destination-path with a file name, '
+                                            'not directory name. Please change to a valid directory name. '
+                                            'If you want to upload to a file, please use '
+                                            '`az storage fs file upload` command.')
+    except AzureException:
+        pass
+
+    if not destination.endswith('/'):
+        destination += '/'
+    url = client.make_blob_url(container, destination)
+    namespace.destination = _add_sas_for_url(cmd, url=url, account_name=namespace.account_name,
+                                             account_key=namespace.account_key, sas_token=namespace.sas_token,
+                                             service='blob', resource_types='co', permissions='rwdlac')
+    del namespace.destination_fs
+    del namespace.destination_path
+
+
+def validate_fs_directory_download_source_url(cmd, namespace):
+    kwargs = {'account_name': namespace.account_name,
+              'account_key': namespace.account_key,
+              'connection_string': namespace.connection_string,
+              'sas_token': namespace.sas_token}
+    client = blob_data_service_factory(cmd.cli_ctx, kwargs)
+    url = client.make_blob_url(namespace.source_fs, namespace.source_path)
+    namespace.source = _add_sas_for_url(cmd, url=url, account_name=namespace.account_name,
+                                        account_key=namespace.account_key, sas_token=namespace.sas_token,
+                                        service='blob', resource_types='co', permissions='rl')
+    del namespace.source_fs
+    del namespace.source_path
+
+
 def validate_text_configuration(cmd, ns):
     DelimitedTextDialect = cmd.get_models('_models#DelimitedTextDialect', resource_type=ResourceType.DATA_STORAGE_BLOB)
     DelimitedJSON = cmd.get_models('_models#DelimitedJSON', resource_type=ResourceType.DATA_STORAGE_BLOB)

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

@@ -779,6 +779,10 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         g.storage_command_oauth('metadata show', 'get_directory_properties', exception_handler=show_exception_handler,
                                 transform=transform_metadata)
 
+    with self.command_group('storage fs directory', custom_command_type=get_custom_sdk('azcopy', None))as g:
+        g.storage_custom_command_oauth('upload', 'storage_fs_directory_copy', is_preview=True)
+        g.storage_custom_command_oauth('download', 'storage_fs_directory_copy', is_preview=True)
+
     with self.command_group('storage fs file', adls_file_sdk, resource_type=ResourceType.DATA_STORAGE_FILEDATALAKE,
                             custom_command_type=get_custom_sdk('fs_file', cf_adls_file), min_api='2018-11-09') as g:
         from ._transformers import transform_storage_list_output, create_boolean_result_output_transformer

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

@@ -65,6 +65,18 @@ def storage_remove(cmd, client, service, target, recursive=None, exclude_pattern
     azcopy.remove(_add_url_sas(target, sas_token), flags=flags)
 
 
+# pylint: disable=unused-argument
+def storage_fs_directory_copy(cmd, source, destination, recursive=None, **kwargs):
+    azcopy = AzCopy()
+    if kwargs.get('token_credential'):
+        azcopy = _azcopy_login_client(cmd)
+
+    flags = []
+    if recursive is not None:
+        flags.append('--recursive')
+    azcopy.copy(source, destination, flags=flags)
+
+
 def storage_blob_sync(cmd, client, source, destination, exclude_pattern=None, include_pattern=None,
                       exclude_path=None):
     azcopy = _azcopy_blob_client(cmd, client)

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

@@ -3,11 +3,12 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 import unittest
+import os
 
-from azure.cli.testsdk import (ScenarioTest, ResourceGroupPreparer, StorageAccountPreparer, JMESPathCheck, NoneCheck,
+from azure.cli.testsdk import (ScenarioTest, LiveScenarioTest, ResourceGroupPreparer, StorageAccountPreparer, JMESPathCheck, NoneCheck,
                                api_version_constraint, RoleBasedServicePrincipalPreparer)
 from azure.cli.core.profiles import ResourceType
-from ..storage_test_util import StorageScenarioMixin
+from ..storage_test_util import StorageScenarioMixin, StorageTestFilesPreparer
 from knack.util import CLIError
 
 
@@ -448,3 +449,97 @@ def test_adls_metadata_scenarios(self, resource_group, storage_account):
                          account_info, file, filesystem)
         self.storage_cmd('storage fs file metadata show -p {} -f {}', account_info, file, filesystem) \
             .assert_with_checks(JMESPathCheck('test', 'beta'), JMESPathCheck('cat', 'file'))
+
+
+class StorageADLSGen2LiveTests(StorageScenarioMixin, LiveScenarioTest):
+    @ResourceGroupPreparer()
+    @StorageAccountPreparer(kind="StorageV2", hns=True)
+    @StorageTestFilesPreparer()
+    def test_adls_directory_upload(self, resource_group, storage_account, test_dir):
+        account_info = self.get_account_info(resource_group, storage_account)
+        connection_string = self.get_connection_string(resource_group, storage_account)
+
+        filesystem = 'testfilesystem'
+        directory = 'testdir'
+        self.storage_cmd('storage fs create -n {}', account_info, filesystem)
+        self.storage_cmd('storage fs directory create -n {} -f {}', account_info, directory, filesystem)
+
+        from datetime import datetime, timedelta
+        expiry = (datetime.utcnow() + timedelta(hours=1)).strftime('%Y-%m-%dT%H:%MZ')
+        sas = self.storage_cmd('storage container generate-sas -n {} --https-only '
+                               '--permissions acdlrw --expiry {} -otsv',
+                               account_info, filesystem, expiry).output.strip()
+
+        # Upload a single file to the fs directory
+        self.storage_cmd('storage fs directory upload -f {} -d {} -s "{}"', account_info, filesystem, directory,
+                         os.path.join(test_dir, 'readme'))
+        self.storage_cmd('storage fs file list -f {} --path {}', account_info, filesystem, directory) \
+            .assert_with_checks(JMESPathCheck('length(@)', 1))
+
+        # Upload a local directory to the fs directory
+        self.oauth_cmd('storage fs directory upload -f {} -d {} -s "{}" --recursive --account-name {}', filesystem,
+                       directory, os.path.join(test_dir, 'apple'), storage_account)
+        self.oauth_cmd('storage fs file list -f {} --path {} --account-name {}',
+                       filesystem, directory, storage_account)\
+            .assert_with_checks(JMESPathCheck('length(@)', 12))
+
+        # Upload files in a local directory to the fs directory
+        self.cmd('storage fs directory upload -f {} -d {} -s "{}" --recursive --connection-string {}'
+                 .format(filesystem, directory, os.path.join(test_dir, 'butter/file_*'), connection_string))
+        self.cmd('storage fs file list -f {} --path {} --connection-string {}'
+                 .format(filesystem, directory, connection_string), checks=[JMESPathCheck('length(@)', 22)])
+
+        # Upload files in a local directory to the fs subdirectory
+        self.cmd('storage fs directory upload -f {} -d {} -s "{}" --recursive --account-name {} --sas-token {}'
+                 .format(filesystem, '/'.join([directory, 'subdir']),
+                         os.path.join(test_dir, 'butter/file_*'), storage_account, sas))
+        self.cmd('storage fs file list -f {} --path {} --account-name {} --sas-token {}'
+                 .format(filesystem, '/'.join([directory, 'subdir']), storage_account, sas),
+                 checks=[JMESPathCheck('length(@)', 10)])
+
+        # Argument validation: Fail when destination path is file name
+        self.cmd('storage fs directory upload -f {} -d {} -s {} --account-name {}'.format(
+            filesystem, '/'.join([directory, 'readme']), test_dir, storage_account), expect_failure=True)
+
+    @ResourceGroupPreparer()
+    @StorageAccountPreparer(kind="StorageV2", hns=True)
+    @StorageTestFilesPreparer()
+    def test_adls_directory_download(self, resource_group, storage_account, test_dir):
+        account_info = self.get_account_info(resource_group, storage_account)
+        connection_string = self.get_connection_string(resource_group, storage_account)
+
+        filesystem = 'testfilesystem'
+        directory = 'testdir'
+        self.storage_cmd('storage fs create -n {}', account_info, filesystem)
+        self.storage_cmd('storage fs directory upload -f {} -d {} -s "{}" --recursive', account_info, filesystem,
+                         directory, os.path.join(test_dir, 'readme'))
+        self.storage_cmd('storage fs directory upload -f {} -d {} -s "{}" --recursive', account_info, filesystem,
+                         directory, os.path.join(test_dir, 'apple'))
+
+        from datetime import datetime, timedelta
+        expiry = (datetime.utcnow() + timedelta(hours=1)).strftime('%Y-%m-%dT%H:%MZ')
+        sas = self.storage_cmd('storage container generate-sas -n {} --https-only '
+                               '--permissions lr --expiry {} -otsv',
+                               account_info, filesystem, expiry).output.strip()
+
+        local_folder = self.create_temp_dir()
+        # Download a single file
+        self.storage_cmd('storage fs directory download -f {} -s "{}" -d "{}" --recursive', account_info, filesystem,
+                         '/'.join([directory, 'readme']), local_folder)
+        self.assertEqual(1, sum(len(f) for r, d, f in os.walk(local_folder)))
+
+        # Download entire directory
+        self.oauth_cmd('storage fs directory download -f {} -s {} -d "{}" --recursive --account-name {}',
+                       filesystem, directory, local_folder, storage_account)
+        self.assertEqual(2, sum(len(d) for r, d, f in os.walk(local_folder)))
+        self.assertEqual(12, sum(len(f) for r, d, f in os.walk(local_folder)))
+
+        # Download an entire subdirectory of a storage blob directory.
+        self.cmd('storage fs directory download -f {} -s {} -d "{}" --recursive --connection-string {}'
+                 .format(filesystem, '/'.join([directory, 'apple']), local_folder, connection_string))
+        self.assertEqual(3, sum(len(d) for r, d, f in os.walk(local_folder)))
+
+        # Download file in a directory
+        self.cmd('storage fs directory download -f {} -s "{}" -d "{}" --recursive --account-name {} --sas-token {}'
+                 .format(filesystem, '/'.join([directory, 'apple/file_0']), local_folder, storage_account, sas))
+        self.assertEqual(23, sum(len(f) for r, d, f in os.walk(local_folder)))

src/azure-cli/azure/cli/command_modules/storage/tests/storage_test_util.py

@@ -26,6 +26,10 @@ def get_account_key(self, group, name):
 
         return self.cmd(template.format(name, group)).output
 
+    def get_connection_string(self, group, name):
+        return self.cmd('storage account show-connection-string -n {} -g {} '
+                        '--query connectionString -otsv'.format(name, group)).output.strip()
+
     def get_account_info(self, group, name):
         """Returns the storage account name and key in a tuple"""
         return name, self.get_account_key(group, name)