upload_file.py

#!/usr/bin/env python3
"""
TriplyDB File Upload - Python Reference Implementation

This script demonstrates how to upload a file to TriplyDB using the TUS protocol.
It replicates the functionality from the triplydb-js client library.

Requirements:
    pip install tuspy requests

Usage:
    python upload_file.py --file path/to/file.ttl --account myaccount --dataset mydataset --token YOUR_API_TOKEN
"""

import argparse
import os
import sys
import time
from typing import Optional

import requests
from tusclient import client as tus_client


class TriplyDBUploader:
    """Handle file uploads to TriplyDB using TUS protocol."""

    def __init__(self, api_url: str, token: str, account: str, dataset: str):
        """
        Initialize the TriplyDB uploader.

        Args:
            api_url: Base URL of the TriplyDB API (e.g., 'https://api.triplydb.com')
            token: TriplyDB API token
            account: Account name
            dataset: Dataset name
        """
        self.api_url = api_url.rstrip("/")
        self.token = token
        self.account = account
        self.dataset = dataset
        self.dataset_path = f"/datasets/{account}/{dataset}"
        self.job_url: Optional[str] = None
        self.job_id: Optional[str] = None

    def _get_headers(self) -> dict:
        """Get HTTP headers with authorization."""
        return {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }

    def ensure_dataset(self) -> None:
        """Create the dataset if it doesn't exist yet."""
        url = f"{self.api_url}{self.dataset_path}"
        response = requests.get(url, headers=self._get_headers())
        if response.status_code == 404:
            create_url = f"{self.api_url}/datasets/{self.account}"
            print(f"Dataset not found, creating: {self.dataset}")
            response = requests.post(
                create_url,
                json={"name": self.dataset},
                headers=self._get_headers(),
            )
            response.raise_for_status()
            print(f"✓ Dataset created: {self.dataset}")
        else:
            response.raise_for_status()

    def create_job(self, **job_config) -> dict:
        """
        Create an upload job in TriplyDB.

        Args:
            **job_config: Optional job configuration (defaultGraphName, overwriteAll, etc.)

        Returns:
            Job information dictionary

        Raises:
            requests.HTTPError: If job creation fails
        """
        url = f"{self.api_url}{self.dataset_path}/jobs"

        # Prepare job data (type is intentionally omitted for file uploads)
        data = {**job_config}

        print(f"Creating job at: {url}")
        response = requests.post(url, json=data, headers=self._get_headers())
        response.raise_for_status()

        job_info = response.json()
        self.job_id = job_info["jobId"]
        self.job_url = f"{self.api_url}{self.dataset_path}/jobs/{self.job_id}"

        print(f"✓ Job created: {self.job_id}")
        return job_info

    def upload_file(self, file_path: str) -> None:
        """
        Upload a file using TUS protocol.

        Args:
            file_path: Path to the file to upload

        Raises:
            ValueError: If job hasn't been created yet
            FileNotFoundError: If file doesn't exist
        """
        if not self.job_url:
            raise ValueError("Job must be created before uploading files")

        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")

        # TUS upload endpoint
        upload_url = f"{self.job_url}/add"

        # Get file info
        file_size = os.path.getsize(file_path)
        file_name = os.path.basename(file_path)

        print(f"Uploading file: {file_name} ({file_size:,} bytes)")

        # Configure TUS client (matching JavaScript implementation)
        # Headers must be passed to TusClient, not to uploader
        tus = tus_client.TusClient(
            upload_url,
            headers={"Authorization": f"Bearer {self.token}"}
        )
        uploader = tus.uploader(
            file_path,
            chunk_size=5 * 1024 * 1024,  # 5MB chunks (same as JS implementation)
            metadata={
                "filename": file_name,
            },
            retries=5,  # Match JS retry delays length
        )

        # Upload with progress tracking
        try:
            uploader.upload()
            print(f"✓ File uploaded: {file_name}")
        except Exception as e:
            print(f"✗ Upload failed: {e}", file=sys.stderr)
            raise

    def start_job(self) -> dict:
        """
        Start the job execution and indexing.

        Returns:
            Updated job information

        Raises:
            ValueError: If job hasn't been created yet
            requests.HTTPError: If starting job fails
        """
        if not self.job_url:
            raise ValueError("Job must be created before starting")

        url = f"{self.job_url}/start"
        print(f"Starting job and indexing: {self.job_id}")

        response = requests.post(url, headers=self._get_headers())
        response.raise_for_status()

        job_info = response.json()
        print("✓ Job started, indexing in progress")
        return job_info

    def wait_for_job(self, poll_interval: float = 0.5, max_wait: float = 10.0) -> dict:
        """
        Poll job status until completion.

        Args:
            poll_interval: Initial polling interval in seconds (increases exponentially)
            max_wait: Maximum polling interval in seconds

        Returns:
            Final job information

        Raises:
            ValueError: If job hasn't been created yet
            RuntimeError: If job fails or is canceled
            requests.HTTPError: If polling fails
        """
        if not self.job_url:
            raise ValueError("Job must be created before polling")

        wait_time = poll_interval
        print("Waiting for indexing to complete...")

        while True:
            time.sleep(wait_time)

            response = requests.get(self.job_url, headers=self._get_headers())
            response.raise_for_status()

            job_info = response.json()
            status = job_info.get("status")

            if status == "error":
                error_msg = job_info.get("error", {}).get("message", "Unknown error")
                raise RuntimeError(f"Job failed: {error_msg}")

            if status in ("canceled", "finished"):
                if status == "finished":
                    print(f"✓ Indexing complete - data is now queryable")
                else:
                    print(f"✓ Job {status}")
                return job_info

            # Exponential backoff (same strategy as JS implementation)
            wait_time = min(max_wait, wait_time * 3)

    def upload_and_import(self, file_path: str, **job_config) -> dict:
        """
        Complete workflow: create job, upload file, start indexing, and wait for completion.

        Args:
            file_path: Path to the file to upload
            **job_config: Optional job configuration

        Returns:
            Final job information
        """
        # Step 1: Ensure dataset exists
        self.ensure_dataset()

        # Step 2: Create job
        self.create_job(**job_config)

        # Step 3: Upload file via TUS
        self.upload_file(file_path)

        # Step 4: Start indexing
        self.start_job()

        # Step 5: Wait for indexing to complete
        return self.wait_for_job()


def main():
    """Command-line interface for the uploader."""
    parser = argparse.ArgumentParser(
        description="Upload a file to TriplyDB using TUS protocol"
    )
    parser.add_argument(
        "--file",
        required=True,
        help="Path to the file to upload"
    )
    parser.add_argument(
        "--account",
        required=True,
        help="TriplyDB account name"
    )
    parser.add_argument(
        "--dataset",
        required=True,
        help="TriplyDB dataset name"
    )
    parser.add_argument(
        "--token",
        required=True,
        help="TriplyDB API token"
    )
    parser.add_argument(
        "--api-url",
        required=True,
        help="TriplyDB API URL (e.g., https://api.triplydb.com)"
    )
    parser.add_argument(
        "--default-graph",
        help="Default graph name for imported triples"
    )
    parser.add_argument(
        "--overwrite-all",
        action="store_true",
        help="Overwrite all existing graphs"
    )
    parser.add_argument(
        "--merge-graphs",
        action="store_true",
        help="Merge data into existing graphs"
    )

    args = parser.parse_args()

    # Build job configuration
    job_config = {}
    if args.default_graph:
        job_config["defaultGraphName"] = args.default_graph
    if args.overwrite_all:
        job_config["overwriteAll"] = True
    if args.merge_graphs:
        job_config["mergeGraphs"] = True

    # Create uploader and run
    uploader = TriplyDBUploader(
        api_url=args.api_url,
        token=args.token,
        account=args.account,
        dataset=args.dataset,
    )

    try:
        result = uploader.upload_and_import(args.file, **job_config)
        print("\n✓ Upload completed successfully!")
        print(f"Final status: {result.get('status')}")
        return 0
    except Exception as e:
        print(f"\n✗ Upload failed: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())