Add profile support

README.rst

@@ -95,6 +95,33 @@ you can try::
     TABLENAMES     Second
     TABLENAMES     Third
 
+Profiles
+--------
+
+The aws-shell supports AWS CLI profiles.  You have two options to use
+profiles.  First, you can provide a profile when you start the aws-shell::
+
+    $ aws-shell --profile prod
+    aws>
+
+When you do this all the server side completion as well as CLI commands
+you run will automatically use the ``prod`` profile.
+
+You can also change the current profile while you're in the aws-shell::
+
+    $ aws-shell
+    aws> .profile demo
+    Current shell profile changed to: demo
+
+You can also check what profile you've configured in the aws-shell using::
+
+    aws> .profile
+    Current shell profile: demo
+
+After changing your profile using the ``.profile`` dot command, all
+server side completion as well as CLI commands will automatically use
+the new profile you've configured.
+
 
 Features
 ========
@@ -218,6 +245,21 @@ variable before defaulting to ``notepad`` on Windows and
     aws> dynamodb list-tables
     aws> .edit
 
+Changing Profiles with .profile
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+You can change the current AWS CLI profile used by the aws-shell
+by using the ``.profile`` dot command.  If you run the ``.profile``
+command with no arguments, the currently configured shell profile
+will be printed.
+
+::
+
+    aws> .profile demo
+    Current shell profile changed to: demo
+    aws> .profile
+    Current shell profile: demo
+
 
 Executing Shell Commands
 ------------------------

awsshell/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import unicode_literals, print_function
 
 import json
+import argparse
 import threading
 
 from awsshell import shellcomplete
@@ -28,6 +29,11 @@ def load_index(filename):
 
 
 def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-p', '--profile', help='The profile name to use '
+                        'when starting the AWS Shell.')
+    args = parser.parse_args()
+
     indexer = completion.CompletionIndex()
     try:
         index_str = indexer.load_index(utils.AWSCLI_VERSION)
@@ -59,6 +65,8 @@ def main():
     model_completer = autocomplete.AWSCLIModelCompleter(index_data)
     completer = shellcomplete.AWSShellCompleter(model_completer)
     shell = app.create_aws_shell(completer, model_completer, doc_data)
+    if args.profile:
+        shell.profile = args.profile
     shell.run()
 
 

awsshell/app.py

@@ -82,9 +82,45 @@ def run(self, command, application):
             p.communicate()
 
 
+class ProfileHandler(object):
+    USAGE = (
+        '.profile           # Print the current profile\n'
+        '.profile <name>    # Change the current profile\n'
+    )
+
+    def __init__(self, output=sys.stdout, err=sys.stderr):
+        self._output = output
+        self._err = err
+
+    def run(self, command, application):
+        """Get or set the profile.
+
+        If .profile is called with no args, the current profile
+        is displayed.  If the .profile command is called with a
+        single arg, then the current profile for the application
+        will be set to the new value.
+        """
+        if len(command) == 1:
+            profile = application.profile
+            if profile is None:
+                self._output.write(
+                    "Current shell profile: no profile configured\n"
+                    "You can change profiles using: .profile profile-name\n")
+            else:
+                self._output.write("Current shell profile: %s\n" % profile)
+        elif len(command) == 2:
+            new_profile_name = command[1]
+            application.profile = new_profile_name
+            self._output.write("Current shell profile changed to: %s\n" %
+                               new_profile_name)
+        else:
+            self._err.write("Usage:\n%s\n" % self.USAGE)
+
+
 class DotCommandHandler(object):
     HANDLER_CLASSES = {
         'edit': EditHandler,
+        'profile': ProfileHandler,
     }
 
     def __init__(self, output=sys.stdout, err=sys.stderr):
@@ -162,6 +198,8 @@ def __init__(self, completer, model_completer, docs):
         self.refresh_cli = False
         self.key_manager = None
         self._dot_cmd = DotCommandHandler()
+        self._env = os.environ.copy()
+        self._profile = None
 
         # These attrs come from the config file.
         self.config_obj = None
@@ -233,7 +271,7 @@ def run(self):
                         initial_document=Document(self.current_docs,
                                                   cursor_position=0))
                     self.cli.request_redraw()
-                    p = subprocess.Popen(full_cmd, shell=True)
+                    p = subprocess.Popen(full_cmd, shell=True, env=self._env)
                     p.communicate()
 
     def stop_input_and_refresh_cli(self):
@@ -386,3 +424,27 @@ def create_cli_interface(self, display_completions_in_columns):
                                       display_completions_in_columns)
         cli = CommandLineInterface(application=app, eventloop=loop)
         return cli
+
+    @property
+    def profile(self):
+        return self._profile
+
+    @profile.setter
+    def profile(self, new_profile_name):
+        # There's only two things that need to know about new profile
+        # changes.
+        #
+        # First, the actual command runner.  If we want
+        # to use a different profile, it should ensure that the CLI
+        # commands that get run use the new profile (via the
+        # AWS_DEFAULT_PROFILE env var).
+        #
+        # Second, we also need to let the server side autocompleter know.
+        #
+        # Given this is easy to manage by hand, I don't think
+        # it's worth adding an event system or observers just yet.
+        # If this gets hard to manage, the complexity of those systems
+        # would be worth it.
+        self._env['AWS_DEFAULT_PROFILE'] = new_profile_name
+        self.completer.change_profile(new_profile_name)
+        self._profile = new_profile_name

awsshell/shellcomplete.py

@@ -11,7 +11,10 @@
 """
 import os
 import logging
+
+import botocore.session
 from prompt_toolkit.completion import Completer, Completion
+
 from awsshell import fuzzy
 
 
@@ -33,7 +36,6 @@ def __init__(self, completer, server_side_completer=None):
         self._server_side_completer = server_side_completer
 
     def _create_server_side_completer(self, session=None):
-        import botocore.session
         from awsshell.resource import index
         if session is None:
             session = botocore.session.Session()
@@ -48,6 +50,11 @@ def _create_server_side_completer(self, session=None):
         completer = index.ServerSideCompleter(client_creator, describer)
         return completer
 
+    def change_profile(self, profile_name):
+        """Change the profile used for server side completions."""
+        self._server_side_completer = self._create_server_side_completer(
+            session=botocore.session.Session(profile=profile_name))
+
     @property
     def completer(self):
         return self._completer

tests/unit/test_app.py

@@ -3,6 +3,7 @@
 
 
 from awsshell import app
+from awsshell import shellcomplete
 from awsshell import compat
 
 
@@ -38,7 +39,57 @@ def test_edit_handler():
     assert command_run[0] == 'my-editor'
 
 
+def test_profile_handler_prints_profile():
+    shell = mock.Mock(spec=app.AWSShell)
+    shell.profile = 'myprofile'
+    stdout = compat.StringIO()
+    handler = app.ProfileHandler(stdout)
+    handler.run(['.profile'], shell)
+    assert stdout.getvalue().strip() == 'Current shell profile: myprofile'
+
+
+def test_profile_handler_when_no_profile_configured():
+    shell = mock.Mock(spec=app.AWSShell)
+    shell.profile = None
+    stdout = compat.StringIO()
+    handler = app.ProfileHandler(stdout)
+    handler.run(['.profile'], shell)
+    assert stdout.getvalue() == (
+        'Current shell profile: no profile configured\n'
+        'You can change profiles using: .profile profile-name\n'
+    )
+
+
+def test_profile_command_changes_profile():
+    shell = mock.Mock(spec=app.AWSShell)
+    shell.profile = 'myprofile'
+    stdout = compat.StringIO()
+    handler = app.ProfileHandler(stdout)
+
+    handler.run(['.profile', 'newprofile'], shell)
+
+    assert shell.profile == 'newprofile'
+
+
+def test_profile_prints_error_on_bad_syntax():
+    stderr = compat.StringIO()
+    handler = app.ProfileHandler(None, stderr)
+    handler.run(['.profile', 'a', 'b', 'c'], None)
+
+    # We don't really care about the exact usage message here,
+    # we just want to ensure usage was written to stderr.
+    assert 'Usage' in stderr.getvalue()
+
+
 def test_prints_error_message_on_unknown_dot_command(errstream):
     handler = app.DotCommandHandler(err=errstream)
     handler.handle_cmd(".unknown foo bar", None)
     assert errstream.getvalue() == "Unknown dot command: .unknown\n"
+
+
+def test_delegates_to_complete_changing_profile():
+    completer = mock.Mock(spec=shellcomplete.AWSShellCompleter)
+    shell = app.AWSShell(completer, mock.Mock(), mock.Mock())
+    shell.profile = 'mynewprofile'
+    assert completer.change_profile.call_args == mock.call('mynewprofile')
+    assert shell.profile == 'mynewprofile'