Merge pull request #1207 from dod-ccpo/logging-improvements

Logging improvements

atst/queue.py

@@ -1,5 +1,6 @@
 from celery import Celery
 
+
 celery = Celery(__name__)
 
 

atst/utils/logging.py

@@ -10,6 +10,7 @@ class RequestContextFilter(logging.Filter):
         if has_request_context():
             if getattr(g, "current_user", None):
                 record.user_id = str(g.current_user.id)
+                record.dod_edipi = g.current_user.dod_id
 
             if request.environ.get("HTTP_X_REQUEST_ID"):
                 record.request_id = request.environ.get("HTTP_X_REQUEST_ID")
@@ -28,16 +29,28 @@ class JsonFormatter(logging.Formatter):
         ("version", lambda r: 1),
         ("request_id", lambda r: r.__dict__.get("request_id")),
         ("user_id", lambda r: r.__dict__.get("user_id")),
+        ("dod_edipi", lambda r: r.__dict__.get("dod_edipi")),
         ("severity", lambda r: r.levelname),
         ("tags", lambda r: r.__dict__.get("tags")),
-        ("message", lambda r: r.msg),
         ("audit_event", lambda r: r.__dict__.get("audit_event")),
     ]
 
-    def format(self, record):
-        message_dict = {"source": "atst"}
+    def __init__(self, *args, source="atst", **kwargs):
+        self.source = source
+        super().__init__(self)
+
+    def format(self, record, *args, **kwargs):
+        message_dict = {"source": self.source}
+
         for field, func in self._DEFAULT_RECORD_FIELDS:
-            message_dict[field] = func(record)
+            result = func(record)
+            if result:
+                message_dict[field] = result
+
+        if record.args:
+            message_dict["message"] = record.msg % record.args
+        else:
+            message_dict["message"] = record.msg
 
         if record.__dict__.get("exc_info") is not None:
             message_dict["details"] = {

celery_worker.py

@@ -1,7 +1,19 @@
 #!/usr/bin/env python
+import logging
 
 from atst.app import celery, make_app, make_config
+from celery.signals import after_setup_task_logger
+
+from atst.utils.logging import JsonFormatter
 
 config = make_config()
 app = make_app(config)
 app.app_context().push()
+
+
+@after_setup_task_logger.connect
+def setup_task_logger(*args, **kwargs):
+    if app.config.get("LOG_JSON"):
+        logger = logging.getLogger()
+        for handler in logger.handlers:
+            handler.setFormatter(JsonFormatter(source="queue"))

tests/utils/test_logging.py

@@ -69,10 +69,12 @@ def test_request_context_filter(logger, log_stream_content, request_ctx, monkeyp
 
     user = Mock(spec=["id"])
     user.id = user_uuid
+    user.dod_id = "5678901234"
 
     monkeypatch.setattr("atst.utils.logging.g", Mock(current_user=user))
     request_ctx.request.environ["HTTP_X_REQUEST_ID"] = request_uuid
     logger.info("this user is doing something")
     log = json.loads(log_stream_content())
     assert log["user_id"] == str(user_uuid)
+    assert log["dod_edipi"] == str(user.dod_id)
     assert log["request_id"] == request_uuid