Skip to content

Commit 7e801a6

Browse files
weirongw23-msftweirongw23-msft
authored
[Storage] Added retry logic in Blob download's process_content for ServiceResponseError (#45055)
1 parent b46abc2 commit 7e801a6

5 files changed

Lines changed: 105 additions & 9 deletions

File tree

‎sdk/storage/azure-storage-blob/assets.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,5 @@
22
"AssetsRepo": "Azure/azure-sdk-assets",
33
"AssetsRepoPrefixPath": "python",
44
"TagPrefix": "python/storage/azure-storage-blob",
5-
"Tag": "python/storage/azure-storage-blob_16c5acad24"
5+
"Tag": "python/storage/azure-storage-blob_89c4f2856e"
66
}

‎sdk/storage/azure-storage-blob/azure/storage/blob/_download.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
overload, Tuple, TypeVar, Union, TYPE_CHECKING
1616
)
1717

18-
from azure.core.exceptions import DecodeError, HttpResponseError, IncompleteReadError
18+
from azure.core.exceptions import DecodeError, HttpResponseError, IncompleteReadError, ServiceResponseError
1919
from azure.core.tracing.common import with_current_context
2020

2121
from ._shared.request_handlers import validate_and_format_range_headers
@@ -235,7 +235,7 @@ def _download_chunk(self, chunk_start: int, chunk_end: int) -> Tuple[bytes, int]
235235
try:
236236
chunk_data = process_content(response, offset[0], offset[1], self.encryption_options)
237237
retry_active = False
238-
except (IncompleteReadError, HttpResponseError, DecodeError) as error:
238+
except (IncompleteReadError, HttpResponseError, DecodeError, ServiceResponseError) as error:
239239
retry_total -= 1
240240
if retry_total <= 0:
241241
raise HttpResponseError(error, error=error) from error
@@ -522,7 +522,7 @@ def _initial_request(self):
522522
self._encryption_options
523523
)
524524
retry_active = False
525-
except (IncompleteReadError, HttpResponseError, DecodeError) as error:
525+
except (IncompleteReadError, HttpResponseError, DecodeError, ServiceResponseError) as error:
526526
retry_total -= 1
527527
if retry_total <= 0:
528528
raise HttpResponseError(error, error=error) from error

‎sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
Tuple, TypeVar, Union, TYPE_CHECKING
2020
)
2121

22-
from azure.core.exceptions import DecodeError, HttpResponseError, IncompleteReadError
22+
from azure.core.exceptions import DecodeError, HttpResponseError, IncompleteReadError, ServiceResponseError
2323

2424
from .._shared.request_handlers import validate_and_format_range_headers
2525
from .._shared.response_handlers import parse_length_from_content_range, process_storage_error
@@ -144,7 +144,7 @@ async def _download_chunk(self, chunk_start: int, chunk_end: int) -> Tuple[bytes
144144
try:
145145
chunk_data = await process_content(response, offset[0], offset[1], self.encryption_options)
146146
retry_active = False
147-
except (IncompleteReadError, HttpResponseError, DecodeError) as error:
147+
except (IncompleteReadError, HttpResponseError, DecodeError, ServiceResponseError) as error:
148148
retry_total -= 1
149149
if retry_total <= 0:
150150
raise HttpResponseError(error, error=error) from error
@@ -432,7 +432,7 @@ async def _initial_request(self):
432432
self._encryption_options
433433
)
434434
retry_active = False
435-
except (IncompleteReadError, HttpResponseError, DecodeError) as error:
435+
except (IncompleteReadError, HttpResponseError, DecodeError, ServiceResponseError) as error:
436436
retry_total -= 1
437437
if retry_total <= 0:
438438
raise HttpResponseError(error, error=error) from error

‎sdk/storage/azure-storage-blob/tests/test_retry.py‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
ClientAuthenticationError,
1414
HttpResponseError,
1515
ResourceExistsError,
16-
ServiceResponseError
16+
ServiceResponseError,
17+
ServiceResponseTimeoutError
1718
)
1819
from azure.core.pipeline.transport import RequestsTransport
1920
from azure.storage.blob._shared.authentication import AzureSigningError
@@ -639,4 +640,51 @@ def assert_exception_retry_hook(**kwargs):
639640

640641
assert retry_counter.count == 3
641642

643+
@BlobPreparer()
644+
@recorded_by_proxy
645+
def test_retry_on_service_response_error(self, **kwargs):
646+
storage_account_name = kwargs.pop("storage_account_name")
647+
storage_account_key = kwargs.pop("storage_account_key")
648+
649+
# Arrange
650+
container_name = self.get_resource_name('utcontainer')
651+
blob_name = self.get_resource_name('blob')
652+
service = self._create_storage_service(
653+
BlobServiceClient, storage_account_name, storage_account_key, max_block_size=4)
654+
container = service.create_container(container_name)
655+
data = b'abcd' * 4
656+
container.upload_blob(blob_name, data, overwrite=True)
657+
658+
retry = LinearRetry(backoff=1, random_jitter_range=1)
659+
retry_counter = RetryCounter()
660+
retry_service = self._create_storage_service(
661+
BlobServiceClient,
662+
storage_account_name,
663+
storage_account_key,
664+
retry_policy=retry,
665+
max_block_size=4
666+
)
667+
blob = retry_service.get_blob_client(container_name, blob_name)
668+
669+
# Mock the internal response to raise ServiceResponseError on first chunk processing
670+
from azure.storage.blob._download import process_content as real_process_content
671+
672+
def mock_process_content_with_error(response, start_offset, end_offset, encryption):
673+
retry_counter.simple_count(retry)
674+
conn_error = AzureError("Connection reset by peer")
675+
if retry_counter.count == 1:
676+
raise ServiceResponseError(conn_error, error=conn_error)
677+
elif retry_counter.count == 2:
678+
raise ServiceResponseTimeoutError(conn_error, error=conn_error)
679+
return real_process_content(response, start_offset, end_offset, encryption)
680+
681+
# Act
682+
try:
683+
with mock.patch('azure.storage.blob._download.process_content', side_effect=mock_process_content_with_error):
684+
downloaded_data = blob.download_blob().readall()
685+
assert downloaded_data == data
686+
assert retry_counter.count >= 3
687+
finally:
688+
service.delete_container(container_name)
689+
642690
# ------------------------------------------------------------------------------

‎sdk/storage/azure-storage-blob/tests/test_retry_async.py‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@
1717
IncompleteReadError,
1818
HttpResponseError,
1919
ResourceExistsError,
20-
ServiceResponseError
20+
ServiceResponseError,
21+
ServiceResponseTimeoutError
2122
)
2223
from azure.core.pipeline.transport import AioHttpTransport
2324
from azure.storage.blob import LocationMode
@@ -618,4 +619,51 @@ def assert_exception_retry_hook(**kwargs):
618619

619620
assert retry_counter.count == 3
620621

622+
@BlobPreparer()
623+
@recorded_by_proxy_async
624+
async def test_retry_on_service_response_error(self, **kwargs):
625+
storage_account_name = kwargs.pop("storage_account_name")
626+
storage_account_key = kwargs.pop("storage_account_key")
627+
628+
# Arrange
629+
container_name = self.get_resource_name('utcontainer')
630+
blob_name = self.get_resource_name('blob')
631+
service = self._create_storage_service(
632+
BlobServiceClient, storage_account_name, storage_account_key, max_block_size=4)
633+
container = await service.create_container(container_name)
634+
data = b'abcd' * 4
635+
await container.upload_blob(blob_name, data, overwrite=True)
636+
637+
retry = LinearRetry(backoff=1, random_jitter_range=1)
638+
retry_counter = RetryCounter()
639+
retry_service = self._create_storage_service(
640+
BlobServiceClient,
641+
storage_account_name,
642+
storage_account_key,
643+
retry_policy=retry,
644+
max_block_size=4
645+
)
646+
blob = retry_service.get_blob_client(container_name, blob_name)
647+
648+
# Mock the internal response to raise ServiceResponseError on first chunk processing
649+
from azure.storage.blob.aio._download_async import process_content as real_process_content
650+
651+
async def mock_process_content_with_error(response, start_offset, end_offset, encryption):
652+
retry_counter.simple_count(retry)
653+
conn_error = AzureError("Connection reset by peer")
654+
if retry_counter.count == 1:
655+
raise ServiceResponseError(conn_error, error=conn_error)
656+
elif retry_counter.count == 2:
657+
raise ServiceResponseTimeoutError(conn_error, error=conn_error)
658+
return await real_process_content(response, start_offset, end_offset, encryption)
659+
660+
# Act
661+
try:
662+
with mock.patch('azure.storage.blob.aio._download_async.process_content', side_effect=mock_process_content_with_error):
663+
downloaded_data = await (await blob.download_blob()).readall()
664+
assert downloaded_data == data
665+
assert retry_counter.count >= 3
666+
finally:
667+
await service.delete_container(container_name)
668+
621669
# ------------------------------------------------------------------------------

0 commit comments

Comments
 (0)