L-1200 Detect and handle circular dependencies (#23)

logtail/frame.py

@@ -41,8 +41,7 @@ def create_frame(record, message, context, include_extra_attributes=False):
     if events:
         frame.update(events)
 
-    return frame
-
+    return _remove_circular_dependencies(frame)
 
 def _parse_custom_events(record, include_extra_attributes):
     default_keys = {
@@ -60,6 +59,22 @@ def _parse_custom_events(record, include_extra_attributes):
         events[key] = val
     return events
 
+def _remove_circular_dependencies(obj, memo=None):
+    if memo is None:
+        memo = set()
+    if id(obj) in memo:
+        return "<omitted circular reference>"
+    memo.add(id(obj))
+    if isinstance(obj, dict):
+        new_dict = {}
+        for key, value in obj.items():
+            new_dict[key] = _remove_circular_dependencies(value, memo)
+        return new_dict
+    elif isinstance(obj, list):
+        new_list = [_remove_circular_dependencies(item, memo) for item in obj]
+        return new_list
+    else:
+        return obj
 
 def _levelname(level):
     return level.lower()

tests/test_handler.py

@@ -159,7 +159,7 @@ def test_can_send_unserializable_extra_data(self, MockWorker):
         self.assertTrue(handler.pipe.empty())
 
     @mock.patch('logtail.handler.FlushWorker')
-    def test_can_send_unserializable_extra_context(self, MockWorker):
+    def test_can_send_unserializable_context(self, MockWorker):
         buffer_capacity = 1
         handler = LogtailHandler(
             source_token=self.source_token,
@@ -178,6 +178,50 @@ def test_can_send_unserializable_extra_context(self, MockWorker):
         self.assertRegex(log_entry['context']['data']['unserializable'], r'^<tests\.test_handler\.UnserializableObject object at 0x[0-f]+>$')
         self.assertTrue(handler.pipe.empty())
 
+    @mock.patch('logtail.handler.FlushWorker')
+    def test_can_send_circular_dependency_in_extra_data(self, MockWorker):
+        buffer_capacity = 1
+        handler = LogtailHandler(
+            source_token=self.source_token,
+            buffer_capacity=buffer_capacity
+        )
+
+        logger = logging.getLogger(__name__)
+        logger.handlers = []
+        logger.addHandler(handler)
+        circular_dependency = {'egg': {}}
+        circular_dependency['egg']['chicken'] = circular_dependency
+        logger.info('hello', extra={'data': circular_dependency})
+
+        log_entry = handler.pipe.get()
+
+        self.assertEqual(log_entry['message'], 'hello')
+        self.assertEqual(log_entry['data']['egg']['chicken'], "<omitted circular reference>")
+        self.assertTrue(handler.pipe.empty())
+
+
+    @mock.patch('logtail.handler.FlushWorker')
+    def test_can_send_circular_dependency_in_context(self, MockWorker):
+        buffer_capacity = 1
+        handler = LogtailHandler(
+            source_token=self.source_token,
+            buffer_capacity=buffer_capacity
+        )
+
+        logger = logging.getLogger(__name__)
+        logger.handlers = []
+        logger.addHandler(handler)
+        circular_dependency = {'egg': {}}
+        circular_dependency['egg']['chicken'] = circular_dependency
+        with context(data=circular_dependency):
+            logger.info('hello')
+
+        log_entry = handler.pipe.get()
+
+        self.assertEqual(log_entry['message'], 'hello')
+        self.assertEqual(log_entry['context']['data']['egg']['chicken']['egg'], "<omitted circular reference>")
+        self.assertTrue(handler.pipe.empty())
+
 
 class UnserializableObject(object):
     """ Because this is a custom class, it cannot be serialized into JSON. """