feat: add new pre_hook_with_context method to Command (#35)

This feature adds a new lifecycle method to the Command base class:

    def pre_hook_with_context(self, context):
        pass

The value the context parameter comes from the new context parameter on AccountRunner.run:

    def run(self, cmd, accounts, key=lambda x: x, context=None):
        pass

This allows users of the awsrun library to pass a runtime context into commands.

As the primary user of the library, the awsrun CLI command uses this new feature to pass in references to the account loader and session provider plug-ins to a command. It allows command authors to lookup metadata associated with other accounts and obtaining sessions for other accounts while processing accounts. Here is an example:

    """Identify shared VPCs from our main account 999999999999 that have
    been shared with other accounts owned by a different Business Unit."""

    import io

    from awsrun.cli import Context
    from awsrun.runner import RegionalCommand


    class CLICommand(RegionalCommand):
        """Display VPCs configured in accounts."""

        def pre_hook_with_context(self, context: Context):
            acct_id = "999999999999"

            # Let's get the account object for our main account. Assume the
            # account loader plug-in used annotates that object with one
            # metadata value called "BU" representing the business unit.
            acct = context.account_loader.accounts(acct_ids=[acct_id])[0]

            # Let's get a session object for our main account so we can get
            # a list of VPCs from it.
            session = context.session_provider.session(acct_id)
            vpc_ids = []
            for region in self.regions:
                ec2 = session.resource("ec2", region_name=region)
                vpc_ids.extend(vpc.id for vpc in ec2.vpcs.all())

            # Save these for use later in regional_execute()
            self.important_acct = acct
            self.important_vpc_ids = vpc_ids

        def regional_execute(self, session, acct, region):
            out = io.StringIO()
            ec2 = session.resource("ec2", region_name=region)

            # Check each VPC in this current account to see if it is in the
            # list of VPCs from our main account. And, then check if the BU
            # for this current account is the same as the BU for our main
            # account. Report when they differ.
            for vpc in ec2.vpcs.all():
                if vpc.id in self.important_vpc_ids and acct.BU != self.important_acct.BU:
                    print(
                        "{acct}/{region}: Shared VPC {vpc.id} owned by other Business Unit!",
                        file=out,
                    )

            return out.getvalue()

The above example demonstrates how you can use these plug-ins in your
commands to inspect metadata of other accounts and to obtain sessions for
any account. This can be useful if you need to access another account while
processing the current account. For example, perhaps you want to accept a
VPC peering request after initiating it.

---------

Co-authored-by: Kazmier, Peter <peter.kazmier@fmr.com>

Author: azurenetworkguy

src/awsrun/cli.py

@@ -1078,7 +1078,12 @@ def _cli(csp):
     # awsrun can be used without the need of the CLI. One only needs a list of
     # accounts, an awsrun.runner.Command, and an awsrun.session.SessionProvider.
     runner = AccountRunner(session_provider, args.threads)
-    elapsed = runner.run(command, accounts, key=account_loader.acct_id)
+    elapsed = runner.run(
+        command,
+        accounts,
+        key=account_loader.acct_id,
+        context=Context(session_provider, account_loader, accounts),
+    )
 
     # Show a quick summary on how long the command took to run.
     pluralize = "s" if len(accounts) != 1 else ""
@@ -1179,5 +1184,19 @@ def default_session_provider(self):
         return "awsrun.plugins.creds." + self.name.lower() + ".Default"
 
 
+class Context:
+    """Used when `Command.pre_hook_with_context` is invoked.
+
+    The `Context` provides a `Command` access to awsrun information/context prior to any
+    accounts being processed by `Runner`.
+
+    """
+
+    def __init__(self, session_provider, account_loader, accounts):
+        self.session_provider = session_provider
+        self.account_loader = account_loader
+        self.accounts = accounts
+
+
 if __name__ == "__main__":
     main()

src/awsrun/commands/__init__.py

@@ -487,6 +487,72 @@ def post_hook(self):
 post-hook to print a summary of the data we collected during the processing of
 accounts.
 
+### Pre-Hook With Context
+
+We can use `awsrun.runner.Command.pre_hook_with_context` to obtain access to
+the CLI's `awsrun.cli.Context` object, which contains a reference to the
+`awsrun.acctload.AccountLoader`, the `awsrun.session.SessionProvider`, and
+the list of accounts being processed by the CLI. This allows command authors
+to use those plug-ins as part of their command. Let's use an example to
+illustrate. Assume that the account loader plug-in annotates the accounts
+with one metadata field called "BU" that represents the business unit that
+owns the account.
+
+    \"\"\"Identify shared VPCs from our main account 999999999999 that have
+    been shared with other accounts owned by a different Business Unit.\"\"\"
+
+    import io
+
+    from awsrun.cli import Context
+    from awsrun.runner import RegionalCommand
+
+
+    class CLICommand(RegionalCommand):
+        \"\"\"Display VPCs configured in accounts.\"\"\"
+
+        def pre_hook_with_context(self, context: Context):
+            acct_id = "999999999999"
+
+            # Let's get the account object for our main account. Assume the
+            # account loader plug-in used annotates that object with one
+            # metadata value called "BU" representing the business unit.
+            acct = context.account_loader.accounts(acct_ids=[acct_id])[0]
+
+            # Let's get a session object for our main account so we can get
+            # a list of VPCs from it.
+            session = context.session_provider.session(acct_id)
+            vpc_ids = []
+            for region in self.regions:
+                ec2 = session.resource("ec2", region_name=region)
+                vpc_ids.extend(vpc.id for vpc in ec2.vpcs.all())
+
+            # Save these for use later in regional_execute()
+            self.important_acct = acct
+            self.important_vpc_ids = vpc_ids
+
+        def regional_execute(self, session, acct, region):
+            out = io.StringIO()
+            ec2 = session.resource("ec2", region_name=region)
+
+            # Check each VPC in this current account to see if it is in the
+            # list of VPCs from our main account. And, then check if the BU
+            # for this current account is the same as the BU for our main
+            # account. Report when they differ.
+            for vpc in ec2.vpcs.all():
+                if vpc.id in self.important_vpc_ids and acct.BU != self.important_acct.BU:
+                    print(
+                        "{acct}/{region}: Shared VPC {vpc.id} owned by other Business Unit!",
+                        file=out,
+                    )
+
+            return out.getvalue()
+
+The above example demonstrates how you can use these plug-ins in your
+commands to inspect metadata of other accounts and to obtain sessions for
+any account. This can be useful if you need to access another account while
+processing the current account. For example, perhaps you want to accept a
+VPC peering request after initiating it.
+
 You should now have a good understanding of the do's and don'ts to keep in mind
 when authoring your own commands. In the next section, we will discuss how to
 install your commands.

src/awsrun/runner.py

@@ -97,6 +97,36 @@ def regional_execute(self, session, acct, region):
     account_runner = AccountRunner(session_provider)
     account_runner.run(cmd, ['111222333444', '222333444111'])
 
+## Context With Pre-Hook
+
+In some cases users may wish to pass additional context to the `Command` that
+can be used during execution. This can be done via the `context` parameter of
+`AccountRunner.run`. That value is passed to `Command.pre_hook_with_context`,
+which command authors can then use. For example, to pass a trivial prefix to be
+used as part of the output:
+
+    from awsrun.runner import AccountRunner, RegionalCommand
+    from awsrun.session import CredsViaProfile
+
+    class VpcInfoCommand(RegionalCommand):
+        # Save our context, a simple string in this case, to self.prefix.
+        def pre_hook_with_context(self, context):
+            self.prefix = context
+
+        def regional_execute(self, session, acct, region):
+            ec2 = session.resource('ec2', region_name=region)
+            vpc_ids = ', '.join(vpc.id for vpc in ec2.vpcs.all())
+            # Prefix the output with the prefix passed via context
+            return f'{self.prefix}/{acct}/{region}: {vpc_ids}\\n'
+
+    cmd = VpcInfoCommand(['us-east-1', 'us-west-2'])
+    session_provider = CredsViaProfile()
+    account_runner = AccountRunner(session_provider)
+    # We pass a simple string prefix as the context in this trivial example
+    account_runner.run(cmd, ['111222333444', '222333444111'], context="Group-A")
+    account_runner.run(cmd, ['333444555666'], context="Group-B")
+    account_runner.run(cmd, ['444555666777', '555666777888'], context="Group-C")
+
 ## Collecting Results
 
 The final example demonstrates how to collect the results from all of the
@@ -173,6 +203,7 @@ def regional_collect_results(self, acct, region, get_result):
     account_runner = AccountRunner(session_provider)
     account_runner.run(cmd, ['111222333444', '222333444111'])
 """
+
 import functools
 import logging
 import sys
@@ -307,12 +338,29 @@ def from_cli(cls, parser, argv, cfg):
         parser.parse_args(argv)
         return cls()
 
-    def pre_hook(self):
-        """Invoked by `AccountRunner.run` before any processing starts.
+    def pre_hook_with_context(self, context):
+        """Invoked by `AccountRunner.run` before any account processing starts.
 
         This method is invoked only once per invocation of `AccountRunner.run`.
-        It is not executed before each account is processed, but rather once
-        before any accounts are processed.
+        The `context` parameter is the opaque object passed as the context
+        parameter to `AccountRunner.run`. It is intended to provide access to
+        additional runtime context information for the command to leverage.
+        The method is not executed before each account is processed, but
+        rather once before any accounts are processed.
+
+        To provide backwards compatibility, the default implementation invokes
+        `Command.pre_hook`.
+        """
+        self.pre_hook()
+
+    def pre_hook(self):
+        """Invoked in the default implementation of `pre_hook_with_context`
+        before any account processing starts.
+
+        This method is invoked in the event a command does not override the
+        default implementation of `pre_hook_with_context`. The method is not
+        executed before each account is processed, but rather once before
+        any accounts are processed.
 
         The default implementation does nothing.
         """
@@ -844,18 +892,23 @@ def __init__(self, session_provider, max_workers=10):
         self.session_provider = session_provider
         self.max_workers = max_workers
 
-    def run(self, cmd, accounts, key=lambda x: x):
+    def run(self, cmd, accounts, key=lambda x: x, context=None):
         """Execute a command concurrently on the specified accounts.
 
         This method will block until all accounts have been processed. The
         return value is the number of seconds it took to process the accounts.
 
-        The `cmd` must be a subclass of `Command`. The runner will invoke the
-        `Command.pre_hook` once before it starts processing any accounts, then
-        accounts are processed concurrently and `Command.execute` is invoked by
-        a worker for each account. As each execute method returns, the main
-        thread will invoke `Command.collect_results`, which ensures results are
-        collected sequentially. Finally, after all accounts have been processed,
+        The `cmd` must be a subclass of `Command`. The runner will invoke
+        `Command.pre_hook_with_context` once before it starts processing any
+        accounts passing it the `context` argument, which is simply passed
+        through as an opaque object. This can be used to pass a runtime context
+        to a `Command`.
+
+        After the pre-hook has been invoked, accounts are then processed
+        concurrently and `Command.execute` is invoked by a worker for each
+        account. As each execute method returns, the main thread will invoke
+        `Command.collect_results`, which ensures results are collected
+        sequentially. Finally, after all accounts have been processed,
         `Command.post_hook` is called.
 
         The specified list of `accounts` can be of any type as long as the
@@ -887,6 +940,12 @@ def run(self, cmd, accounts, key=lambda x: x):
         exception, then an `InvalidAccountIDError` is raised in the worker
         thread processing the account, which will then propagate to the
         `Command.collect_results`.
+
+        The optional `context` parameter can be any value. It is passed
+        as-is to `Command.pre_hook_with_context`. It provides a mechanism
+        for the caller to pass a runtime context to the command being
+        executed. See the [Context With Pre-Hook](#context-with-pre-hook)
+        section of the use guide for an example.
         """
         # This will ensure v1 users of awsrun aren't mixing v1 Command's with
         # the v2 framework.
@@ -899,7 +958,7 @@ def run(self, cmd, accounts, key=lambda x: x):
         key = _valid_key_fn(key)
 
         start = time.time()
-        cmd.pre_hook()
+        cmd.pre_hook_with_context(context)
 
         with ThreadPoolExecutor(max_workers=self.max_workers) as pool:
             # The worker task processes a single account. The worker task takes