Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions ayon_api/server_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -1660,7 +1660,7 @@ def _download_file_to_stream(
url = self._endpoint_to_url(endpoint, use_rest=False)
progress.set_source_url(url)

retries = self.get_default_max_retries()
retries = max(self.max_retries, 1)
api_prepended = False
for attempt in range(retries):
# Continue in download
Expand All @@ -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"]
Expand All @@ -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
Expand Down Expand Up @@ -2095,7 +2113,7 @@ def _upload_file(
headers.pop(orig_key)
headers[key] = value

retries = self.get_default_max_retries()
retries = max(self.max_retries, 1)
response = None

# Get size of file
Expand Down
4 changes: 2 additions & 2 deletions ayon_api/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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.
Expand Down
60 changes: 60 additions & 0 deletions tests/test_download_resume.py
Original file line number Diff line number Diff line change
@@ -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
Loading