Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for creation-defer-length extension #96

Merged
merged 5 commits into from
Dec 11, 2024
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
40 changes: 40 additions & 0 deletions tests/test_uploader.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from unittest import mock

import responses
from responses import matchers
from parametrize import parametrize
import pytest

Expand Down Expand Up @@ -209,3 +210,42 @@ def test_upload_checksum(self, request_mock):
self.uploader.upload_checksum = True
self.uploader.upload()
self.assertEqual(self.uploader.offset, self.uploader.get_file_size())

@parametrize("chunk_size", [1, 2, 3, 4, 5, 6])
@responses.activate
def test_upload_length_deferred(self, chunk_size: int):
upload_url = f"{self.client.url}test_upload_length_deferred"

responses.head(
upload_url,
adding_headers={"upload-offset": "0", "Upload-Defer-Length": "1"},
)
uploader = self.client.uploader(
file_stream=io.BytesIO(b"hello"),
url=upload_url,
chunk_size=chunk_size,
upload_length_deferred=True,
)
self.assertTrue(uploader.upload_length_deferred)
self.assertTrue(uploader.stop_at is None)

offset = 0
while not (offset + chunk_size > 5):
next_offset = min(offset + chunk_size, 5)
responses.patch(
upload_url,
adding_headers={"upload-offset": str(next_offset)},
match=[matchers.header_matcher({"upload-offset": str(offset)})],
)
offset = next_offset
last_req_headers = {"upload-offset": str(offset)}
last_req_headers["upload-length"] = "5"
responses.patch(
upload_url,
adding_headers={"upload-offset": "5"},
match=[matchers.header_matcher(last_req_headers)],
)

uploader.upload()
self.assertEqual(uploader.offset, 5)
self.assertEqual(uploader.stop_at, 5)
12 changes: 10 additions & 2 deletions tusclient/request.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,10 @@ class BaseTusRequest:

def __init__(self, uploader):
self._url = uploader.url
self.response_headers = {}
self.status_code = None
self.response_headers = {}
self.response_content = None
self.stream_eof = False
self.verify_tls_cert = bool(uploader.verify_tls_cert)
self.file = uploader.get_file_stream()
self.file.seek(uploader.offset)
Expand All @@ -51,6 +52,8 @@ def __init__(self, uploader):
"upload-offset": str(uploader.offset),
"Content-Type": "application/offset+octet-stream",
}
self._offset = uploader.offset
self._upload_length_deferred = uploader.upload_length_deferred
self._request_headers.update(uploader.get_headers())
self._content_length = uploader.get_request_length()
self._upload_checksum = uploader.upload_checksum
Expand Down Expand Up @@ -78,16 +81,21 @@ def perform(self):
"""
try:
chunk = self.file.read(self._content_length)
stream_eof = len(chunk) < self._content_length
self.add_checksum(chunk)
headers = self._request_headers
if stream_eof and self._upload_length_deferred:
headers["upload-length"] = str(self._offset + len(chunk))
resp = requests.patch(
self._url,
data=chunk,
headers=self._request_headers,
headers=headers,
verify=self.verify_tls_cert,
)
self.status_code = resp.status_code
self.response_content = resp.content
self.response_headers = {k.lower(): v for k, v in resp.headers.items()}
self.stream_eof = stream_eof
except requests.exceptions.RequestException as error:
raise TusUploadFailed(error)

Expand Down
20 changes: 16 additions & 4 deletions tusclient/uploader/baseuploader.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,10 @@ class BaseUploader:
- upload_checksum (bool):
Whether or not to supply the Upload-Checksum header along with each
chunk. Defaults to False.
- upload_length_deferred (bool):
Whether or not to declare the upload length when finished reading the file stream instead of when the upload is started. This is useful
when uploading from a streaming resource, where the total file size isn't available when the upload is created
but only becomes known when the stream finishes. The server must support the `creation-defer-length` extension.

:Constructor Args:
- file_path (str)
Expand All @@ -89,6 +93,7 @@ class BaseUploader:
- url_storage (Optinal [<tusclient.storage.interface.Storage>])
- fingerprinter (Optional [<tusclient.fingerprint.interface.Fingerprint>])
- upload_checksum (Optional[bool])
- upload_length_deferred (Optional[bool])
"""

DEFAULT_HEADERS = {"Tus-Resumable": "1.0.0"}
Expand All @@ -114,6 +119,7 @@ def __init__(
url_storage: Optional[Storage] = None,
fingerprinter: Optional[interface.Fingerprint] = None,
upload_checksum=False,
upload_length_deferred=False,
):
if file_path is None and file_stream is None:
raise ValueError("Either 'file_path' or 'file_stream' cannot be None.")
Expand All @@ -129,7 +135,8 @@ def __init__(
self.verify_tls_cert = verify_tls_cert
self.file_path = file_path
self.file_stream = file_stream
self.stop_at = self.get_file_size()
self.file_size = self.get_file_size() if not upload_length_deferred else None
self.stop_at = self.file_size
self.client = client
self.metadata = metadata or {}
self.metadata_encoding = metadata_encoding
Expand All @@ -145,6 +152,7 @@ def __init__(
self._retried = 0
self.retry_delay = retry_delay
self.upload_checksum = upload_checksum
self.upload_length_deferred = upload_length_deferred
(
self.__checksum_algorithm_name,
self.__checksum_algorithm,
Expand All @@ -161,7 +169,10 @@ def get_headers(self):
def get_url_creation_headers(self):
"""Return headers required to create upload url"""
headers = self.get_headers()
headers["upload-length"] = str(self.get_file_size())
if self.upload_length_deferred:
headers['upload-defer-length'] = '1'
else:
headers["upload-length"] = str(self.file_size)
headers["upload-metadata"] = ",".join(self.encode_metadata())
return headers

Expand Down Expand Up @@ -247,8 +258,9 @@ def get_request_length(self):
"""
Return length of next chunk upload.
"""
remainder = self.stop_at - self.offset
return self.chunk_size if remainder > self.chunk_size else remainder
if self.stop_at is None:
return self.chunk_size
return min(self.chunk_size, self.stop_at - self.offset)

def get_file_stream(self):
"""
Expand Down
12 changes: 8 additions & 4 deletions tusclient/uploader/uploader.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def upload(self, stop_at: Optional[int] = None):
Determines at what offset value the upload should stop. If not specified this
defaults to the file size.
"""
self.stop_at = stop_at or self.get_file_size()
self.stop_at = stop_at or self.file_size

if not self.url:
# Ensure the POST request is performed even for empty files.
Expand All @@ -41,7 +41,7 @@ def upload(self, stop_at: Optional[int] = None):
self.set_url(self.create_url())
self.offset = 0

while self.offset < self.stop_at:
while self.stop_at is None or (self.offset < self.stop_at):
self.upload_chunk()

def upload_chunk(self):
Expand All @@ -58,6 +58,8 @@ def upload_chunk(self):

self._do_request()
self.offset = int(self.request.response_headers.get("upload-offset"))
if self.upload_length_deferred and self.request.stream_eof:
self.stop_at = self.offset

@catch_requests_error
def create_url(self):
Expand Down Expand Up @@ -118,13 +120,13 @@ async def upload(self, stop_at: Optional[int] = None):
Determines at what offset value the upload should stop. If not specified this
defaults to the file size.
"""
self.stop_at = stop_at or self.get_file_size()
self.stop_at = stop_at or self.file_size

if not self.url:
self.set_url(await self.create_url())
self.offset = 0

while self.offset < self.stop_at:
while self.stop_at is None or (self.offset < self.stop_at):
await self.upload_chunk()

async def upload_chunk(self):
Expand All @@ -141,6 +143,8 @@ async def upload_chunk(self):

await self._do_request()
self.offset = int(self.request.response_headers.get("upload-offset"))
if self.upload_length_deferred and self.request.stream_eof:
self.stop_at = self.offset

async def create_url(self):
"""
Expand Down
Loading