From 2f729bfdbcdd22556f8128ead11a2565c293b6ab Mon Sep 17 00:00:00 2001
From: James Turk <james.p.turk@gmail.com>
Date: Tue, 15 May 2012 23:47:45 -0400
Subject: [PATCH] add mongolog

---
 oyster/mongolog.py            | 54 +++++++++++++++++++++++++++++++++++
 oyster/tests/test_mongolog.py | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+)
 create mode 100644 oyster/mongolog.py
 create mode 100644 oyster/tests/test_mongolog.py

diff --git a/oyster/mongolog.py b/oyster/mongolog.py
new file mode 100644
index 0000000..f99cbb4
--- /dev/null
+++ b/oyster/mongolog.py
@@ -0,0 +1,54 @@
+"""
+    MongoDB handler for Python Logging
+
+    inspired by https://github.com/andreisavu/mongodb-log
+"""
+
+import logging
+import datetime
+import socket
+import pymongo
+
+
+class MongoFormatter(logging.Formatter):
+
+    def format(self, record):
+        """ turn a LogRecord into something mongo can store """
+        data = record.__dict__.copy()
+
+        data.update(
+            # format message
+            message=record.getMessage(),
+            # overwrite created (float) w/ a mongo-compatible datetime
+            created=datetime.datetime.utcnow(),
+            host=socket.gethostname(),
+            args=tuple(unicode(arg) for arg in record.args)
+        )
+        data.pop('msecs')   # not needed, stored in created
+
+        # TODO: ensure everything in 'extra' is MongoDB-ready
+        exc_info = data.get('exc_info')
+        if exc_info:
+            data['exc_info'] = self.formatException(exc_info)
+        return data
+
+
+class MongoHandler(logging.Handler):
+    def __init__(self, db, collection='logs', host='localhost', port=None,
+                 capped_size=100000000, level=logging.NOTSET, async=True):
+        db = pymongo.connection.Connection(host, port)[db]
+        # try and create the capped log collection
+        if capped_size:
+            try:
+                db.create_collection(collection, capped=True, size=capped_size)
+            except pymongo.errors.CollectionInvalid:
+                pass
+        self.collection = db[collection]
+        self.async = async
+        super(MongoHandler, self).__init__(level)
+        self.formatter = MongoFormatter()
+
+    def emit(self, record):
+        # explicitly set safe=False to get async insert
+        # TODO: what to do if an error occurs? not safe to log-- ignore?
+        self.collection.save(self.format(record), safe=not self.async)
diff --git a/oyster/tests/test_mongolog.py b/oyster/tests/test_mongolog.py
new file mode 100644
index 0000000..3cdde1d
--- /dev/null
+++ b/oyster/tests/test_mongolog.py
@@ -0,0 +1,53 @@
+import unittest
+import logging
+import datetime
+
+import pymongo
+from ..mongolog import MongoHandler
+
+class TestMongoLog(unittest.TestCase):
+
+    DB_NAME = 'oyster_test'
+
+    def setUp(self):
+        pymongo.Connection().drop_database(self.DB_NAME)
+        self.log = logging.getLogger('mongotest')
+        self.log.setLevel(logging.DEBUG)
+        self.logs = pymongo.Connection()[self.DB_NAME]['logs']
+        # clear handlers upon each setup 
+        self.log.handlers = []
+        # async = False for testing
+        self.log.addHandler(MongoHandler(self.DB_NAME, capped_size=4000,
+                                         async=False))
+
+    def tearDown(self):
+        pymongo.Connection().drop_database(self.DB_NAME)
+
+    def test_basic_write(self):
+        self.log.debug('test')
+        self.assertEqual(self.logs.count(), 1)
+        self.log.debug('test')
+        self.assertEqual(self.logs.count(), 2)
+        # capped_size will limit these
+        self.log.debug('test'*200)
+        self.log.debug('test'*200)
+        self.assertEqual(self.logs.count(), 1)
+
+    def test_attributes(self):
+        self.log.debug('pi=%s', 3.14, extra={'pie':'pizza'})
+        logged = self.logs.find_one()
+        self.assertEqual(logged['message'], 'pi=3.14')
+        self.assertTrue(isinstance(logged['created'], datetime.datetime))
+        self.assertTrue('host' in logged)
+        self.assertEqual(logged['name'], 'mongotest')
+        self.assertEqual(logged['levelname'], 'DEBUG')
+        self.assertEqual(logged['pie'], 'pizza')
+
+        # and exc_info
+        try:
+            raise Exception('error!')
+        except:
+            self.log.warning('oh no!', exc_info=True)
+            logged = self.logs.find_one(sort=[('$natural', -1)])
+        self.assertEqual(logged['levelname'], 'WARNING')
+        self.assertTrue('error!' in logged['exc_info'])