Coverage for app/core/middleware.py: 100%

1import time 

2from typing import Awaitable, Callable 

3import uuid 

4from fastapi import Request, Response 

5from starlette.middleware.base import BaseHTTPMiddleware 

6from app.core.context import correlation_id_context 

7from app.core.logging import Logger 

8 

9 

10class CorrelationIDMiddleware(BaseHTTPMiddleware): 

11 """ 

12 Middleware to generate and manage a unique `correlation` ID for each request. 

13 This correlation ID is stored in a context variable, allowing it to be accessed 

14 throughout the request lifecycle for logging purposes. 

15 """ 

16 

17 async def dispatch( 

18 self, request: Request, call_next: Callable[[Request], Awaitable[Response]] 

19 ) -> Response: 

20 correlation_id = request.headers.get("X-Correlation-Id", str(uuid.uuid4())) 

21 token = correlation_id_context.set(correlation_id) 

22 

23 try: 

24 response = await call_next(request) 

25 response.headers["X-Correlation-Id"] = correlation_id 

26 return response 

27 finally: 

28 correlation_id_context.reset(token) 

29 

30 

31class RequestLoggingMiddleware(BaseHTTPMiddleware): 

32 """ 

33 Middleware to log incoming HTTP requests and their corresponding responses. 

34 It captures the HTTP method, request path, response status code, and the duration of the request. 

35 """ 

36 

37 async def dispatch( 

38 self, request: Request, call_next: Callable[[Request], Awaitable[Response]] 

39 ) -> Response: 

40 start = time.perf_counter() 

41 response = await call_next(request) 

42 duration_ms = round((time.perf_counter() - start) * 1000, 2) 

43 

44 Logger.log( 

45 "completed_request", 

46 method=request.method, 

47 path=request.url.path, 

48 status_code=response.status_code, 

49 duration_ms=duration_ms, 

50 ) 

51 

52 return response