[Fix niteoweb/pyramid_heroku#22] Add request_id...

Pipfile

@@ -18,6 +18,7 @@ structlog = "*"
 black = "==18.9b0"
 mypy = "*"
 typecov = "*"
+sentry-sdk = "*"
 
 [requires]
 python_version = "3.7"

pyramid_heroku/request_id.py

@@ -0,0 +1,60 @@
+"""Set request ID in Sentry logs if sentry is imported, and in structlog logs
+if structlog is imported.
+"""
+
+IS_SENTRY_INSTALLED: bool
+IS_STRUCTLOG_INSTALLED: bool
+
+try:
+    import sentry_sdk
+
+    IS_SENTRY_INSTALLED = True
+except ImportError:
+    IS_SENTRY_INSTALLED = False
+
+try:
+    import structlog
+
+    IS_STRUCTLOG_INSTALLED = True
+except ImportError:
+    IS_STRUCTLOG_INSTALLED = 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 None:
+            return self.handler(request)
+
+        if IS_SENTRY_INSTALLED:
+            self.set_sentry_scope(request_id)
+
+        if IS_STRUCTLOG_INSTALLED:
+            self.set_structlog_context(request_id)
+
+        return self.handler(request)
+
+    def set_sentry_scope(self, request_id: int) -> None:
+        with sentry_sdk.configure_scope() as scope:
+            scope.set_tag("request_id", request_id)
+
+    def set_structlog_context(self, request_id: int) -> None:
+        WrappedDictClass = structlog.threadlocal.wrap_dict(dict)
+        _ = WrappedDictClass({"request_id": request_id})
+        structlog.configure(context_class=WrappedDictClass)

pyramid_heroku/tests/__init__.py

@@ -0,0 +1 @@
+

pyramid_heroku/tests/test_request_id.py

@@ -0,0 +1,120 @@
+"""Tests for RequestIDLogger tween"""
+
+from pyramid import request
+from pyramid import testing
+
+from importlib import reload
+import unittest
+from unittest import mock
+import structlog
+import sentry_sdk
+import sys
+
+
+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"
+
+        self.some_random_dsn = "https://[email protected]/12345"
+
+    def tearDown(self):
+        testing.tearDown()
+
+    def test_request_id_in_sentry(self):
+        sentry_sdk.Hub.current._stack[-1][1].clear()
+        sentry_sdk.init(self.some_random_dsn)
+        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_sdk.Hub.current._stack[-1][1]._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_request_id_not_in_header_sentry(self):
+        sentry_sdk.Hub.current._stack[-1][1].clear()
+        sentry_sdk.init(self.some_random_dsn)
+        from pyramid_heroku.request_id import RequestIDLogger
+
+        RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+        self.assertEqual(sentry_sdk.Hub.current._stack[-1][1]._tags, {})
+
+    def test_request_id_not_in_header_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)
+
+    def test_smooth_run_without_sentry(self):
+
+        structlog.reset_defaults()
+
+        # remove sentry_sdk from modules to disable its import,
+        # in order to cover the code that runs without sentry
+        from pyramid_heroku import request_id
+
+        sys.modules["sentry_sdk"] = None
+        reload(request_id)
+
+        self.request.headers["X-Request-ID"] = self.request_id
+        request_id.RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+
+        # structlog is still there
+        context_class = structlog.get_config().get("context_class")
+        self.assertEqual(context_class._tl.dict_, {"request_id": self.request_id})
+
+        # revert sys.modules to its original state
+        sys.modules["sentry_sdk"] = sentry_sdk
+
+    def test_smooth_run_without_structlog(self):
+        sentry_sdk.Hub.current._stack[-1][1].clear()
+        sentry_sdk.init(self.some_random_dsn)
+
+        # remove structlog from modules to disable its import,
+        # in order to cover the code that runs without structlog
+        from pyramid_heroku import request_id
+
+        sys.modules["structlog"] = None
+        reload(request_id)
+
+        from pyramid_heroku import request_id
+
+        self.request.headers["X-Request-ID"] = self.request_id
+        request_id.RequestIDLogger(self.handler, self.registry)(self.request)
+        self.handler.assert_called_with(self.request)
+        self.assertEqual(
+            sentry_sdk.Hub.current._stack[-1][1]._tags, {"request_id": self.request_id}
+        )
+
+        # revert sys.modules to its original state
+        sys.modules["structlog"] = structlog