From 44089aaf6305ca9820472ee61a24c0b3336d9084 Mon Sep 17 00:00:00 2001
From: TomasKuznetsov <tomas.k@logz.io>
Date: Thu, 30 Nov 2017 16:14:34 +0200
Subject: [PATCH] Added extra fields support

---
 logzio/handler.py           | 30 ++++++++++++++++++++++++++++++
 tests/test_logzioHandler.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/logzio/handler.py b/logzio/handler.py
index 5df4b12..e5bb314 100644
--- a/logzio/handler.py
+++ b/logzio/handler.py
@@ -3,6 +3,7 @@ import json
 import logging
 import logging.handlers
 import traceback
+import sys
 
 from .sender import LogzioSender
 
@@ -19,6 +20,31 @@ class LogzioHandler(logging.Handler):
         self.logzio_sender = LogzioSender(token=token, url=url, logs_drain_timeout=logs_drain_timeout, debug=debug)
         logging.Handler.__init__(self)
 
+    def extra_fields(self, message):
+
+        not_allowed_keys = (
+            'args', 'asctime', 'created', 'exc_info', 'exc_text', 'filename',
+            'funcName', 'levelname', 'levelno', 'lineno', 'module',
+            'msecs', 'msecs', 'message', 'msg', 'name', 'pathname', 'process',
+            'processName', 'relativeCreated', 'thread', 'threadName')
+
+
+        if sys.version_info < (3, 0):
+            var_type = (basestring, bool, dict, float, int, long, list, type(None))
+        else:
+            var_type = (str, bool, dict, float, int, list, type(None))
+
+        extra_fields = {}
+
+        for key, value in message.__dict__.items():
+            if key not in not_allowed_keys:
+                if isinstance(value, var_type):
+                    extra_fields[key] = value
+                else:
+                    extra_fields[key] = repr(value)
+
+        return extra_fields
+
     def format(self, record):
         message = super(LogzioHandler, self).format(record)
         try:
@@ -47,6 +73,9 @@ class LogzioHandler(logging.Handler):
             return_json["exception"] = self.format_exception(message.exc_info)
         else:
             formatted_message = self.format(message)
+
+            return_json.update(self.extra_fields(message))
+
             if isinstance(formatted_message, dict):
                 return_json.update(formatted_message)
             else:
@@ -54,5 +83,6 @@ class LogzioHandler(logging.Handler):
 
         return return_json
 
+
     def emit(self, record):
         self.logzio_sender.append(self.format_message(record))
diff --git a/tests/test_logzioHandler.py b/tests/test_logzioHandler.py
index c4e0194..88689a7 100644
--- a/tests/test_logzioHandler.py
+++ b/tests/test_logzioHandler.py
@@ -77,6 +77,38 @@ class TestLogzioHandler(TestCase):
              }
         )
 
+    def test_extra_formatting(self):
+        record = logging.LogRecord(
+            name='my-logger',
+            level=0,
+            pathname='handler_test.py',
+            lineno=10,
+            msg="this is a test: moo.",
+            args=[],
+            exc_info=None,
+            func='test_json'
+        )
+
+        record.__dict__["extra_key"] = "extra_value"
+        record.__dict__["module"] = "testing"
+
+        formatted_message = self.handler.format_message(record)
+        formatted_message["@timestamp"] = None
+
+        self.assertDictEqual(
+            formatted_message,
+            {
+                '@timestamp': None,
+                'line_number': 10,
+                'log_level': 'NOTSET',
+                'logger': 'my-logger',
+                'message': 'this is a test: moo.',
+                'path_name': 'handler_test.py',
+                'type': 'python',
+                'extra_key': 'extra_value'
+             }
+        )
+
     def test_format_string_message(self):
         record = logging.LogRecord(
             name='my-logger',