{Storage} Make storage batch tests non-live so CI can catch issues (#32143)

Author: calvinhzy

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

@@ -374,6 +374,9 @@ def storage_blob_copy_batch(cmd, client, source_client, container_name=None, des
     if source_container:
         # copy blobs for blob container, skip empty dir
         # pylint: disable=inconsistent-return-statements
+        if source_client is None:
+            source_client = client
+
         def action_blob_copy(blob_name):
             if dryrun:
                 logger.warning('  - copy blob %s', blob_name)
@@ -392,12 +395,25 @@ def action_blob_copy(blob_name):
     if source_share:
         # copy blob from file share, skip empty dir
         # pylint: disable=inconsistent-return-statements
+        if source_client is None:
+            t_share_service = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_FILESHARE,
+                                      '_share_service_client#ShareServiceClient')
+            account_url = client.url.replace('blob', 'file')
+            if client.credential and client.credential.account_key:
+                credential = {
+                    "account_name": client.credential.account_name,
+                    "account_key": client.credential.account_key
+                }
+                source_client = t_share_service(account_url=account_url, credential=credential)
+            else:
+                source_client = t_share_service(account_url=account_url, credential=client.credential)
+
         def action_file_copy(file_info):
             dir_name, file_name = file_info
             if dryrun:
-                logger.warning('  - copy file %s', os.path.join(dir_name, file_name))
+                logger.warning('  - copy file %s', '/'.join(dir_name, file_name))
             else:
-                return _copy_file_to_blob_container(client, source_client, container_name, destination_path,
+                return _copy_file_to_blob_container(cmd, client, source_client, container_name, destination_path,
                                                     source_share, source_sas, dir_name, file_name)
 
         return list(filter_none(action_file_copy(file) for file in collect_files_track2(source_client,
@@ -908,16 +924,21 @@ def create_blob_url(client, container_name, blob_name, snapshot, protocol='https
 def _copy_blob_to_blob_container(cmd, blob_service, source_blob_service, destination_container, destination_path,
                                  source_container, source_blob_name, source_sas, **kwargs):
     t_blob_client = cmd.get_models('_blob_client#BlobClient')
-    # generate sas for oauth copy source
-    if not source_sas:
-        from ..util import create_short_lived_blob_sas_v2
-        start = datetime.utcnow()
-        expiry = datetime.utcnow() + timedelta(hours=1)
-        source_user_delegation_key = source_blob_service.get_user_delegation_key(start, expiry)
-        source_sas = create_short_lived_blob_sas_v2(cmd, source_blob_service.account_name, source_container,
-                                                    source_blob_name, user_delegation_key=source_user_delegation_key)
-    source_client = t_blob_client(account_url=source_blob_service.url, container_name=source_container,
-                                  blob_name=source_blob_name, credential=source_sas)
+    # if blob_service and source_blob_service are the same
+    if blob_service == source_blob_service:
+        source_client = source_blob_service.get_blob_client(container=source_container, blob=source_blob_name)
+    else:
+        # generate sas for oauth copy source
+        if not source_sas:
+            from ..util import create_short_lived_blob_sas_v2
+            start = datetime.utcnow()
+            expiry = datetime.utcnow() + timedelta(days=1)
+            source_user_delegation_key = source_blob_service.get_user_delegation_key(start, expiry)
+            source_sas = create_short_lived_blob_sas_v2(cmd, source_blob_service.account_name, source_container,
+                                                        source_blob_name,
+                                                        user_delegation_key=source_user_delegation_key)
+        source_client = t_blob_client(account_url=source_blob_service.url, container_name=source_container,
+                                      blob_name=source_blob_name, credential=source_sas)
     source_blob_url = source_client.url
 
     destination_blob_name = normalize_blob_file_path(destination_path, source_blob_name)
@@ -935,16 +956,30 @@ def _copy_blob_to_blob_container(cmd, blob_service, source_blob_service, destina
             raise CLIError(error_template.format(source_blob_name, destination_container, ex))
 
 
-def _copy_file_to_blob_container(blob_service, source_file_service, destination_container, destination_path,
+def _copy_file_to_blob_container(cmd, blob_service, source_file_service, destination_container, destination_path,
                                  source_share, source_sas, source_file_dir, source_file_name):
     t_share_client = source_file_service.get_share_client(source_share)
-    t_file_client = t_share_client.get_file_client(os.path.join(source_file_dir, source_file_name))
+    source_path = os.path.join(source_file_dir, source_file_name) if source_file_dir else source_file_name
+    source_path = normalize_blob_file_path(None, source_path)
+    t_file_client = t_share_client.get_file_client(source_path)
+    if source_sas is None:
+        t_generate_share_sas = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_FILESHARE,
+                                       '_shared_access_signature#generate_share_sas')
+        t_file_permissions = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_FILESHARE,
+                                     '_models#FileSasPermissions')
+        start = datetime.utcnow()
+        expiry = datetime.utcnow() + timedelta(days=1)
+        source_sas = t_generate_share_sas(account_name=t_file_client.account_name, share_name=source_share,
+                                          account_key=t_file_client.credential.account_key,
+                                          permission=t_file_permissions(read=True),
+                                          expiry=expiry, start=start)
+        from urllib.parse import quote
+        source_sas = quote(source_sas, safe='&%()$=\',~')
     if '?' not in t_file_client.url:
         source_file_url = '{}?{}'.format(t_file_client.url, source_sas)
     else:
         source_file_url = t_file_client.url
 
-    source_path = os.path.join(source_file_dir, source_file_name) if source_file_dir else source_file_name
     destination_blob_name = normalize_blob_file_path(destination_path, source_path)
     try:
         blob_client = blob_service.get_blob_client(container=destination_container, blob=destination_blob_name)

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

@@ -184,7 +184,7 @@ def _upload_action(src, dst2):
         dir_name = os.path.dirname(dst2)
         file_name = os.path.basename(dst2)
 
-        _make_directory_in_files_share(client, destination, dir_name, V2=True)
+        _make_directory_in_files_share(client, dir_name)
 
         logger.warning('uploading %s', src)
         storage_file_upload(client.get_file_client(dst2), src, content_settings, metadata, validate_content,
@@ -233,7 +233,7 @@ def storage_file_download_batch(client, source, destination, pattern=None, dryru
     Download files from file share to local directory in batch
     """
 
-    from azure.cli.command_modules.storage.util import glob_files_remotely_track2
+    from azure.cli.command_modules.storage.util import glob_files_remotely_track2, normalize_blob_file_path
 
     source_files = glob_files_remotely_track2(client, source, pattern, is_share_client=True)
 
@@ -255,12 +255,13 @@ def storage_file_download_batch(client, source, destination, pattern=None, dryru
     def _download_action(pair):
         path = os.path.join(*pair)
         local_path = os.path.join(destination, *pair)
+        path = normalize_blob_file_path(None, path)
         file_client = client.get_file_client(path)
 
         download_file(file_client, destination_path=local_path, max_connections=max_connections,
                       progress_callback=progress_callback, validate_content=validate_content)
 
-        return file_client.url.replace('%5C', '/')
+        return file_client.url
 
     return list(_download_action(f) for f in source_files)
 
@@ -298,6 +299,19 @@ def storage_file_copy_batch(cmd, client, source_client, share_name=None, destina
         # repeatedly create existing directory so as to optimize the performance.
         existing_dirs = set()
 
+        if source_client is None:
+            t_blob_service_client = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_BLOB,
+                                            '_blob_service_client#BlobServiceClient')
+            account_url = '/'.join(client.url.replace('file', 'blob').split('/')[:-1])
+            if client.credential and client.credential.account_key:
+                credential = {
+                    "account_name": client.credential.account_name,
+                    "account_key": client.credential.account_key
+                }
+                source_client = t_blob_service_client(account_url=account_url, credential=credential)
+            else:
+                source_client = t_blob_service_client(account_url=account_url, credential=client.credential)
+
         # pylint: disable=inconsistent-return-statements
         def action_blob_copy(blob_name):
             if dryrun:
@@ -318,6 +332,12 @@ def action_blob_copy(blob_name):
         # repeatedly create existing directory so as to optimize the performance.
         existing_dirs = set()
 
+        if source_client is None:
+            t_share_service = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_FILESHARE,
+                                      '_share_service_client#ShareServiceClient')
+            account_url = '/'.join(client.url.split('/')[:-1])
+            source_client = t_share_service(account_url=account_url, credential=client.credential)
+
         # pylint: disable=inconsistent-return-statements
         def action_file_copy(file_info):
             dir_name, file_name = file_info
@@ -339,9 +359,11 @@ def storage_file_delete_batch(client, source, pattern=None, dryrun=False, timeou
     """
     Delete files from file share in batch
     """
+    from azure.cli.command_modules.storage.util import normalize_blob_file_path
 
     def delete_action(pair):
         path = os.path.join(*pair)
+        path = normalize_blob_file_path(None, path)
         file_client = client.get_file_client(path)
         return file_client.delete_file(timeout=timeout)
 
@@ -362,6 +384,7 @@ def delete_action(pair):
         delete_action(f)
 
 
+# pylint: disable=too-many-locals
 def _create_file_and_directory_from_blob(cmd, file_service, blob_service, share, container, sas, blob_name,
                                          destination_dir=None, metadata=None, timeout=None, existing_dirs=None):
     """
@@ -371,13 +394,29 @@ def _create_file_and_directory_from_blob(cmd, file_service, blob_service, share,
     from azure.cli.command_modules.storage.util import normalize_blob_file_path
 
     t_blob_client = cmd.get_models('_blob_client#BlobClient', resource_type=ResourceType.DATA_STORAGE_BLOB)
+    if sas is None:
+        t_generate_blob_sas = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_BLOB,
+                                      '_shared_access_signature#generate_blob_sas')
+        t_blob_permissions = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_BLOB,
+                                     '_models#BlobSasPermissions')
+        from datetime import datetime, timedelta
+        start = datetime.utcnow()
+        expiry = datetime.utcnow() + timedelta(days=1)
+        source_sas = t_generate_blob_sas(account_name=blob_service.account_name, container_name=container,
+                                         blob_name=blob_name,
+                                         account_key=blob_service.credential.account_key,
+                                         permission=t_blob_permissions(read=True),
+                                         expiry=expiry, start=start)
+        from urllib.parse import quote
+        sas = quote(source_sas, safe='&%()$=\',~')
+
     source_client = t_blob_client(account_url=blob_service.url, container_name=container,
                                   blob_name=blob_name, credential=sas)
     blob_url = source_client.url
 
     full_path = normalize_blob_file_path(destination_dir, blob_name)
     dir_name = os.path.dirname(full_path)
-    _make_directory_in_files_share(file_service, share, dir_name, existing_dirs, V2=True)
+    _make_directory_in_files_share(file_service, dir_name, existing_dirs)
 
     try:
         file_client = file_service.get_file_client(full_path)
@@ -390,6 +429,7 @@ def _create_file_and_directory_from_blob(cmd, file_service, blob_service, share,
         raise CLIError(error_template.format(blob_name, share))
 
 
+# pylint: disable=too-many-locals
 def _create_file_and_directory_from_file(cmd, file_service, source_file_service, share, source_share, sas,
                                          source_file_dir, source_file_name, destination_dir=None, metadata=None,
                                          timeout=None, existing_dirs=None):
@@ -403,14 +443,29 @@ def _create_file_and_directory_from_file(cmd, file_service, source_file_service,
     if source_file_dir:
         file_path = source_file_dir + '/' + file_path
     t_file_client = cmd.get_models('_file_client#ShareFileClient', resource_type=ResourceType.DATA_STORAGE_FILESHARE)
+    if sas is None:
+        t_generate_share_sas = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_FILESHARE,
+                                       '_shared_access_signature#generate_share_sas')
+        t_file_permissions = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_FILESHARE,
+                                     '_models#FileSasPermissions')
+        from datetime import datetime, timedelta
+        start = datetime.utcnow()
+        expiry = datetime.utcnow() + timedelta(days=1)
+        source_sas = t_generate_share_sas(account_name=source_file_service.account_name, share_name=source_share,
+                                          account_key=source_file_service.credential.account_key,
+                                          permission=t_file_permissions(read=True),
+                                          expiry=expiry, start=start)
+        from urllib.parse import quote
+        sas = quote(source_sas, safe='&%()$=\',~')
+
     source_client = t_file_client(account_url=source_file_service.url, share_name=source_share, file_path=file_path,
                                   credential=sas)
     file_url = source_client.url
 
     full_path = normalize_blob_file_path(destination_dir, os.path.join(source_file_dir, source_file_name))
     file_name = os.path.basename(full_path)
     dir_name = os.path.dirname(full_path)
-    _make_directory_in_files_share(file_service, share, dir_name, existing_dirs, V2=True)
+    _make_directory_in_files_share(file_service, dir_name, existing_dirs)
 
     try:
         file_client = file_service.get_file_client(full_path)
@@ -423,7 +478,7 @@ def _create_file_and_directory_from_file(cmd, file_service, source_file_service,
         raise CLIError(error_template.format(file_name, source_share, share))
 
 
-def _make_directory_in_files_share(file_service, file_share, directory_path, existing_dirs=None, V2=False):
+def _make_directory_in_files_share(file_service, directory_path, existing_dirs=None):
     """
     Create directories recursively.
     This method accept a existing_dirs set which serves as the cache of existing directory. If the
@@ -446,10 +501,7 @@ def _make_directory_in_files_share(file_service, file_share, directory_path, exi
             continue
 
         try:
-            if V2:
-                file_service.create_directory(directory_name=dir_name)
-            else:
-                file_service.create_directory(share_name=file_share, directory_name=dir_name, fail_on_exist=False)
+            file_service.create_directory(directory_name=dir_name)
         except ResourceExistsError:
             pass
         except AzureHttpError: