Merge pull request #559 from kkedziak-splunk/feature/improve_logging

feat: Add exception logging in Script

splunklib/modularinput/event_writer.py

@@ -13,6 +13,7 @@
 # under the License.
 
 import sys
+import traceback
 
 from splunklib.utils import ensure_str
 from .event import ET
@@ -66,6 +67,25 @@ def log(self, severity, message):
         self._err.write(f"{severity} {message}\n")
         self._err.flush()
 
+    def log_exception(self, message, exception=None, severity=None):
+        """Logs messages about the exception thrown by this modular input to Splunk.
+        These messages will show up in Splunk's internal logs.
+
+        :param message: ``string``, message to log.
+        :param exception: ``Exception``, exception thrown by this modular input; if none, sys.exc_info() is used
+        :param severity: ``string``, severity of message, see severities defined as class constants. Default severity: ERROR
+        """
+        if exception is not None:
+            tb_str = traceback.format_exception(type(exception), exception, exception.__traceback__)
+        else:
+            tb_str = traceback.format_exc()
+
+        if severity is None:
+            severity = EventWriter.ERROR
+
+        self._err.write(("%s %s - %s" % (severity, message, tb_str)).replace("\n", " "))
+        self._err.flush()
+
     def write_xml_document(self, document):
         """Writes a string representation of an
         ``ElementTree`` object to the output stream.

splunklib/modularinput/script.py

@@ -91,13 +91,12 @@ def run_script(self, args, event_writer, input_stream):
                     event_writer.write_xml_document(root)
 
                     return 1
-            err_string = "ERROR Invalid arguments to modular input script:" + ' '.join(
-                args)
-            event_writer._err.write(err_string)
+            event_writer.log(EventWriter.ERROR, "Invalid arguments to modular input script:" + ' '.join(
+                args))
             return 1
 
         except Exception as e:
-            event_writer.log(EventWriter.ERROR, str(e))
+            event_writer.log_exception(str(e))
             return 1
 
     @property

tests/modularinput/test_event.py

@@ -15,7 +15,9 @@
 # under the License.
 
 
+import re
 import sys
+from io import StringIO
 
 import pytest
 
@@ -150,3 +152,29 @@ def test_write_xml_is_sane(capsys):
         found_xml = ET.fromstring(captured.out)
 
         assert xml_compare(expected_xml, found_xml)
+
+
+def test_log_exception():
+    out, err = StringIO(), StringIO()
+    ew = EventWriter(out, err)
+
+    exc = Exception("Something happened!")
+
+    try:
+        raise exc
+    except Exception:
+        ew.log_exception("ex1")
+
+    assert out.getvalue() == ""
+
+    # Remove paths and line
+    err = re.sub(r'File "[^"]+', 'File "...', err.getvalue())
+    err = re.sub(r'line \d+', 'line 123', err)
+
+    # One line
+    assert err == (
+        'ERROR ex1 - Traceback (most recent call last): '
+        '  File "...", line 123, in test_log_exception '
+        '    raise exc '
+        'Exception: Something happened! '
+    )

tests/modularinput/test_script.py

@@ -1,6 +1,7 @@
 import sys
 
 import io
+import re
 import xml.etree.ElementTree as ET
 from splunklib.client import Service
 from splunklib.modularinput import Script, EventWriter, Scheme, Argument, Event
@@ -228,3 +229,37 @@ def stream_events(self, inputs, ew):
         assert output.err == ""
         assert isinstance(script.service, Service)
         assert script.service.authority == script.authority_uri
+
+
+def test_log_script_exception(monkeypatch):
+    out, err = io.StringIO(), io.StringIO()
+
+    # Override abstract methods
+    class NewScript(Script):
+        def get_scheme(self):
+            return None
+
+        def stream_events(self, inputs, ew):
+            raise RuntimeError("Some error")
+
+    script = NewScript()
+    input_configuration = data_open("data/conf_with_2_inputs.xml")
+
+    ew = EventWriter(out, err)
+
+    assert script.run_script([TEST_SCRIPT_PATH], ew, input_configuration) == 1
+
+    # Remove paths and line numbers
+    err = re.sub(r'File "[^"]+', 'File "...', err.getvalue())
+    err = re.sub(r'line \d+', 'line 123', err)
+
+    assert out.getvalue() == ""
+    assert err == (
+        'ERROR Some error - '
+        'Traceback (most recent call last): '
+        '  File "...", line 123, in run_script '
+        '    self.stream_events(self._input_definition, event_writer) '
+        '  File "...", line 123, in stream_events '
+        '    raise RuntimeError("Some error") '
+        'RuntimeError: Some error '
+    )