Add exception logging in Script

Author: kkedziak-splunk

splunklib/modularinput/event_writer.py

@@ -14,6 +14,7 @@
 
 from __future__ import absolute_import
 import sys
+import traceback
 
 from splunklib.six import ensure_str
 from .event import ET
@@ -23,6 +24,7 @@
 except ImportError:
     from splunklib.six import StringIO
 
+
 class EventWriter(object):
     """``EventWriter`` writes events and error messages to Splunk from a modular input.
     Its two important methods are ``writeEvent``, which takes an ``Event`` object,
@@ -71,6 +73,25 @@ def log(self, severity, message):
         self._err.write("%s %s\n" % (severity, message))
         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: 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

@@ -99,13 +99,12 @@ def run_script(self, args, event_writer, input_stream):
 
                     return 1
             else:
-                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

@@ -16,7 +16,9 @@
 
 from __future__ import absolute_import
 
+import re
 import sys
+from io import StringIO
 
 import pytest
 
@@ -151,3 +153,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:
+        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

@@ -3,6 +3,7 @@
 from splunklib.client import Service
 from splunklib.modularinput import Script, EventWriter, Scheme, Argument, Event
 import io
+import re
 
 from splunklib.modularinput.utils import xml_compare
 from tests.modularinput.modularinput_testlib import data_open
@@ -231,3 +232,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 '
+    )