| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import logging |
| 5 | import re |
| 6 | from contextvars import ContextVar, Token |
| 7 | from datetime import UTC, datetime |
| 8 | |
| 9 | |
| 10 | _REQUEST_ID: ContextVar[str | None] = ContextVar("aip_request_id", default=None) |
| 11 | _SENSITIVE = re.compile( |
| 12 | r"(?i)\b(authorization|cookie|password|token|api[_-]?(?:key|token)|" |
| 13 | r"access[_-]?token|refresh[_-]?token|request[_-]?token|client[_-]?secret)" |
| 14 | r"\s*[:=]\s*(?:bearer\s+)?([^\s,;&]+)" |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | class StructuredFormatter(logging.Formatter): |
| 19 | def __init__(self, service: str, environment: str) -> None: |
| 20 | super().__init__() |
| 21 | self.service = service |
| 22 | self.environment = environment |
| 23 | |
| 24 | def format(self, record: logging.LogRecord) -> str: |
| 25 | message = _SENSITIVE.sub(r"\1=<redacted>", record.getMessage()).replace("\r", " ").replace("\n", " ") |
| 26 | return json.dumps({ |
| 27 | "timestamp": datetime.fromtimestamp(record.created, UTC).isoformat(), |
| 28 | "service": self.service, |
| 29 | "environment": self.environment, |
| 30 | "level": record.levelname, |
| 31 | "event": record.name, |
| 32 | "requestId": _REQUEST_ID.get(), |
| 33 | "message": message, |
| 34 | }, separators=(",", ":"), default=str) |
| 35 | |
| 36 | |
| 37 | def configure_logging(service: str, environment: str) -> None: |
| 38 | handler = logging.StreamHandler() |
| 39 | handler.setFormatter(StructuredFormatter(service, environment)) |
| 40 | logging.basicConfig(level=logging.INFO, handlers=[handler], force=True) |
| 41 | |
| 42 | |
| 43 | def set_request_id(value: str) -> Token: |
| 44 | return _REQUEST_ID.set(value) |
| 45 | |
| 46 | |
| 47 | def reset_request_id(token: Token) -> None: |
| 48 | _REQUEST_ID.reset(token) |