Add limit kwarg to repo functions (#261)

* Add `limit` kwarg to repo functions

Allow users to specify a `limit` that will limit how many roles are
repoed at one time. This is an optional kwarg that defaults to a value
of -1, which results in not limiting the repo action.

* fix hook calls in command test

* change unlimited value to < 0

* update defaults to reflect unlimited change

README.md

@@ -140,7 +140,8 @@ Repokid is extensible via hooks that are called before, during, and after variou
 
 | Hook name | Context |
 |-----------|---------|
-| `AFTER_REPO` | role |
+| `AFTER_REPO` | role, errors |
+| `AFTER_REPO_ROLES` | roles, errors |
 | `BEFORE_REPO_ROLES` | account_number, roles |
 | `AFTER_SCHEDULE_REPO` | roles |
 | `DURING_REPOABLE_CALCULATION` | account_number, role_name, potentially_repoable_permissions, minimum_age |

repokid/commands/repo.py

@@ -193,7 +193,7 @@ def _repo_role(
         dynamo_table, role.role_id, {"RepoScheduled": 0, "ScheduledPerms": []}
     )
 
-    repokid.hooks.call_hooks(hooks, "AFTER_REPO", {"role": role})
+    repokid.hooks.call_hooks(hooks, "AFTER_REPO", {"role": role, "errors": errors})
 
     if not errors:
         # repos will stay scheduled until they are successful
@@ -367,7 +367,7 @@ def _rollback_role(
 
 
 def _repo_all_roles(
-    account_number, dynamo_table, config, hooks, commit=False, scheduled=True
+    account_number, dynamo_table, config, hooks, commit=False, scheduled=True, limit=-1
 ):
     """
     Repo all scheduled or eligible roles in an account.  Collect any errors and display them at the end.
@@ -378,6 +378,7 @@ def _repo_all_roles(
         config
         commit (bool): actually make the changes
         scheduled (bool): if True only repo the scheduled roles, if False repo all the (eligible) roles
+        limit (int): limit number of roles to be repoed per run (< 0 is unlimited)
 
     Returns:
         None
@@ -420,7 +421,11 @@ def _repo_all_roles(
         hooks, "BEFORE_REPO_ROLES", {"account_number": account_number, "roles": roles}
     )
 
+    count = 0
+    repoed = Roles([])
     for role in roles:
+        if limit >= 0 and count == limit:
+            break
         error = _repo_role(
             account_number,
             role.role_name,
@@ -432,13 +437,19 @@ def _repo_all_roles(
         )
         if error:
             errors.append(error)
+        repoed.append(role)
+        count += 1
 
     if errors:
-        LOGGER.error(
-            "Error(s) during repo: \n{} (account: {})".format(errors, account_number)
-        )
+        LOGGER.error(f"Error(s) during repo: \n{errors} (account: {account_number})")
     else:
-        LOGGER.info("Successfully repoed roles in account {}".format(account_number))
+        LOGGER.info(f"Successfully repoed {count} roles in account {account_number}")
+
+    repokid.hooks.call_hooks(
+        hooks,
+        "AFTER_REPO_ROLES",
+        {"account_number": account_number, "roles": repoed, "errors": errors},
+    )
 
 
 def _repo_stats(output_file, dynamo_table, account_number=None):

repokid/lib/__init__.py

@@ -199,7 +199,9 @@ def schedule_repo(account_number: str):
     return _schedule_repo(account_number, dynamo_table, CONFIG, hooks)
 
 
-def repo_all_roles(account_number: str, commit: bool = False, update: bool = True):
+def repo_all_roles(
+    account_number: str, commit: bool = False, update: bool = True, limit: int = -1
+):
     """
     Convenience wrapper for repo_roles() with scheduled=False.
 
@@ -209,15 +211,18 @@ def repo_all_roles(account_number: str, commit: bool = False, update: bool = Tru
         account_number (string): The current account number Repokid is being run against
         commit (bool): actually make the changes
         update (bool): if True run update_role_cache before repoing
+        limit (int): limit number of roles to be repoed per run (< 0 is unlimited)
 
     Returns:
         None
     """
-    return repo_roles(account_number, commit=commit, scheduled=False, update=update)
+    return repo_roles(
+        account_number, commit=commit, scheduled=False, update=update, limit=limit
+    )
 
 
 def repo_scheduled_roles(
-    account_number: str, commit: bool = False, update: bool = True
+    account_number: str, commit: bool = False, update: bool = True, limit: int = -1
 ):
     """
     Convenience wrapper for repo_roles() with scheduled=True.
@@ -228,18 +233,22 @@ def repo_scheduled_roles(
         account_number (string): The current account number Repokid is being run against
         commit (bool): actually make the changes
         update (bool): if True run update_role_cache before repoing
+        limit (int): limit number of roles to be repoed per run (< 0 is unlimited)
 
     Returns:
         None
     """
-    return repo_roles(account_number, commit=commit, scheduled=True, update=update)
+    return repo_roles(
+        account_number, commit=commit, scheduled=True, update=update, limit=limit
+    )
 
 
 def repo_roles(
     account_number: str,
     commit: bool = False,
     scheduled: bool = False,
     update: bool = True,
+    limit: int = -1,
 ):
     """
     Library wrapper to repo all scheduled or eligible roles in an account. Collect any errors and display them at the
@@ -252,14 +261,21 @@ def repo_roles(
         commit (bool): actually make the changes
         scheduled (bool): if True only repo the scheduled roles, if False repo all the (eligible) roles
         update (bool): if True run update_role_cache before repoing
+        limit (int): limit number of roles to be repoed per run (< 0 is unlimited)
 
     Returns:
         None
     """
     if update:
         _update_role_cache(account_number, dynamo_table, CONFIG, hooks)
     return _repo_all_roles(
-        account_number, dynamo_table, CONFIG, hooks, commit=commit, scheduled=scheduled
+        account_number,
+        dynamo_table,
+        CONFIG,
+        hooks,
+        commit=commit,
+        scheduled=scheduled,
+        limit=limit,
     )
 
 

repokid/tests/test_commands.py

@@ -536,11 +536,21 @@ def test_repo_all_roles(
                 "BEFORE_REPO_ROLES",
                 {"account_number": None, "roles": roles_items},
             ),
+            call(
+                hooks,
+                "AFTER_REPO_ROLES",
+                {"account_number": None, "roles": roles_items, "errors": []},
+            ),
             call(
                 hooks,
                 "BEFORE_REPO_ROLES",
                 {"account_number": None, "roles": [roles_items[2]]},
             ),
+            call(
+                hooks,
+                "AFTER_REPO_ROLES",
+                {"account_number": None, "roles": [roles_items[2]], "errors": []},
+            ),
         ]
 
     @patch("repokid.commands.schedule.find_role_in_cache")