From 94b2cbfe5272b8795110511a3f2fce01086c5494 Mon Sep 17 00:00:00 2001 From: James Saryerwinnie Date: Tue, 29 Dec 2015 16:35:11 -0800 Subject: [PATCH 1/4] Add --profile argument --- awsshell/__init__.py | 8 ++++++++ awsshell/app.py | 25 ++++++++++++++++++++++++- awsshell/shellcomplete.py | 9 ++++++++- tests/unit/test_app.py | 8 ++++++++ 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/awsshell/__init__.py b/awsshell/__init__.py index 5e7ddf4..ca5da0b 100644 --- a/awsshell/__init__.py +++ b/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.on_profile_change(args.profile) shell.run() diff --git a/awsshell/app.py b/awsshell/app.py index 694ebf8..0631833 100644 --- a/awsshell/app.py +++ b/awsshell/app.py @@ -162,6 +162,7 @@ def __init__(self, completer, model_completer, docs): self.refresh_cli = False self.key_manager = None self._dot_cmd = DotCommandHandler() + self._env = os.environ.copy() # These attrs come from the config file. self.config_obj = None @@ -233,7 +234,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 +387,25 @@ def create_cli_interface(self, display_completions_in_columns): display_completions_in_columns) cli = CommandLineInterface(application=app, eventloop=loop) return cli + + def on_profile_change(self, new_profile_name): + """Handler invoked when the user requests to use a new profile. + + This will inform the necessary objects to use the new profile. + """ + # 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) diff --git a/awsshell/shellcomplete.py b/awsshell/shellcomplete.py index 34eab75..8488e53 100644 --- a/awsshell/shellcomplete.py +++ b/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 diff --git a/tests/unit/test_app.py b/tests/unit/test_app.py index 9702545..42e95a0 100644 --- a/tests/unit/test_app.py +++ b/tests/unit/test_app.py @@ -3,6 +3,7 @@ from awsshell import app +from awsshell import shellcomplete from awsshell import compat @@ -42,3 +43,10 @@ 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.on_profile_change('mynewprofile') + assert completer.change_profile.call_args == mock.call('mynewprofile') From fd7400d7442e7bd918fc27fa500c29b4b8afa5a5 Mon Sep 17 00:00:00 2001 From: James Saryerwinnie Date: Tue, 29 Dec 2015 16:52:30 -0800 Subject: [PATCH 2/4] Add .profile dot command Closes #9. --- awsshell/__init__.py | 2 +- awsshell/app.py | 47 ++++++++++++++++++++++++++++++++++++++---- tests/unit/test_app.py | 45 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 88 insertions(+), 6 deletions(-) diff --git a/awsshell/__init__.py b/awsshell/__init__.py index ca5da0b..f668344 100644 --- a/awsshell/__init__.py +++ b/awsshell/__init__.py @@ -66,7 +66,7 @@ def main(): completer = shellcomplete.AWSShellCompleter(model_completer) shell = app.create_aws_shell(completer, model_completer, doc_data) if args.profile: - shell.on_profile_change(args.profile) + shell.profile = args.profile shell.run() diff --git a/awsshell/app.py b/awsshell/app.py index 0631833..680a43c 100644 --- a/awsshell/app.py +++ b/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 # 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): @@ -163,6 +199,7 @@ def __init__(self, completer, model_completer, docs): 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 @@ -388,11 +425,12 @@ def create_cli_interface(self, display_completions_in_columns): cli = CommandLineInterface(application=app, eventloop=loop) return cli - def on_profile_change(self, new_profile_name): - """Handler invoked when the user requests to use a new profile. + @property + def profile(self): + return self._profile - This will inform the necessary objects to use the new profile. - """ + @profile.setter + def profile(self, new_profile_name): # There's only two things that need to know about new profile # changes. # @@ -409,3 +447,4 @@ def on_profile_change(self, new_profile_name): # would be worth it. self._env['AWS_DEFAULT_PROFILE'] = new_profile_name self.completer.change_profile(new_profile_name) + self._profile = new_profile_name diff --git a/tests/unit/test_app.py b/tests/unit/test_app.py index 42e95a0..55fa697 100644 --- a/tests/unit/test_app.py +++ b/tests/unit/test_app.py @@ -39,6 +39,48 @@ 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) @@ -48,5 +90,6 @@ def test_prints_error_message_on_unknown_dot_command(errstream): def test_delegates_to_complete_changing_profile(): completer = mock.Mock(spec=shellcomplete.AWSShellCompleter) shell = app.AWSShell(completer, mock.Mock(), mock.Mock()) - shell.on_profile_change('mynewprofile') + shell.profile = 'mynewprofile' assert completer.change_profile.call_args == mock.call('mynewprofile') + assert shell.profile == 'mynewprofile' From dab0d38f8b6c2561d250b89bdfda3e2ca970e5ae Mon Sep 17 00:00:00 2001 From: James Saryerwinnie Date: Tue, 29 Dec 2015 17:02:41 -0800 Subject: [PATCH 3/4] Update README with new profile functionality --- README.rst | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/README.rst b/README.rst index 223e1f1..63e1926 100644 --- a/README.rst +++ b/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 ======== From a4f68de769742a966b49a38fa9551918a2a355f2 Mon Sep 17 00:00:00 2001 From: James Saryerwinnie Date: Tue, 29 Dec 2015 17:06:13 -0800 Subject: [PATCH 4/4] Also add .profile under docs for dot commands --- README.rst | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/README.rst b/README.rst index 63e1926..ba4f52b 100644 --- a/README.rst +++ b/README.rst @@ -245,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 ------------------------