[Fix teamniteo#22] Add request_id...

... to sentry and structlog logs

Signed-off-by: Sourya Vatsyayan <[email protected]>

pyramid_heroku/request_id.py

@@ -0,0 +1,50 @@
+"""Set request ID in Sentry logs if sentry is imported, and in structlog logs
+if structlog is imported.
+"""
+
+
+try:
+    import sentry_sdk
+    IS_SENTRY_ENABLED = True
+except ImportError:
+    IS_SENTRY_ENABLED = False
+
+try:
+    import structlog
+    IS_STRUCTLOG_ENABLED = True
+except ImportError:
+    IS_STRUCTLOG_ENABLED = False
+
+
+def includeme(config):
+    config.add_tween("pyramid_heroku.request_id.RequestIDLogger")
+
+
+class RequestIDLogger(object):
+    """Set request ID in sentry and structlog logs.
+
+    If the request headers contain a request ID (X-Request-ID), it should be
+    logged for better debugging. In case sentry is being used, log the request
+    ID along with other info. Similarly, if structlog is being used, add the
+    request ID in the logs.
+    """
+
+    def __init__(self, handler, _):
+        self.handler = handler
+
+    def __call__(self, request):
+        request_id = request.headers.get("X-Request-ID")
+
+        if request_id is not None:
+            if IS_SENTRY_ENABLED:
+                with sentry_sdk.configure_scope() as scope:
+                    scope.set_tag("request_id", request_id)
+
+            if IS_STRUCTLOG_ENABLED:
+                WrappedDictClass = structlog.threadlocal.wrap_dict(dict)
+                wrapped_request_id = WrappedDictClass({
+                    "request_id": request_id
+                })
+                structlog.configure(context_class=WrappedDictClass)
+
+        return self.handler(request)

pyramid_heroku/tests/__init__.py

@@ -0,0 +1,6 @@
+import unittest
+
+from .test_request_id import TestRequestIDLogger
+
+if __name__ == "__main__":
+    unittest.main()

pyramid_heroku/tests/test_client_addr.py

@@ -3,8 +3,8 @@
 from pyramid import request
 from pyramid import testing
 
-import mock
 import unittest
+from unittest import mock
 
 
 class TestClientAddrTween(unittest.TestCase):

pyramid_heroku/tests/test_request_id.py

@@ -0,0 +1,92 @@
+"""Tests for RequestIDLogger tween"""
+
+from contextlib import contextmanager
+from pyramid import request
+from pyramid import testing
+
+import unittest
+from unittest import mock
+import structlog
+import sys
+
+
+SENTRY_TAGS = {}
+
+class MockScope(object):
+
+    def __init__(self):
+        # reset the tags for every new test
+        SENTRY_TAGS = {}
+
+    def set_tag(self, key, value):
+        SENTRY_TAGS[key] = value
+
+
+class MockSentrySDK:
+    @classmethod
+    def configure_scope(cls):
+        @contextmanager
+        def inner():
+            yield MockScope()
+        return inner()
+
+
+sys.modules['sentry_sdk'] = MockSentrySDK
+
+
+class TestRequestIDLogger(unittest.TestCase):
+    def setUp(self):
+        self.config = testing.setUp()
+        self.request = request.Request({})
+
+        self.handler = mock.Mock()
+        self.registry = None
+
+        self.request_id = "some-random-requestid"
+
+    def tearDown(self):
+        testing.tearDown()
+
+    def test_request_id_in_sentry(self):
+        from pyramid_heroku.request_id import RequestIDLogger
+
+        self.request.headers["X-Request-ID"] = self.request_id
+        RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+        self.assertEqual(SENTRY_TAGS, {"request_id": self.request_id})
+
+    def test_request_id_in_structlog(self):
+        structlog.reset_defaults()
+        from pyramid_heroku.request_id import RequestIDLogger
+
+        context_class = structlog.get_config().get('context_class')
+        self.assertEqual(context_class, dict)
+
+        self.request.headers["X-Request-ID"] = self.request_id
+        RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+
+        context_class = structlog.get_config().get('context_class')
+        self.assertEqual(
+            context_class._tl.dict_, {"request_id": self.request_id}
+        )
+
+    def test_no_request_id_in_sentry(self):
+        from pyramid_heroku.request_id import RequestIDLogger
+
+        RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+        self.assertEqual(SENTRY_TAGS, {})
+
+    def test_no_request_id_in_structlog(self):
+        structlog.reset_defaults()
+        from pyramid_heroku.request_id import RequestIDLogger
+
+        context_class = structlog.get_config().get('context_class')
+        self.assertEqual(context_class, dict)
+
+        RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+
+        context_class = structlog.get_config().get('context_class')
+        self.assertEqual(context_class, dict)