Generate TUS assembly upload helper

Author: kvz

examples/api2-devdock-tus-assembly/main.py

@@ -7,12 +7,9 @@
 
 import json
 import os
-from io import BytesIO
 from pathlib import Path
-from urllib.parse import quote
 
 from transloadit.client import Transloadit
-from tusclient import client as tus
 
 
 def required_env(name):
@@ -35,45 +32,6 @@ def load_scenario():
         return json.load(scenario_file)
 
 
-def read_path(value, path_parts, label):
-    current = value
-    for part in path_parts:
-        if isinstance(current, list) and isinstance(part, int):
-            if part >= len(current):
-                fail(f"{label} path {path_parts!r} index {part} is out of range")
-            current = current[part]
-            continue
-
-        if isinstance(current, dict) and isinstance(part, str):
-            if part not in current:
-                fail(f"{label} path {path_parts!r} is missing key {part!r}")
-            current = current[part]
-            continue
-
-        fail(f"{label} path {path_parts!r} cannot read {part!r} from {current!r}")
-
-    return current
-
-
-def resolve_value(value_spec, context, label):
-    if "value" in value_spec:
-        return value_spec["value"]
-
-    source = value_spec.get("source")
-    if not isinstance(source, dict):
-        fail(f"{label} value spec has no literal value or source")
-
-    root = source.get("root")
-    if root not in context:
-        fail(f"{label} value source root {root!r} is unavailable")
-
-    path_parts = source.get("path") or []
-    if not isinstance(path_parts, list):
-        fail(f"{label} value source path must be a list")
-
-    return read_path(context[root], path_parts, label)
-
-
 def response_data(response, operation):
     data = response.data
     if not isinstance(data, dict):
@@ -100,19 +58,13 @@ def feature_step(scenario, collection_name, feature_id, kind):
     fail(f"scenario has no {collection_name} step for feature {feature_id!r}")
 
 
-def create_assembly(client, scenario):
+def file_count(scenario):
     feature = feature_step(scenario, "preparations", "createTusAssembly", "feature-call")
     input_values = list(feature["input"].values())
     if len(input_values) != 1:
         fail(f"{feature['featureId']} expected exactly one input value")
 
-    response = client.create_tus_assembly(input_values[0])
-    data = response_data(response, feature["featureId"])
-    for required_path in feature["requiredResponsePaths"]:
-        value = read_path(data, required_path, feature["featureId"])
-        if value is None or value == "":
-            fail(f"{feature['featureId']} returned empty value at {required_path!r}")
-    return data
+    return input_values[0]
 
 
 def scenario_bytes(scenario):
@@ -124,53 +76,14 @@ def scenario_bytes(scenario):
     return source["value"].encode("utf-8")
 
 
-def upload_metadata(scenario, create_response):
-    context = {"createResponse": create_response, "scenario": scenario}
-    metadata = {}
-    for field in scenario["upload"]["metadata"]:
-        metadata[field["name"]] = str(resolve_value(field["value"], context, field["name"]))
-    return metadata
-
-
-def upload_with_tus(scenario, create_response):
-    context = {"createResponse": create_response, "scenario": scenario}
-    endpoint_url = str(resolve_value(scenario["upload"]["tusUrl"], context, "tusUrl"))
-    content = scenario_bytes(scenario)
-    chunk_size = len(content) if scenario["upload"]["chunkSize"] == "full-file" else None
-    if chunk_size is None:
-        fail(f"unsupported chunk size policy {scenario['upload']['chunkSize']!r}")
-
-    uploader = tus.TusClient(endpoint_url).uploader(
-        file_stream=BytesIO(content),
-        chunk_size=chunk_size,
-        metadata=upload_metadata(scenario, create_response),
-        retries=scenario["upload"]["retries"],
-    )
-    uploader.upload()
-    if not uploader.url:
-        fail("TUS upload did not expose an upload URL")
-    if uploader.offset != len(content):
-        fail(f"TUS upload offset {uploader.offset}, expected {len(content)}")
-    return uploader.url
-
-
-def render_path_template(template_config, context, label):
-    rendered = template_config["template"]
-    for name, value_spec in template_config["replacements"].items():
-        value = resolve_value(value_spec, context, f"{label}.{name}")
-        rendered = rendered.replace("{" + name + "}", quote(str(value), safe=""))
-
-    if "{" in rendered or "}" in rendered:
-        fail(f"{label} still has unresolved placeholders: {rendered}")
-
-    return rendered
-
-
-def wait_for_assembly(client, scenario, create_response):
-    feature = feature_step(scenario, "observations", "waitForAssembly", "feature-poll")
-    context = {"createResponse": create_response, "scenario": scenario}
-    wait_input = render_path_template(feature["input"], context, feature["featureId"])
-    return response_data(client.wait_for_assembly(wait_input), feature["featureId"])
+def upload_config(scenario):
+    upload = scenario["upload"]
+    return {
+        "content": scenario_bytes(scenario),
+        "fieldname": upload["fieldName"],
+        "filename": upload["fileName"],
+        "user_meta": upload.get("userMeta") or {},
+    }
 
 
 def write_result(create_response, status, upload_url):
@@ -200,10 +113,16 @@ def main():
         service=endpoint,
     )
 
-    create_response = create_assembly(client, scenario)
-    upload_url = upload_with_tus(scenario, create_response)
-    status = wait_for_assembly(client, scenario, create_response)
-    write_result(create_response, status, upload_url)
+    upload = upload_config(scenario)
+    completed_assembly, upload_url = client.upload_tus_assembly(
+        file_count(scenario),
+        upload["content"],
+        upload["fieldname"],
+        upload["filename"],
+        upload["user_meta"],
+    )
+    status = response_data(completed_assembly, "uploadTusAssembly")
+    write_result(status, status, upload_url)
 
     print(f"Python Transloadit SDK devdock scenario {scenario['scenarioId']} passed for {endpoint}")
 

transloadit/async_client.py

@@ -267,6 +267,74 @@ async def create_tus_assembly(self, file_count: int):
 
         return assembly
 
+    async def upload_tus_assembly(self, file_count: int, content: bytes, fieldname: str, filename: str, user_meta: Optional[dict] = None):
+        """
+        Create a TUS-ready Assembly, upload one file with the TUS protocol, and wait for the Assembly to finish.
+        """
+        createdAssembly = await self.create_tus_assembly(file_count)
+
+        import base64
+        from urllib.parse import urljoin
+
+        endpointUrl = createdAssembly.data.get("tus_url")
+        if not endpointUrl:
+            raise RuntimeError("TUS singleUploadLifecycle needs input.endpointUrl")
+
+        metadataMap = {}
+        if user_meta:
+            metadataMap.update({str(key): str(value) for key, value in user_meta.items()})
+        metadataMap["assembly_url"] = str(createdAssembly.data.get("assembly_url"))
+        metadataMap["fieldname"] = str(fieldname)
+        metadataMap["filename"] = str(filename)
+
+        session = await self.request._ensure_session()
+
+        createHeaders = {}
+        createHeaders["Tus-Resumable"] = "1.0.0"
+        createHeaders["Upload-Length"] = str(len(content))
+        createMetadataParts = []
+        for key, value in metadataMap.items():
+            encoded_value = base64.b64encode(str(value).encode("utf-8")).decode("ascii")
+            createMetadataParts.append(f"{key} {encoded_value}")
+        createHeaders["Upload-Metadata"] = ",".join(createMetadataParts)
+        async with session.request(
+            "POST",
+            endpointUrl,
+            data=b"",
+            headers=createHeaders,
+            timeout=self.request._timeout(),
+        ) as createResponse:
+            if createResponse.status != 201:
+                raise RuntimeError(f"TUS create returned HTTP {createResponse.status}, expected 201")
+            uploadUrlLocation = createResponse.headers.get("Location")
+            if not uploadUrlLocation:
+                raise RuntimeError("TUS create did not return a Location header")
+            uploadUrlText = urljoin(endpointUrl, uploadUrlLocation)
+
+        uploadHeaders = {}
+        uploadHeaders["Tus-Resumable"] = "1.0.0"
+        uploadHeaders["Upload-Offset"] = "0"
+        uploadHeaders["Content-Type"] = "application/offset+octet-stream"
+        async with session.request(
+            "PATCH",
+            uploadUrlText,
+            data=content,
+            headers=uploadHeaders,
+            timeout=self.request._timeout(),
+        ) as uploadResponse:
+            if uploadResponse.status != 204:
+                raise RuntimeError(f"TUS upload returned HTTP {uploadResponse.status}, expected 204")
+            try:
+                remote_offset = int(uploadResponse.headers.get("Upload-Offset", ""))
+            except ValueError as error:
+                raise RuntimeError("TUS upload returned an invalid Upload-Offset header") from error
+            if remote_offset != len(content):
+                raise RuntimeError(f"TUS upload offset {remote_offset}, expected {len(content)}")
+
+        completedAssembly = await self.wait_for_assembly(createdAssembly.data.get("assembly_ssl_url"))
+
+        return completedAssembly, uploadUrlText
+
     async def wait_for_assembly(self, assembly_url: str):
         """
         Wait for an Assembly to finish uploading and executing.

transloadit/client.py

@@ -275,6 +275,74 @@ def create_tus_assembly(self, file_count: int):
 
         return assembly
 
+    def upload_tus_assembly(self, file_count: int, content: bytes, fieldname: str, filename: str, user_meta: Optional[dict] = None):
+        """
+        Create a TUS-ready Assembly, upload one file with the TUS protocol, and wait for the Assembly to finish.
+        """
+        createdAssembly = self.create_tus_assembly(file_count)
+
+        import base64
+        from urllib.parse import urljoin
+
+        import requests
+
+        endpointUrl = createdAssembly.data.get("tus_url")
+        if not endpointUrl:
+            raise RuntimeError("TUS singleUploadLifecycle needs input.endpointUrl")
+
+        metadataMap = {}
+        if user_meta:
+            metadataMap.update({str(key): str(value) for key, value in user_meta.items()})
+        metadataMap["assembly_url"] = str(createdAssembly.data.get("assembly_url"))
+        metadataMap["fieldname"] = str(fieldname)
+        metadataMap["filename"] = str(filename)
+
+        createHeaders = {}
+        createHeaders["Tus-Resumable"] = "1.0.0"
+        createHeaders["Upload-Length"] = str(len(content))
+        createMetadataParts = []
+        for key, value in metadataMap.items():
+            encoded_value = base64.b64encode(str(value).encode("utf-8")).decode("ascii")
+            createMetadataParts.append(f"{key} {encoded_value}")
+        createHeaders["Upload-Metadata"] = ",".join(createMetadataParts)
+        createResponse = requests.request(
+            "POST",
+            endpointUrl,
+            data=b"",
+            headers=createHeaders,
+            timeout=request.TIMEOUT,
+        )
+        if createResponse.status_code != 201:
+            raise RuntimeError(f"TUS create returned HTTP {createResponse.status_code}, expected 201")
+        uploadUrlLocation = createResponse.headers.get("Location")
+        if not uploadUrlLocation:
+            raise RuntimeError("TUS create did not return a Location header")
+        uploadUrlText = urljoin(endpointUrl, uploadUrlLocation)
+
+        uploadHeaders = {}
+        uploadHeaders["Tus-Resumable"] = "1.0.0"
+        uploadHeaders["Upload-Offset"] = "0"
+        uploadHeaders["Content-Type"] = "application/offset+octet-stream"
+        uploadResponse = requests.request(
+            "PATCH",
+            uploadUrlText,
+            data=content,
+            headers=uploadHeaders,
+            timeout=request.TIMEOUT,
+        )
+        if uploadResponse.status_code != 204:
+            raise RuntimeError(f"TUS upload returned HTTP {uploadResponse.status_code}, expected 204")
+        try:
+            remote_offset = int(uploadResponse.headers.get("Upload-Offset", ""))
+        except ValueError as error:
+            raise RuntimeError("TUS upload returned an invalid Upload-Offset header") from error
+        if remote_offset != len(content):
+            raise RuntimeError(f"TUS upload offset {remote_offset}, expected {len(content)}")
+
+        completedAssembly = self.wait_for_assembly(createdAssembly.data.get("assembly_ssl_url"))
+
+        return completedAssembly, uploadUrlText
+
     def wait_for_assembly(self, assembly_url: str):
         """
         Wait for an Assembly to finish uploading and executing.