From 5d05c146d6a084713262b73a2d397d42320bdddd Mon Sep 17 00:00:00 2001
From: dandds <dan-ctr@friends.dds.mil>
Date: Fri, 29 Mar 2019 15:47:44 -0400
Subject: [PATCH] context filter for adding additional data to logs

---
 atst/utils/logging.py       | 14 +++++++++++---
 tests/utils/test_logging.py | 19 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/atst/utils/logging.py b/atst/utils/logging.py
index 0f097b0b..9c165610 100644
--- a/atst/utils/logging.py
+++ b/atst/utils/logging.py
@@ -2,10 +2,18 @@ import datetime
 import json
 import logging
 
+from flask import g, request
 
-class ContextFilter(logging.Filter):
-    # this should impart the request_id and user_id if available
-    pass
+
+class RequestContextFilter(logging.Filter):
+    def filter(self, record):
+        if getattr(g, "current_user", None):
+            record.user_id = str(g.current_user.id)
+
+        if request.environ.get("HTTP_X_REQUEST_ID"):
+            record.request_id = request.environ.get("HTTP_X_REQUEST_ID")
+
+        return True
 
 
 def epoch_to_iso8601(ts):
diff --git a/tests/utils/test_logging.py b/tests/utils/test_logging.py
index 79a30c8c..cb8198d3 100644
--- a/tests/utils/test_logging.py
+++ b/tests/utils/test_logging.py
@@ -1,10 +1,13 @@
 from io import StringIO
 import json
 import logging
+from uuid import uuid4
 
 import pytest
 
-from atst.utils.logging import JsonFormatter
+from atst.utils.logging import JsonFormatter, RequestContextFilter
+
+from tests.factories import UserFactory
 
 
 @pytest.fixture
@@ -30,6 +33,8 @@ def logger(log_stream):
     logHandler = logging.StreamHandler(log_stream)
     formatter = JsonFormatter()
     logHandler.setFormatter(formatter)
+    logger.setLevel(logging.INFO)
+    logger.addFilter(RequestContextFilter())
     logger.addHandler(logHandler)
 
     return logger
@@ -55,3 +60,15 @@ def test_json_formatter_for_exceptions(logger, log_stream_content):
     log = json.loads(log_stream_content())
     assert log["severity"] == "ERROR"
     assert log.get("details")
+
+
+def test_request_context_filter(logger, log_stream_content, request_ctx):
+    user = UserFactory.create()
+    uuid = str(uuid4())
+
+    request_ctx.g.current_user = user
+    request_ctx.request.environ["HTTP_X_REQUEST_ID"] = uuid
+    logger.info("this user is doing something")
+    log = json.loads(log_stream_content())
+    assert log["user_id"] == str(user.id)
+    assert log["request_id"] == uuid