cesnet invenio multipart

Author: martenson

lib/galaxy/files/sources/_rdm.py

@@ -25,11 +25,15 @@
 class RDMFileSourceTemplateConfiguration(BaseFileSourceTemplateConfiguration):
     token: Optional[Union[str, TemplateExpansion]] = None
     public_name: Optional[Union[str, TemplateExpansion]] = None
+    multipart_threshold: Optional[Union[int, TemplateExpansion]] = None  # bytes
+    multipart_chunk_size: Optional[Union[int, TemplateExpansion]] = None  # bytes
 
 
 class RDMFileSourceConfiguration(BaseFileSourceConfiguration):
     token: Optional[str] = None
     public_name: Optional[str] = None
+    multipart_threshold: Optional[int] = None  # bytes
+    multipart_chunk_size: Optional[int] = None  # bytes
 
 
 class ContainerAndFileIdentifier(NamedTuple):

lib/galaxy/files/sources/invenio.py

@@ -1,7 +1,14 @@
 import datetime
 import json
+import logging
+import math
+import os
 import re
 import urllib.request
+from concurrent.futures import (
+    as_completed,
+    ThreadPoolExecutor,
+)
 from typing import (
     Any,
     cast,
@@ -10,6 +17,8 @@
 from urllib.error import HTTPError
 from urllib.parse import quote
 
+log = logging.getLogger(__name__)
+
 from typing_extensions import (
     Literal,
     TypedDict,
@@ -104,6 +113,56 @@ class RecordLinks(TypedDict):
     reserve_doi: str
 
 
+# AWS S3 multipart limits (used by Invenio RDM)
+MIN_UPLOAD_PART_SIZE = 50 * 1024 * 1024  # 50 MiB
+MAX_UPLOAD_PART_SIZE = 5 * 1024**3  # 5 GiB
+MAX_UPLOAD_PARTS = 10_000
+
+# Default threshold for using multipart upload (100 MiB)
+DEFAULT_MULTIPART_THRESHOLD = 100 * 1024 * 1024
+
+
+def calculate_multipart_params(file_size: int, preferred_part_size: int | None = None) -> tuple[int, int]:
+    """Calculate optimal parts count and part size for multipart upload.
+
+    Args:
+        file_size: Total file size in bytes
+        preferred_part_size: Preferred part size in bytes (optional)
+
+    Returns:
+        Tuple of (parts_count, part_size)
+
+    Note:
+        Maximum uploadable file size is MAX_UPLOAD_PARTS * MAX_UPLOAD_PART_SIZE (~48.8 TiB).
+        Files larger than this will still return valid params but would fail server-side.
+    """
+    if file_size == 0:
+        return 1, 0
+
+    # Start with preferred or minimum part size
+    part_size = preferred_part_size or MIN_UPLOAD_PART_SIZE
+
+    # Ensure part_size is within bounds
+    part_size = max(part_size, MIN_UPLOAD_PART_SIZE)
+    part_size = min(part_size, MAX_UPLOAD_PART_SIZE)
+
+    # Calculate parts needed
+    parts = math.ceil(file_size / part_size)
+
+    # If too many parts, increase part size (up to max)
+    while parts > MAX_UPLOAD_PARTS and part_size < MAX_UPLOAD_PART_SIZE:
+        part_size = min(part_size * 2, MAX_UPLOAD_PART_SIZE)
+        parts = math.ceil(file_size / part_size)
+
+    # For extremely large files, cap parts at MAX_UPLOAD_PARTS
+    # This means part_size may effectively be larger than calculated
+    # but such files would likely fail server-side anyway
+    if parts > MAX_UPLOAD_PARTS:
+        parts = MAX_UPLOAD_PARTS
+
+    return parts, part_size
+
+
 class InvenioRecord(TypedDict):
     id: str
     title: str
@@ -310,7 +369,7 @@ def create_draft_file_container(
             "metadata": {
                 "title": title,
                 "publication_date": today,
-                "resource_type": {"id": "dataset"},
+                "resource_type": {"id": "c_393c"},
                 "creators": [
                     creator,
                 ],
@@ -331,6 +390,33 @@ def upload_file_to_draft_container(
         file_path: str,
         context: FilesSourceRuntimeContext[RDMFileSourceConfiguration],
     ):
+        file_size = os.path.getsize(file_path)
+        threshold = context.config.multipart_threshold
+
+        # Use default threshold if not configured
+        if threshold is None or threshold <= 0:
+            threshold = DEFAULT_MULTIPART_THRESHOLD
+
+        use_multipart = file_size >= threshold
+
+        if use_multipart:
+            log.info(f"Using multipart upload for file '{filename}' ({file_size} bytes >= threshold {threshold})")
+            self._upload_file_multipart(record_id, filename, file_path, file_size, context)
+        else:
+            self._upload_file_single(record_id, filename, file_path, context, file_size)
+
+    def _upload_file_single(
+        self,
+        record_id: str,
+        filename: str,
+        file_path: str,
+        context: FilesSourceRuntimeContext[RDMFileSourceConfiguration],
+        file_size: Optional[int] = None,
+    ):
+        """Upload a file using single PUT request."""
+        if file_size is None:
+            file_size = os.path.getsize(file_path)
+
         record = self._get_draft_record(record_id, context)
         upload_file_url = record["links"]["files"]
         headers = self._get_request_headers(context, auth_required=True)
@@ -346,12 +432,148 @@ def upload_file_to_draft_container(
         commit_file_upload_url = file_entry["links"]["commit"]
         with open(file_path, "rb") as file:
             response = requests.put(upload_file_content_url, data=file, headers=headers)
+            # Handle 413 (Payload Too Large) - suggest using multipart upload
+            if response.status_code == 413:
+                threshold_mb = DEFAULT_MULTIPART_THRESHOLD / (1024 * 1024)
+                raise Exception(
+                    f"Failed to upload file '{filename}' ({file_size} bytes): HTTP 413 Payload Too Large. "
+                    f"The server rejected the upload because the file is too large for a single request. "
+                    f"Please configure 'multipart_threshold' to {threshold_mb}MB or lower to enable multipart upload for files of this size."
+                )
             self._ensure_response_has_expected_status_code(response, 200)
 
         # Commit file upload
         response = requests.post(commit_file_upload_url, headers=headers)
         self._ensure_response_has_expected_status_code(response, 200)
 
+    def _upload_file_multipart(
+        self,
+        record_id: str,
+        filename: str,
+        file_path: str,
+        file_size: int,
+        context: FilesSourceRuntimeContext[RDMFileSourceConfiguration],
+    ):
+        """Upload a file using multipart upload.
+
+        Flow:
+        1. Calculate parts/part_size
+        2. POST with transfer metadata
+        3. Server returns links.parts[] with URL for each part
+        4. Upload parts (parallel for > 2 parts)
+        5. POST to commit URL
+        """
+        preferred_part_size = context.config.multipart_chunk_size
+        num_parts, part_size = calculate_multipart_params(file_size, preferred_part_size)
+
+        log.info(f"Multipart upload: {num_parts} parts of {part_size} bytes each for '{filename}'")
+
+        record = self._get_draft_record(record_id, context)
+        upload_file_url = record["links"]["files"]
+        headers = self._get_request_headers(context, auth_required=True)
+
+        # Initialize multipart upload with transfer metadata
+        file_metadata = {
+            "key": filename,
+            "size": file_size,
+            "transfer": {
+                "type": "M",
+                "parts": num_parts,
+                "part_size": part_size,
+            },
+        }
+        response = requests.post(upload_file_url, json=[file_metadata], headers=headers)
+        self._ensure_response_has_expected_status_code(response, 201)
+
+        # Get part upload URLs from response
+        entries = response.json()["entries"]
+        file_entry = next(entry for entry in entries if entry["key"] == filename)
+        commit_url = file_entry["links"]["commit"]
+        part_links = file_entry.get("links", {}).get("parts", [])
+
+        if len(part_links) != num_parts:
+            raise Exception(
+                f"Server returned {len(part_links)} part URLs but expected {num_parts} for file '{filename}'"
+            )
+
+        # Sort part links by part number to ensure correct ordering
+        # Invenio uses 'part' key, not 'part_number'
+        part_links = sorted(part_links, key=lambda p: p.get("part", 0))
+
+        # Upload parts
+        self._upload_parts(file_path, file_size, part_size, part_links, headers)
+
+        # Commit multipart upload
+        response = requests.post(commit_url, json={}, headers=headers)
+        self._ensure_response_has_expected_status_code(response, 200)
+        log.info(f"Multipart upload completed for '{filename}'")
+
+    def _upload_parts(
+        self,
+        file_path: str,
+        file_size: int,
+        part_size: int,
+        part_links: list[dict],
+        headers: dict,
+    ):
+        """Upload all parts, sequentially for <=2 parts, parallel otherwise."""
+        num_parts = len(part_links)
+
+        if num_parts <= 2:
+            for part_index, part_info in enumerate(part_links):
+                self._upload_single_part(file_path, file_size, part_size, part_index, part_info)
+        else:
+            max_workers = min(4, num_parts)
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                futures = {}
+                for part_index, part_info in enumerate(part_links):
+                    future = executor.submit(
+                        self._upload_single_part,
+                        file_path,
+                        file_size,
+                        part_size,
+                        part_index,
+                        part_info,
+                    )
+                    futures[future] = part_index
+
+                for future in as_completed(futures):
+                    part_index = futures[future]
+                    try:
+                        future.result()
+                    except Exception as e:
+                        log.error(f"Failed to upload part {part_index}: {e}")
+                        raise
+
+    def _upload_single_part(
+        self,
+        file_path: str,
+        file_size: int,
+        part_size: int,
+        part_index: int,
+        part_info: dict,
+    ):
+        """Upload a single part of a multipart upload."""
+        part_url = part_info.get("url")
+        if not part_url:
+            raise Exception(f"No URL provided for part {part_index}")
+
+        start_byte = part_index * part_size
+        end_byte = min(start_byte + part_size, file_size)
+        part_content_length = end_byte - start_byte
+
+        log.debug(f"Uploading part {part_index}: bytes {start_byte}-{end_byte-1} ({part_content_length} bytes)")
+
+        # Read the entire part into memory and upload
+        with open(file_path, "rb") as f:
+            f.seek(start_byte)
+            part_data = f.read(part_content_length)
+
+        # Use empty headers - presigned URLs are authenticated via query parameters
+        # Adding Authorization or other headers would invalidate the signature
+        response = requests.put(part_url, data=part_data)
+        self._ensure_response_has_expected_status_code(response, 200)
+
     def download_file_from_container(
         self,
         container_id: str,
@@ -544,7 +766,12 @@ def _raise_auth_required(self):
         )
 
     def _get_response_error_message(self, response):
-        response_json = response.json()
+        try:
+            response_json = response.json()
+        except Exception:
+            # Response is not JSON, return raw text or status info
+            return response.text or f"HTTP {response.status_code} error"
+
         error_message = response_json.get("message") if response.status_code == 400 else response.text
         errors = response_json.get("errors", [])
         for error in errors: