From 4d000b63cb356daae4401d557029abed5fadad11 Mon Sep 17 00:00:00 2001
From: TannedCung <dx.tan@samsung.com>
Date: Fri, 5 Jul 2024 13:32:00 +0700
Subject: [PATCH] Add: TraceID for every logging

---
 cope2n-api/fwd/settings.py                            | 10 ++++++++--
 cope2n-api/fwd_api/celery_worker/client_connector.py  |  7 +++++--
 cope2n-api/fwd_api/middleware/local_storage.py        |  9 +++++++++
 .../middleware/logging_request_response_middleware.py | 11 ++++++++++-
 4 files changed, 32 insertions(+), 5 deletions(-)
 create mode 100644 cope2n-api/fwd_api/middleware/local_storage.py

diff --git a/cope2n-api/fwd/settings.py b/cope2n-api/fwd/settings.py
index 8548f92..a53ce27 100755
--- a/cope2n-api/fwd/settings.py
+++ b/cope2n-api/fwd/settings.py
@@ -13,6 +13,7 @@ import os
 from pathlib import Path
 import environ
 from django.urls import reverse_lazy
+from fwd_api.middleware.logging_request_response_middleware import TraceIDLogFilter
 
 # Build paths inside the project like this: BASE_DIR / 'subdir'.
 
@@ -262,13 +263,18 @@ CACHES = {
 LOGGING = {
     'version': 1,
     'disable_existing_loggers': False,
+    'filters': {
+        'trace_id': {
+            '()': TraceIDLogFilter,
+        },
+    },
     'formatters': {
         'verbose': {
-            'format': '{levelname} {asctime} {module} {message}',
+            'format': f'{{asctime}} - {{levelname}} - Trace_ID: {{trace_id}} - {{module}} - {{message}}',
             'style': '{',
         },
         'simple': {
-            'format': '{levelname} {message}',
+            'format': f'{{asctime}} - {{levelname}} - Trace_ID: {{trace_id}} - {{message}}',
             'style': '{',
         },
     },
diff --git a/cope2n-api/fwd_api/celery_worker/client_connector.py b/cope2n-api/fwd_api/celery_worker/client_connector.py
index 76e90df..f736c4b 100755
--- a/cope2n-api/fwd_api/celery_worker/client_connector.py
+++ b/cope2n-api/fwd_api/celery_worker/client_connector.py
@@ -2,6 +2,7 @@ from celery import Celery
 
 from fwd import settings
 from fwd_api.exception.exceptions import GeneralException
+from fwd_api.middleware.local_storage import get_current_request
 from kombu.utils.uuid import uuid
 from celery.utils.log import get_task_logger
 logger = get_task_logger(__name__)
@@ -127,8 +128,10 @@ class CeleryConnector:
     def send_task(self, name=None, args=None, countdown=None):
         if name not in self.task_routes or 'queue' not in self.task_routes[name]:
             raise GeneralException("System")
-        task_id = args[0] + "_" + uuid()[:4] if isinstance(args, tuple) and is_it_an_index(args[0]) else uuid()
+        # task_id = args[0] + "_" + uuid()[:4] if isinstance(args, tuple) and is_it_an_index(args[0]) else uuid()
+        request = get_current_request()
+        task_id = request.META.get('X-Trace-ID', uuid()) + "_" + uuid()[:4] if request else uuid()
         logger.info(f"SEND task name: {name} - {task_id} | args: {args} | countdown: {countdown}")
         return self.app.send_task(name, args, queue=self.task_routes[name]['queue'], expires=300, countdown=countdown, task_id=task_id)
 
-c_connector = CeleryConnector()
+c_connector = CeleryConnector()
\ No newline at end of file
diff --git a/cope2n-api/fwd_api/middleware/local_storage.py b/cope2n-api/fwd_api/middleware/local_storage.py
new file mode 100644
index 0000000..d87e795
--- /dev/null
+++ b/cope2n-api/fwd_api/middleware/local_storage.py
@@ -0,0 +1,9 @@
+from threading import local
+
+_thread_locals = local()
+
+def get_current_request():
+    return getattr(_thread_locals, 'request', None)
+
+def set_current_request(request):
+    _thread_locals.request = request
diff --git a/cope2n-api/fwd_api/middleware/logging_request_response_middleware.py b/cope2n-api/fwd_api/middleware/logging_request_response_middleware.py
index 3ec0730..3614fe2 100644
--- a/cope2n-api/fwd_api/middleware/logging_request_response_middleware.py
+++ b/cope2n-api/fwd_api/middleware/logging_request_response_middleware.py
@@ -2,6 +2,7 @@ import logging
 import uuid
 
 from django.utils.deprecation import MiddlewareMixin
+from .local_storage import set_current_request, get_current_request
 
 logger = logging.getLogger(__name__)
 
@@ -9,6 +10,7 @@ class LoggingMiddleware(MiddlewareMixin):
     def process_request(self, request):
         trace_id = request.headers.get('X-Trace-ID', str(uuid.uuid4()))
         request.META['X-Trace-ID'] = trace_id
+        set_current_request(request)
 
         request_body = ""
         content_type = request.headers.get("Content-Type", "")
@@ -35,4 +37,11 @@ class LoggingMiddleware(MiddlewareMixin):
     def process_exception(self, request, exception):
         trace_id = request.META.get('X-Trace-ID', str(uuid.uuid4()))
         logger.error(f"Exception: {request.method} {request.path} | Trace ID: {trace_id} | "
-                     f"Error: {exception}")
\ No newline at end of file
+                     f"Error: {exception}")
+        
+class TraceIDLogFilter(logging.Filter):
+    def filter(self, record):
+        request = get_current_request()
+        trace_id = request.META.get('X-Trace-ID', 'unknown') if request else 'unknown'
+        record.trace_id = trace_id
+        return True
\ No newline at end of file