add_collaborators_to_dataset_from_csv.py

"""Bulk-add dataset collaborators to a Segments.ai dataset from a CSV of usernames and roles.

Usage:
    python scripts/add_dataset_collaborators.py --api-key <API_KEY> --dataset jane/flowers --csv collaborators.csv

The CSV is expected to have "username" and "role" columns, e.g.:

    username,role
    john,reviewer
    jane,labeler

Valid roles: labeler, reviewer, manager, admin. Rows with an empty role default to "labeler".
"""

import argparse
import csv
import sys

from segments import SegmentsClient
from segments.exceptions import SegmentsError


VALID_ROLES = ["labeler", "reviewer", "manager", "admin"]
DEFAULT_ROLE = "labeler"


def format_error_message(error: SegmentsError) -> str:
    return " ".join(str(error).split())


def read_collaborators(csv_path: str) -> tuple[list[tuple[str, str]], list[str]]:
    with open(csv_path, newline="") as csv_file:
        reader = csv.DictReader(csv_file)
        if reader.fieldnames is None or "username" not in reader.fieldnames:
            raise ValueError('CSV must have a "username" column')

        collaborators = []
        skipped_rows = []
        for row_index, row in enumerate(reader, start=2):
            username = row["username"].strip()
            if username == "":
                skipped_rows.append(f"Row {row_index}: skipped, missing username")
                continue

            role = (row.get("role") or "").strip() or DEFAULT_ROLE
            if role not in VALID_ROLES:
                raise ValueError(f"Row {row_index}: invalid role '{role}' for username '{username}'. Must be one of {VALID_ROLES}")

            collaborators.append((username, role))

        return collaborators, skipped_rows


def add_collaborators(client: SegmentsClient, dataset_identifier: str, collaborators: list[tuple[str, str]]) -> tuple[list[str], bool]:
    logs = []
    failed_usernames = []
    for username, role in collaborators:
        try:
            client.add_dataset_collaborator(dataset_identifier, username, role)
            logs.append(f"Added {username} to {dataset_identifier} as {role}")
        except SegmentsError as error:
            logs.append(f"Failed to add {username}: {format_error_message(error)}")
            failed_usernames.append(username)

    if failed_usernames:
        logs.append(f"\n{len(failed_usernames)} of {len(collaborators)} usernames failed: {', '.join(failed_usernames)}")

    return logs, len(failed_usernames) > 0


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--api-key", required=True, help="Segments.ai API key")
    parser.add_argument("--dataset", required=True, help="Dataset identifier, e.g. jane/flowers")
    parser.add_argument("--csv", required=True, help="Path to the CSV file")
    parser.add_argument("--api-url", default="https://api.segments.ai/api/v1/", help="Segments API URL")
    args = parser.parse_args()

    try:
        client = SegmentsClient(args.api_key, api_url=args.api_url)
    except SegmentsError as error:
        print(format_error_message(error), file=sys.stderr)
        sys.exit(1)

    try:
        collaborators, skipped_rows = read_collaborators(args.csv)
    except ValueError as error:
        print(error, file=sys.stderr)
        sys.exit(1)

    for skipped_row in skipped_rows:
        print(skipped_row, file=sys.stderr)

    if not collaborators:
        print("No usernames found in CSV", file=sys.stderr)
        sys.exit(1)

    logs, has_failures = add_collaborators(client, args.dataset, collaborators)
    for log in logs:
        print(log)

    if has_failures or skipped_rows:
        sys.exit(1)

if __name__ == "__main__":
    main()