diff --git a/poetry.lock b/poetry.lock index 575e11aa..8abf7ffc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "alabaster" @@ -11,6 +11,28 @@ files = [ {file = "alabaster-0.7.13.tar.gz", hash = "sha256:a27a4a084d5e690e16e01e03ad2b2e552c61a65469419b907243193de1a84ae2"}, ] +[[package]] +name = "anyio" +version = "4.3.0" +description = "High level compatibility layer for multiple asynchronous event loop implementations" +optional = false +python-versions = ">=3.8" +files = [ + {file = "anyio-4.3.0-py3-none-any.whl", hash = "sha256:048e05d0f6caeed70d731f3db756d35dcc1f35747c8c403364a8332c630441b8"}, + {file = "anyio-4.3.0.tar.gz", hash = "sha256:f75253795a87df48568485fd18cdd2a3fa5c4f7c5be8e5e36637733fce06fed6"}, +] + +[package.dependencies] +exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""} +idna = ">=2.8" +sniffio = ">=1.1" +typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""} + +[package.extras] +doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"] +test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"] +trio = ["trio (>=0.23)"] + [[package]] name = "arrow" version = "0.17.0" @@ -431,6 +453,62 @@ files = [ {file = "flaky-3.7.0.tar.gz", hash = "sha256:3ad100780721a1911f57a165809b7ea265a7863305acb66708220820caf8aa0d"}, ] +[[package]] +name = "h11" +version = "0.14.0" +description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1" +optional = false +python-versions = ">=3.7" +files = [ + {file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"}, + {file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"}, +] + +[[package]] +name = "httpcore" +version = "1.0.5" +description = "A minimal low-level HTTP client." +optional = false +python-versions = ">=3.8" +files = [ + {file = "httpcore-1.0.5-py3-none-any.whl", hash = "sha256:421f18bac248b25d310f3cacd198d55b8e6125c107797b609ff9b7a6ba7991b5"}, + {file = "httpcore-1.0.5.tar.gz", hash = "sha256:34a38e2f9291467ee3b44e89dd52615370e152954ba21721378a87b2960f7a61"}, +] + +[package.dependencies] +certifi = "*" +h11 = ">=0.13,<0.15" + +[package.extras] +asyncio = ["anyio (>=4.0,<5.0)"] +http2 = ["h2 (>=3,<5)"] +socks = ["socksio (==1.*)"] +trio = ["trio (>=0.22.0,<0.26.0)"] + +[[package]] +name = "httpx" +version = "0.27.0" +description = "The next generation HTTP client." +optional = false +python-versions = ">=3.8" +files = [ + {file = "httpx-0.27.0-py3-none-any.whl", hash = "sha256:71d5465162c13681bff01ad59b2cc68dd838ea1f10e51574bac27103f00c91a5"}, + {file = "httpx-0.27.0.tar.gz", hash = "sha256:a0cb88a46f32dc874e04ee956e4c2764aba2aa228f650b06788ba6bda2962ab5"}, +] + +[package.dependencies] +anyio = "*" +certifi = "*" +httpcore = "==1.*" +idna = "*" +sniffio = "*" + +[package.extras] +brotli = ["brotli", "brotlicffi"] +cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"] +http2 = ["h2 (>=3,<5)"] +socks = ["socksio (==1.*)"] + [[package]] name = "identify" version = "2.5.35" @@ -993,6 +1071,7 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, + {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -1000,8 +1079,16 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, + {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, + {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, + {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, + {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -1018,6 +1105,7 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, + {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -1025,6 +1113,7 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, + {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -1159,6 +1248,17 @@ files = [ {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"}, ] +[[package]] +name = "sniffio" +version = "1.3.1" +description = "Sniff out which async library your code is running under" +optional = false +python-versions = ">=3.7" +files = [ + {file = "sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2"}, + {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, +] + [[package]] name = "snowballstemmer" version = "2.2.0" @@ -1464,4 +1564,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "0d38d200cf62fa0c09b90693d3f862e262e56ef8f7fca56162f10d828628abc5" +content-hash = "ae1b53ae512540bd171fd27573ba403cbb0305600894fcde428fcc629ecf04c8" diff --git a/pyproject.toml b/pyproject.toml index 885e142f..1b7dfd9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,8 @@ pyparsing = ">=2.0,<3" clique = "==1.6.1" websocket-client = ">=0.40.0,<1" platformdirs = ">=4.0.0,<5" +httpx = "^0.27.0" +anyio = "^4.3.0" [tool.poetry.group.dev.dependencies] black = "^23.7.0" diff --git a/source/ftrack_api/_http.py b/source/ftrack_api/_http.py new file mode 100644 index 00000000..56cd9cfa --- /dev/null +++ b/source/ftrack_api/_http.py @@ -0,0 +1,26 @@ +# :coding: utf-8 +# :copyright: Copyright (c) 2024 ftrack + +import httpx +import os +from pathlib import Path + + +def _get_ssl_context(): + ssl_context = httpx.create_ssl_context() + + requests_ca_env = os.environ.get("REQUESTS_CA_BUNDLE") + if not requests_ca_env: + return + + ca_path = Path(requests_ca_env) + + if ca_path.is_file(): + ssl_context.load_verify_locations(cafile=str(ca_path)) + elif ca_path.is_dir(): + ssl_context.load_verify_locations(capath=str(ca_path)) + + return ssl_context + + +ssl_context = _get_ssl_context() diff --git a/source/ftrack_api/accessor/server.py b/source/ftrack_api/accessor/server.py index da6b0262..a8f5a188 100644 --- a/source/ftrack_api/accessor/server.py +++ b/source/ftrack_api/accessor/server.py @@ -4,13 +4,13 @@ import os import hashlib import base64 -import json import requests from .base import Accessor from ..data import String import ftrack_api.exception +from ftrack_api.uploader import Uploader import ftrack_api.symbol @@ -72,7 +72,6 @@ def _read(self): def _write(self): """Write current data to remote key.""" position = self.tell() - self.seek(0) # Retrieve component from cache to construct a filename. component = self._session.get("FileComponent", self.resource_identifier) @@ -89,28 +88,16 @@ def _write(self): name = "{0}.{1}".format(name, component["file_type"].lstrip(".")) try: - metadata = self._session.get_upload_metadata( + uploader = Uploader( + self._session, component_id=self.resource_identifier, file_name=name, file_size=self._get_size(), + file=self.wrapped_file, checksum=self._compute_checksum(), ) + uploader.start() except Exception as error: - raise ftrack_api.exception.AccessorOperationFailedError( - "Failed to get put metadata: {0}.".format(error) - ) - - # Ensure at beginning of file before put. - self.seek(0) - - # Put the file based on the metadata. - response = requests.put( - metadata["url"], data=self.wrapped_file, headers=metadata["headers"] - ) - - try: - response.raise_for_status() - except requests.exceptions.HTTPError as error: raise ftrack_api.exception.AccessorOperationFailedError( "Failed to put file to server: {0}.".format(error) ) @@ -120,8 +107,7 @@ def _write(self): def _get_size(self): """Return size of file in bytes.""" position = self.tell() - self.seek(0, os.SEEK_END) - length = self.tell() + length = self.seek(0, os.SEEK_END) self.seek(position) return length diff --git a/source/ftrack_api/data.py b/source/ftrack_api/data.py index 8d05d166..9978beb9 100644 --- a/source/ftrack_api/data.py +++ b/source/ftrack_api/data.py @@ -5,6 +5,7 @@ import os from abc import ABCMeta, abstractmethod import tempfile +import typing class Data(metaclass=ABCMeta): @@ -25,7 +26,7 @@ def write(self, content): def flush(self): """Flush buffers ensuring data written.""" - def seek(self, offset, whence=os.SEEK_SET): + def seek(self, offset, whence=os.SEEK_SET) -> int: """Move internal pointer by *offset*. The *whence* argument is optional and defaults to os.SEEK_SET or 0 @@ -33,6 +34,7 @@ def seek(self, offset, whence=os.SEEK_SET): (seek relative to the current position) and os.SEEK_END or 2 (seek relative to the file's end). + Return the new absolute position. """ raise NotImplementedError("Seek not supported.") @@ -49,7 +51,7 @@ def close(self): class FileWrapper(Data): """Data wrapper for Python file objects.""" - def __init__(self, wrapped_file): + def __init__(self, wrapped_file: typing.IO): """Initialise access to *wrapped_file*.""" self.wrapped_file = wrapped_file self._read_since_last_write = False @@ -81,7 +83,7 @@ def flush(self): def seek(self, offset, whence=os.SEEK_SET): """Move internal pointer by *offset*.""" - self.wrapped_file.seek(offset, whence) + return self.wrapped_file.seek(offset, whence) def tell(self): """Return current position of internal pointer.""" diff --git a/source/ftrack_api/session.py b/source/ftrack_api/session.py index ff6f2988..998d9789 100644 --- a/source/ftrack_api/session.py +++ b/source/ftrack_api/session.py @@ -2249,7 +2249,9 @@ def encode_media(self, media, version_id=None, keep_original="auto"): return self.get("Job", result[0]["job_id"]) - def get_upload_metadata(self, component_id, file_name, file_size, checksum=None): + def get_upload_metadata( + self, component_id, file_name, file_size, checksum=None, parts=None + ): """Return URL and headers used to upload data for *component_id*. *file_name* and *file_size* should match the components details. @@ -2268,6 +2270,7 @@ def get_upload_metadata(self, component_id, file_name, file_size, checksum=None) "file_name": file_name, "file_size": file_size, "checksum": checksum, + "parts": parts, } try: @@ -2286,6 +2289,29 @@ def get_upload_metadata(self, component_id, file_name, file_size, checksum=None) return result[0] + def complete_multipart_upload(self, component_id, upload_id, parts): + operation = { + "action": "complete_multipart_upload", + "component_id": component_id, + "upload_id": upload_id, + "parts": parts, + } + + try: + result = self.call([operation]) + except ftrack_api.exception.ServerError as error: + # Raise informative error if the action is not supported. + if "Invalid action u'complete_multipart_upload'" in error.message: + raise ftrack_api.exception.ServerCompatibilityError( + "Server version {0!r} does not support " + '"complete_multipart_upload", please update server and try ' + "again.".format(self.server_information.get("version")) + ) + else: + raise + + return result[0] + def send_user_invite(self, user): """Send a invitation to the provided *user*. diff --git a/source/ftrack_api/uploader.py b/source/ftrack_api/uploader.py new file mode 100644 index 00000000..e672c33a --- /dev/null +++ b/source/ftrack_api/uploader.py @@ -0,0 +1,157 @@ +# :coding: utf-8 +# :copyright: Copyright (c) 2024 ftrack + +import logging +import math +import os +from typing import IO, Awaitable, Callable, TYPE_CHECKING, List, Optional +import typing +import anyio + +import anyio.from_thread +import httpx +from ._http import ssl_context + +if TYPE_CHECKING: + from ftrack_api.session import Session + + +SIZE_MEGABYTES = 1024**2 +SIZE_GIGABYTES = 1024**3 +MAX_PARTS = 10000 + +logger = logging.getLogger(__name__) + + +def get_chunk_size(file_size: int) -> int: + chunk_profiles = [ + (0, 8), + (0.1, 16), + (1, 32), + (8, 64), + (64, 128), + ] + + for min_size_in_gb, chunk_size_in_mb in reversed(chunk_profiles): + if file_size >= min_size_in_gb * SIZE_GIGABYTES: + return chunk_size_in_mb * SIZE_MEGABYTES + + raise ValueError("Invalid file size.") + + +async def back_off( + func: Callable[..., Awaitable], + *args: typing.Any, + retries: int = 5, + delay: int = 5, +): + for i in range(retries): + try: + return await func(*args) + except httpx.HTTPError as e: + if i == retries - 1: + raise e + + sleep_time = delay * 2**i + logger.warn(f"Retrying in {sleep_time}s") + await anyio.sleep(sleep_time) + + +class Uploader: + max_concurrency: int = int(os.environ.get("FTRACK_UPLOAD_MAX_CONCURRENCY", 5)) + + def __init__( + self, + session: "Session", + component_id: str, + file_name: str, + file_size: int, + file: "IO", + checksum: Optional[str], + ): + self.session = session + self.component_id = component_id + self.file = file + self.file_name = file_name + self.file_size = file_size + self.checksum = checksum + + self.chunk_size = get_chunk_size(self.file_size) + + self.upload_id: Optional[str] = None + self.upload_urls: List[dict] = [] + self.uploaded_parts: List[dict] = [] + + def _single_upload(self, url, headers): + self.file.seek(0) + + response = httpx.put( + url, + verify=ssl_context, + content=self.file, + headers=headers, + ) + + response.raise_for_status() + + async def _upload_part_task(self, http: httpx.AsyncClient): + async def send_data(part_num, url, content): + resp = await http.put(url=url, content=content) + resp.raise_for_status() + + return { + "part_number": part_num, + "e_tag": resp.headers["ETag"].strip('"'), + } + + while True: + url_info = self.upload_urls.pop(0) if self.upload_urls else None + if not url_info: + break + + url = url_info["signed_url"] + part_num = url_info["part_number"] + + startPos = (part_num - 1) * self.chunk_size + self.file.seek(startPos) + content = self.file.read(self.chunk_size) + + uploaded_part = await back_off(send_data, part_num, url, content) + self.uploaded_parts.append(uploaded_part) + + async def _multi_upload(self): + async with httpx.AsyncClient(verify=ssl_context) as http: + async with anyio.create_task_group() as tg: + for _ in range(self.max_concurrency): + tg.start_soon(self._upload_part_task, http) + + def start(self): + parts_count = math.ceil(self.file_size / self.chunk_size) + + if parts_count > MAX_PARTS: + raise ValueError("File is too big.") + + metadata = self.session.get_upload_metadata( + self.component_id, + self.file_name, + self.file_size, + self.checksum, + parts_count if parts_count > 1 else None, + ) + + if "urls" in metadata: + self.upload_id = metadata["upload_id"] + self.upload_urls = metadata["urls"] + + # start upload in a separate thread, with anyio + with anyio.from_thread.start_blocking_portal() as portal: + portal.call(self._multi_upload) + + uploaded = sorted(self.uploaded_parts, key=lambda x: x["part_number"]) + self.session.complete_multipart_upload( + self.component_id, + self.upload_id, + uploaded, + ) + else: + self._single_upload(metadata["url"], metadata["headers"]) diff --git a/test/unit/test_uploader.py b/test/unit/test_uploader.py new file mode 100644 index 00000000..35b23e29 --- /dev/null +++ b/test/unit/test_uploader.py @@ -0,0 +1,47 @@ +# :coding: utf-8 +# :copyright: Copyright (c) 2024 ftrack + +import io +import pytest +from ftrack_api.uploader import Uploader, get_chunk_size, SIZE_GIGABYTES, SIZE_MEGABYTES + + +@pytest.mark.parametrize( + "file_size, expected_chunk_size", + [ + (0, 8), + (3 * SIZE_MEGABYTES, 8), + (255 * SIZE_MEGABYTES, 16), + (2 * SIZE_GIGABYTES, 32), + (70 * SIZE_GIGABYTES, 128), + ], +) +def test_get_chunk_size(file_size, expected_chunk_size): + chunk_size_in_mb = get_chunk_size(file_size) // SIZE_MEGABYTES + assert chunk_size_in_mb == expected_chunk_size + + +@pytest.fixture() +def make_dummy_file(): + chunk = b"DEADBEEF" + files = [] + + def _make_dummy_file(size): + chunk_count = size // len(chunk) + f = io.BytesIO(chunk * chunk_count) + files.append(f) + return f + + yield _make_dummy_file + + for f in files: + f.close() + + +@pytest.mark.parametrize("size", [1 * SIZE_MEGABYTES, 10 * SIZE_MEGABYTES]) +def test_uploader(session, make_dummy_file, size): + file = make_dummy_file(size) + uploader = Uploader( + session, "17db4ccc-dd37-49c9-8be5-9afc4abf7c2c", "test.jpg", size, file, None + ) + uploader.start()