Skip to content

Commit

Permalink
Fix GH #16723
Browse files Browse the repository at this point in the history
  • Loading branch information
jochen-ott-by committed Feb 12, 2021
1 parent 5dee530 commit a198cf7
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 53 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -100,14 +100,11 @@ class StreamDownloadGenerator(object):
:param pipeline: The pipeline object
:param response: The response object.
"""
def __init__(self, pipeline, response):
self.pipeline = pipeline
self.request = response.request
def __init__(self, response):
self.response = response
self.block_size = response.block_size
self.iter_content_func = self.response.internal_response.iter_content(self.block_size)
self.content_length = int(response.headers.get('Content-Length', 0))
self.downloaded = 0

def __len__(self):
return self.content_length
Expand All @@ -116,42 +113,20 @@ def __iter__(self):
return self

def __next__(self):
retry_active = True
retry_total = 3
retry_interval = 1 # 1 second
while retry_active:
try:
chunk = next(self.iter_content_func)
if not chunk:
raise StopIteration()
self.downloaded += self.block_size
return chunk
except StopIteration:
self.response.internal_response.close()
try:
chunk = next(self.iter_content_func)
if not chunk:
raise StopIteration()
except (requests.exceptions.ChunkedEncodingError,
requests.exceptions.ConnectionError):
retry_total -= 1
if retry_total <= 0:
retry_active = False
else:
time.sleep(retry_interval)
headers = {'range': 'bytes=' + str(self.downloaded) + '-'}
resp = self.pipeline.run(self.request, stream=True, headers=headers)
if resp.http_response.status_code == 416:
raise
chunk = next(self.iter_content_func)
if not chunk:
raise StopIteration()
self.downloaded += len(chunk)
return chunk
continue
except requests.exceptions.StreamConsumedError:
raise
except Exception as err:
_LOGGER.warning("Unable to stream download: %s", err)
self.response.internal_response.close()
raise
return chunk
except StopIteration:
self.response.internal_response.close()
raise StopIteration()
except requests.exceptions.StreamConsumedError:
raise
except Exception as err:
_LOGGER.warning("Unable to stream download: %s", err)
self.response.internal_response.close()
raise
next = __next__ # Python 2 compatibility.


Expand All @@ -161,7 +136,7 @@ class RequestsTransportResponse(HttpResponse, _RequestsTransportResponseBase):
def stream_download(self, pipeline):
# type: (PipelineType) -> Iterator[bytes]
"""Generator for streaming request body data."""
return StreamDownloadGenerator(pipeline, self)
return StreamDownloadGenerator(self)


class RequestsTransport(HttpTransport):
Expand Down
16 changes: 3 additions & 13 deletions sdk/core/azure-core/tests/test_requests_universal.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,16 +117,6 @@ def close(self):
block_size,
)

def mock_run(self, *args, **kwargs):
return PipelineResponse(
None,
requests.Response(),
None,
)

transport = RequestsTransport()
pipeline = Pipeline(transport)
pipeline.run = mock_run
downloader = response.stream_download(pipeline)
full_response = b"".join(downloader)
assert len(full_response) == total_response_size
downloader = response.stream_download(None)
with pytest.raises(requests.exceptions.ConnectionError):
list(downloader)

0 comments on commit a198cf7

Please sign in to comment.