Feature/add bulk user role commands (#331)

* bulk user role commands

* bulk user roles

* bulk cmds

* fixing unit tests

* renaming variable

* add cache decorate to _get_role_id()

* changed to lru_cache for <3.9 compatability

* fix

Author: tora-kozic

CHANGELOG.md

@@ -16,6 +16,10 @@ how a consumer would use the library (e.g. adding unit tests, updating documenta
 
 ### Added
 
+- New bulk commands to manage user roles
+    - `code42 users bulk add-roles`
+    - `code42 users bulk remove-roles`
+
 - New option `--include-roles` on `code42 users list` that includes the roles for all users.
 
 - New command `code42 users show <username>` that prints all the details of that user.

src/code42cli/cmds/users.py

@@ -1,3 +1,5 @@
+import functools
+
 import click
 from pandas import DataFrame
 from pandas import json_normalize
@@ -217,6 +219,8 @@ def reactivate(state, username):
 
 _bulk_user_move_headers = ["username", "org_id"]
 
+_bulk_user_roles_headers = ["username", "role_name"]
+
 
 @users.command(name="move")
 @username_option("The username of the user to move.", required=True)
@@ -461,6 +465,86 @@ def handle_row(**row):
     formatter.echo_formatted_list(result_rows)
 
 
+@bulk.command(
+    name="add-roles",
+    help=f"Add roles to a list of users from the provided CSV in format: {','.join(_bulk_user_roles_headers)}",
+)
+@read_csv_arg(headers=_bulk_user_roles_headers)
+@format_option
+@sdk_options()
+def bulk_add_roles(state, csv_rows, format):
+    """Bulk add roles to a list of users."""
+
+    # Initialize the SDK before starting any bulk processes
+    # to prevent multiple instances and having to enter 2fa multiple times.
+    sdk = state.sdk
+    status_header = "role added"
+
+    csv_rows[0][status_header] = "False"
+    formatter = OutputFormatter(format, {key: key for key in csv_rows[0].keys()})
+    stats = create_worker_stats(len(csv_rows))
+
+    def handle_row(**row):
+        try:
+            _add_user_role(
+                sdk, **{key: row[key] for key in row.keys() if key != status_header}
+            )
+            row[status_header] = "True"
+        except Exception as err:
+            row[status_header] = f"False: {err}"
+            stats.increment_total_errors()
+        return row
+
+    result_rows = run_bulk_process(
+        handle_row,
+        csv_rows,
+        progress_label="Adding roles to users:",
+        stats=stats,
+        raise_global_error=False,
+    )
+    formatter.echo_formatted_list(result_rows)
+
+
+@bulk.command(
+    name="remove-roles",
+    help=f"Remove roles from a list of users from the provided CSV in format: {','.join(_bulk_user_roles_headers)}",
+)
+@read_csv_arg(headers=_bulk_user_roles_headers)
+@format_option
+@sdk_options()
+def bulk_remove_roles(state, csv_rows, format):
+    """Bulk remove roles from a list of users."""
+
+    # Initialize the SDK before starting any bulk processes
+    # to prevent multiple instances and having to enter 2fa multiple times.
+    sdk = state.sdk
+    success_header = "role removed"
+
+    csv_rows[0][success_header] = "False"
+    formatter = OutputFormatter(format, {key: key for key in csv_rows[0].keys()})
+    stats = create_worker_stats(len(csv_rows))
+
+    def handle_row(**row):
+        try:
+            _remove_user_role(
+                sdk, **{key: row[key] for key in row.keys() if key != success_header}
+            )
+            row[success_header] = "True"
+        except Exception as err:
+            row[success_header] = f"False: {err}"
+            stats.increment_total_errors()
+        return row
+
+    result_rows = run_bulk_process(
+        handle_row,
+        csv_rows,
+        progress_label="Removing roles from users:",
+        stats=stats,
+        raise_global_error=False,
+    )
+    formatter.echo_formatted_list(result_rows)
+
+
 def _add_user_role(sdk, username, role_name):
     user_id = _get_legacy_user_id(sdk, username)
     _get_role_id(sdk, role_name)  # function provides role name validation
@@ -484,6 +568,7 @@ def _get_legacy_user_id(sdk, username):
     return user_id
 
 
+@functools.lru_cache()
 def _get_role_id(sdk, role_name):
     try:
         roles_dataframe = DataFrame.from_records(

tests/cmds/test_users.py

@@ -1108,6 +1108,154 @@ def _get(username, *args, **kwargs):
     assert worker_stats.increment_total_errors.call_count == 1
 
 
+def test_bulk_add_roles_uses_expected_arguments(runner, mocker, cli_state_with_user):
+    bulk_processor = mocker.patch(f"{_NAMESPACE}.run_bulk_process")
+    with runner.isolated_filesystem():
+        with open("test_bulk_add_roles.csv", "w") as csv:
+            csv.writelines(
+                ["username,role_name\n", f"{TEST_USERNAME},{TEST_ROLE_NAME}\n"]
+            )
+        command = ["users", "bulk", "add-roles", "test_bulk_add_roles.csv"]
+        runner.invoke(
+            cli, command, obj=cli_state_with_user,
+        )
+    assert bulk_processor.call_args[0][1] == [
+        {"username": TEST_USERNAME, "role_name": TEST_ROLE_NAME, "role added": "False"},
+    ]
+    bulk_processor.assert_called_once()
+
+
+def test_bulk_add_roles_ignores_blank_lines(runner, mocker, cli_state):
+    bulk_processor = mocker.patch(f"{_NAMESPACE}.run_bulk_process")
+    with runner.isolated_filesystem():
+        with open("test_bulk_add_roles.csv", "w") as csv:
+            csv.writelines(
+                ["username,role_name\n\n\n", f"{TEST_USERNAME},{TEST_ROLE_NAME}\n\n\n"]
+            )
+        runner.invoke(
+            cli,
+            ["users", "bulk", "add-roles", "test_bulk_add_roles.csv"],
+            obj=cli_state,
+        )
+    assert bulk_processor.call_args[0][1] == [
+        {"username": TEST_USERNAME, "role_name": TEST_ROLE_NAME, "role added": "False"},
+    ]
+    bulk_processor.assert_called_once()
+
+
+def test_bulk_add_roles_uses_handler_that_when_encounters_error_increments_total_errors(
+    runner,
+    mocker,
+    cli_state,
+    worker_stats,
+    get_users_response,
+    get_available_roles_success,
+):
+    def _get(username, *args, **kwargs):
+        if username == "[email protected]":
+            raise Exception("TEST")
+        return get_users_response
+
+    cli_state.sdk.users.get_by_username.side_effect = _get
+    bulk_processor = mocker.patch(f"{_NAMESPACE}.run_bulk_process")
+    with runner.isolated_filesystem():
+        with open("test_bulk_add_roles.csv", "w") as csv:
+            csv.writelines(
+                ["username,role_name\n", f"{TEST_USERNAME},{TEST_ROLE_NAME}\n"]
+            )
+
+        runner.invoke(
+            cli,
+            ["users", "bulk", "add-roles", "test_bulk_add_roles.csv"],
+            obj=cli_state,
+        )
+    handler = bulk_processor.call_args[0][0]
+
+    handler(
+        username="[email protected]", role_name=TEST_ROLE_NAME,
+    )
+    handler(username="[email protected]", role_name=TEST_ROLE_NAME)
+    assert worker_stats.increment_total_errors.call_count == 1
+
+
+def test_bulk_remove_roles_uses_expected_arguments(runner, mocker, cli_state_with_user):
+    bulk_processor = mocker.patch(f"{_NAMESPACE}.run_bulk_process")
+    with runner.isolated_filesystem():
+        with open("test_bulk_remove_roles.csv", "w") as csv:
+            csv.writelines(
+                ["username,role_name\n", f"{TEST_USERNAME},{TEST_ROLE_NAME}\n"]
+            )
+        command = ["users", "bulk", "remove-roles", "test_bulk_remove_roles.csv"]
+        runner.invoke(
+            cli, command, obj=cli_state_with_user,
+        )
+    assert bulk_processor.call_args[0][1] == [
+        {
+            "username": TEST_USERNAME,
+            "role_name": TEST_ROLE_NAME,
+            "role removed": "False",
+        },
+    ]
+    bulk_processor.assert_called_once()
+
+
+def test_bulk_remove_roles_ignores_blank_lines(runner, mocker, cli_state):
+    bulk_processor = mocker.patch(f"{_NAMESPACE}.run_bulk_process")
+    with runner.isolated_filesystem():
+        with open("test_bulk_remove_roles.csv", "w") as csv:
+            csv.writelines(
+                ["username,role_name\n\n\n", f"{TEST_USERNAME},{TEST_ROLE_NAME}\n\n\n"]
+            )
+        runner.invoke(
+            cli,
+            ["users", "bulk", "remove-roles", "test_bulk_remove_roles.csv"],
+            obj=cli_state,
+        )
+    assert bulk_processor.call_args[0][1] == [
+        {
+            "username": TEST_USERNAME,
+            "role_name": TEST_ROLE_NAME,
+            "role removed": "False",
+        },
+    ]
+    bulk_processor.assert_called_once()
+
+
+def test_bulk_remove_roles_uses_handler_that_when_encounters_error_increments_total_errors(
+    runner,
+    mocker,
+    cli_state,
+    worker_stats,
+    get_users_response,
+    get_available_roles_success,
+):
+    def _get(username, *args, **kwargs):
+        if username == "[email protected]":
+            raise Exception("TEST")
+
+        return get_users_response
+
+    cli_state.sdk.users.get_by_username.side_effect = _get
+    bulk_processor = mocker.patch(f"{_NAMESPACE}.run_bulk_process")
+    with runner.isolated_filesystem():
+        with open("test_bulk_remove_roles.csv", "w") as csv:
+            csv.writelines(
+                ["username,role_name\n", f"{TEST_USERNAME},{TEST_ROLE_NAME}\n"]
+            )
+
+        runner.invoke(
+            cli,
+            ["users", "bulk", "remove-roles", "test_bulk_remove_roles.csv"],
+            obj=cli_state,
+        )
+    handler = bulk_processor.call_args[0][0]
+    handler(
+        username="[email protected]", role_name=TEST_ROLE_NAME,
+    )
+    handler(username="[email protected]", role_name=TEST_ROLE_NAME)
+    assert worker_stats.increment_total_errors.call_count == 1
+
+
 def test_orgs_list_calls_orgs_get_all_with_expected_params(runner, cli_state):
     runner.invoke(cli, ["users", "orgs", "list"], obj=cli_state)
     assert cli_state.sdk.orgs.get_all.call_count == 1