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

@@ -218,38 +218,50 @@ async def __anext__(self):
         retry_active = True
         retry_total = 3
         retry_interval = 1  # 1 second
-        while retry_active:
-            try:
-                chunk = await self.response.internal_response.content.read(self.block_size)
-                if not chunk:
-                    raise _ResponseStopIteration()
-                self.downloaded += self.block_size
-                return chunk
-            except _ResponseStopIteration:
-                self.response.internal_response.close()
-                raise StopAsyncIteration()
-            except (ChunkedEncodingError, ConnectionError):
+        try:
+            chunk = await self.response.internal_response.content.read(self.block_size)
+            if not chunk:
+                raise _ResponseStopIteration()
+            self.downloaded += self.block_size
+            return chunk
+        except _ResponseStopIteration:
+            self.response.internal_response.close()
+            raise StopAsyncIteration()
+        except (ChunkedEncodingError, ConnectionError) as ex:
+            while retry_active:
                 retry_total -= 1
                 if retry_total <= 0:
-                    retry_active = False
+                    _LOGGER.warning("Unable to stream download: %s", ex)
+                    raise ex
                 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)
+                    # todo handle pre-set range & x-ms-range
+                    headers = self.request.headers
+                    headers.update({'range': 'bytes=' + str(self.downloaded) + '-'})
+                    try:
+                        resp = await self.pipeline.run(self.request, stream=True, headers=headers)
+                        if not resp.http_response:
+                            continue
+                        if resp.http_response.status_code == 416:
+                            continue
+                        self.response = resp
+                        chunk = await self.response.http_response.content.read(self.block_size)
+                    except StopIteration:
+                        self.response.internal_response.close()
+                        raise StopIteration()
+                    except Exception:  # pylint: disable=broad-except
+                        continue
                     if not chunk:
+                        self.response.internal_response.close()
                         raise StopAsyncIteration()
-                    self.downloaded += len(chunk)
+                    self.downloaded += self.block_size
                     return chunk
-                continue
-            except StreamConsumedError:
-                raise
-            except Exception as err:
-                _LOGGER.warning("Unable to stream download: %s", err)
-                self.response.internal_response.close()
-                raise
+        except StreamConsumedError:
+            raise
+        except Exception as err:
+            _LOGGER.warning("Unable to stream download: %s", err)
+            self.response.internal_response.close()
+            raise
 
 class AioHttpTransportResponse(AsyncHttpResponse):
     """Methods for accessing response body data.

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

@@ -27,6 +27,7 @@
 import logging
 from typing import Iterator, Optional, Any, Union, TypeVar
 import time
+import six
 import urllib3 # type: ignore
 from urllib3.util.retry import Retry # type: ignore
 import requests
@@ -49,6 +50,27 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+def parse_range_header(header_value):
+    range_value = header_value.strip()
+    if not range_value.startswith("bytes="):
+        raise ValueError("Invalid header")
+    range = range_value[6:]
+    ret = range.split("-")
+    if len(ret) < 2:
+        raise ValueError("Invalid header")
+    start = int(ret[0]) if ret[0] else -1
+    end = int(ret[1]) if ret[1] else -1
+    return (start, end)
+
+def make_range_header(original_range, downloaded_size=0):
+    if original_range[0] == -1:
+        end = original_range[1] - downloaded_size
+        return "bytes=-" + str(end)
+    start = original_range[0] + downloaded_size
+    if original_range[1] == -1:
+        return "bytes=" + str(start) + "-"
+    return "bytes=" + str(start) + "-" + str(original_range[1])
+
 
 class _RequestsTransportResponseBase(_HttpResponseBase):
     """Base class for accessing response data.
@@ -108,6 +130,19 @@ def __init__(self, pipeline, response):
         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
+        headers = response.request.headers
+        if "x-ms-range" in headers:
+            self.range_header = "x-ms-range"
+            self.range = headers["x-ms-range"]
+        elif "Range" in headers:
+            self.range_header = "Range"
+            self.range = headers["Range"]
+        else:
+            self.range_header = None
+        if 'etag' in response.headers:
+            self.etag = response.headers['etag']
+        else:
+            self.etag = None
 
     def __len__(self):
         return self.content_length
@@ -119,39 +154,58 @@ 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):
+            self.downloaded += self.block_size
+            return chunk
+        except StopIteration:
+            self.response.internal_response.close()
+            raise StopIteration()
+        except (requests.exceptions.ChunkedEncodingError,
+                requests.exceptions.ConnectionError) as ex:
+            while retry_active:
                 retry_total -= 1
                 if retry_total <= 0:
-                    retry_active = False
+                    _LOGGER.warning("Unable to stream download: %s", ex)
+                    raise ex
                 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)
+                    # todo handle pre-set range & x-ms-range
+                    headers = self.request.headers
+                    if not self.range_header:
+                        range_header = {'range': 'bytes=' + str(self.downloaded) + '-'}
+                    else:
+                        range_header = {self.range_header: make_range_header(self.range, self.downloaded)}
+                    if self.etag:
+                        range_header.update({'If-Match': self.etag})
+                    headers.update(range_header)
+                    try:
+                        resp = self.pipeline.run(self.request, stream=True, headers=headers)
+                        if not resp.http_response:
+                            continue
+                        if resp.http_response.status_code == 412:
+                            raise ex
+                        self.response = resp
+                        self.iter_content_func = resp.http_response.iter_content(self.block_size)
+                        chunk = next(self.iter_content_func)
+                        self.downloaded += self.block_size
+                        return chunk
+                    except StopIteration:
+                        self.response.internal_response.close()
+                        raise StopIteration()
+                    except Exception:   # pylint: disable=broad-except
+                        continue
                     if not chunk:
+                        self.response.internal_response.close()
                         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
+        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.
 
 

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

@@ -2,12 +2,15 @@
 # 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 import AsyncPipeline, PipelineResponse
 from azure.core.pipeline.transport._aiohttp import AioHttpStreamDownloadGenerator
 from unittest import mock
 import pytest
@@ -90,7 +93,7 @@ def __init__(self):
             self.headers = {}
             self.content = MockContent()
 
-        async def close(self):
+        def close(self):
             pass
 
     class AsyncMock(mock.MagicMock):
@@ -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_range_header.py

@@ -0,0 +1,60 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+from azure.core.pipeline.transport._requests_basic import parse_range_header, make_range_header
+
+def test_basic_range_header():
+    range_header = "bytes=0-500"
+    range = parse_range_header(range_header)
+    assert range == (0,500)
+
+def test_no_start_range_header():
+    range_header = "bytes=-500"
+    range = parse_range_header(range_header)
+    assert range == (-1,500)
+
+def test_no_end_range_header():
+    range_header = "bytes=0-"
+    range = parse_range_header(range_header)
+    assert range == (0,-1)
+
+def test_make_basic_range_header():
+    range_header = "bytes=0-500"
+    range = parse_range_header(range_header)
+    header = make_range_header(range, 100)
+    assert header == "bytes=100-500"
+
+def test_make_no_start_range_header():
+    range_header = "bytes=-500"
+    range = parse_range_header(range_header)
+    header = make_range_header(range, 100)
+    assert header == "bytes=-400"
+
+def test_make_no_end_range_header():
+    range_header = "bytes=0-"
+    range = parse_range_header(range_header)
+    header = make_range_header(range, 100)
+    assert header == "bytes=100-"