Raise exception rather than swallowing it if there is something wrong…

sdk/core/azure-core/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## 1.11.1 (Unreleased)
 
+### Bug Fixes
+
+- Raise exception rather than swallowing it if there is something wrong in retry stream downloading  #16723
 
 ## 1.11.0 (2021-02-08)
 

sdk/core/azure-core/azure/core/pipeline/transport/_aiohttp.py

@@ -228,17 +228,24 @@ async def __anext__(self):
             except _ResponseStopIteration:
                 self.response.internal_response.close()
                 raise StopAsyncIteration()
-            except (ChunkedEncodingError, ConnectionError):
+            except (ChunkedEncodingError, ConnectionError) as ex:
                 retry_total -= 1
                 if retry_total <= 0:
                     retry_active = False
                 else:
                     await asyncio.sleep(retry_interval)
                     headers = {'range': 'bytes=' + str(self.downloaded) + '-'}
-                    resp = await self.pipeline.run(self.request, stream=True, headers=headers)
-                    if resp.http_response.status_code == 416:
-                        raise
-                    chunk = await self.response.internal_response.content.read(self.block_size)
+                    try:
+                        resp = await self.pipeline.run(self.request, stream=True, headers=headers)
+                        if not resp.http_response:
+                            raise
+                        if resp.http_response.status_code == 416:
+                            raise
+                        chunk = await self.response.internal_response.content.read(self.block_size)
+                    except Exception as err:  # pylint: disable=broad-except
+                        _LOGGER.warning("Unable to stream download: %s", err)
+                        self.response.internal_response.close()
+                        raise ex
                     if not chunk:
                         raise StopAsyncIteration()
                     self.downloaded += len(chunk)

sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py

@@ -130,17 +130,24 @@ def __next__(self):
                 self.response.internal_response.close()
                 raise StopIteration()
             except (requests.exceptions.ChunkedEncodingError,
-                    requests.exceptions.ConnectionError):
+                    requests.exceptions.ConnectionError) as ex:
                 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)
+                    try:
+                        resp = self.pipeline.run(self.request, stream=True, headers=headers)
+                        if not resp.http_response:
+                            raise
+                        if resp.http_response.status_code == 416:
+                            raise
+                        chunk = next(self.iter_content_func)
+                    except Exception as err:   # pylint: disable=broad-except
+                        _LOGGER.warning("Unable to stream download: %s", err)
+                        self.response.internal_response.close()
+                        raise ex
                     if not chunk:
                         raise StopIteration()
                     self.downloaded += len(chunk)

sdk/core/azure-core/tests/async_tests/test_stream_generator_async.py

@@ -2,10 +2,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import requests
 from azure.core.pipeline.transport import (
     HttpRequest,
     AsyncHttpResponse,
     AsyncHttpTransport,
+    AsyncioRequestsTransportResponse,
+    AioHttpTransport,
 )
 from azure.core.pipeline import AsyncPipeline
 from azure.core.pipeline.transport._aiohttp import AioHttpStreamDownloadGenerator
@@ -105,3 +108,50 @@ async def __call__(self, *args, **kwargs):
     with mock.patch('asyncio.sleep', new_callable=AsyncMock):
         with pytest.raises(ConnectionError):
             await stream.__anext__()
+
+@pytest.mark.asyncio
+async def test_response_streaming_error_behavior():
+    # Test to reproduce https://github.com/Azure/azure-sdk-for-python/issues/16723
+    block_size = 103
+    total_response_size = 500
+    req_response = requests.Response()
+    req_request = requests.Request()
+
+    class FakeStreamWithConnectionError:
+        # fake object for urllib3.response.HTTPResponse
+
+        def stream(self, chunk_size, decode_content=False):
+            assert chunk_size == block_size
+            left = total_response_size
+            while left > 0:
+                if left <= block_size:
+                    raise ConnectionError()
+                data = b"X" * min(chunk_size, left)
+                left -= len(data)
+                yield data
+
+        def close(self):
+            pass
+
+    req_response.raw = FakeStreamWithConnectionError()
+
+    response = AsyncioRequestsTransportResponse(
+        req_request,
+        req_response,
+        block_size,
+    )
+
+    async def mock_run(self, *args, **kwargs):
+        return PipelineResponse(
+            None,
+            requests.Response(),
+            None,
+        )
+
+    transport = AioHttpTransport()
+    pipeline = AsyncPipeline(transport)
+    pipeline.run = mock_run
+    downloader = response.stream_download(pipeline)
+    with pytest.raises(ConnectionError):
+        while True:
+            await downloader.__anext__()

sdk/core/azure-core/tests/test_stream_generator.py

@@ -7,6 +7,8 @@
     HttpRequest,
     HttpResponse,
     HttpTransport,
+    RequestsTransport,
+    RequestsTransportResponse,
 )
 from azure.core.pipeline import Pipeline, PipelineResponse
 from azure.core.pipeline.transport._requests_basic import StreamDownloadGenerator
@@ -98,4 +100,49 @@ def close(self):
     stream = StreamDownloadGenerator(pipeline, http_response)
     with mock.patch('time.sleep', return_value=None):
         with pytest.raises(requests.exceptions.ConnectionError):
-            stream.__next__()
+            stream.__next__()
+
+def test_response_streaming_error_behavior():
+    # Test to reproduce https://github.com/Azure/azure-sdk-for-python/issues/16723
+    block_size = 103
+    total_response_size = 500
+    req_response = requests.Response()
+    req_request = requests.Request()
+
+    class FakeStreamWithConnectionError:
+        # fake object for urllib3.response.HTTPResponse
+
+        def stream(self, chunk_size, decode_content=False):
+            assert chunk_size == block_size
+            left = total_response_size
+            while left > 0:
+                if left <= block_size:
+                    raise requests.exceptions.ConnectionError()
+                data = b"X" * min(chunk_size, left)
+                left -= len(data)
+                yield data
+
+        def close(self):
+            pass
+
+    req_response.raw = FakeStreamWithConnectionError()
+
+    response = RequestsTransportResponse(
+        req_request,
+        req_response,
+        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)
+    with pytest.raises(requests.exceptions.ConnectionError):
+        full_response = b"".join(downloader)