As applications grow and move into production, robust error handling and comprehensive logging become indispensable. This chapter focuses on setting up structured logging, handling custom exceptions, and providing graceful error responses in our FastAPI chat application.

Purpose of this Chapter

By the end of this chapter, you will:

  • Configure Python’s logging module for structured output.
  • Implement custom exception handlers for specific application errors.
  • Ensure that unhandled exceptions are caught and logged appropriately.
  • Understand best practices for logging sensitive information.

Concepts Explained: Structured Logging & Custom Exception Handling

Structured Logging

Traditional logging often outputs plain text messages. Structured logging outputs logs in a consistent, machine-readable format, typically JSON. This makes logs much easier to parse, filter, and analyze with log management tools (e.g., ELK Stack, Splunk, DataDog).

Key benefits:

  • Easier Searching & Filtering: Query logs by specific fields (e.g., user_id, endpoint).
  • Better Analytics: Generate metrics and dashboards from log data.
  • Context Rich: Attach relevant context (user ID, request ID, timestamp) to every log entry.

Custom Exception Handling

FastAPI automatically handles HTTPExceptions. However, you might have custom exceptions specific to your application’s business logic (e.g., RoomNotFoundError, UserAlreadyExistsError). By implementing custom exception handlers, you can control the exact HTTP response given for these specific errors, ensuring consistent client feedback.

Step-by-Step Tasks

1. Configure Structured Logging

Let’s create a new file app/logger_config.py to set up our logging. We’ll use the standard logging module and configure a custom formatter.

# app/logger_config.py

import logging
import json
from datetime import datetime

class JsonFormatter(logging.Formatter):
    """
    A custom logging formatter that outputs logs as JSON.
    """
    def format(self, record):
        log_entry = {
            "timestamp": datetime.fromtimestamp(record.created).isoformat(),
            "level": record.levelname,
            "message": record.getMessage(),
            "logger_name": record.name,
            "pathname": record.pathname,
            "lineno": record.lineno,
            "funcName": record.funcName,
            "process": record.process,
            "thread": record.thread,
        }

        # Add any extra attributes passed to the logger
        if hasattr(record, 'extra_data') and isinstance(record.extra_data, dict):
            log_entry.update(record.extra_data)

        # Include exception information if present
        if record.exc_info:
            log_entry["exc_info"] = self.formatException(record.exc_info)

        return json.dumps(log_entry)

def configure_logging():
    """
    Configures the root logger and optionally adds a file handler.
    """
    # Remove any existing handlers from the root logger
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    # Set up console handler with JSON formatter
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(JsonFormatter())
    logging.basicConfig(level=logging.INFO, handlers=[console_handler])

    # Optionally, add a file handler for persistent logs
    # file_handler = logging.FileHandler("app.log")
    # file_handler.setFormatter(JsonFormatter())
    # logging.root.addHandler(file_handler)

    # Disable default uvicorn access logs if you want full control
    # logging.getLogger("uvicorn.access").handlers = []
    # logging.getLogger("uvicorn.access").propagate = False

    # Example logger for our app
    logger = logging.getLogger("chat_app")
    logger.setLevel(logging.INFO)
    logger.info("Logging configured successfully.")
    return logger

# Call configure_logging once when this module is imported
app_logger = configure_logging()

Now, import this app_logger into app/main.py and use it.

# app/main.py (imports section)

# ... (other imports) ...
from .logger_config import app_logger # New import

# ... (rest of the code) ...

And use app_logger.info(...) or app_logger.error(...) instead of print():

# app/main.py (relevant parts)

# In register_user
    app_logger.info(f"User '{user.username}' registered successfully.", extra_data={"user_id": db_user.id})
    return db_user

# In login_for_access_token
    if not user or not Hasher.verify_password(form_data.password, user.hashed_password):
        app_logger.warning(f"Failed login attempt for username: {form_data.username}")
        raise HTTPException(...)
    app_logger.info(f"User '{user.username}' logged in successfully.", extra_data={"user_id": user.id})
    return {"access_token": access_token, "token_type": "bearer"}

# In websocket_room_endpoint connect
    app_logger.info(f"User '{current_user.username}' connected to room '{room_name}'.", extra_data={"user_id": current_user.id, "room_name": room_name})
    await manager.broadcast(welcome_message, room_name)

# In websocket_room_endpoint receive_text (when a message is sent)
            app_logger.info(
                f"User '{current_user.username}' sent message in room '{room_name}'.",
                extra_data={
                    "user_id": current_user.id,
                    "room_name": room_name,
                    "message_content_len": len(data) # Avoid logging full content if sensitive
                }
            )
            # ... (save message to DB) ...
            await manager.broadcast(full_message, room_name)
    except WebSocketDisconnect:
        manager.disconnect(websocket, room_name)
        app_logger.info(f"User '{current_user.username}' disconnected from room '{room_name}'.", extra_data={"user_id": current_user.id, "room_name": room_name})
        await manager.broadcast(f"User '{current_user.username}' left room '{room_name}'.", room_name)
    except HTTPException as e:
        app_logger.warning(f"WebSocket authentication failed for user {current_user.username if current_user else 'unknown'} in room '{room_name}': {e.detail}")
    except Exception as e:
        app_logger.error(f"Unexpected error in WebSocket for user '{current_user.username}' in room '{room_name}': {e}", exc_info=True)
        manager.disconnect(websocket, room_name)
        await manager.broadcast(f"User '{current_user.username}' left room '{room_name}' due to an error.", room_name)

2. Implement Custom Exception Handlers

Let’s define a custom exception for when a chat room is not found, and then register an exception handler for it.

# app/exceptions.py (new file)

class ChatRoomNotFoundException(HTTPException):
    def __init__(self, room_name: str):
        super().__init__(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Chat room '{room_name}' not found."
        )

class WebSocketAuthException(HTTPException):
    def __init__(self):
        super().__init__(
            status_code=status.WS_1008_POLICY_VIOLATION, # Or WS_1000_NORMAL_CLOSURE with a reason
            detail="WebSocket authentication failed. Please provide a valid token."
        )

Now, register this custom handler in app/main.py.

# app/main.py (updated with exception handlers)

# ... (imports) ...
from .exceptions import ChatRoomNotFoundException, WebSocketAuthException # New import
from starlette.requests import Request # New import
from starlette.responses import JSONResponse # New import

# ... (app = FastAPI() and manager = ConnectionManager() lines) ...


# Register custom exception handler for ChatRoomNotFoundException
@app.exception_handler(ChatRoomNotFoundException)
async def chat_room_not_found_exception_handler(request: Request, exc: ChatRoomNotFoundException):
    app_logger.warning(f"ChatRoomNotFoundException: {exc.detail} for path: {request.url.path}")
    return JSONResponse(
        status_code=exc.status_code,
        content={"message": exc.detail}
    )

# Generic handler for unexpected errors (to return JSON instead of default Starlette HTML error)
@app.exception_handler(500)
async def generic_server_error_handler(request: Request, exc: Exception):
    app_logger.error(f"Unhandled server error for path: {request.url.path}", exc_info=True)
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"message": "An unexpected server error occurred."}
    )

# ... (replace existing HTTPException for room not found with ChatRoomNotFoundException) ...
@app.get("/messages/{room_name}", response_model=List[schemas.MessageResponse])
async def get_chat_history_for_room(
    room_name: str,
    db: Session = Depends(get_db),
    current_user: models.User = Depends(get_current_user_db),
    skip: int = 0,
    limit: int = 100
):
    if room_name not in AVAILABLE_ROOMS:
        raise ChatRoomNotFoundException(room_name=room_name) # Raise custom exception

    # ... (rest of the code) ...

# ... (adjust get_current_active_user_ws to raise WebSocketAuthException on failure) ...
async def get_current_active_user_ws(websocket: WebSocket, db: Session = Depends(get_db)):
    # ... (token extraction) ...
    try:
        # ... (JWT decode and user lookup) ...
        if username is None:
            raise WebSocketAuthException() # Raise custom exception
        user = db.query(models.User).filter(models.User.username == username).first()
        if user is None:
            raise WebSocketAuthException() # Raise custom exception
        return user
    except JWTError:
        raise WebSocketAuthException() # Raise custom exception
    except HTTPException:
        raise
    except Exception as e:
        app_logger.error(f"Internal WebSocket authentication error: {e}", exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Internal authentication error")

# ... (in websocket_room_endpoint, catch WebSocketAuthException specifically for logging)
@app.websocket("/ws/{room_name}")
async def websocket_room_endpoint(
    websocket: WebSocket,
    room_name: str,
    db: Session = Depends(get_db),
    current_user: models.User = Depends(get_current_active_user_ws)
):
    if room_name not in AVAILABLE_ROOMS:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid chat room")
        app_logger.warning(f"Attempted connection to invalid room '{room_name}' by '{current_user.username if current_user else 'unknown'}'")
        return
    # ... (rest of websocket_room_endpoint) ...
    except WebSocketAuthException as e: # Catch the custom WebSocket auth exception
        app_logger.warning(f"WebSocket authentication failed for connection to room '{room_name}': {e.detail}")
        # FastAPI's dependency injection will close the WebSocket with the right code
        # if HTTPException is raised. No need for explicit close here.
    except Exception as e:
        app_logger.error(f"Unexpected error in WebSocket for user '{current_user.username}' in room '{room_name}': {e}", exc_info=True)
        manager.disconnect(websocket, room_name)
        await manager.broadcast(f"User '{current_user.username}' left room '{room_name}' due to an error.", room_name)

3. Test Logging and Error Handling

  1. Start the server: pipenv shell then uvicorn app.main:app --reload Observe your console. You should now see JSON-formatted log messages for server startup.

  2. Test HTTP Endpoint Error:

    • Try to access http://localhost:8000/messages/nonexistent-room.
    • You should get a 404 Not Found response with {"message": "Chat room 'nonexistent-room' not found."}.
    • Check your server console; you should see a WARNING level JSON log entry for ChatRoomNotFoundException.

  3. Test WebSocket Authentication Error:

    • Open client.html. Don’t log in. Try to connect to a room.
    • The WebSocket connection should fail. Check your browser’s developer console for WebSocket errors.
    • Check your server console; you should see a WARNING level JSON log entry for WebSocket authentication failed....

  4. Test Generic 500 Error (Optional):

    • To trigger a generic 500 error, you might temporarily introduce a bug that causes an unhandled exception in an HTTP endpoint (e.g., divide by zero in /health).
    • Access that endpoint. You should get a 500 Internal Server Error with {"message": "An unexpected server error occurred."}.
    • Crucially, your server console should show an ERROR level JSON log with exc_info (stack trace) included, which is invaluable for debugging.

Pitfalls and Best Practices

  • Sensitive Data in Logs: Never log raw passwords, API keys, or full JWTs. Only log metadata or masked values. In app_logger.info("User sent message...", extra_data={"message_content_len": len(data)}), we log only the length, not the content.
  • Logging Levels: Use appropriate logging levels (DEBUG, INFO, WARNING, ERROR, CRITICAL). DEBUG is verbose for development, INFO for general operations, ERROR for critical issues.
  • Centralized Logging: For production, integrate with a centralized logging solution (e.g., send logs to Logstash, Grafana Loki, cloud logging services) instead of just writing to a local file or console.
  • Uvicorn’s Default Logs: Uvicorn has its own access logs. You might want to disable them or reconfigure them to use your structured logging for consistency. (The configure_logging function includes commented-out lines for this).

Summary/Key Takeaways

You’ve successfully implemented structured logging and custom exception handling in your FastAPI chat application. Your server now produces rich, machine-readable logs and provides consistent, user-friendly error responses for both expected and unexpected issues. This makes your application significantly more robust, maintainable, and observable, moving it closer to a production-ready state.

The next step in making our application production-ready is to secure the WebSocket communication itself, building on our authentication.