Merge pull request #88 from sj26/query

Allow JMES path queries for JSON logs

awslogs/bin.py

@@ -128,6 +128,12 @@ def add_date_range_arguments(parser):
                             dest='color_enabled',
                             help="Do not color output")
 
+    get_parser.add_argument("-q",
+                            "--query",
+                            action="store",
+                            dest="query",
+                            help="JMESPath query to use in filtering the response data")
+
     # groups
     groups_parser = subparsers.add_parser('groups', description='List groups')
     groups_parser.set_defaults(func="list_groups")

awslogs/core.py

@@ -6,7 +6,9 @@
 from collections import deque
 
 import boto3
-from botocore.compat import total_seconds
+from botocore.compat import json, six, total_seconds
+
+import jmespath
 
 from termcolor import colored
 from dateutil.parser import parse
@@ -47,6 +49,9 @@ def __init__(self, **kwargs):
             'output_ingestion_time_enabled')
         self.start = self.parse_datetime(kwargs.get('start'))
         self.end = self.parse_datetime(kwargs.get('end'))
+        self.query = kwargs.get('query')
+        if self.query is not None:
+            self.query_expression = jmespath.compile(self.query)
 
         self.client = boto3.client(
             'logs',
@@ -167,7 +172,15 @@ def consumer():
                             'blue'
                         )
                     )
-                output.append(event['message'].rstrip())
+
+                message = event['message']
+                if self.query is not None and message[0] == '{':
+                    parsed = json.loads(event['message'])
+                    message = self.query_expression.search(parsed)
+                    if not isinstance(message, six.string_types):
+                        message = json.dumps(message)
+                output.append(message)
+
                 print(' '.join(output))
                 sys.stdout.flush()
         try:

setup.py

@@ -5,6 +5,7 @@
 
 install_requires = [
     'boto3>=1.2.1',
+    'jmespath>=0.7.1,<1.0.0',
     'termcolor>=1.1.0',
     'python-dateutil>=2.4.0'
 ]

tests/test_it.py

@@ -137,6 +137,43 @@ def paginator(value):
         client.get_paginator.side_effect = paginator
         client.filter_log_events.side_effect = logs
 
+    def set_json_logs(self, botoclient):
+        client = Mock()
+        botoclient.return_value = client
+
+        event_keys = ["eventId", "timestamp", "ingestionTime",
+                      "message", "logStreamName"]
+        logs = [
+            {'events': mapkeys(event_keys,
+                               [[1, 0, 5000, '{"foo": "bar"}', "DDD"],
+                                [2, 0, 5000, '{"foo": {"bar": "baz"}}', "EEE"],
+                                [3, 0, 5006, "Hello 3", "DDD"]]),
+             'nextToken': 'token'},
+            {'events': []}
+        ]
+
+        groups = [
+            {'logGroups': [{'logGroupName': 'AAA'},
+                           {'logGroupName': 'BBB'},
+                           {'logGroupName': 'CCC'}]},
+        ]
+
+        streams = [
+            {'logStreams': [self._stream('DDD'),
+                            self._stream('EEE')]}
+        ]
+
+        def paginator(value):
+            mock = Mock()
+            mock.paginate.return_value = {
+                'describe_log_groups': groups,
+                'describe_log_streams': streams
+            }.get(value)
+            return mock
+
+        client.get_paginator.side_effect = paginator
+        client.filter_log_events.side_effect = logs
+
     @patch('boto3.client')
     def test_get_groups(self, botoclient):
         client = Mock()
@@ -258,6 +295,19 @@ def test_main_get_with_color(self, mock_stdout, botoclient):
 
         assert output == expected
 
+    @patch('boto3.client')
+    @patch('sys.stdout', new_callable=StringIO)
+    def test_main_get_query(self, mock_stdout, botoclient):
+        self.set_json_logs(botoclient)
+        main("awslogs get AAA DDD --query foo".split())
+        output = mock_stdout.getvalue()
+        expected = ("\x1b[32mAAA\x1b[0m \x1b[36mDDD\x1b[0m bar\n"
+                    "\x1b[32mAAA\x1b[0m \x1b[36mEEE\x1b[0m {\"bar\": \"baz\"}\n"
+                    "\x1b[32mAAA\x1b[0m \x1b[36mDDD\x1b[0m Hello 3\n"
+                    )
+
+        assert output == expected
+
     @patch('boto3.client')
     @patch('sys.stdout', new_callable=StringIO)
     def test_get_nogroup(self, mock_stdout, botoclient):