middleware.py

Author: MehrazRumman

.gitignore

@@ -205,4 +205,3 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
-.pre-commit-config.yaml

.pre-commit-config.yaml

@@ -0,0 +1,14 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-toml
+      - id: check-merge-conflict
+
+  - repo: https://github.com/psf/black
+    rev: 24.10.0
+    hooks:
+      - id: black

fastapi_observer/__init__.py

@@ -1,5 +1,6 @@
 from .config import ObserverConfig
 from .logger import build_logger, log_event
+from .middleware import ObserverMiddleware
 from .models import LogEvent
 
-__all__ = ["ObserverConfig", "LogEvent", "build_logger", "log_event"]
+__all__ = ["ObserverConfig", "ObserverMiddleware", "LogEvent", "build_logger", "log_event"]

fastapi_observer/middleware.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import json
+import logging
+import time
+import uuid
+from typing import Any, Callable
+
+from .config import ObserverConfig
+from .logger import build_logger, log_event
+from .models import LogEvent
+
+try:
+    from starlette.middleware.base import BaseHTTPMiddleware
+    from starlette.requests import Request
+    from starlette.responses import Response
+except ModuleNotFoundError as exc:  # pragma: no cover - guarded runtime fallback
+    _IMPORT_ERROR = exc
+else:
+    _IMPORT_ERROR = None
+
+
+def _redact_sensitive(value: Any, redact_fields: set[str]) -> Any:
+    if isinstance(value, dict):
+        redacted: dict[str, Any] = {}
+        for key, nested_value in value.items():
+            if str(key).lower() in redact_fields:
+                redacted[key] = "***"
+            else:
+                redacted[key] = _redact_sensitive(nested_value, redact_fields)
+        return redacted
+    if isinstance(value, list):
+        return [_redact_sensitive(item, redact_fields) for item in value]
+    return value
+
+
+def _redact_headers(headers: dict[str, str], redact_headers: set[str]) -> dict[str, str]:
+    result: dict[str, str] = {}
+    for key, value in headers.items():
+        if key.lower() in redact_headers:
+            result[key] = "***"
+        else:
+            result[key] = value
+    return result
+
+
+def _to_log_level(status_code: int) -> str:
+    if status_code >= 500:
+        return "ERROR"
+    if status_code >= 400:
+        return "WARNING"
+    return "INFO"
+
+
+if _IMPORT_ERROR is None:
+
+    class ObserverMiddleware(BaseHTTPMiddleware):
+        def __init__(
+            self,
+            app: Any,
+            *,
+            config: ObserverConfig | None = None,
+            logger: logging.Logger | None = None,
+        ) -> None:
+            super().__init__(app)
+            self.config = config or ObserverConfig()
+            self.logger = logger or build_logger(self.config)
+
+        async def dispatch(
+            self,
+            request: Request,
+            call_next: Callable[[Request], Any],
+        ) -> Response:
+            if not self.config.enabled:
+                return await call_next(request)
+
+            path = request.url.path
+            method = request.method
+            if not self.config.should_log_method(method) or not self.config.should_log_path(path):
+                return await call_next(request)
+
+            correlation_id = self._resolve_correlation_id(request)
+            request_body_for_log: Any | None = None
+            request_for_next = request
+            if self.config.log_request_body:
+                body = await request.body()
+                request_body_for_log = self._format_body_for_log(body)
+                request_for_next = Request(request.scope, receive=_build_body_replay_receive(body))
+
+            start = time.perf_counter()
+            try:
+                response = await call_next(request_for_next)
+            except Exception as exc:
+                duration_ms = round((time.perf_counter() - start) * 1000, 3)
+                log_event(
+                    self.logger,
+                    self.config,
+                    LogEvent(
+                        level="ERROR",
+                        message="HTTP request failed",
+                        method=method,
+                        path=path,
+                        status_code=500,
+                        duration_ms=duration_ms,
+                        correlation_id=correlation_id,
+                        error=str(exc),
+                        metadata=self._build_metadata(
+                            request,
+                            request_body=request_body_for_log,
+                        ),
+                    ),
+                )
+                raise
+
+            duration_ms = round((time.perf_counter() - start) * 1000, 3)
+            response_body_for_log = (
+                self._extract_response_body_for_log(response)
+                if self.config.log_response_body
+                else None
+            )
+
+            log_event(
+                self.logger,
+                self.config,
+                LogEvent(
+                    level=_to_log_level(response.status_code),
+                    message="HTTP request completed",
+                    method=method,
+                    path=path,
+                    status_code=response.status_code,
+                    duration_ms=duration_ms,
+                    correlation_id=correlation_id,
+                    metadata=self._build_metadata(
+                        request,
+                        request_body=request_body_for_log,
+                        response_body=response_body_for_log,
+                    ),
+                ),
+            )
+
+            if correlation_id:
+                response.headers.setdefault(self.config.correlation_id_header, correlation_id)
+            return response
+
+        def _resolve_correlation_id(self, request: Request) -> str | None:
+            value = request.headers.get(self.config.correlation_id_header)
+            if value:
+                request.state.correlation_id = value
+                return value
+            if self.config.generate_correlation_id:
+                generated = str(uuid.uuid4())
+                request.state.correlation_id = generated
+                return generated
+            return None
+
+        def _build_metadata(
+            self,
+            request: Request,
+            *,
+            request_body: Any | None = None,
+            response_body: Any | None = None,
+        ) -> dict[str, Any]:
+            metadata: dict[str, Any] = {
+                "query_params": dict(request.query_params),
+            }
+            if request.client:
+                metadata["client"] = {
+                    "host": request.client.host,
+                    "port": request.client.port,
+                }
+            if self.config.log_headers:
+                metadata["headers"] = _redact_headers(
+                    dict(request.headers),
+                    self.config.redact_headers,
+                )
+            if request_body is not None:
+                metadata["request_body"] = request_body
+            if response_body is not None:
+                metadata["response_body"] = response_body
+            return metadata
+
+        def _extract_response_body_for_log(self, response: Response) -> Any | None:
+            body = getattr(response, "body", None)
+            if body is None:
+                return None
+            if isinstance(body, bytes):
+                return self._format_body_for_log(body)
+            if isinstance(body, str):
+                return self._format_body_for_log(body.encode("utf-8"))
+            return None
+
+        def _format_body_for_log(self, body: bytes) -> Any:
+            body_size = len(body)
+            body_for_parse = body[: self.config.max_body_bytes]
+            parsed: Any
+            try:
+                parsed = json.loads(body_for_parse.decode("utf-8"))
+            except (UnicodeDecodeError, json.JSONDecodeError):
+                parsed = body_for_parse.decode("utf-8", errors="replace")
+
+            redacted = _redact_sensitive(parsed, self.config.redact_fields)
+            if body_size <= self.config.max_body_bytes:
+                return redacted
+            return {
+                "truncated": True,
+                "original_size": body_size,
+                "content": redacted,
+            }
+
+
+    def _build_body_replay_receive(body: bytes) -> Callable[[], Any]:
+        sent = False
+
+        async def receive() -> dict[str, Any]:
+            nonlocal sent
+            if not sent:
+                sent = True
+                return {"type": "http.request", "body": body, "more_body": False}
+            return {"type": "http.request", "body": b"", "more_body": False}
+
+        return receive
+
+
+else:
+
+    class ObserverMiddleware:  # pragma: no cover - import-time guard only
+        def __init__(self, *_args: Any, **_kwargs: Any) -> None:
+            raise RuntimeError(
+                "ObserverMiddleware requires FastAPI/Starlette. "
+                "Install optional runtime dependencies first."
+            ) from _IMPORT_ERROR
+
+
+__all__ = ["ObserverMiddleware"]

pyproject.toml

@@ -11,11 +11,13 @@ requires-python = ">=3.10,<3.15"
 license = { file = "LICENSE" }
 authors = [{ name = "fastapi-observer maintainers" }]
 dependencies = [
+  "fastapi>=0.100,<1.0",
   "pydantic>=2.0,<3.0",
 ]
 
 [project.optional-dependencies]
 test = [
+  "httpx>=0.24",
   "pytest>=8.0",
   "pytest-cov>=5.0",
 ]

tests/test_middleware.py

@@ -0,0 +1,120 @@
+import logging
+
+import pytest
+
+fastapi = pytest.importorskip("fastapi")
+testclient = pytest.importorskip("fastapi.testclient")
+
+from fastapi_observer.config import ObserverConfig
+from fastapi_observer.middleware import ObserverMiddleware
+
+FastAPI = fastapi.FastAPI
+TestClient = testclient.TestClient
+
+
+class InMemoryHandler(logging.Handler):
+    def __init__(self) -> None:
+        super().__init__()
+        self.records: list[logging.LogRecord] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.records.append(record)
+
+
+def _build_memory_logger(name: str) -> tuple[logging.Logger, InMemoryHandler]:
+    handler = InMemoryHandler()
+    logger = logging.getLogger(name)
+    logger.handlers.clear()
+    logger.setLevel(logging.INFO)
+    logger.propagate = False
+    logger.addHandler(handler)
+    return logger, handler
+
+
+def test_middleware_logs_completed_request_and_sets_correlation_id():
+    app = FastAPI()
+
+    @app.get("/items")
+    async def read_items():
+        return {"ok": True}
+
+    logger, memory = _build_memory_logger("fastapi_observer.test.middleware.success")
+    config = ObserverConfig(log_headers=True)
+    app.add_middleware(ObserverMiddleware, config=config, logger=logger)
+
+    client = TestClient(app)
+    response = client.get("/items?limit=10")
+
+    assert response.status_code == 200
+    assert config.correlation_id_header in response.headers
+    assert len(memory.records) == 1
+
+    event = memory.records[0].event
+    assert event["message"] == "HTTP request completed"
+    assert event["path"] == "/items"
+    assert event["status_code"] == 200
+    assert event["correlation_id"] == response.headers[config.correlation_id_header]
+    assert event["metadata"]["query_params"] == {"limit": "10"}
+
+
+def test_middleware_skips_excluded_path():
+    app = FastAPI()
+
+    @app.get("/health")
+    async def health():
+        return {"ok": True}
+
+    logger, memory = _build_memory_logger("fastapi_observer.test.middleware.exclude")
+    config = ObserverConfig()
+    app.add_middleware(ObserverMiddleware, config=config, logger=logger)
+
+    client = TestClient(app)
+    response = client.get("/health")
+
+    assert response.status_code == 200
+    assert memory.records == []
+
+
+def test_middleware_logs_request_and_response_body_with_redaction():
+    app = FastAPI()
+
+    @app.post("/echo")
+    async def echo(payload: dict):
+        return payload
+
+    logger, memory = _build_memory_logger("fastapi_observer.test.middleware.body")
+    config = ObserverConfig(log_request_body=True, log_response_body=True)
+    app.add_middleware(ObserverMiddleware, config=config, logger=logger)
+
+    client = TestClient(app)
+    response = client.post("/echo", json={"password": "secret", "name": "alice"})
+
+    assert response.status_code == 200
+    assert len(memory.records) == 1
+
+    event = memory.records[0].event
+    assert event["metadata"]["request_body"]["password"] == "***"
+    assert event["metadata"]["response_body"]["password"] == "***"
+    assert event["metadata"]["request_body"]["name"] == "alice"
+
+
+def test_middleware_logs_exceptions():
+    app = FastAPI()
+
+    @app.get("/boom")
+    async def boom():
+        raise RuntimeError("something failed")
+
+    logger, memory = _build_memory_logger("fastapi_observer.test.middleware.error")
+    config = ObserverConfig()
+    app.add_middleware(ObserverMiddleware, config=config, logger=logger)
+
+    client = TestClient(app, raise_server_exceptions=False)
+    response = client.get("/boom")
+
+    assert response.status_code == 500
+    assert len(memory.records) == 1
+    event = memory.records[0].event
+    assert event["message"] == "HTTP request failed"
+    assert event["status_code"] == 500
+    assert "something failed" in event["error"]