uv add tableauserverclient>=v0.40

import argparse
from collections.abc import Iterable, Sequence
import csv
from collections import defaultdict
from dataclasses import dataclass
from itertools import batched
import os
from pathlib import Path
import sys

from dotenv import load_dotenv
import tableauserverclient as TSC

load_dotenv()


@dataclass
class AddUsersArgs:
    filename: Path


def parse_args(args: Sequence[str] | None = None) -> AddUsersArgs:
    if args is None:
        args = sys.argv[1:]
    parser = argparse.ArgumentParser(
        description="Bulk add users to Tableau Server from a CSV file."
    )
    parser.add_argument(
        "filename", type=Path, help="Path to the CSV file containing user information."
    )
    parsed_args = parser.parse_args(args)
    return AddUsersArgs(filename=parsed_args.filename)


def read_users_csv(
    filename: str | Path,
) -> tuple[list[TSC.UserItem], dict[str, list[str]]]:
    users = []
    groups: dict[str, list[str]] = defaultdict(list)
    filepath = Path(filename)
    with open(filepath, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            # Name and site_role are required, everything else is optional.
            name = row.pop("name")
            site_role = row.pop("site_role")
            default_auth_setting = "TableauIDWithMFA"  # Cloud
            # default_auth_setting = "ServerDefault" # Server

            user = TSC.UserItem(
                name=name,
                site_role=site_role,
                auth_setting=row.get("auth_setting") or default_auth_setting,
            )

            if (group_names := row.pop("group_names", None)) is not None:
                for group_name in group_names.split(";"):
                    groups[group_name].append(name)

            for key, value in row.items():
                if value:  # Only add non-empty values
                    setattr(user, key, value)

            users.append(user)

    return users, groups


def bulk_add_users(
    server: TSC.Server, users: Iterable[TSC.UserItem]
) -> dict[str, TSC.UserItem]:

    job = server.users.bulk_add(users)
    server.jobs.wait_for_job(job)
    names = [user.name for user in users]
    site_users: dict[str, TSC.UserItem] = {}
    for batch in batched(names, 100):
        site_users |= {u.name: u for u in server.users.filter(name__in=batch)}
    if len(site_users) != len(users):  # Sanity check that all users were added
        raise ValueError("Not all users were added successfully")
    return site_users


def add_users(
    server: TSC.Server, users: Iterable[TSC.UserItem]
) -> dict[str, TSC.UserItem]:
    site_users: dict[str, TSC.UserItem] = {}
    for user in users:
        try:
            created_user = server.users.add(user)
            site_users[created_user.name] = created_user
        except TSC.ServerResponseError as e:
            print(f"Error creating user '{user.name}': {e}")
    return site_users


def add_users_to_groups(
    server: TSC.Server,
    groups: dict[str, list[str]],
    site_users: dict[str, TSC.UserItem],
) -> None:
    for group_name, user_names in groups.items():
        try:
            group = server.groups.filter(name=group_name)[0]
        except IndexError:
            group = TSC.GroupItem(name=group_name)
            group = server.groups.create(group)
        added_users = server.groups.add_users(
            group, [site_users[name] for name in user_names]
        )

        requested_users = set(user_names)
        added_user_names = {u.name for u in added_users}
        if added_user_names != requested_users:
            missing_users = requested_users - added_user_names
            print(
                f"Warning: The following users were not added to group '{group_name}': {', '.join(missing_users)}"
            )


def main():
    args = parse_args()
    users, groups = read_users_csv(args.filename)

    server = TSC.Server(os.environ["TABLEAU_SERVER"], use_server_version=True)
    auth = TSC.PersonalAccessTokenAuth(
        os.environ["TABLEAU_TOKEN_NAME"],
        os.environ["TABLEAU_TOKEN_SECRET"],
        site_id=os.getenv("TABLEAU_SITE", ""),
    )

    with server.auth.sign_in(auth):
        if len(users) > 1_000:
            site_users = bulk_add_users(server, users)
        else:
            site_users = add_users(server, users)

        add_users_to_groups(server, groups, site_users)


if __name__ == "__main__":
    main()