diff --git a/ayon_api/server_api.py b/ayon_api/server_api.py index 2fce18fc0..6c98ecdf9 100644 --- a/ayon_api/server_api.py +++ b/ayon_api/server_api.py @@ -1686,6 +1686,14 @@ def _download_file_to_stream( progress.set_destination_url(url) continue response.raise_for_status() + if offset > 0 and response.status_code != 206: + # Server ignored 'Range' and sends whole file again, + # already downloaded content must be discarded + stream.seek(0) + stream.truncate() + progress.reset_transferred() + headers.pop("Range", None) + if progress.get_content_size() is None: progress.set_content_size( response.headers["Content-length"] @@ -1694,11 +1702,21 @@ def _download_file_to_stream( for chunk in response.iter_content(chunk_size=chunk_size): stream.write(chunk) progress.add_transferred_chunk(len(chunk)) + + content_size = progress.get_content_size() + transferred = progress.get_transferred_size() + if content_size is not None and transferred != content_size: + # Connection was closed before all content was received + raise requests.exceptions.ConnectionError( + f"Downloaded {transferred} out of {content_size}" + f" bytes from '{url}'." + ) break except ( requests.exceptions.Timeout, requests.exceptions.ConnectionError, + requests.exceptions.ChunkedEncodingError, ): if attempt == retries - 1: raise diff --git a/ayon_api/utils.py b/ayon_api/utils.py index ab2e4c7dd..f5769e7d8 100644 --- a/ayon_api/utils.py +++ b/ayon_api/utils.py @@ -955,7 +955,7 @@ def get_content_size(self) -> int | None: """ return self._content_size - def set_content_size(self, content_size: int) -> None: + def set_content_size(self, content_size: int | str) -> None: """Set content size in bytes. Args: @@ -967,7 +967,7 @@ def set_content_size(self, content_size: int) -> None: """ if self._content_size is not None: raise ValueError("Content size was set more then once") - self._content_size = content_size + self._content_size = int(content_size) def get_started(self) -> bool: """Transfer was started. diff --git a/tests/test_download_resume.py b/tests/test_download_resume.py new file mode 100644 index 000000000..9c353cbcb --- /dev/null +++ b/tests/test_download_resume.py @@ -0,0 +1,60 @@ +"""Interrupted downloads. Does not require running AYON server.""" +import io + +import pytest + +from ayon_api.server_api import ServerAPI +from ayon_api.utils import RequestTypes + +from .fake_transfer import FakeResponse + +CONTENT = b"0123456789" +SIZE_HEADER = {"Content-Length": str(len(CONTENT))} + + +@pytest.fixture +def con(monkeypatch): + monkeypatch.setattr("time.sleep", lambda *args, **kwargs: None) + return ServerAPI("http://localhost:0", create_session=False, max_retries=3) + + +def _download(con, get_func): + con._base_functions_mapping[RequestTypes.get] = get_func + stream = io.BytesIO() + progress = con.download_file_to_stream("api/file", stream) + return stream.getvalue(), progress + + +def test_incomplete_download_is_continued(con): + def get_func(url, **kwargs): + if "Range" not in kwargs["headers"]: + # Connection closed after 4 bytes without an exception + return FakeResponse(200, CONTENT[:4], SIZE_HEADER) + assert kwargs["headers"]["Range"] == "bytes=4-" + return FakeResponse(206, CONTENT[4:], {"Content-Length": "6"}) + + content, progress = _download(con, get_func) + assert content == CONTENT + assert progress.transferred_size == len(CONTENT) + + +def test_resume_without_range_support_does_not_duplicate(con): + calls = [] + + def get_func(url, **kwargs): + calls.append(url) + if len(calls) == 1: + return FakeResponse(200, CONTENT[:4], SIZE_HEADER) + # Server ignores 'Range' header and sends whole file again + return FakeResponse(200, CONTENT, SIZE_HEADER) + + content, progress = _download(con, get_func) + assert content == CONTENT + assert progress.transferred_size == len(CONTENT) + + +def test_download_without_content_length(con): + content, _ = _download( + con, lambda url, **kwargs: FakeResponse(200, CONTENT) + ) + assert content == CONTENT