ai4mats-mcp-tools/.opencode/skills/dataset/create_dataset/scripts/upload_file.py

import asyncio
import hashlib
import json
import os
import sys
import time

import httpx
from common.config import API_BASE_URL

DEFAULT_SIZE = 10 * 1024 * 1024  # 10MB

def compute_md5(file_path: str) -> str:
    file_size = os.path.getsize(file_path)
    chunk_size = min(DEFAULT_SIZE, file_size)
    with open(file_path, "rb") as f:
        data = f.read(chunk_size)
    md5 = hashlib.md5(data).hexdigest()
    filename = os.path.basename(file_path)
    name_bytes = filename.encode('utf-8')
    combined = md5.encode('utf-8') + name_bytes
    return hashlib.md5(combined).hexdigest()

async def get_upload_task(token: str, params: dict) -> dict:
    url = f"{API_BASE_URL}/api/mmp/uploader/chunk"
    headers = {"Authorization": f"Bearer {token}"}
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.get(url, params=params, headers=headers)
        response.raise_for_status()
        return response.json()

async def upload_chunk(token: str, file_path: str, part_number: int, total_chunks: int, identifier: str) -> dict:
    url = f"{API_BASE_URL}/api/mmp/uploader/chunk"
    file_size = os.path.getsize(file_path)
    filename = os.path.basename(file_path)

    start = DEFAULT_SIZE * (part_number - 1)
    end = min(start + DEFAULT_SIZE, file_size)
    current_chunk_size = end - start

    with open(file_path, "rb") as f:
        f.seek(start)
        blob_data = f.read(current_chunk_size)

    files = {
        "chunkNumber": (None, str(part_number)),
        "chunkSize": (None, str(DEFAULT_SIZE)),
        "currentChunkSize": (None, str(current_chunk_size)),
        "filename": (None, filename),
        "relativePath": (None, filename),
        "identifier": (None, identifier),
        "totalChunks": (None, str(total_chunks)),
        "totalSize": (None, str(file_size)),
        "upfile": (filename, blob_data),
    }

    headers = {"Authorization": f"Bearer {token}"}
    async with httpx.AsyncClient(timeout=600.0) as client:
        response = await client.post(url, files=files, headers=headers)
        response.raise_for_status()
        return response.json()

async def merge_chunks(token: str, file_path: str, identifier: str) -> dict:
    url = f"{API_BASE_URL}/api/mmp/uploader/mergeFile"
    file_size = os.path.getsize(file_path)
    filename = os.path.basename(file_path)
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json; charset=UTF-8",
    }
    payload = {
        "fileType": "application/zip",
        "name": filename,
        "relativePath": filename,
        "size": file_size,
        "uniqueIdentifier": identifier,
        "refProjectId": "123456789",
    }
    async with httpx.AsyncClient(timeout=600.0) as client:
        response = await client.post(url, json=payload, headers=headers)
        response.raise_for_status()
        return response.json()

async def get_merge_status(token: str, filename: str, identifier: str) -> dict:
    url = f"{API_BASE_URL}/api/mmp/uploader/selectFile"
    headers = {"Authorization": f"Bearer {token}"}
    params = {"filename": filename, "identifier": identifier}
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.get(url, params=params, headers=headers)
        response.raise_for_status()
        return response.json()

async def upload_file(token: str, file_path: str) -> dict:
    file_size = os.path.getsize(file_path)
    filename = os.path.basename(file_path)
    total_chunks = max(1, (file_size + DEFAULT_SIZE - 1) // DEFAULT_SIZE)

    identifier = compute_md5(file_path)
    print(f"File: {filename}, Size: {file_size}, Chunks: {total_chunks}, MD5: {identifier}")

    task_params = {
        "chunkNumber": 1,
        "chunkSize": DEFAULT_SIZE,
        "currentChunkSize": min(DEFAULT_SIZE, file_size),
        "totalSize": file_size,
        "identifier": identifier,
        "filename": filename,
        "relativePath": filename,
        "totalChunks": total_chunks,
    }
    task_result = await get_upload_task(token, task_params)
    print(f"Task result: {json.dumps(task_result, ensure_ascii=False)}")

    if task_result.get("code") != 200:
        raise Exception(f"Get upload task failed: {task_result}")

    task = task_result.get("data", {})
    if task.get("skip_upload"):
        print("File already uploaded, skipping upload")
        return task

    for part in range(1, total_chunks + 1):
        print(f"Uploading chunk {part}/{total_chunks}...")
        result = await upload_chunk(token, file_path, part, total_chunks, identifier)
        print(f"  Chunk {part} result: {json.dumps(result, ensure_ascii=False)}")

    print("Merging chunks...")
    merge_result = await merge_chunks(token, file_path, identifier)
    print(f"Merge result: {json.dumps(merge_result, ensure_ascii=False)}")

    if merge_result.get("code") != 200:
        raise Exception(f"Merge failed: {merge_result}")

    merge_data = merge_result.get("data", {})
    if merge_data.get("state") == "Succeeded":
        print("Merge succeeded immediately!")
        return merge_data
    if merge_data.get("location"):
        print("Merge has location, returning immediately despite state:", merge_data.get("state"))
        return merge_data

    for i in range(30):
        await asyncio.sleep(3)
        status_result = await get_merge_status(token, filename, identifier)
        status_data = status_result.get("data", {})
        state = status_data.get("state")
        print(f"  Merge status poll #{i+1}: {state}")
        if state == "Succeeded":
            print("Merge succeeded!")
            return status_data
        elif state == "Failed":
            if status_data.get("location"):
                print("Merge has location despite Failed state, using it")
                return status_data
            raise Exception(f"Merge failed: {status_result}")

    raise Exception("Merge status polling timed out")

async def main():
    token = sys.argv[1]
    file_path = sys.argv[2]
    result = await upload_file(token, file_path)
    print(f"\nFinal result:\n{json.dumps(result, ensure_ascii=False, indent=2)}")

if __name__ == "__main__":
    asyncio.run(main())