Proof of concept of server side completion

Author: jamesls

awsshell/makeindex.py

@@ -39,6 +39,9 @@ def index_command(index_dict, help_command):
             'type_name': arg_obj.cli_type_name,
             'minidoc': '',
             'example': '',
+            # The name used in the API call/botocore,
+            # typically CamelCased.
+            'api_name': getattr(arg_obj, '_serialized_name', '')
         }
         if arg_obj.documentation:
             metadata['minidoc'] = remove_html(arg_obj.documentation.split('\n')[0])
@@ -132,5 +135,9 @@ def main():
     args = parser.parse_args()
     if args.output is None:
         args.output = determine_index_filename()
-    #write_index(args.output)
-    write_doc_index()
+    write_index(args.output)
+    #write_doc_index()
+
+
+if __name__ == '__main__':
+    main()

awsshell/resource/index.py

@@ -1,7 +1,27 @@
-"""Index and retrive information from the resource JSON."""
-import jmespath
+"""Index and retrive information from the resource JSON.
+
+The classes are organized as follows:
+
+* ResourceIndexBuilder - Takes a boto3 resource and converts into the
+  index format we need to do server side completions.
+* CompleterQuery - Takes the index from ResourceIndexBuilder and looks
+  up how to perform the autocompletion.  Note that this class does
+  *not* actually do the autocompletion.  It merely tells you how
+  you _would_ do the autocompletion if you made the appropriate
+  service calls.
+* ServerSideCompleter - The thing that does the actual autocompletion.
+  You tell it the command/operation/param you're on, and it will
+  return a list of completions for you.
+
+"""
+import logging
 from collections import namedtuple
 
+import jmespath
+from botocore import xform_name
+
+LOG = logging.getLogger(__name__)
+
 # service - The name of the AWS service
 # operation - The name of the AWS operation
 # params - A dict of params to send in the request (not implemented yet)
@@ -90,7 +110,9 @@ def describe_autocomplete(self, service, operation, param):
 
         """
         service_index = self._index[service]
+        LOG.debug(service_index)
         if param not in service_index.get('operations', {}).get(operation, {}):
+            LOG.debug("param not in index: %s", param)
             return None
         p = service_index['operations'][operation][param]
         resource_name = p['resourceName']
@@ -103,6 +125,69 @@ def describe_autocomplete(self, service, operation, param):
                                 params={}, path=path)
 
 
+class ServerSideCompleter(object):
+    def __init__(self, session, builder):
+        # session is a boto3 session.
+        # It is a public attribute as it is intended to be
+        # changed if the profile changes.
+        self.session = session
+        self._loader = session._loader
+        self._builder = builder
+        self._client_cache = {}
+        self._completer_cache = {}
+
+    def _get_completer_for_service(self, service_name, resource_model):
+        if service_name not in self._completer_cache:
+            index = self._builder.build_index(resource_model)
+            cq = CompleterQuery({service_name: index})
+            self._completer_cache[service_name] = cq
+        return self._completer_cache[service_name]
+
+    def _get_client(self, service_name):
+        if service_name in self._client_cache:
+            return self._client_cache[service_name]
+        client = self.session.client(service_name)
+        self._client_cache[service_name] = client
+        return client
+
+    def autocomplete(self, service, operation, param):
+        # Example call:
+        # service='ec2', operation='terminate-instances',
+        # param='--instance-ids'.
+        # We need to convert this to botocore syntax.
+        # First try to load the resource model.
+        LOG.debug("Called with: %s, %s, %s", service, operation, param)
+        try:
+            resource_model = self._loader.load_service_model(
+                service, 'resources-1')
+        except Exception as e:
+            # No resource == no server side completion.
+            return
+        # Now convert operation to the name used by botocore.
+        client = self._get_client(service)
+        api_operation_name = client.meta.method_to_api_mapping.get(
+            operation.replace('-', '_'))
+        if api_operation_name is None:
+            return
+        # Now we need to convert the param name to the
+        # casing used by the API.
+        completer = self._get_completer_for_service(service,
+                                                    resource_model)
+        result = completer.describe_autocomplete(
+            service, api_operation_name, param)
+        if result is None:
+            return
+        # DEBUG:awsshell.resource.index:RESULTS:
+            # ServerCompletion(service=u'ec2', operation=u'DescribeInstances',
+            # params={}, path=u'Reservations[].Instances[].InstanceId')
+        try:
+            response = getattr(client, xform_name(result.operation, '_'))()
+        except Exception as e:
+            return
+        results = jmespath.search(result.path, response)
+        return results
+
+
 def main():
     import sys
     import json

awsshell/shellcomplete.py

@@ -9,9 +9,14 @@
 logic, see awsshell.autocomplete.
 
 """
+import logging
 from prompt_toolkit.completion import Completer, Completion
 
 
+logging.basicConfig(filename='/tmp/completions', level=logging.DEBUG)
+LOG = logging.getLogger(__name__)
+
+
 class AWSShellCompleter(Completer):
     """Completer class for the aws-shell.
 
@@ -20,8 +25,19 @@ class AWSShellCompleter(Completer):
     low level, and can be reused in contexts other than the
     aws shell.
     """
-    def __init__(self, completer):
+    def __init__(self, completer, server_side_completer=None):
         self._completer = completer
+        if server_side_completer is None:
+            server_side_completer = self._create_server_side_completer()
+        self._server_side_completer = server_side_completer
+
+    def _create_server_side_completer(self):
+        import boto3.session
+        from awsshell.resource import index
+        session = boto3.session.Session()
+        builder = index.ResourceIndexBuilder()
+        completer = index.ServerSideCompleter(session, builder)
+        return completer
 
     @property
     def completer(self):
@@ -39,14 +55,17 @@ def last_option(self):
     def current_command(self):
         return u' '.join(self._completer.cmd_path)
 
-    def get_completions(self, document, complete_event):
-        text_before_cursor = document.text_before_cursor
+    def _convert_to_prompt_completions(self, low_level_completions,
+                                       text_before_cursor):
+        # Converts the low level completions from the model autocompleter
+        # and converts them to Completion() objects used by
+        # prompt_toolkit.  We also try to enhance the metadata of the
+        # completion by including docs and marking required fields.
+        arg_meta = self._completer.arg_metadata
         word_before_cursor = ''
         if text_before_cursor.strip():
             word_before_cursor = text_before_cursor.strip().split()[-1]
-        completions = self._completer.autocomplete(text_before_cursor)
-        arg_meta = self._completer.arg_metadata
-        for completion in completions:
+        for completion in low_level_completions:
             # Go through the completions and add inline docs and
             # mark which options are required.
             if completion.startswith('--') and completion in arg_meta:
@@ -69,4 +88,41 @@ def get_completions(self, document, complete_event):
                              display=display_text, display_meta=display_meta)
 
 
-
+    def get_completions(self, document, complete_event):
+        text_before_cursor = document.text_before_cursor
+        completions = self._completer.autocomplete(text_before_cursor)
+        prompt_completions = list(self._convert_to_prompt_completions(
+            completions, text_before_cursor))
+        LOG.debug("num_completions: %s, text: %s", len(prompt_completions),
+                  text_before_cursor)
+        if (not prompt_completions and self._completer.last_option and
+                len(self._completer.cmd_path) == 3):
+            # If we couldn't complete anything from the JSON model
+            # completer and we're on a cli option (e.g --foo), we
+            # can ask the server side completer if it knows anything
+            # about this resource.
+            LOG.debug("No local autocompletions found, trying "
+                      "server side completion.")
+            command = self._completer.cmd_path
+            service = command[1]
+            if service == 's3api':
+                # TODO: we need a more generic way to capture renames
+                # of commands.  This currently lives in the CLI
+                # customization code.
+                service = 's3'
+            operation = command[2]
+            param = self._completer.arg_metadata.get(
+                self._completer.last_option, {}).get('api_name')
+            if param is not None:
+                results = self._server_side_completer.autocomplete(
+                    service, operation, param)
+                if results is not None:
+                    for result in results:
+                        # Insert at the end
+                        location = 0
+                        yield Completion(result, location,
+                                         display=result,
+                                         display_meta='')
+        else:
+            for c in prompt_completions:
+                yield c