diff --git a/README.md b/README.md index 34f349b46..3962ab23b 100644 --- a/README.md +++ b/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 | diff --git a/repokid/commands/repo.py b/repokid/commands/repo.py index 62aac9d18..f0b6f4d3a 100644 --- a/repokid/commands/repo.py +++ b/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): diff --git a/repokid/lib/__init__.py b/repokid/lib/__init__.py index 513bbb867..1bf3791b1 100644 --- a/repokid/lib/__init__.py +++ b/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,11 +233,14 @@ 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( @@ -240,6 +248,7 @@ def repo_roles( 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,6 +261,7 @@ 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 @@ -259,7 +269,13 @@ def repo_roles( 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, ) diff --git a/repokid/tests/test_commands.py b/repokid/tests/test_commands.py index 95db78565..1ed6cb477 100644 --- a/repokid/tests/test_commands.py +++ b/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")