From d0054012d02f9207d698e10614aa303f3c05d72a Mon Sep 17 00:00:00 2001 From: Davide Depau Date: Fri, 13 Sep 2024 15:11:01 +0200 Subject: [PATCH 1/2] Add support for jq queries --- awslogs/bin.py | 20 ++++++++++++++++++ awslogs/core.py | 55 ++++++++++++++++++++++++++++++++++-------------- setup.py | 1 + tests/test_it.py | 16 ++++++++++++++ 4 files changed, 76 insertions(+), 16 deletions(-) diff --git a/awslogs/bin.py b/awslogs/bin.py index c88c688..d4058d1 100644 --- a/awslogs/bin.py +++ b/awslogs/bin.py @@ -178,6 +178,22 @@ def add_date_range_arguments(parser, default_start="5m"): help="JMESPath query to use in filtering the response data", ) + get_parser.add_argument( + "-j", + "--jq", + action="store", + dest="jq", + help="jq query to use in filtering the response data", + ) + + get_parser.add_argument( + "--jq-all", + action="store_true", + dest="jq_all", + default=False, + help="Use all results from the jq query (by default only the first result is used)", + ) + # groups groups_parser = subparsers.add_parser("groups", description="List groups") groups_parser.set_defaults(func="list_groups") @@ -202,6 +218,10 @@ def add_date_range_arguments(parser, default_start="5m"): # Parse input options, _ = parser.parse_known_args(argv) + if getattr(options, "query", None) and getattr(options, "jq", None): + sys.stderr.write("Cannot use both --query and --jq at the same time\n") + return 1 + try: logs = AWSLogs(**vars(options)) if not hasattr(options, "func"): diff --git a/awslogs/core.py b/awslogs/core.py index 748896e..d86601c 100644 --- a/awslogs/core.py +++ b/awslogs/core.py @@ -1,20 +1,20 @@ +import errno +import os import re import sys -import os import time -import errno -from datetime import datetime, timedelta from collections import deque +from datetime import datetime, timedelta +from typing import TypeAlias import boto3 import botocore -from botocore.compat import json, total_seconds - import jmespath - -from termcolor import colored +import jq +from botocore.compat import json, total_seconds from dateutil.parser import parse from dateutil.tz import tzutc +from termcolor import colored from . import exceptions @@ -25,12 +25,12 @@ def milis2iso(milis): def boto3_client( - aws_profile, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region, - aws_endpoint_url, + aws_profile, + aws_access_key_id, + aws_secret_access_key, + aws_session_token, + aws_region, + aws_endpoint_url, ): core_session = botocore.session.get_session() core_session.set_config_variable("profile", aws_profile) @@ -52,8 +52,26 @@ def boto3_client( ) -class AWSLogs(object): +# noinspection PyProtectedMember +JQProgram: TypeAlias = jq._Program +# noinspection PyProtectedMember +JQProgramWithInput: TypeAlias = jq._ProgramWithInput + +class JQAdapter: + # noinspection PyShadowingBuiltins + def __init__(self, jq_query: JQProgram, all: bool = False): + self.jq = jq_query + self.all = all + + def search(self, data): + j: JQProgramWithInput = self.jq.input_value(data) + if self.all: + return j.all() + return j.first() + + +class AWSLogs(object): ACTIVE = 1 EXHAUSTED = 2 WATCH_SLEEP = 2 @@ -82,8 +100,13 @@ def __init__(self, **kwargs): self.start = self.parse_datetime(kwargs.get("start")) self.end = self.parse_datetime(kwargs.get("end")) self.query = kwargs.get("query") + self.query_expression = None + self.jq = kwargs.get("jq") + self.jq_all = kwargs.get("jq_all") if self.query is not None: self.query_expression = jmespath.compile(self.query) + elif self.jq is not None: + self.query_expression = JQAdapter(jq.compile(self.jq), all=self.jq_all) self.log_group_prefix = kwargs.get("log_group_prefix") self.client = boto3_client( self.aws_profile, @@ -197,7 +220,7 @@ def consumer(): output.append(self.color(milis2iso(event["ingestionTime"]), "blue")) message = event["message"] - if self.query is not None and message[0] == "{": + if self.query_expression is not None and message[0] == "{": parsed = json.loads(event["message"]) message = self.query_expression.search(parsed) if not isinstance(message, str): @@ -256,7 +279,7 @@ def get_streams(self, log_group_name=None): # no firstEventTimestamp. yield stream["logStreamName"] elif max(stream["firstEventTimestamp"], window_start) <= min( - stream["lastIngestionTime"], window_end + stream["lastIngestionTime"], window_end ): yield stream["logStreamName"] diff --git a/setup.py b/setup.py index cdbd8dd..20143ca 100644 --- a/setup.py +++ b/setup.py @@ -6,6 +6,7 @@ "jmespath>=1.0.1", "termcolor>=2.4.0", "python-dateutil>=2.9.0", + "jq>=1.8.0", ] diff --git a/tests/test_it.py b/tests/test_it.py index 53411aa..0ef88e9 100644 --- a/tests/test_it.py +++ b/tests/test_it.py @@ -380,6 +380,22 @@ def test_main_get_query(self, mock_stdout, botoclient): self.assertEqual(output, expected) assert exit_code == 0 + @patch("awslogs.core.boto3_client") + @patch("sys.stdout", new_callable=StringIO) + @patch("os.isatty", lambda fd: True) + @patch("sys.stdout.isatty", lambda: True) + def test_main_get_jq(self, mock_stdout, botoclient): + self.set_json_logs(botoclient) + exit_code = main("awslogs get AAA DDD --jq .foo --color=always".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" + ) + self.assertEqual(output, expected) + assert exit_code == 0 + @patch("awslogs.core.boto3_client") @patch("sys.stdout", new_callable=StringIO) def test_get_nogroup(self, mock_stdout, botoclient): From 35f9081b9c94b4223bbcf6ac1f6d5c0646f449bb Mon Sep 17 00:00:00 2001 From: Davide Depau Date: Fri, 13 Sep 2024 15:40:21 +0200 Subject: [PATCH 2/2] Allow outputting JSON logs This allows post-processing this tool's output using other tools. --- awslogs/bin.py | 8 +++++ awslogs/core.py | 81 ++++++++++++++++++++++++++++++++++-------------- tests/test_it.py | 17 +++++++++- 3 files changed, 81 insertions(+), 25 deletions(-) diff --git a/awslogs/bin.py b/awslogs/bin.py index d4058d1..434f25a 100644 --- a/awslogs/bin.py +++ b/awslogs/bin.py @@ -194,6 +194,14 @@ def add_date_range_arguments(parser, default_start="5m"): help="Use all results from the jq query (by default only the first result is used)", ) + get_parser.add_argument( + "--json", + action="store_true", + dest="json", + default=False, + help="Output logs in JSON format for processing by other tools", + ) + # groups groups_parser = subparsers.add_parser("groups", description="List groups") groups_parser.set_defaults(func="list_groups") diff --git a/awslogs/core.py b/awslogs/core.py index d86601c..07099ad 100644 --- a/awslogs/core.py +++ b/awslogs/core.py @@ -107,6 +107,7 @@ def __init__(self, **kwargs): self.query_expression = jmespath.compile(self.query) elif self.jq is not None: self.query_expression = JQAdapter(jq.compile(self.jq), all=self.jq_all) + self.json_out = kwargs.get("json") self.log_group_prefix = kwargs.get("log_group_prefix") self.client = boto3_client( self.aws_profile, @@ -201,32 +202,64 @@ def consumer(): else: return - output = [] - if self.output_group_enabled: - output.append( - self.color( - self.log_group_name.ljust(group_length, " "), "green" + if self.json_out: + message = event["message"] + if message[0] == "{": + try: + message = json.loads(event["message"]) + except json.JSONDecodeError: + pass + if self.query_expression is not None and not isinstance(message, str): + # noinspection PyBroadException + try: + message = self.query_expression.search(message) + except Exception as _: + pass + + metadata = { + "log_group": self.log_group_name, + "log_stream": event["logStreamName"], + "timestamp": event["timestamp"], + "ingestion_time": event["ingestionTime"], + } + + output = { + "awslogs": metadata, + } + if isinstance(message, dict): + output.update(message) + else: + output["message"] = message + + print(json.dumps(output)) + + else: + output = [] + if self.output_group_enabled: + output.append( + self.color( + self.log_group_name.ljust(group_length, " "), "green" + ) ) - ) - if self.output_stream_enabled: - output.append( - self.color( - event["logStreamName"].ljust(max_stream_length, " "), "cyan" + if self.output_stream_enabled: + output.append( + self.color( + event["logStreamName"].ljust(max_stream_length, " "), "cyan" + ) ) - ) - if self.output_timestamp_enabled: - output.append(self.color(milis2iso(event["timestamp"]), "yellow")) - if self.output_ingestion_time_enabled: - output.append(self.color(milis2iso(event["ingestionTime"]), "blue")) - - message = event["message"] - if self.query_expression is not None and message[0] == "{": - parsed = json.loads(event["message"]) - message = self.query_expression.search(parsed) - if not isinstance(message, str): - message = json.dumps(message) - output.append(message.rstrip()) - print(" ".join(output)) + if self.output_timestamp_enabled: + output.append(self.color(milis2iso(event["timestamp"]), "yellow")) + if self.output_ingestion_time_enabled: + output.append(self.color(milis2iso(event["ingestionTime"]), "blue")) + + message = event["message"] + if self.query_expression is not None and message[0] == "{": + parsed = json.loads(event["message"]) + message = self.query_expression.search(parsed) + if not isinstance(message, str): + message = json.dumps(message) + output.append(message.rstrip()) + print(" ".join(output)) try: sys.stdout.flush() diff --git a/tests/test_it.py b/tests/test_it.py index 0ef88e9..542fd6b 100644 --- a/tests/test_it.py +++ b/tests/test_it.py @@ -34,7 +34,6 @@ class TestAWSLogsDatetimeParse(unittest.TestCase): @patch("awslogs.core.boto3_client") @patch("awslogs.core.datetime") def test_parse_datetime(self, datetime_mock, botoclient): - awslogs = AWSLogs() datetime_mock.utcnow.return_value = datetime(2015, 1, 1, 3, 0, 0, 0) datetime_mock.return_value = datetime(1970, 1, 1) @@ -396,6 +395,22 @@ def test_main_get_jq(self, mock_stdout, botoclient): self.assertEqual(output, expected) assert exit_code == 0 + @patch("awslogs.core.boto3_client") + @patch("sys.stdout", new_callable=StringIO) + @patch("os.isatty", lambda fd: True) + @patch("sys.stdout.isatty", lambda: True) + def test_main_get_json(self, mock_stdout, botoclient): + self.set_json_logs(botoclient) + exit_code = main("awslogs get AAA DDD --json".split()) + output = mock_stdout.getvalue() + expected = ( + '{"awslogs": {"log_group": "AAA", "log_stream": "DDD", "timestamp": 0, "ingestion_time": 5000}, "foo": "bar"}\n' + '{"awslogs": {"log_group": "AAA", "log_stream": "EEE", "timestamp": 0, "ingestion_time": 5000}, "foo": {"bar": "baz"}}\n' + '{"awslogs": {"log_group": "AAA", "log_stream": "DDD", "timestamp": 0, "ingestion_time": 5006}, "message": "Hello 3 \\ud83d\\udc4e"}\n' + ) + self.assertEqual(output, expected) + assert exit_code == 0 + @patch("awslogs.core.boto3_client") @patch("sys.stdout", new_callable=StringIO) def test_get_nogroup(self, mock_stdout, botoclient):