Drive Image Reader

1import os2import sys3import json4import traceback5import requests67EMBLEMA_API_BASE_URL = os.getenv("EMBLEMA_API_BASE_URL")8USER_TOKEN = os.getenv("USER_TOKEN")9OUTPUT_DIR = "/data/output"1011if not EMBLEMA_API_BASE_URL:12    raise ValueError("EMBLEMA_API_BASE_URL environment variable not set")131415def api_request_json(method, path, **kwargs):16    """Make authenticated API request and return JSON response."""17    if not USER_TOKEN:18        raise ValueError("USER_TOKEN environment variable not set")1920    headers = {21        "Authorization": f"Bearer {USER_TOKEN}",22    }23    url = f"{EMBLEMA_API_BASE_URL}{path}"24    response = requests.request(method, url, headers=headers, **kwargs)25    response.raise_for_status()26    return response.json()272829def fetch_image(drive_item_id_input) -> str:30    """Fetch an image from Drive and save it to the output directory.3132    Args:33        drive_item_id_input: Single Drive Item ID (string). Arrays are rejected.3435    Returns:36        The output filename saved to /data/output/37    """38    # Normalize input — reject arrays, this is a single-image reader39    if isinstance(drive_item_id_input, list):40        if len(drive_item_id_input) == 1:41            drive_item_id = drive_item_id_input[0]42        else:43            raise ValueError(44                f"Drive Image Reader accepts a single image, got {len(drive_item_id_input)} items. "45                "Use one Drive Image Reader per image."46            )47    elif isinstance(drive_item_id_input, str):48        drive_item_id = drive_item_id_input49    else:50        raise ValueError(f"driveItemId must be a string, got {type(drive_item_id_input).__name__}")5152    if not drive_item_id:53        raise ValueError("No valid Drive Item ID provided")5455    # Get presigned download URL + metadata56    download_info = api_request_json("GET", f"/api/v2/drive/items/{drive_item_id}/download")5758    presigned_url = download_info.get("url")59    file_name = download_info.get("name", "image.png")60    mime_type = download_info.get("mimeType", "")6162    if not presigned_url:63        raise ValueError("Download API did not return a presigned URL")6465    # Validate that the file is an image66    if mime_type and not mime_type.startswith("image/"):67        raise ValueError(68            f"Selected Drive item is not an image (mimeType: {mime_type}). "69            "Please select an image file."70        )7172    # Download the image binary from presigned URL73    img_response = requests.get(presigned_url, timeout=120)74    img_response.raise_for_status()7576    # Save to output directory77    os.makedirs(OUTPUT_DIR, exist_ok=True)78    out_path = os.path.join(OUTPUT_DIR, file_name)79    with open(out_path, "wb") as f:80        f.write(img_response.content)8182    return file_name838485def main():86    """Main execution function"""87    try:88        input_json = sys.stdin.read()89        execution_input = json.loads(input_json)9091        inputs = execution_input.get("inputs", {})92        drive_item_id = inputs.get("driveItemId")9394        if not drive_item_id:95            raise ValueError("Required input 'driveItemId' not provided")9697        filename = fetch_image(drive_item_id)9899        output = {100            "image": filename,101        }102103        print(json.dumps(output, indent=2))104105    except requests.HTTPError as e:106        error_output = {107            "error": f"API request failed: {e}",108            "errorType": "HTTPError",109            "status": e.response.status_code,110            "detail": e.response.text,111        }112        print(json.dumps(error_output), file=sys.stderr)113        sys.exit(1)114    except Exception as e:115        error_output = {116            "error": str(e),117            "errorType": type(e).__name__,118            "traceback": traceback.format_exc(),119        }120        print(json.dumps(error_output), file=sys.stderr)121        sys.exit(1)122123124if __name__ == "__main__":125    main()