[Storage] Add progress tracking callback for blob upload (Azure#23965)

Author: jalauzon-msft

sdk/storage/azure-storage-blob/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 12.12.0 (Unreleased)
 
 ### Features Added
+- Added an optional callback `progress_hook` to `upload_blob()` that can be used to track progress of a large upload.
 
 ### Bugs Fixed:
 - Fixed a bug in `BlobClient.from_blob_url()` such that users will receive a more helpful error

sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py

@@ -697,6 +697,10 @@ def upload_blob(  # pylint: disable=too-many-locals
 
         :keyword str encoding:
             Defaults to UTF-8.
+        :keyword Callable[[int, Optional[int]], None] progress_hook:
+            A callback to track the progress of a long running upload. The signature is
+            function(current: int, total: Optional[int]) where current is the number of bytes transfered
+            so far, and total is the size of the blob or None if the size is unknown.
         :keyword int timeout:
             The timeout parameter is expressed in seconds. This method may make
             multiple calls to the Azure service and the timeout will apply to

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/uploads.py

@@ -53,6 +53,7 @@ def upload_data_chunks(
         stream=None,
         validate_content=None,
         encryption_options=None,
+        progress_hook=None,
         **kwargs):
 
     if encryption_options:
@@ -75,6 +76,7 @@ def upload_data_chunks(
         stream=stream,
         parallel=parallel,
         validate_content=validate_content,
+        progress_hook=progress_hook,
         **kwargs)
     if parallel:
         with futures.ThreadPoolExecutor(max_concurrency) as executor:
@@ -98,6 +100,7 @@ def upload_substream_blocks(
         chunk_size=None,
         max_concurrency=None,
         stream=None,
+        progress_hook=None,
         **kwargs):
     parallel = max_concurrency > 1
     if parallel and 'modified_access_conditions' in kwargs:
@@ -109,6 +112,7 @@ def upload_substream_blocks(
         chunk_size=chunk_size,
         stream=stream,
         parallel=parallel,
+        progress_hook=progress_hook,
         **kwargs)
 
     if parallel:
@@ -128,7 +132,16 @@ def upload_substream_blocks(
 
 class _ChunkUploader(object):  # pylint: disable=too-many-instance-attributes
 
-    def __init__(self, service, total_size, chunk_size, stream, parallel, encryptor=None, padder=None, **kwargs):
+    def __init__(
+            self, service,
+            total_size,
+            chunk_size,
+            stream,
+            parallel,
+            encryptor=None,
+            padder=None,
+            progress_hook=None,
+            **kwargs):
         self.service = service
         self.total_size = total_size
         self.chunk_size = chunk_size
@@ -142,6 +155,7 @@ def __init__(self, service, total_size, chunk_size, stream, parallel, encryptor=
         # Progress feedback
         self.progress_total = 0
         self.progress_lock = Lock() if parallel else None
+        self.progress_hook = progress_hook
 
         # Encryption
         self.encryptor = encryptor
@@ -199,6 +213,9 @@ def _update_progress(self, length):
         else:
             self.progress_total += length
 
+        if self.progress_hook:
+            self.progress_hook(self.progress_total, self.total_size)
+
     def _upload_chunk(self, chunk_offset, chunk_data):
         raise NotImplementedError("Must be implemented by child class.")
 

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/uploads_async.py

@@ -53,6 +53,7 @@ async def upload_data_chunks(
         max_concurrency=None,
         stream=None,
         encryption_options=None,
+        progress_hook=None,
         **kwargs):
 
     if encryption_options:
@@ -74,6 +75,7 @@ async def upload_data_chunks(
         chunk_size=chunk_size,
         stream=stream,
         parallel=parallel,
+        progress_hook=progress_hook,
         **kwargs)
 
     if parallel:
@@ -100,6 +102,7 @@ async def upload_substream_blocks(
         chunk_size=None,
         max_concurrency=None,
         stream=None,
+        progress_hook=None,
         **kwargs):
     parallel = max_concurrency > 1
     if parallel and 'modified_access_conditions' in kwargs:
@@ -111,6 +114,7 @@ async def upload_substream_blocks(
         chunk_size=chunk_size,
         stream=stream,
         parallel=parallel,
+        progress_hook=progress_hook,
         **kwargs)
 
     if parallel:
@@ -131,7 +135,16 @@ async def upload_substream_blocks(
 
 class _ChunkUploader(object):  # pylint: disable=too-many-instance-attributes
 
-    def __init__(self, service, total_size, chunk_size, stream, parallel, encryptor=None, padder=None, **kwargs):
+    def __init__(
+            self, service,
+            total_size,
+            chunk_size,
+            stream,
+            parallel,
+            encryptor=None,
+            padder=None,
+            progress_hook=None,
+            **kwargs):
         self.service = service
         self.total_size = total_size
         self.chunk_size = chunk_size
@@ -145,6 +158,7 @@ def __init__(self, service, total_size, chunk_size, stream, parallel, encryptor=
         # Progress feedback
         self.progress_total = 0
         self.progress_lock = Lock() if parallel else None
+        self.progress_hook = progress_hook
 
         # Encryption
         self.encryptor = encryptor
@@ -202,6 +216,9 @@ async def _update_progress(self, length):
         else:
             self.progress_total += length
 
+        if self.progress_hook:
+            self.progress_hook(self.progress_total, self.total_size)
+
     async def _upload_chunk(self, chunk_offset, chunk_data):
         raise NotImplementedError("Must be implemented by child class.")
 

sdk/storage/azure-storage-blob/azure/storage/blob/_upload_helpers.py

@@ -84,6 +84,7 @@ def upload_block_blob(  # pylint: disable=too-many-locals
         immutability_policy_expiry = None if immutability_policy is None else immutability_policy.expiry_time
         immutability_policy_mode = None if immutability_policy is None else immutability_policy.policy_mode
         legal_hold = kwargs.pop('legal_hold', None)
+        progress_hook = kwargs.pop('progress_hook', None)
 
         # Do single put if the size is smaller than or equal config.max_single_put_size
         if adjusted_count is not None and (adjusted_count <= blob_settings.max_single_put_size):
@@ -96,7 +97,8 @@ def upload_block_blob(  # pylint: disable=too-many-locals
             if encryption_options.get('key'):
                 encryption_data, data = encrypt_blob(data, encryption_options['key'])
                 headers['x-ms-meta-encryptiondata'] = encryption_data
-            return client.upload(
+
+            response = client.upload(
                 body=data,
                 content_length=adjusted_count,
                 blob_http_headers=blob_headers,
@@ -112,6 +114,11 @@ def upload_block_blob(  # pylint: disable=too-many-locals
                 legal_hold=legal_hold,
                 **kwargs)
 
+            if progress_hook:
+                progress_hook(adjusted_count, adjusted_count)
+
+            return response
+
         use_original_upload_path = blob_settings.use_byte_buffer or \
             validate_content or encryption_options.get('required') or \
             blob_settings.max_block_size < blob_settings.min_large_block_upload_threshold or \
@@ -133,6 +140,7 @@ def upload_block_blob(  # pylint: disable=too-many-locals
                 stream=stream,
                 validate_content=validate_content,
                 encryption_options=encryption_options,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs
             )
@@ -145,6 +153,7 @@ def upload_block_blob(  # pylint: disable=too-many-locals
                 max_concurrency=max_concurrency,
                 stream=stream,
                 validate_content=validate_content,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs
             )
@@ -200,6 +209,7 @@ def upload_page_blob(
         if encryption_options and encryption_options.get('data'):
             headers['x-ms-meta-encryptiondata'] = encryption_options['data']
         blob_tags_string = kwargs.pop('blob_tags_string', None)
+        progress_hook = kwargs.pop('progress_hook', None)
 
         response = client.create(
             content_length=0,
@@ -223,6 +233,7 @@ def upload_page_blob(
             max_concurrency=max_concurrency,
             validate_content=validate_content,
             encryption_options=encryption_options,
+            progress_hook=progress_hook,
             headers=headers,
             **kwargs)
 
@@ -254,6 +265,7 @@ def upload_append_blob(  # pylint: disable=unused-argument
             max_size=kwargs.pop('maxsize_condition', None),
             append_position=None)
         blob_tags_string = kwargs.pop('blob_tags_string', None)
+        progress_hook = kwargs.pop('progress_hook', None)
 
         try:
             if overwrite:
@@ -272,6 +284,7 @@ def upload_append_blob(  # pylint: disable=unused-argument
                 max_concurrency=max_concurrency,
                 validate_content=validate_content,
                 append_position_access_conditions=append_conditions,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs)
         except HttpResponseError as error:
@@ -300,6 +313,7 @@ def upload_append_blob(  # pylint: disable=unused-argument
                 max_concurrency=max_concurrency,
                 validate_content=validate_content,
                 append_position_access_conditions=append_conditions,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs)
     except HttpResponseError as error:

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_client_async.py

@@ -366,6 +366,10 @@ async def upload_blob(
 
         :keyword str encoding:
             Defaults to UTF-8.
+        :keyword Callable[[int, Optional[int]], None] progress_hook:
+            A callback to track the progress of a long running upload. The signature is
+            function(current: int, total: Optional[int]) where current is the number of bytes transfered
+            so far, and total is the size of the blob or None if the size is unknown.
         :keyword int timeout:
             The timeout parameter is expressed in seconds. This method may make
             multiple calls to the Azure service and the timeout will apply to

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_upload_helpers.py

@@ -59,6 +59,7 @@ async def upload_block_blob(  # pylint: disable=too-many-locals
         immutability_policy_expiry = None if immutability_policy is None else immutability_policy.expiry_time
         immutability_policy_mode = None if immutability_policy is None else immutability_policy.policy_mode
         legal_hold = kwargs.pop('legal_hold', None)
+        progress_hook = kwargs.pop('progress_hook', None)
 
         # Do single put if the size is smaller than config.max_single_put_size
         if adjusted_count is not None and (adjusted_count <= blob_settings.max_single_put_size):
@@ -71,7 +72,7 @@ async def upload_block_blob(  # pylint: disable=too-many-locals
             if encryption_options.get('key'):
                 encryption_data, data = encrypt_blob(data, encryption_options['key'])
                 headers['x-ms-meta-encryptiondata'] = encryption_data
-            return await client.upload(
+            response = await client.upload(
                 body=data,
                 content_length=adjusted_count,
                 blob_http_headers=blob_headers,
@@ -87,6 +88,11 @@ async def upload_block_blob(  # pylint: disable=too-many-locals
                 legal_hold=legal_hold,
                 **kwargs)
 
+            if progress_hook:
+                progress_hook(adjusted_count, adjusted_count)
+
+            return response
+
         use_original_upload_path = blob_settings.use_byte_buffer or \
             validate_content or encryption_options.get('required') or \
             blob_settings.max_block_size < blob_settings.min_large_block_upload_threshold or \
@@ -108,6 +114,7 @@ async def upload_block_blob(  # pylint: disable=too-many-locals
                 stream=stream,
                 validate_content=validate_content,
                 encryption_options=encryption_options,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs
             )
@@ -120,6 +127,7 @@ async def upload_block_blob(  # pylint: disable=too-many-locals
                 max_concurrency=max_concurrency,
                 stream=stream,
                 validate_content=validate_content,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs
             )
@@ -175,6 +183,7 @@ async def upload_page_blob(
         if encryption_options and encryption_options.get('data'):
             headers['x-ms-meta-encryptiondata'] = encryption_options['data']
         blob_tags_string = kwargs.pop('blob_tags_string', None)
+        progress_hook = kwargs.pop('progress_hook', None)
 
         response = await client.create(
             content_length=0,
@@ -198,6 +207,7 @@ async def upload_page_blob(
             max_concurrency=max_concurrency,
             validate_content=validate_content,
             encryption_options=encryption_options,
+            progress_hook=progress_hook,
             headers=headers,
             **kwargs)
 
@@ -229,6 +239,7 @@ async def upload_append_blob(  # pylint: disable=unused-argument
             max_size=kwargs.pop('maxsize_condition', None),
             append_position=None)
         blob_tags_string = kwargs.pop('blob_tags_string', None)
+        progress_hook = kwargs.pop('progress_hook', None)
 
         try:
             if overwrite:
@@ -247,6 +258,7 @@ async def upload_append_blob(  # pylint: disable=unused-argument
                 max_concurrency=max_concurrency,
                 validate_content=validate_content,
                 append_position_access_conditions=append_conditions,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs)
         except HttpResponseError as error:
@@ -275,6 +287,7 @@ async def upload_append_blob(  # pylint: disable=unused-argument
                 max_concurrency=max_concurrency,
                 validate_content=validate_content,
                 append_position_access_conditions=append_conditions,
+                progress_hook=progress_hook,
                 headers=headers,
                 **kwargs)
     except HttpResponseError as error: