Drive File Reader

1"""drive-file-reader — read one Drive file as it was uploaded, output it as a File.23The generic sibling of drive-image-reader / drive-3d-model-loader: no type4check, so any file (a PDF, .pptx, .docx, image, archive, ...) reaches a File5port byte for byte. Drive Content Reader is the one for a file's processed6text. Privileged Drive IO is delegated to the gais.drive SDK (USER_TOKEN +7Drive API live there); this node imports no HTTP client and uses no os/shutil8(both are lint-blocked). The downloaded file is stream-copied into9/data/output with plain open() so the executor uploads it as the File output.10"""1112from __future__ import annotations1314import json15import re16import sys17import traceback18from typing import Any1920from gais import Gais2122OUTPUT_DIR = "/data/output"23DEFAULT_NAME = "drive-file"24# The next node stages an input as "<input>_<i>_<j>_<name>" and a file name25# stops at 255 bytes, so the output name stays well under that.26MAX_NAME_BYTES = 20027MAX_EXTENSION_BYTES = 162829_UUID_RE = re.compile(30    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",31    re.IGNORECASE,32)33# The Drive mimeType is whatever the uploader's client sent, and it reaches the34# agent's context through getNodeOutput: only a plain type/subtype is passed on.35_MIME_TYPE_RE = re.compile(r"[A-Za-z0-9.+_-]{1,64}/[A-Za-z0-9.+_-]{1,64}")36# What the workspace refuses in an output file name: a control character, or a37# %XX escape (the upload check decodes it first, so "%2F" is a slash). A38# leading "." is refused at the next edge. Either way the node would report39# success and the downstream node would get no file.40_UNSAFE_NAME_CHARS = re.compile(r"[\x00-\x1f\x7f]|%(?=[0-9A-Fa-f]{2})")414243def _normalize_drive_item_id(value: Any) -> str:44    """Accept a bare UUID or a mention://drive_{file,folder}/<uuid> URI."""45    if not isinstance(value, str):46        raise ValueError(f"Drive item id must be a string, got {type(value).__name__}")47    s = value.strip()48    for prefix in ("mention://drive_file/", "mention://drive_folder/"):49        if s.startswith(prefix):50            s = s[len(prefix):]51            break52    if not _UUID_RE.match(s):53        raise ValueError(f"Drive item id is not a valid UUID: {value!r}")54    return s555657def _cap_name_bytes(name: str) -> str:58    """Cut the stem so the UTF-8 name fits MAX_NAME_BYTES, keeping a short59    extension and never splitting a character."""60    if len(name.encode("utf-8")) <= MAX_NAME_BYTES:61        return name62    stem, dot, ext = name.rpartition(".")63    if not dot or not stem or len(ext.encode("utf-8")) > MAX_EXTENSION_BYTES:64        stem, ext = name, ""65    suffix = f".{ext}" if ext else ""66    budget = MAX_NAME_BYTES - len(suffix.encode("utf-8"))67    return stem.encode("utf-8")[:budget].decode("utf-8", "ignore") + suffix686970def _safe_output_name(name: Any, default: str) -> str:71    """Reduce a Drive-supplied name to a bare filename the workspace accepts:72    no directory part (it cannot escape OUTPUT_DIR), no control character or73    %XX escape (refused at upload), no leading dot (refused at the next edge),74    at most MAX_NAME_BYTES (the next edge prefixes it when staging)."""75    base = str(name or "").replace("\\", "/").rsplit("/", 1)[-1].strip()76    return _cap_name_bytes(_UNSAFE_NAME_CHARS.sub("_", base).lstrip(". ")) or default777879def _copy_to_output(src_path: str, name: str) -> int:80    """Stream-copy the downloaded file into /data/output in 1 MiB chunks; return bytes copied."""81    total = 082    with open(src_path, "rb") as src, open(f"{OUTPUT_DIR}/{name}", "wb") as dst:83        while True:84            chunk = src.read(1024 * 1024)85            if not chunk:86                break87            dst.write(chunk)88            total += len(chunk)89    return total909192def process(inputs: dict[str, Any]) -> dict[str, Any]:93    """Validate the input, download via gais.drive, stage the File output."""94    raw = inputs.get("driveItemId")95    if not raw:96        raise ValueError("Required input 'driveItemId' not provided")9798    if isinstance(raw, list):99        if len(raw) == 1:100            raw = raw[0]101        else:102            raise ValueError(103                f"Drive File Reader accepts a single file, got {len(raw)} items. "104                "Use one Drive File Reader per file."105            )106    elif not isinstance(raw, str):107        raise ValueError(f"driveItemId must be a string, got {type(raw).__name__}")108109    item_id = _normalize_drive_item_id(raw)110111    result = Gais.drive.download(item_id=item_id)112    file_name = _safe_output_name(result.metadata.get("name"), DEFAULT_NAME)113    raw_mime = result.metadata.get("mimeType", "") or ""114    mime_type = raw_mime if isinstance(raw_mime, str) and _MIME_TYPE_RE.fullmatch(raw_mime) else ""115116    size_bytes = _copy_to_output(str(result.file("file")), file_name)117    return {"file": file_name, "mime_type": mime_type, "size_bytes": size_bytes}118119120def main() -> None:121    try:122        envelope = json.loads(sys.stdin.read() or "{}")123        if not isinstance(envelope, dict):124            envelope = {}125        inputs = envelope.get("inputs", {}) or {}126        json.dump(process(inputs), sys.stdout)127    except Exception as e:128        print(129            json.dumps({130                "error": str(e),131                "errorType": type(e).__name__,132                "traceback": traceback.format_exc(),133            }),134            file=sys.stderr,135        )136        sys.exit(1)137138139if __name__ == "__main__":140    main()141