Adding unique trace id (#45)

omni_rpc/custom_logging.py

@@ -3,13 +3,16 @@
 
 from pythonjsonlogger import jsonlogger
 
+trace_id = None
+
+
 def get_trace_id():
     return uuid.uuid4()
 
 
 class ContextFilter(logging.Filter):
     def filter(self, record):
-        record.trace_id = get_trace_id()
+        record.trace_id = trace_id
         return True
 
 

omni_rpc/main.py

@@ -1,10 +1,17 @@
 from fastapi import FastAPI
+from fastapi import Request
 
-from omni_rpc.custom_logging import logger
+import omni_rpc.custom_logging as custom_logging
 from omni_rpc.routers import proxy_rpc
 
 app = FastAPI()
 
 
-logger.info("Starting Omni RPC server")
+@app.middleware("http")
+async def add_process_time_header(request: Request, call_next):
+    custom_logging.trace_id = custom_logging.get_trace_id()
+    response = await call_next(request)
+    return response
+
+
 app.include_router(proxy_rpc.router)