From 3dd040837054af8fb24432b8d3bead80991694f1 Mon Sep 17 00:00:00 2001
From: Stephen Leong Koan <sleongkoan@gmail.com>
Date: Fri, 20 Jan 2017 14:07:38 -0500
Subject: [PATCH 1/2] introduce argument allowing the formatter to ignore key
 errors

---
 fluent/handler.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/fluent/handler.py b/fluent/handler.py
index 1e06e08..af44856 100644
--- a/fluent/handler.py
+++ b/fluent/handler.py
@@ -23,8 +23,12 @@ class FluentRecordFormatter(logging.Formatter, object):
     Best used with server storing data in an ElasticSearch cluster for example.
 
     :param fmt: a dict with format string as values to map to provided keys.
+    :param datefmt: strftime()-compatible date/time format string.
+    :param style: (NOT USED)
+    :param fill_missing_fmt_key: if True, do not raise a KeyError if the format
+        key is not found. Put None if not found.s
     """
-    def __init__(self, fmt=None, datefmt=None, style='%'):
+    def __init__(self, fmt=None, datefmt=None, style='%', fill_missing_fmt_key=False):
         super(FluentRecordFormatter, self).__init__(None, datefmt)
 
         if not fmt:
@@ -38,6 +42,8 @@ def __init__(self, fmt=None, datefmt=None, style='%'):
 
         self.hostname = socket.gethostname()
 
+        self.fill_missing_fmt_key = fill_missing_fmt_key
+
     def format(self, record):
         # Only needed for python2.6
         if sys.version_info[0:2] <= (2, 6) and self.usesTime():
@@ -47,9 +53,18 @@ def format(self, record):
         super(FluentRecordFormatter, self).format(record)
         # Add ours
         record.hostname = self.hostname
+
         # Apply format
-        data = dict([(key, value % record.__dict__)
-                     for key, value in self._fmt_dict.items()])
+        data = {}
+        for key, value in self._fmt_dict.items():
+            try:
+                value = value % record.__dict__
+            except KeyError as exc:
+                value = None
+                if not self.fill_missing_fmt_key:
+                    raise exc
+
+            data[key] = value
 
         self._structuring(data, record)
         return data

From 268c6a7831915bb8558e9b0c9772996d263e8ea4 Mon Sep 17 00:00:00 2001
From: Stephen Leong Koan <sleongkoan@gmail.com>
Date: Fri, 20 Jan 2017 14:47:44 -0500
Subject: [PATCH 2/2] add test for new feature

---
 tests/test_handler.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/tests/test_handler.py b/tests/test_handler.py
index dfb3d29..ea44a98 100644
--- a/tests/test_handler.py
+++ b/tests/test_handler.py
@@ -62,6 +62,48 @@ def test_custom_fmt(self):
         self.assertTrue('lineno' in data[0][2])
         self.assertTrue('emitted_at' in data[0][2])
 
+    def test_custom_field_raise_exception(self):
+        handler = fluent.handler.FluentHandler('app.follow', port=self._port)
+
+        logging.basicConfig(level=logging.INFO)
+        log = logging.getLogger('fluent.test')
+        handler.setFormatter(
+            fluent.handler.FluentRecordFormatter(fmt={
+                'name': '%(name)s',
+                'custom_field': '%(custom_field)s'
+            })
+        )
+        log.addHandler(handler)
+        with self.assertRaises(KeyError):
+            log.info({'sample': 'value'})
+        log.removeHandler(handler)
+        handler.close()
+
+    def test_custom_field_fill_missing_fmt_key_is_true(self):
+        handler = fluent.handler.FluentHandler('app.follow', port=self._port)
+
+        logging.basicConfig(level=logging.INFO)
+        log = logging.getLogger('fluent.test')
+        handler.setFormatter(
+            fluent.handler.FluentRecordFormatter(fmt={
+                'name': '%(name)s',
+                'custom_field': '%(custom_field)s'
+                },
+                fill_missing_fmt_key=True
+            )
+        )
+        log.addHandler(handler)
+        log.info({'sample': 'value'})
+        log.removeHandler(handler)
+        handler.close()
+
+        data = self.get_data()
+        self.assertTrue('name' in data[0][2])
+        self.assertEqual('fluent.test', data[0][2]['name'])
+        self.assertTrue('custom_field' in data[0][2])
+        # field defaults to none if not in log record
+        self.assertIsNone(data[0][2]['custom_field'])
+
     def test_json_encoded_message(self):
         handler = fluent.handler.FluentHandler('app.follow', port=self._port)