[Storage] Add progress callback to download_blob methods (Azure#24276)

Author: jalauzon-msft

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

@@ -840,6 +840,11 @@ def download_blob(self, offset=None, length=None, **kwargs):
             The number of parallel connections with which to download.
         :keyword str encoding:
             Encoding to decode the downloaded bytes. Default is None, i.e. no decoding.
+        :keyword progress_hook:
+            A callback to track the progress of a long running download. The signature is
+            function(current: int, total: int) where current is the number of bytes transfered
+            so far, and total is the total size of the download.
+        :paramtype progress_hook: Callable[[int, int], None]
         :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/_container_client.py

@@ -1134,6 +1134,11 @@ def download_blob(self, blob, offset=None, length=None, **kwargs):
             The number of parallel connections with which to download.
         :keyword str encoding:
             Encoding to decode the downloaded bytes. Default is None, i.e. no decoding.
+        :keyword progress_hook:
+            A callback to track the progress of a long running download. The signature is
+            function(current: int, total: int) where current is the number of bytes transfered
+            so far, and total is the total size of the download.
+        :paramtype progress_hook: Callable[[int, int], None]
         :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/_download.py

@@ -10,7 +10,7 @@
 
 import warnings
 from io import BytesIO
-from typing import Iterator
+from typing import Iterator, Union
 
 import requests
 from azure.core.exceptions import HttpResponseError, ServiceResponseError
@@ -81,6 +81,7 @@ def __init__(
         parallel=None,
         validate_content=None,
         encryption_options=None,
+        progress_hook=None,
         **kwargs
     ):
         self.client = client
@@ -96,6 +97,7 @@ def __init__(
         self.stream = stream
         self.stream_lock = threading.Lock() if parallel else None
         self.progress_lock = threading.Lock() if parallel else None
+        self.progress_hook = progress_hook
 
         # For a parallel download, the stream is always seekable, so we note down the current position
         # in order to seek to the right place when out-of-order chunks come in
@@ -143,6 +145,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 _write_to_stream(self, chunk_data, chunk_start):
         if self.stream_lock:
             with self.stream_lock:  # pylint: disable=not-context-manager
@@ -322,6 +327,7 @@ def __init__(
         self._encoding = encoding
         self._validate_content = validate_content
         self._encryption_options = encryption_options or {}
+        self._progress_hook = kwargs.pop('progress_hook', None)
         self._request_options = kwargs
         self._location_mode = None
         self._download_complete = False
@@ -514,7 +520,6 @@ def readall(self):
         """Download the contents of this blob.
 
         This operation is blocking until all data is downloaded.
-
         :rtype: bytes or str
         """
         stream = BytesIO()
@@ -583,6 +588,9 @@ def readinto(self, stream):
 
         # Write the content to the user stream
         stream.write(self._current_content)
+        if self._progress_hook:
+            self._progress_hook(len(self._current_content), self.size)
+
         if self._download_complete:
             return self.size
 
@@ -604,6 +612,7 @@ def readinto(self, stream):
             validate_content=self._validate_content,
             encryption_options=self._encryption_options,
             use_location=self._location_mode,
+            progress_hook=self._progress_hook,
             **self._request_options
         )
         if parallel:

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

@@ -465,6 +465,11 @@ async def download_blob(self, offset=None, length=None, **kwargs):
             The number of parallel connections with which to download.
         :keyword str encoding:
             Encoding to decode the downloaded bytes. Default is None, i.e. no decoding.
+        :keyword progress_hook:
+            An async callback to track the progress of a long running download. The signature is
+            function(current: int, total: int) where current is the number of bytes transfered
+            so far, and total is the total size of the download.
+        :paramtype progress_hook: Callable[[int, int], Awaitable[None]]
         :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/_container_client_async.py

@@ -994,6 +994,11 @@ async def download_blob(self, blob, offset=None, length=None, **kwargs):
             The number of parallel connections with which to download.
         :keyword str encoding:
             Encoding to decode the downloaded bytes. Default is None, i.e. no decoding.
+        :keyword progress_hook:
+            An async callback to track the progress of a long running download. The signature is
+            function(current: int, total: int) where current is the number of bytes transfered
+            so far, and total is the total size of the download.
+        :paramtype progress_hook: Callable[[int, int], Awaitable[None]]
         :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/_download_async.py

@@ -67,6 +67,9 @@ async def _update_progress(self, length):
         else:
             self.progress_total += length
 
+        if self.progress_hook:
+            await self.progress_hook(self.progress_total, self.total_size)
+
     async def _write_to_stream(self, chunk_data, chunk_start):
         if self.stream_lock:
             async with self.stream_lock:  # pylint: disable=not-async-context-manager
@@ -220,6 +223,7 @@ def __init__(
         self._encoding = encoding
         self._validate_content = validate_content
         self._encryption_options = encryption_options or {}
+        self._progress_hook = kwargs.pop('progress_hook', None)
         self._request_options = kwargs
         self._location_mode = None
         self._download_complete = False
@@ -472,6 +476,9 @@ async def readinto(self, stream):
 
         # Write the content to the user stream
         stream.write(self._current_content)
+        if self._progress_hook:
+            await self._progress_hook(len(self._current_content), self.size)
+
         if self._download_complete:
             return self.size
 
@@ -493,6 +500,7 @@ async def readinto(self, stream):
             validate_content=self._validate_content,
             encryption_options=self._encryption_options,
             use_location=self._location_mode,
+            progress_hook=self._progress_hook,
             **self._request_options)
 
         dl_tasks = downloader.get_chunk_offsets()