Add support to collapse JavaStackTrace into a single event

.github/workflows/test.yml

@@ -0,0 +1,27 @@
+name: Test Package
+
+on:
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  id-token: write
+  contents: write
+  packages: write
+  deployments: write
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up environment
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Test
+        run: ./test.sh

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 **/.idea/
+
+# test python env
+test_env/

README.md

@@ -15,6 +15,10 @@ This lambda function can be configured with the following attributes:
 where keys and values should only contain alphanumerical character or `-` or `_`.
 - `max_batch_size`: the maximum size of an uncompressed payload sent to BrontoBytes. This lambda function compresses 
 the data with gzip. As a rule of thumb, a compression ratio of about 90% can be expected.
+- `aggregator`: the name of an aggregator to use: either `java_stack_trace` or `default`. This property defaults to 
+`default` if not set. Aggregators aggregate multiline log entries into a single entry. The `default` aggregator is a 
+noop (no entries get aggregated), while the `java_stack_trace` aggregator aggregates Java stack trace entries into a 
+single one.
 - `destination_config`: a base64 encoded map representing the configuration to where each log should be sent to.
 - `paths_regex`: `paths_regex` is a base64-encoded list of objects, each containing a regular expression pattern with a 
 named capture group called `dest_config_id`. This is used for log data delivered to S3 when the S3 object key does not 

log_forwarder/aggregator.py

@@ -0,0 +1,85 @@
+from typing import Union
+from queue import Queue
+
+AGGREGATED_TERMINATION_MARKER = None
+
+
+class Aggregator:
+
+    _NO_AGGREGATED_LINE = None
+
+    def add_line(self, line: Union[str, None]):
+        raise NotImplementedError()
+
+    def get_complete_aggregated_line(self):
+        raise NotImplementedError()
+
+    def complete(self):
+        self.add_line(AGGREGATED_TERMINATION_MARKER)
+
+    def has_complete_aggregated_line(self):
+        raise NotImplementedError()
+
+
+class NoopAggregator(Aggregator):
+
+    def __init__(self):
+        self.lines = Queue(maxsize=2)
+
+    def add_line(self, line: Union[str, None]):
+        if self.has_complete_aggregated_line():
+            raise Exception('Complete aggregated lines must be retrieved first')
+        if line is not None and line == '':
+            return Aggregator._NO_AGGREGATED_LINE
+        self.lines.put(line)
+        return
+
+    def get_complete_aggregated_line(self):
+        if self.has_complete_aggregated_line():
+            return self.lines.get()
+        return Aggregator._NO_AGGREGATED_LINE
+
+    def has_complete_aggregated_line(self):
+        return self.lines.full()
+
+
+class JavaStackTraceAggregator(Aggregator):
+    """ Inspired from https://www.elastic.co/docs/reference/beats/filebeat/multiline-examples#_java_stack_traces """
+
+    def __init__(self):
+        self.lines = Queue(maxsize=2)
+
+    def add_line(self, line):
+        if self.lines.full():
+            raise Exception('Aggregator queue is full')
+        if line is None:
+            current_line = line
+        else:
+            stripped_line = line.lstrip()
+            if stripped_line == '':
+                return
+            if (stripped_line.startswith('at ') or stripped_line.startswith('Caused by: ') or
+                    stripped_line.startswith('...')):
+                current_line = self.lines.get().rstrip('\n')
+                current_line += f'\\n{line}'
+            else:
+                current_line = line
+        self.lines.put(current_line)
+        return
+
+    def get_complete_aggregated_line(self):
+        if self.has_complete_aggregated_line():
+            return self.lines.get()
+        return Aggregator._NO_AGGREGATED_LINE
+
+    def has_complete_aggregated_line(self):
+        return self.lines.full()
+
+
+class AggregatorFactory:
+
+    @staticmethod
+    def get_aggregator(aggregator_name) -> Aggregator:
+        if aggregator_name == 'java_stack_trace':
+            return JavaStackTraceAggregator()
+        return NoopAggregator()

log_forwarder/clients.py

@@ -23,9 +23,10 @@ def download(self, bucket, object_key):
 
 class Batch:
 
-    def __init__(self, no_formatting=False):
+    def __init__(self, max_size: int, no_formatting=False):
         self.batch = []
         self.size = 0
+        self.max_size = max_size
         self.no_formatting = no_formatting
 
     def add(self, line):
@@ -46,6 +47,10 @@ def get_formatted_data(self, attributes: Dict[str, str]):
             log_message.update(attributes)
         return '\n'.join([json.dumps(log_message) for log_message in log_messages])
 
+    def reset(self):
+        self.batch = []
+        self.size = 0
+
 
 class BrontoClient:
 

log_forwarder/config.py

@@ -38,6 +38,7 @@ def __init__(self, event, filepath):
         self.filepath = filepath
         self.event = event
         self.path_regexes = os.environ.get('path_regexes')
+        self.aggregator = os.environ.get('aggregator', 'default')
         raw_tags = os.environ.get('tags')
         self.tags = Config._extract_kvps(raw_tags)
         raw_attributes = os.environ.get('attributes')

log_forwarder/exporter.py

@@ -0,0 +1,32 @@
+from typing import Dict, List
+from aggregator import Aggregator
+from clients import BrontoClient, Batch
+from parser import Parser
+
+
+class BrontoExporter:
+
+    def __init__(self, client: BrontoClient, parser: Parser, batch: Batch, aggregator: Aggregator,
+                 attributes: Dict[str, str]):
+        self.client = client
+        self.parser = parser
+        self.batch = batch
+        self.aggregator = aggregator
+        self.attributes = attributes
+
+    def export(self):
+        for line in self.parser.get_parsed_lines():
+            self.aggregator.add_line(line)
+            if not self.aggregator.has_complete_aggregated_line():
+                continue
+            _line = self.aggregator.get_complete_aggregated_line()
+            self.batch.add(_line)
+            if self.batch.get_batch_size() > self.batch.max_size:
+                self.client.send_data(self.batch, self.attributes)
+                self.batch.reset()
+        self.aggregator.complete()
+        if self.aggregator.has_complete_aggregated_line():
+            _line = self.aggregator.get_complete_aggregated_line()
+            self.batch.add(_line)
+        if self.batch.get_batch_size() > 0:
+            self.client.send_data(self.batch, self.attributes)

log_forwarder/forward.py

@@ -2,9 +2,11 @@
 import tempfile
 import shutil
 
+from aggregator import JavaStackTraceAggregator, AggregatorFactory
 from config import Config, DestinationConfig
 from data_retriever import DataRetrieverFactory
 from destination_provider import DestinationProvider
+from exporter import BrontoExporter
 from parser import ParserFactory
 from clients import BrontoClient, Batch
 from logfile import LogFileFactory
@@ -58,14 +60,10 @@ def process(event):
             bronto_client = BrontoClient(dest_config.bronto_api_key, dest_config.bronto_endpoint, dataset, collection,
                 client_type, config.tags)
             no_formatting = client_type is not None
-            batch = Batch(no_formatting)
-            for line in parser.get_parsed_lines():
-                batch.add(line)
-                if batch.get_batch_size() > dest_config.max_batch_size:
-                    bronto_client.send_data(batch, attributes)
-                    batch = Batch()
-            if batch.get_batch_size() > 0:
-                bronto_client.send_data(batch, config.get_resource_attributes())
+            batch = Batch(dest_config.max_batch_size, no_formatting)
+            aggregator = AggregatorFactory.get_aggregator(config.aggregator)
+            exporter = BrontoExporter(bronto_client, parser, batch, aggregator, attributes)
+            exporter.export()
 
 
 def forward_logs(_event, _):

log_forwarder/logfile.py

@@ -46,7 +46,7 @@ def get_file(self):
     def get_lines(self):
         with self.file as f:
             for line in f.readlines():
-                yield line.strip()
+                yield line.rstrip()  # we keep spaces at the start as they may be indicative of a stack trace
 
 
 class LogFileFactory:

test.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+set -e
+
+PROJECT_ROOT_DIR="${PWD}"
+TEST_REQUIREMENT_FILE="${PROJECT_ROOT_DIR}/test_requirements.txt"
+
+rm -rf test_env
+python3 -m venv test_env
+"${PROJECT_ROOT_DIR}/test_env/bin/python" -m pip install --upgrade pip
+test_env/bin/pip install -r test_requirements.txt
+
+cd "${PROJECT_ROOT_DIR}/tests" || exit 1
+
+echo "Running Tests"
+PYTHONPATH="${PROJECT_ROOT_DIR}/log_forwarder" "${PROJECT_ROOT_DIR}/test_env/bin/python" -m pytest -v -s

test_requirements.txt

@@ -0,0 +1,2 @@
+pytest
+boto3