Skip to content

Commit 155eb8b

Browse files
[Storage] Add progress_hook to file-share upload/download (#24997)
1 parent eb4afb3 commit 155eb8b

19 files changed

+2868
-29
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
## 12.9.0 (Unreleased)
44

55
### Features Added
6+
- Added support for progress tracking to `upload_file()` and `download_file()` via a new optional callback, `progress_hook`.
67

78
### Bugs Fixed
89

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -845,6 +845,11 @@ def upload_file(
845845
file.
846846
:keyword int max_concurrency:
847847
Maximum number of parallel connections to use.
848+
:keyword progress_hook:
849+
A callback to track the progress of a long running upload. The signature is
850+
function(current: int, total: Optional[int]) where current is the number of bytes transferred
851+
so far, and total is the size of the blob or None if the size is unknown.
852+
:paramtype progress_hook: Callable[[int, Optional[int]], None]
848853
:keyword int timeout:
849854
The timeout parameter is expressed in seconds.
850855
:keyword str encoding:

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ def __init__(
3838
stream=None,
3939
parallel=None,
4040
validate_content=None,
41+
progress_hook=None,
4142
etag=None,
4243
**kwargs
4344
):
@@ -53,6 +54,7 @@ def __init__(
5354
self.stream = stream
5455
self.stream_lock = threading.Lock() if parallel else None
5556
self.progress_lock = threading.Lock() if parallel else None
57+
self.progress_hook = progress_hook
5658

5759
# For a parallel download, the stream is always seekable, so we note down the current position
5860
# in order to seek to the right place when out-of-order chunks come in
@@ -97,6 +99,9 @@ def _update_progress(self, length):
9799
else:
98100
self.progress_total += length
99101

102+
if self.progress_hook:
103+
self.progress_hook(self.progress_total, self.total_size)
104+
100105
def _write_to_stream(self, chunk_data, chunk_start):
101106
if self.stream_lock:
102107
with self.stream_lock: # pylint: disable=not-context-manager
@@ -227,6 +232,7 @@ def __init__(
227232
self._max_concurrency = max_concurrency
228233
self._encoding = encoding
229234
self._validate_content = validate_content
235+
self._progress_hook = kwargs.pop('progress_hook', None)
230236
self._request_options = kwargs
231237
self._location_mode = None
232238
self._download_complete = False
@@ -446,6 +452,9 @@ def readinto(self, stream):
446452

447453
# Write the content to the user stream
448454
stream.write(self._current_content)
455+
if self._progress_hook:
456+
self._progress_hook(len(self._current_content), self.size)
457+
449458
if self._download_complete:
450459
return self.size
451460

@@ -465,6 +474,7 @@ def readinto(self, stream):
465474
parallel=parallel,
466475
validate_content=self._validate_content,
467476
use_location=self._location_mode,
477+
progress_hook=self._progress_hook,
468478
etag=self._etag,
469479
**self._request_options
470480
)

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ def _upload_file_helper(
6666
file_last_write_time="now",
6767
file_permission=None,
6868
file_permission_key=None,
69+
progress_hook=None,
6970
**kwargs):
7071
try:
7172
if size is None or size < 0:
@@ -93,6 +94,7 @@ def _upload_file_helper(
9394
stream=stream,
9495
max_concurrency=max_concurrency,
9596
validate_content=validate_content,
97+
progress_hook=progress_hook,
9698
timeout=timeout,
9799
**kwargs
98100
)
@@ -489,6 +491,11 @@ def upload_file(
489491
.. versionadded:: 12.1.0
490492
491493
:paramtype lease: ~azure.storage.fileshare.ShareLeaseClient or str
494+
:keyword progress_hook:
495+
A callback to track the progress of a long running upload. The signature is
496+
function(current: int, total: Optional[int]) where current is the number of bytes transferred
497+
so far, and total is the size of the blob or None if the size is unknown.
498+
:paramtype progress_hook: Callable[[int, Optional[int]], None]
492499
:keyword int timeout:
493500
The timeout parameter is expressed in seconds.
494501
:keyword str encoding:
@@ -509,6 +516,7 @@ def upload_file(
509516
content_settings = kwargs.pop('content_settings', None)
510517
max_concurrency = kwargs.pop('max_concurrency', 1)
511518
validate_content = kwargs.pop('validate_content', False)
519+
progress_hook = kwargs.pop('progress_hook', None)
512520
timeout = kwargs.pop('timeout', None)
513521
encoding = kwargs.pop('encoding', 'UTF-8')
514522

@@ -542,6 +550,7 @@ def upload_file(
542550
file_last_write_time=file_last_write_time,
543551
file_permission=file_permission,
544552
file_permission_key=permission_key,
553+
progress_hook=progress_hook,
545554
**kwargs)
546555

547556
@distributed_trace
@@ -740,6 +749,11 @@ def download_file(
740749
.. versionadded:: 12.1.0
741750
742751
:paramtype lease: ~azure.storage.fileshare.ShareLeaseClient or str
752+
:keyword progress_hook:
753+
A callback to track the progress of a long running download. The signature is
754+
function(current: int, total: int) where current is the number of bytes transferred
755+
so far, and total is the total size of the download.
756+
:paramtype progress_hook: Callable[[int, int], None]
743757
:keyword int timeout:
744758
The timeout parameter is expressed in seconds.
745759
:returns: A streaming object (StorageStreamDownloader)

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -718,6 +718,11 @@ async def upload_file(
718718
file.
719719
:keyword int max_concurrency:
720720
Maximum number of parallel connections to use.
721+
:keyword progress_hook:
722+
An async callback to track the progress of a long running upload. The signature is
723+
function(current: int, total: Optional[int]) where current is the number of bytes transferred
724+
so far, and total is the size of the blob or None if the size is unknown.
725+
:paramtype progress_hook: Callable[[int, Optional[int]], Awaitable[None]]
721726
:keyword int timeout:
722727
The timeout parameter is expressed in seconds.
723728
:keyword str encoding:

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@ async def _update_progress(self, length):
5252
else:
5353
self.progress_total += length
5454

55+
if self.progress_hook:
56+
await self.progress_hook(self.progress_total, self.total_size)
57+
5558
async def _write_to_stream(self, chunk_data, chunk_start):
5659
if self.stream_lock:
5760
async with self.stream_lock: # pylint: disable=not-async-context-manager
@@ -184,6 +187,7 @@ def __init__(
184187
self._max_concurrency = max_concurrency
185188
self._encoding = encoding
186189
self._validate_content = validate_content
190+
self._progress_hook = kwargs.pop('progress_hook', None)
187191
self._request_options = kwargs
188192
self._location_mode = None
189193
self._download_complete = False
@@ -400,6 +404,9 @@ async def readinto(self, stream):
400404

401405
# Write the content to the user stream
402406
stream.write(self._current_content)
407+
if self._progress_hook:
408+
await self._progress_hook(len(self._current_content), self.size)
409+
403410
if self._download_complete:
404411
return self.size
405412

@@ -419,6 +426,7 @@ async def readinto(self, stream):
419426
parallel=parallel,
420427
validate_content=self._validate_content,
421428
use_location=self._location_mode,
429+
progress_hook=self._progress_hook,
422430
etag=self._etag,
423431
**self._request_options)
424432

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ async def _upload_file_helper(
5959
file_last_write_time="now",
6060
file_permission=None,
6161
file_permission_key=None,
62+
progress_hook=None,
6263
**kwargs
6364
):
6465
try:
@@ -85,6 +86,7 @@ async def _upload_file_helper(
8586
stream=stream,
8687
max_concurrency=max_concurrency,
8788
validate_content=validate_content,
89+
progress_hook=progress_hook,
8890
timeout=timeout,
8991
**kwargs
9092
)
@@ -360,6 +362,11 @@ async def upload_file(
360362
.. versionadded:: 12.1.0
361363
362364
:paramtype lease: ~azure.storage.fileshare.aio.ShareLeaseClient or str
365+
:keyword progress_hook:
366+
An async callback to track the progress of a long running upload. The signature is
367+
function(current: int, total: Optional[int]) where current is the number of bytes transferred
368+
so far, and total is the size of the blob or None if the size is unknown.
369+
:paramtype progress_hook: Callable[[int, Optional[int]], Awaitable[None]]
363370
:keyword int timeout:
364371
The timeout parameter is expressed in seconds.
365372
:returns: File-updated property dict (Etag and last modified).
@@ -378,6 +385,7 @@ async def upload_file(
378385
content_settings = kwargs.pop('content_settings', None)
379386
max_concurrency = kwargs.pop('max_concurrency', 1)
380387
validate_content = kwargs.pop('validate_content', False)
388+
progress_hook = kwargs.pop('progress_hook', None)
381389
timeout = kwargs.pop('timeout', None)
382390
encoding = kwargs.pop('encoding', 'UTF-8')
383391

@@ -411,6 +419,7 @@ async def upload_file(
411419
file_last_write_time=file_last_write_time,
412420
file_permission=file_permission,
413421
file_permission_key=permission_key,
422+
progress_hook=progress_hook,
414423
**kwargs
415424
)
416425

@@ -613,6 +622,11 @@ async def download_file(
613622
.. versionadded:: 12.1.0
614623
615624
:paramtype lease: ~azure.storage.fileshare.aio.ShareLeaseClient or str
625+
:keyword progress_hook:
626+
An async callback to track the progress of a long running download. The signature is
627+
function(current: int, total: int) where current is the number of bytes transferred
628+
so far, and total is the size of the blob or None if the size is unknown.
629+
:paramtype progress_hook: Callable[[int, int], Awaitable[None]]
616630
:keyword int timeout:
617631
The timeout parameter is expressed in seconds.
618632
:returns: A streaming object (StorageStreamDownloader)

0 commit comments

Comments
 (0)