Add request id middleware

Author: Raekkeri

README.md

@@ -60,7 +60,20 @@ MIDDLEWARE = [
 ]
 ```
 
-This will start storing the request logs using the default `STORAGE_CLASS`, which in fact just uses Python logger named `requestlogs`. Now you can, for example, redirect these logs to a file with the following `LOGGING` configuration:
+Set `'requestlogs.views.exception_handler'` as rest_framework's exception handler
+(this will make sure requestlog entry has all possible data available about the
+request in case of a 500 error):
+
+```python
+REST_FRAMEWORK={
+    ...
+    'EXCEPTION_HANDLER': 'requestlogs.views.exception_handler',
+}
+```
+
+The middleware is now ready to start storing requestlog entries using the default
+`STORAGE_CLASS`, which in fact just uses Python logger named `requestlogs`. Now you can,
+for example, redirect these logs to a file with the following `LOGGING` configuration:
 
 ```python
 LOGGING = {
@@ -77,6 +90,7 @@ LOGGING = {
         'requestlogs': {
             'handlers': ['requestlogs_to_file'],
             'level': 'INFO',
+            'propagate': False,
         },
     },
 }
@@ -103,3 +117,95 @@ REQUESTLOGS = {
   - List of keys in request/response data which will be replaced with `'***'` in the stored entry.
 - **ATTRIBUTE_NAME**
   - django-requestlogs internally attaches the entry object to the Django request object, and uses this attribute name. Override if it causes collisions.
+
+
+# Logging with Request ID
+
+django-requestlogs also contains a middleware and logging helpers to associate a unique
+identifier (uuid) for each request (implemented using `threading.local()`).
+The request id can be added to the standard logging messages (Django application logs)
+by specifying a custom formatter and using the provided logging filter.
+The request id can be stored to requestlog entries as well.
+The middleware to enable the request id logging can be used separately from the
+core requestlogs middleware.
+
+The feature is enabled by adding `requestlogs.middleware.RequestLogsMiddleware`
+to the `MIDDLEWARE` setting:
+
+```python
+MIDDLEWARE = [
+    ...
+    'requestlogs.middleware.RequestLogsMiddleware',
+    'requestlogs.middleware.RequestIdMiddleware',
+]
+```
+
+Once installed, the application logs should start showing messages with a format such as
+the following:
+
+```
+2019-07-18 11:56:07,261 INFO 954fb004fb404751a2fa33326101442c urls:31 Handling GET request
+2019-07-18 11:56:07,262 DEBUG 954fb004fb404751a2fa33326101442c urls:32 No parameters given
+2019-07-18 11:56:07,262 INFO 954fb004fb404751a2fa33326101442c urls:33 All good
+```
+
+The middleware has some configuration possiblities:
+
+
+```python
+REQUESTLOGS = {
+    ...
+    'REQUEST_ID_HTTP_HEADER': 'X_DJANGO_REQUEST_ID',
+    'REQUEST_ID_ATTRIBUTE_NAME': 'request_id',
+}
+```
+- **REQUEST_ID_HTTP_HEADER**
+  - If set, the value of this request header is used as request id (instead of it being
+    randomly generated). Must be a valid uuid.
+- **REQUEST_ID_ATTRIBUTE_NAME**
+  - The attribute name which is used internally to attach request id to
+    `threading.locals()`. Override if it causes collisions.
+
+To add the request id to logging messages of your Django application, use the provided
+logging filter and include `request_id` to the log formatter.
+Here is the complete logging configuration:
+
+```python
+LOGGING = {
+    'version': 1,
+    'disable_existing_loggers': False,
+    'handlers': {
+        'requestlogs_to_file': {
+            'level': 'INFO',
+            'class': 'logging.FileHandler',
+            'filename': '/tmp/requestlogs.log',
+        },
+        'root': {
+            'class': 'logging.StreamHandler',
+            'filters': ['request_id_context'],
+            'formatter': 'verbose',
+        },
+    },
+    'loggers': {
+        '': {
+            'handlers': ['root'],
+            'level': 'DEBUG',
+        },
+        'requestlogs': {
+            'handlers': ['requestlogs_to_file'],
+            'level': 'INFO',
+            'propagate': False,
+        },
+    },
+    'filters': {
+        'request_id_context': {
+            '()': 'requestlogs.logging.RequestIdContext',
+        },
+    },
+    'formatters': {
+        'verbose': {
+            'format': '%(asctime)s %(levelname)s %(request_id)s %(module)s:%(lineno)s %(message)s'
+        },
+    },
+}
+```

requestlogs/base.py

@@ -8,6 +8,8 @@
     'ENTRY_CLASS': 'requestlogs.entries.RequestLogEntry',
     'STORAGE_CLASS': 'requestlogs.storages.LoggingStorage',
     'SECRETS': ['password', 'token'],
+    'REQUEST_ID_ATTRIBUTE_NAME': 'request_id',
+    'REQUEST_ID_HTTP_HEADER': None,
 }
 
 

requestlogs/entries.py

@@ -5,6 +5,7 @@
 from rest_framework.request import Request
 
 from .base import SETTINGS
+from .logging import get_request_id
 from .utils import remove_secrets, get_client_ip
 
 
@@ -28,6 +29,10 @@ def query_params(self):
     def full_path(self):
         return self.request.get_full_path()
 
+    @property
+    def request_id(self):
+        return get_request_id()
+
 
 class DRFRequestHandler(RequestHandler):
     @property

requestlogs/logging.py

@@ -0,0 +1,33 @@
+import logging
+import uuid
+import threading
+
+from .base import SETTINGS
+
+local = threading.local()
+
+
+def get_request_id():
+    return getattr(local, SETTINGS['REQUEST_ID_ATTRIBUTE_NAME'], '')
+
+
+def set_request_id(_uuid=None):
+    _uuid = _uuid or uuid.uuid4().hex
+    setattr(local, SETTINGS['REQUEST_ID_ATTRIBUTE_NAME'], _uuid)
+    return _uuid
+
+
+def validate_uuid(_uuid):
+    try:
+        val = uuid.UUID(_uuid, version=4)
+    except (ValueError, TypeError):
+        return None
+
+    return _uuid if val.hex == _uuid else None
+
+
+class RequestIdContext(logging.Filter):
+    def filter(self, record):
+        setattr(record, SETTINGS['REQUEST_ID_ATTRIBUTE_NAME'],
+                get_request_id())
+        return True

requestlogs/middleware.py

@@ -1,3 +1,5 @@
+from .base import SETTINGS
+from .logging import set_request_id, validate_uuid
 from . import get_requestlog_entry
 
 
@@ -14,3 +16,13 @@ def __call__(self, request):
         response = self.get_response(request)
         get_requestlog_entry(request).finalize(response)
         return response
+
+
+class RequestIdMiddleware(object):
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        reuse_request_id = request.META.get(SETTINGS['REQUEST_ID_HTTP_HEADER'])
+        set_request_id(validate_uuid(reuse_request_id))
+        return self.get_response(request)

tests/test_views.py

@@ -1,4 +1,6 @@
-from unittest.mock import patch
+import io
+import logging
+from unittest.mock import patch, Mock
 
 import django
 from django.contrib.auth import get_user_model
@@ -17,6 +19,7 @@
 from rest_framework.views import APIView
 
 from requestlogs import get_requestlog_entry
+from requestlogs.logging import RequestIdContext
 from requestlogs.storages import BaseEntrySerializer
 
 
@@ -52,6 +55,11 @@ def retrieve(self, request):
     def create(self, request):
         return Response({})
 
+    def w_logging(self, request):
+        logger = logging.getLogger('request_id_test')
+        logger.info('GET with logging ({})'.format(request.GET.get('q')))
+        return Response({})
+
 
 class SimpleAuth(BaseAuthentication):
     def authenticate(self, request):
@@ -100,6 +108,7 @@ def api_view_function(request):
     url(r'^viewset/1/?$', ViewSet.as_view({'get': 'retrieve'})),
     url(r'^func/?$', api_view_function),
     url(r'^error/?$', ServerErrorView.as_view()),
+    url(r'^logging/?$', ViewSet.as_view({'get': 'w_logging'})),
 ]
 
 
@@ -389,3 +398,127 @@ def test_500_response_without_custom_exception_handler(self):
                 KeyError, self.client.post, '/error', {'pay': 'load'})
             self.expected['request']['data'] = None
             self.assert_stored(mocked_store, self.expected)
+
+
+class RequestIdStorage(TestStorage):
+    class serializer_class(serializers.Serializer):
+        request_id = serializers.CharField(source='request.request_id')
+
+
+class LoggingMixin(object):
+    def _setup_logging(self):
+        log_format = '{levelname} {request_id} {message}'
+        stream = io.StringIO('')
+        handler = logging.StreamHandler(stream)
+        handler.setFormatter(logging.Formatter(log_format, style='{'))
+        logger = logging.getLogger('request_id_test')
+        logger.setLevel(logging.INFO)
+        logger.addHandler(handler)
+        logger.addFilter(RequestIdContext())
+        self._log_stream = stream
+
+    def _assert_logged_lines(self, lines):
+        self._log_stream.seek(0)
+        raw = self._log_stream.read()
+        assert raw.endswith('\n')
+        logged_lines = [i for i in raw.split('\n')][:-1]
+        assert logged_lines == lines
+
+
+@override_settings(
+    ROOT_URLCONF=__name__,
+    REQUESTLOGS={'STORAGE_CLASS': 'tests.test_views.RequestIdStorage'},
+)
+@modify_settings(MIDDLEWARE={
+    'append': [
+        'requestlogs.middleware.RequestLogsMiddleware',
+        'requestlogs.middleware.RequestIdMiddleware',
+    ],
+})
+class TestRequestIdMiddleware(LoggingMixin, APITestCase):
+    def test_request_id_generated(self):
+        with patch('tests.test_views.RequestIdStorage.do_store') \
+                as mocked_store, patch('uuid.uuid4') as mocked_uuid:
+            mocked_uuid.side_effect = [Mock(hex='12345dcba')]
+            response = self.client.get('/')
+            assert mocked_store.call_args[0][0] == {'request_id': '12345dcba'}
+
+    def test_python_logging_with_request_id(self):
+        # First build logging setup which outputs entries with request_id.
+        # We cannot use `self.assertLogs`, because it uses the default
+        # formatter, and so request_id wouldn't be seen in log output.
+        self._setup_logging()
+
+        # Now that logging is set up to capture formatted log entries into the
+        # `stream`, we can do the request and finally check the logged result.
+        with patch('uuid.uuid4') as mocked_uuid:
+            mocked_uuid.side_effect = [Mock(hex='12345dcba')]
+            self.client.get('/logging?q=yes')
+
+        self._assert_logged_lines(['INFO 12345dcba GET with logging (yes)'])
+
+    def test_python_logging_and_requestlogs_entry(self):
+        self._setup_logging()
+
+        with patch('tests.test_views.RequestIdStorage.do_store') \
+                as mocked_store, patch('uuid.uuid4') as mocked_uuid, \
+                patch('uuid.uuid4') as mocked_uuid:
+            mocked_uuid.side_effect = [
+                Mock(hex='12345dcba'), Mock(hex='ffc999123')]
+            self.client.get('/logging?q=1')
+            self.client.get('/logging?q=2')
+
+            call_args1, call_args2 = mocked_store.call_args_list
+            assert call_args1[0][0] == {'request_id': '12345dcba'}
+            assert call_args2[0][0] == {'request_id': 'ffc999123'}
+
+        self._assert_logged_lines([
+            'INFO 12345dcba GET with logging (1)',
+            'INFO ffc999123 GET with logging (2)',
+        ])
+
+
+@override_settings(
+    ROOT_URLCONF=__name__,
+    REQUESTLOGS={
+        'STORAGE_CLASS': 'tests.test_views.RequestIdStorage',
+        'REQUEST_ID_HTTP_HEADER': 'X_DJANGO_REQUEST_ID',
+    },
+)
+@modify_settings(MIDDLEWARE={
+    'append': [
+        'requestlogs.middleware.RequestLogsMiddleware',
+        'requestlogs.middleware.RequestIdMiddleware',
+    ],
+})
+class TestReuseRequestId(LoggingMixin, APITestCase):
+    def test_reuse_request_id(self):
+        self._setup_logging()
+
+        uuid = '6359abe9f7d849e09a324791c6a6c976'
+        self.client.credentials(X_DJANGO_REQUEST_ID=uuid)
+
+        with patch('tests.test_views.RequestIdStorage.do_store') \
+                as mocked_store, patch('uuid.uuid4') as mocked_uuid, \
+                patch('uuid.uuid4') as mocked_uuid:
+            mocked_uuid.side_effect = []
+            response = self.client.get('/logging?q=p')
+            assert mocked_store.call_args[0][0] == {'request_id': uuid}
+
+        self._assert_logged_lines([
+            'INFO 6359abe9f7d849e09a324791c6a6c976 GET with logging (p)'])
+
+    def test_bad_request_id(self):
+        self.client.credentials(X_DJANGO_REQUEST_ID='BAD')
+        self._test_request_id_generated()
+
+    def test_no_request_id_present(self):
+        self._test_request_id_generated()
+
+    def _test_request_id_generated(self):
+        with patch('tests.test_views.RequestIdStorage.do_store') \
+                as mocked_store, patch('uuid.uuid4') as mocked_uuid, \
+                patch('uuid.uuid4') as mocked_uuid:
+            mocked_uuid.side_effect = [Mock(hex='12345dcba')]
+            response = self.client.get('/')
+            assert mocked_store.call_args[0][0] == {'request_id': '12345dcba'}