| 1 | from fastapi import HTTPException |
| 2 | from fastapi import Request |
| 3 | from fastapi.responses import JSONResponse |
| 4 | from sqlalchemy.ext.asyncio import AsyncSession |
| 5 | |
| 6 | from app.auth.utils import AuthHandler |
| 7 | from app.db.db_session import async_engine |
| 8 | from app.utils import Logger |
| 9 | |
| 10 | EXCLUDED_PATHS = ["/auth/token", "/auth/register", "/auth/refresh", "/influxdb/alerts?days=1&status=active&exclude_ok=true"] |
| 11 | INTERNAL_SERVER_ERROR = 500 |
| 12 | |
| 13 | |
| 14 | async def process_request(request: Request, call_next, session, logger_instance): |
| 15 | """ |
| 16 | Process the incoming request and log the user ID. |
| 17 | |
| 18 | Args: |
| 19 | request (Request): The incoming request. |
| 20 | call_next (Callable): The next middleware or endpoint to call. |
| 21 | session: The session object. |
| 22 | logger_instance: The logger instance. |
| 23 | |
| 24 | Returns: |
| 25 | Tuple: A tuple containing the response and the user ID. |
| 26 | """ |
| 27 | response = await call_next(request) |
| 28 | user_id = await logger_instance.get_user_id_from_request(request) |
| 29 | return response, user_id |
| 30 | |
| 31 | |
| 32 | def is_excluded_path(path: str) -> bool: |
| 33 | """ |
| 34 | Check if the given path is in the list of excluded paths. |
| 35 | |
| 36 | Args: |
| 37 | path (str): The path to check. |
| 38 | |
| 39 | Returns: |
| 40 | bool: True if the path is in the excluded paths list, False otherwise. |
| 41 | """ |
| 42 | return path in EXCLUDED_PATHS |
| 43 | |
| 44 | |
| 45 | async def handle_exception(e, user_id, request, logger_instance): |
| 46 | """ |
| 47 | Handles exceptions that occur during request processing. |
| 48 | |
| 49 | Args: |
| 50 | e (Exception): The exception that occurred. |
| 51 | user_id (int): The ID of the user associated with the request. |
| 52 | request (Request): The request object. |
| 53 | logger_instance (Logger): An instance of the logger. |
| 54 | |
| 55 | Returns: |
| 56 | JSONResponse: The response containing the error message. |
| 57 | """ |
| 58 | try: |
| 59 | user_id = await logger_instance.get_user_id_from_request(request) if user_id is None else user_id |
| 60 | except HTTPException as http_exc: |
| 61 | return JSONResponse( |
| 62 | status_code=http_exc.status_code, |
| 63 | content={"message": str(http_exc), "success": False}, |
| 64 | ) |
| 65 | await logger_instance.log_error(user_id, request, e) |
| 66 | status_code = e.status_code if isinstance(e, HTTPException) else INTERNAL_SERVER_ERROR |
| 67 | return JSONResponse( |
| 68 | status_code=status_code, |
| 69 | content={"message": str(e), "success": False}, |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | async def log_requests(request: Request, call_next): |
| 74 | """ |
| 75 | Middleware function to log incoming requests and their responses. |
| 76 | |
| 77 | Args: |
| 78 | request (Request): The incoming request object. |
| 79 | call_next (Callable): The next middleware or endpoint to call. |
| 80 | |
| 81 | Returns: |
| 82 | The response generated by the next middleware or endpoint. |
| 83 | """ |
| 84 | if request.method == "OPTIONS": |
| 85 | return await call_next(request) |
| 86 | |
| 87 | async with AsyncSession(async_engine) as session: |
| 88 | logger_instance = Logger(session, AuthHandler()) |
| 89 | user_id = None |
| 90 | |
| 91 | try: |
| 92 | if not is_excluded_path(request.url.path): |
| 93 | response, user_id = await process_request( |
| 94 | request, |
| 95 | call_next, |
| 96 | session, |
| 97 | logger_instance, |
| 98 | ) |
| 99 | else: |
| 100 | response = await call_next(request) |
| 101 | except Exception as e: |
| 102 | return await handle_exception(e, user_id, request, logger_instance) |
| 103 | |
| 104 | await logger_instance.log_route_access(user_id, request, response) |
| 105 | |
| 106 | return response |