[Storage] Add progress_hook to file-share upload/download (#24997)

sdk/storage/azure-storage-file-share/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 12.9.0 (Unreleased)
 
 ### Features Added
+- Added support for progress tracking to `upload_file()` and `download_file()` via a new optional callback, `progress_hook`.
 
 ### Bugs Fixed
 

sdk/storage/azure-storage-file-share/azure/storage/fileshare/_directory_client.py

@@ -845,6 +845,11 @@ def upload_file(
             file.
         :keyword int max_concurrency:
             Maximum number of parallel connections to use.
+        :keyword 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 transferred
+            so far, and total is the size of the blob or None if the size is unknown.
+        :paramtype progress_hook: Callable[[int, Optional[int]], None]
         :keyword int timeout:
             The timeout parameter is expressed in seconds.
         :keyword str encoding:

sdk/storage/azure-storage-file-share/azure/storage/fileshare/_download.py

@@ -38,6 +38,7 @@ def __init__(
         stream=None,
         parallel=None,
         validate_content=None,
+        progress_hook=None,
         etag=None,
         **kwargs
     ):
@@ -53,6 +54,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
@@ -97,6 +99,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
@@ -227,6 +232,7 @@ def __init__(
         self._max_concurrency = max_concurrency
         self._encoding = encoding
         self._validate_content = validate_content
+        self._progress_hook = kwargs.pop('progress_hook', None)
         self._request_options = kwargs
         self._location_mode = None
         self._download_complete = False
@@ -446,6 +452,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
 
@@ -465,6 +474,7 @@ def readinto(self, stream):
             parallel=parallel,
             validate_content=self._validate_content,
             use_location=self._location_mode,
+            progress_hook=self._progress_hook,
             etag=self._etag,
             **self._request_options
         )

sdk/storage/azure-storage-file-share/azure/storage/fileshare/_file_client.py

@@ -66,6 +66,7 @@ def _upload_file_helper(
         file_last_write_time="now",
         file_permission=None,
         file_permission_key=None,
+        progress_hook=None,
         **kwargs):
     try:
         if size is None or size < 0:
@@ -93,6 +94,7 @@ def _upload_file_helper(
             stream=stream,
             max_concurrency=max_concurrency,
             validate_content=validate_content,
+            progress_hook=progress_hook,
             timeout=timeout,
             **kwargs
         )
@@ -489,6 +491,11 @@ def upload_file(
             .. versionadded:: 12.1.0
 
         :paramtype lease: ~azure.storage.fileshare.ShareLeaseClient or str
+        :keyword 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 transferred
+            so far, and total is the size of the blob or None if the size is unknown.
+        :paramtype progress_hook: Callable[[int, Optional[int]], None]
         :keyword int timeout:
             The timeout parameter is expressed in seconds.
         :keyword str encoding:
@@ -509,6 +516,7 @@ def upload_file(
         content_settings = kwargs.pop('content_settings', None)
         max_concurrency = kwargs.pop('max_concurrency', 1)
         validate_content = kwargs.pop('validate_content', False)
+        progress_hook = kwargs.pop('progress_hook', None)
         timeout = kwargs.pop('timeout', None)
         encoding = kwargs.pop('encoding', 'UTF-8')
 
@@ -542,6 +550,7 @@ def upload_file(
             file_last_write_time=file_last_write_time,
             file_permission=file_permission,
             file_permission_key=permission_key,
+            progress_hook=progress_hook,
             **kwargs)
 
     @distributed_trace
@@ -740,6 +749,11 @@ def download_file(
             .. versionadded:: 12.1.0
 
         :paramtype lease: ~azure.storage.fileshare.ShareLeaseClient or str
+        :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 transferred
+            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.
         :returns: A streaming object (StorageStreamDownloader)

sdk/storage/azure-storage-file-share/azure/storage/fileshare/aio/_directory_client_async.py

@@ -718,6 +718,11 @@ async def upload_file(
             file.
         :keyword int max_concurrency:
             Maximum number of parallel connections to use.
+        :keyword progress_hook:
+            An async 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 transferred
+            so far, and total is the size of the blob or None if the size is unknown.
+        :paramtype progress_hook: Callable[[int, Optional[int]], Awaitable[None]]
         :keyword int timeout:
             The timeout parameter is expressed in seconds.
         :keyword str encoding:

sdk/storage/azure-storage-file-share/azure/storage/fileshare/aio/_download_async.py

@@ -52,6 +52,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
@@ -184,6 +187,7 @@ def __init__(
         self._max_concurrency = max_concurrency
         self._encoding = encoding
         self._validate_content = validate_content
+        self._progress_hook = kwargs.pop('progress_hook', None)
         self._request_options = kwargs
         self._location_mode = None
         self._download_complete = False
@@ -400,6 +404,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
 
@@ -419,6 +426,7 @@ async def readinto(self, stream):
             parallel=parallel,
             validate_content=self._validate_content,
             use_location=self._location_mode,
+            progress_hook=self._progress_hook,
             etag=self._etag,
             **self._request_options)
 

sdk/storage/azure-storage-file-share/azure/storage/fileshare/aio/_file_client_async.py

@@ -59,6 +59,7 @@ async def _upload_file_helper(
     file_last_write_time="now",
     file_permission=None,
     file_permission_key=None,
+    progress_hook=None,
     **kwargs
 ):
     try:
@@ -85,6 +86,7 @@ async def _upload_file_helper(
             stream=stream,
             max_concurrency=max_concurrency,
             validate_content=validate_content,
+            progress_hook=progress_hook,
             timeout=timeout,
             **kwargs
         )
@@ -360,6 +362,11 @@ async def upload_file(
             .. versionadded:: 12.1.0
 
         :paramtype lease: ~azure.storage.fileshare.aio.ShareLeaseClient or str
+        :keyword progress_hook:
+            An async 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 transferred
+            so far, and total is the size of the blob or None if the size is unknown.
+        :paramtype progress_hook: Callable[[int, Optional[int]], Awaitable[None]]
         :keyword int timeout:
             The timeout parameter is expressed in seconds.
         :returns: File-updated property dict (Etag and last modified).
@@ -378,6 +385,7 @@ async def upload_file(
         content_settings = kwargs.pop('content_settings', None)
         max_concurrency = kwargs.pop('max_concurrency', 1)
         validate_content = kwargs.pop('validate_content', False)
+        progress_hook = kwargs.pop('progress_hook', None)
         timeout = kwargs.pop('timeout', None)
         encoding = kwargs.pop('encoding', 'UTF-8')
 
@@ -411,6 +419,7 @@ async def upload_file(
             file_last_write_time=file_last_write_time,
             file_permission=file_permission,
             file_permission_key=permission_key,
+            progress_hook=progress_hook,
             **kwargs
         )
 
@@ -613,6 +622,11 @@ async def download_file(
             .. versionadded:: 12.1.0
 
         :paramtype lease: ~azure.storage.fileshare.aio.ShareLeaseClient or str
+        :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 transferred
+            so far, and total is the size of the blob or None if the size is unknown.
+        :paramtype progress_hook: Callable[[int, int], Awaitable[None]]
         :keyword int timeout:
             The timeout parameter is expressed in seconds.
         :returns: A streaming object (StorageStreamDownloader)