95 lines
4.2 KiB
Python
95 lines
4.2 KiB
Python
"""Multipart transport for release files larger than a single upload part."""
|
|
from __future__ import annotations
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import hashlib
|
|
import http.client
|
|
from typing import Iterable
|
|
from xml.etree import ElementTree as ET
|
|
|
|
|
|
def _file_range(path: str, offset: int, size: int) -> Iterable[bytes]:
|
|
with open(path, "rb") as handle:
|
|
handle.seek(offset)
|
|
while size:
|
|
chunk = handle.read(min(size, 1024 * 1024))
|
|
if not chunk:
|
|
raise OSError(f"Upload source was truncated: {path}")
|
|
size -= len(chunk)
|
|
yield chunk
|
|
|
|
|
|
def upload_file(
|
|
url: str, path: str, size: int, creds: dict[str, str],
|
|
content_type: str | None, headers: dict[str, str], part_size: int,
|
|
) -> None:
|
|
from . import r2
|
|
|
|
if not 5 * 1024 * 1024 <= part_size <= 5 * 1024**3:
|
|
raise ValueError("Multipart parts must be between 5 MiB and 5 GiB")
|
|
if (size + part_size - 1) // part_size > 10000:
|
|
raise ValueError("Multipart upload exceeds 10,000 parts")
|
|
response = r2.signed_request(
|
|
"POST", url + "?uploads=", creds=creds, now=r2.amz_timestamp(),
|
|
content_length=0, content_type=content_type, extra_headers=headers,
|
|
)
|
|
upload_id = ET.fromstring(response.text()).findtext("{*}UploadId")
|
|
if not upload_id:
|
|
raise r2.R2RequestError("POST", url, 502, "Missing multipart upload ID")
|
|
upload_url = url + "?" + r2.canonical_query({"uploadId": upload_id})
|
|
|
|
def upload_part(offset: int) -> tuple[int, str]:
|
|
number = offset // part_size + 1
|
|
length = min(part_size, size - offset)
|
|
digest = hashlib.sha256()
|
|
for chunk in _file_range(path, offset, length):
|
|
digest.update(chunk)
|
|
part_url = url + "?" + r2.canonical_query({
|
|
"partNumber": str(number), "uploadId": upload_id,
|
|
})
|
|
part = r2.signed_request(
|
|
"PUT", part_url, creds=creds, now=r2.amz_timestamp(),
|
|
body_iter_factory=lambda: _file_range(path, offset, length),
|
|
body_hash=digest.hexdigest(), content_length=length,
|
|
)
|
|
etag = part.header("etag")
|
|
if not etag:
|
|
raise r2.R2RequestError("PUT", part_url, 502, "Missing part ETag")
|
|
print(f"r2: uploaded part {number} ({length} bytes)", flush=True)
|
|
return number, etag
|
|
|
|
try:
|
|
# Each worker reopens its range on retries; no whole-artifact buffer.
|
|
with ThreadPoolExecutor() as pool:
|
|
parts = list(pool.map(upload_part, range(0, size, part_size)))
|
|
document = ET.Element("CompleteMultipartUpload")
|
|
for number, etag in parts:
|
|
part = ET.SubElement(document, "Part")
|
|
ET.SubElement(part, "PartNumber").text = str(number)
|
|
ET.SubElement(part, "ETag").text = etag
|
|
body = ET.tostring(document, encoding="utf-8")
|
|
conditions = {key: value for key, value in headers.items() if key.lower().startswith("if-")}
|
|
try:
|
|
response = r2.signed_request(
|
|
"POST", upload_url, creds=creds, now=r2.amz_timestamp(),
|
|
body=body, body_hash=hashlib.sha256(body).hexdigest(),
|
|
content_length=len(body), content_type="application/xml",
|
|
extra_headers=conditions, tries=1,
|
|
)
|
|
except (OSError, http.client.HTTPException) as error:
|
|
# Completion may already have committed. A fresh conditional upload
|
|
# either succeeds or reaches the caller's immutable digest check.
|
|
raise r2.R2RequestError("POST", upload_url, 503, str(error)) from error
|
|
# S3 can return an error XML document after sending HTTP 200 headers.
|
|
result = ET.fromstring(response.text())
|
|
if result.tag.rsplit("}", 1)[-1] != "CompleteMultipartUploadResult":
|
|
code = result.findtext("{*}Code")
|
|
status = 503 if code in {"InternalError", "SlowDown", "ServiceUnavailable", "RequestTimeout"} else 400
|
|
raise r2.R2RequestError("POST", upload_url, status, response.text())
|
|
except BaseException as error:
|
|
try:
|
|
r2.signed_request("DELETE", upload_url, creds=creds, now=r2.amz_timestamp())
|
|
except Exception as cleanup_error:
|
|
error.add_note(f"Multipart abort failed: {cleanup_error}")
|
|
raise
|