From f9136e696d25ca69b44d8e93501291f0a95a9287 Mon Sep 17 00:00:00 2001 From: Dan Barrett Date: Fri, 23 May 2025 14:40:20 +1000 Subject: [PATCH] [Storage] `az storage file upload-batch/download-batch`: Add OAuth support --- .../cli/command_modules/storage/commands.py | 6 +- .../latest/test_storage_batch_operations.py | 159 ++++++++++++++++++ 2 files changed, 162 insertions(+), 3 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/storage/commands.py b/src/azure-cli/azure/cli/command_modules/storage/commands.py index edb737433e4..b376fe78ebd 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/commands.py +++ b/src/azure-cli/azure/cli/command_modules/storage/commands.py @@ -714,12 +714,12 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT g.storage_custom_command('copy start-batch', 'storage_file_copy_batch', client_factory=cf_share_client) g.storage_custom_command_oauth('upload', 'storage_file_upload', exception_handler=file_related_exception_handler) - g.storage_custom_command('upload-batch', 'storage_file_upload_batch', - custom_command_type=get_custom_sdk('file', client_factory=cf_share_client)) + g.storage_custom_command_oauth('upload-batch', 'storage_file_upload_batch', + custom_command_type=get_custom_sdk('file', client_factory=cf_share_client)) g.storage_custom_command_oauth('download', 'download_file', exception_handler=file_related_exception_handler, transform=transform_file_show_result) - g.storage_custom_command('download-batch', 'storage_file_download_batch', client_factory=cf_share_client) + g.storage_custom_command_oauth('download-batch', 'storage_file_download_batch', client_factory=cf_share_client) g.storage_command_oauth('hard-link create', 'create_hardlink') with self.command_group('storage cors', get_custom_sdk('cors', multi_service_properties_factory)) as g: diff --git a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_batch_operations.py b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_batch_operations.py index d8e2daa45d7..3993ab29654 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_batch_operations.py +++ b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_batch_operations.py @@ -626,6 +626,165 @@ def test_storage_blob_batch_sas_scenarios(self, test_dir, storage_account_info): self.storage_cmd('storage blob list -c {}', storage_account_info, container_name).assert_with_checks( JMESPathCheck('length(@)', 0)) + @ResourceGroupPreparer() + @StorageAccountPreparer() + @StorageTestFilesPreparer() + def test_storage_file_batch_download_scenarios_oauth(self, resource_group, test_dir, storage_account): + account_info = self.get_account_info(resource_group, storage_account) + src_share = self.create_share(account_info) + # Prepare files + snapshot = self.storage_cmd('storage share snapshot -n {} ', + account_info, src_share).get_output_in_json()["snapshot"] + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --max-connections 3 --account-name {}', + test_dir, src_share, storage_account) + + # download without pattern + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', src_share, local_folder, storage_account) + self.assertEqual(41, sum(len(f) for r, d, f in os.walk(local_folder))) + + # download with pattern apple/* + local_folder = self.create_temp_dir() + share_url = self.file_oauth_cmd('storage file url -s {} -p \'\' -otsv --account-name {}', + src_share, storage_account).output.strip()[:-1] + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --pattern apple/* --account-name {}', + share_url, local_folder, storage_account) + self.assertEqual(10, sum(len(f) for r, d, f in os.walk(local_folder))) + + # download with pattern */file0 + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --pattern */file_0 --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(4, sum(len(f) for r, d, f in os.walk(local_folder))) + + # download with pattern nonexsits/* + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --pattern nonexists/* --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(0, sum(len(f) for r, d, f in os.walk(local_folder))) + + # download with snapshot + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --snapshot {} --account-name {}', + src_share, local_folder, snapshot, storage_account) + self.assertEqual(0, sum(len(f) for r, d, f in os.walk(local_folder))) + + snapshot = self.storage_cmd('storage share snapshot -n {} ', + account_info, src_share).get_output_in_json()["snapshot"] + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --snapshot {} --account-name {}', + src_share, local_folder, snapshot, storage_account) + self.assertEqual(41, sum(len(f) for r, d, f in os.walk(local_folder))) + + local_folder = self.create_temp_dir() + share_url = self.file_oauth_cmd('storage file url -s {} -p \'\' -otsv --account-name {}', + src_share, storage_account).output.strip()[:-1] + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --pattern apple/* --snapshot {} --account-name {} ', + share_url, local_folder, snapshot, storage_account) + self.assertEqual(10, sum(len(f) for r, d, f in os.walk(local_folder))) + + + @ResourceGroupPreparer() + @StorageAccountPreparer() + @StorageTestFilesPreparer() + def test_storage_file_batch_upload_scenarios_oauth(self, resource_group, test_dir, storage_account): + account_info = self.get_account_info(resource_group, storage_account) + # upload without pattern + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --max-connections 3 --account-name {}', + test_dir, src_share, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(41, sum(len(f) for r, d, f in os.walk(local_folder))) + + # upload with pattern apple/* + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --pattern apple/* --account-name {}', + test_dir, src_share, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(10, sum(len(f) for r, d, f in os.walk(local_folder))) + + # upload with pattern */file_0 + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + share_url = self.file_oauth_cmd('storage file url -s {} -p \'\' -otsv --account-name {}' + , src_share, storage_account).output.strip()[:-1] + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --pattern */file_0 --account-name {}', + test_dir, share_url, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(4, sum(len(f) for r, d, f in os.walk(local_folder))) + + # upload with pattern nonexists/* + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --pattern nonexists/* --account-name {}', + test_dir, src_share, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(0, sum(len(f) for r, d, f in os.walk(local_folder))) + + # upload while specifying share path + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + share_url = self.file_oauth_cmd('storage file url -s {} -p \'\' -otsv --account-name {}', + src_share, storage_account).output.strip()[:-1] + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --pattern */file_0 --destination-path some_dir --account-name {}', + test_dir, share_url, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --pattern some_dir* --account-name {}', + src_share, local_folder, storage_account) + self.assertEqual(4, sum(len(f) for r, d, f in os.walk(local_folder))) + + # upload to specifying share path + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + sub_dir = 'test_dir/sub_dir' + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --pattern */file_0 --destination-path {} --account-name {} ', + test_dir, src_share, sub_dir, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', + src_share + "/" + sub_dir, local_folder, storage_account) + self.assertEqual(4, sum(len(f) for r, d, f in os.walk(local_folder))) + + # upload with content settings + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --pattern apple/file_0 ' + '--content-cache-control no-cache ' + '--content-disposition attachment ' + '--content-encoding compress ' + '--content-language en-US ' + '--content-type "multipart/form-data;" ' + '--metadata key=val ' + '--account-name {}', test_dir, src_share, storage_account) + self.file_oauth_cmd('storage file show -s {} -p "{}" --account-name {}', src_share, 'apple/file_0', storage_account). \ + assert_with_checks(JMESPathCheck('name', 'file_0'), + JMESPathCheck('properties.contentSettings.cacheControl', 'no-cache'), + JMESPathCheck('properties.contentSettings.contentDisposition', 'attachment'), + JMESPathCheck('properties.contentSettings.contentEncoding', 'compress'), + JMESPathCheck('properties.contentSettings.contentLanguage', 'en-US'), + JMESPathCheck('properties.contentSettings.contentType', 'multipart/form-data;'), + JMESPathCheck('metadata', {'key': 'val'})) + + # parallel upload with max-connections + import time + start_time = time.time() + src_share = self.create_share(account_info) + local_folder = self.create_temp_dir() + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --max-connections 3 --account-name {}', + test_dir, src_share, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', src_share, local_folder, storage_account) + self.assertEqual(41, sum(len(f) for _, _, f in os.walk(local_folder))) + multi_thread_time = time.time() - start_time + start_time = time.time() + src_share = self.create_share(account_info) + self.file_oauth_cmd('storage file upload-batch -s "{}" -d {} --max-connections 1 --account-name {}', + test_dir, src_share, storage_account) + self.file_oauth_cmd('storage file download-batch -s {} -d "{}" --account-name {}', src_share, local_folder, storage_account) + self.assertEqual(41, sum(len(f) for _, _, f in os.walk(local_folder))) + single_thread_time = time.time() - start_time + self.assertGreater(single_thread_time, multi_thread_time) if __name__ == '__main__': import unittest