Coverage for src/taipanstack/utils/logging.py: 100%

1""" 

2Structured logging with context. 

3 

4Provides a configured logger with support for structured output, 

5context propagation, and proper formatting. 

6""" 

7 

8import logging 

9import logging as _std_logging 

10import re 

11import sys 

12from collections.abc import Iterator, MutableMapping 

13from contextlib import AbstractContextManager, contextmanager 

14from datetime import UTC, datetime 

15from functools import lru_cache 

16from typing import Literal 

17 

18from taipanstack.utils.context import get_correlation_id 

19 

20try: 

21 import structlog 

22 

23 HAS_STRUCTLOG = True 

24except ImportError: 

25 HAS_STRUCTLOG = False 

26 

27 

28# Default log format 

29DEFAULT_FORMAT = "%(asctime)s | %(levelname)-8s | %(name)s | %(message)s" 

30JSON_FORMAT = ( 

31 '{"timestamp": "%(asctime)s", "level": "%(levelname)s", ' 

32 '"logger": "%(name)s", "message": "%(message)s"}' 

33) 

34 

35SENSITIVE_KEY_PATTERNS: tuple[str, ...] = ( 

36 "password", 

37 "secret", 

38 "token", 

39 "authorization", 

40 "api_key", 

41) 

42 

43# Pre-compile the regex for O(1) checks per key instead of O(N*M) 

44_SENSITIVE_KEY_REGEX = ( 

45 re.compile("|".join(map(re.escape, SENSITIVE_KEY_PATTERNS)), re.IGNORECASE) 

46 if SENSITIVE_KEY_PATTERNS 

47 else None 

48) 

49 

50REDACTED_VALUE = "***REDACTED***" 

51 

52 

53@lru_cache(maxsize=1024) 

54def _is_sensitive(key: object, regex: re.Pattern[str] | None) -> bool: 

55 """Check if a key is sensitive using cached regex matching. 

56 

57 Args: 

58 key: The key to check. 

59 regex: The regex pattern to use for matching. 

60 

61 Returns: 

62 True if the key is a string and sensitive, False otherwise. 

63 

64 """ 

65 if regex is None: 

66 return False 

67 if not isinstance(key, str): 

68 return False 

69 return bool(regex.search(key)) 

70 

71 

72def _redact_mapping( 

73 obj: MutableMapping[object, object], 

74 seen: set[int], 

75) -> dict[object, object]: 

76 """Redact a mapping object.""" 

77 seen.add(id(obj)) 

78 return { 

79 k: ( 

80 REDACTED_VALUE 

81 if _is_sensitive(k, _SENSITIVE_KEY_REGEX) 

82 else _redact(v, seen) 

83 ) 

84 for k, v in obj.items() 

85 } 

86 

87 

88def _redact_list(obj: list[object], seen: set[int]) -> list[object]: 

89 """Redact a list object.""" 

90 seen.add(id(obj)) 

91 return [_redact(item, seen) for item in obj] 

92 

93 

94def _redact_tuple(obj: tuple[object, ...], seen: set[int]) -> tuple[object, ...]: 

95 """Redact a tuple object.""" 

96 seen.add(id(obj)) 

97 return tuple(_redact(item, seen) for item in obj) 

98 

99 

100def _redact_set(obj: set[object], seen: set[int]) -> set[object]: 

101 """Redact a set object.""" 

102 seen.add(id(obj)) 

103 redacted = set() 

104 for item in obj: 

105 redacted_item = _redact(item, seen) 

106 try: 

107 redacted.add(redacted_item) 

108 except TypeError: 

109 # If the redacted item is unhashable (e.g. a nested mutable object 

110 # that was partially redacted or changed type), fallback to string 

111 redacted.add(str(redacted_item)) 

112 return redacted 

113 

114 

115def _redact_iterable(obj: object, seen: set[int]) -> object: 

116 """Dispatch redaction based on iterable type.""" 

117 if isinstance(obj, list): 

118 return _redact_list(obj, seen) 

119 if isinstance(obj, tuple): 

120 return _redact_tuple(obj, seen) 

121 if isinstance(obj, set): 

122 return _redact_set(obj, seen) 

123 return obj 

124 

125 

126def _redact_collection(obj: object, seen: set[int]) -> object: 

127 """Dispatch redaction based on collection type.""" 

128 if isinstance(obj, (dict, MutableMapping)): 

129 return _redact_mapping(obj, seen) 

130 

131 return _redact_iterable(obj, seen) 

132 

133 

134def _redact(obj: object, seen: set[int] | None = None) -> object: 

135 """Redact sensitive data recursively with circular reference protection. 

136 

137 Args: 

138 obj: The object to redact. 

139 seen: Set of IDs of already processed containers. 

140 

141 Returns: 

142 A new object with sensitive data redacted. 

143 

144 """ 

145 if _SENSITIVE_KEY_REGEX is None: 

146 return obj 

147 

148 if seen is None: 

149 seen = set() 

150 

151 # Detect circular references 

152 if id(obj) in seen: 

153 return REDACTED_VALUE 

154 

155 return _redact_collection(obj, seen) 

156 

157 

158def _redact_dict(d: MutableMapping[str, object]) -> None: 

159 """Redact sensitive keys in a dictionary in-place (top-level only). 

160 

161 Note: Nested structures are replaced with redacted copies to prevent 

162 shared state mutation. 

163 

164 Args: 

165 d: The dictionary to redact. 

166 

167 """ 

168 if _SENSITIVE_KEY_REGEX is None: 

169 return 

170 

171 for key, value in list(d.items()): 

172 if _is_sensitive(key, _SENSITIVE_KEY_REGEX): 

173 d[key] = REDACTED_VALUE 

174 else: 

175 d[key] = _redact(value) 

176 

177 

178def mask_sensitive_data_processor( 

179 _logger: object, 

180 _method: str, 

181 event_dict: MutableMapping[str, object], 

182) -> MutableMapping[str, object]: 

183 """Mask sensitive data in structlog event dictionaries. 

184 

185 Intercept the *event_dict* produced by structlog and replace the 

186 value of any key whose name contains a sensitive substring with 

187 ``"***REDACTED***"``. Matching is case-insensitive. 

188 

189 Args: 

190 _logger: The wrapped logger object (unused, required by structlog). 

191 _method: The name of the log method called (unused, required by structlog). 

192 event_dict: The structured event dictionary. 

193 

194 Returns: 

195 The event dictionary with sensitive values masked. 

196 

197 """ 

198 _redact_dict(event_dict) 

199 return event_dict 

200 

201 

202def correlation_id_processor( 

203 _logger: object, 

204 _method: str, 

205 event_dict: MutableMapping[str, object], 

206) -> MutableMapping[str, object]: 

207 """Structlog processor to inject correlation ID into events. 

208 

209 Args: 

210 _logger: The wrapped logger object. 

211 _method: The name of the log method called. 

212 event_dict: The structured event dictionary. 

213 

214 Returns: 

215 The event dictionary with correlation_id injected if sets. 

216 

217 """ 

218 correlation_id = get_correlation_id() 

219 if correlation_id is not None: 

220 event_dict["correlation_id"] = correlation_id 

221 return event_dict 

222 

223 

224class StackLogger: 

225 """Enhanced logger with context support. 

226 

227 Provides a wrapper around standard logging with additional features 

228 like context propagation and structured output support. 

229 

230 Attributes: 

231 name: Logger name. 

232 level: Current log level. 

233 

234 """ 

235 

236 def __init__( 

237 self, 

238 name: str = "stack", 

239 level: str = "INFO", 

240 *, 

241 use_structured: bool = False, 

242 ) -> None: 

243 """Initialize the logger. 

244 

245 Args: 

246 name: Logger name. 

247 level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL). 

248 use_structured: Use structured logging if structlog is available. 

249 

250 """ 

251 self.name = name 

252 self.level = level 

253 self._context: dict[str, object] = {} 

254 

255 if use_structured and HAS_STRUCTLOG: 

256 self._logger = structlog.get_logger(name) 

257 self._structured = True 

258 else: 

259 self._logger = logging.getLogger(name) 

260 self._logger.setLevel(getattr(logging, level.upper())) 

261 self._structured = False 

262 

263 def bind(self, **context: object) -> "StackLogger": 

264 """Add context to logger. 

265 

266 Args: 

267 **context: Key-value pairs to add to context. 

268 

269 Returns: 

270 Self for chaining. 

271 

272 """ 

273 self._context.update(context) 

274 if self._structured and HAS_STRUCTLOG: 

275 self._logger = self._logger.bind(**context) 

276 return self 

277 

278 def unbind(self, *keys: str) -> "StackLogger": 

279 """Remove context keys. 

280 

281 Args: 

282 *keys: Keys to remove from context. 

283 

284 Returns: 

285 Self for chaining. 

286 

287 """ 

288 for key in keys: 

289 self._context.pop(key, None) 

290 if self._structured and HAS_STRUCTLOG: 

291 self._logger = self._logger.unbind(*keys) 

292 return self 

293 

294 def _format_message(self, message: str, **kwargs: object) -> str: 

295 """Format message with context. 

296 

297 Args: 

298 message: The log message. 

299 **kwargs: Additional context for this message. 

300 

301 Returns: 

302 Formatted message string. 

303 

304 """ 

305 if not kwargs and not self._context: 

306 return message 

307 

308 context = {**self._context, **kwargs} 

309 _redact_dict(context) 

310 

311 context_str = " ".join(f"{k}={v}" for k, v in context.items()) 

312 return f"{message} | {context_str}" 

313 

314 def debug(self, message: str, **kwargs: object) -> None: 

315 """Log a debug message. 

316 

317 Args: 

318 message: The message to log. 

319 **kwargs: Additional context. 

320 

321 """ 

322 if self._structured: 

323 self._logger.debug(message, **kwargs) 

324 else: 

325 self._logger.debug(self._format_message(message, **kwargs)) 

326 

327 def info(self, message: str, **kwargs: object) -> None: 

328 """Log an info message. 

329 

330 Args: 

331 message: The message to log. 

332 **kwargs: Additional context. 

333 

334 """ 

335 if self._structured: 

336 self._logger.info(message, **kwargs) 

337 else: 

338 self._logger.info(self._format_message(message, **kwargs)) 

339 

340 def warning(self, message: str, **kwargs: object) -> None: 

341 """Log a warning message. 

342 

343 Args: 

344 message: The message to log. 

345 **kwargs: Additional context. 

346 

347 """ 

348 if self._structured: 

349 self._logger.warning(message, **kwargs) 

350 else: 

351 self._logger.warning(self._format_message(message, **kwargs)) 

352 

353 def error(self, message: str, **kwargs: object) -> None: 

354 """Log an error message. 

355 

356 Args: 

357 message: The message to log. 

358 **kwargs: Additional context. 

359 

360 """ 

361 if self._structured: 

362 self._logger.error(message, **kwargs) 

363 else: 

364 self._logger.error(self._format_message(message, **kwargs)) 

365 

366 def critical(self, message: str, **kwargs: object) -> None: 

367 """Log a critical message. 

368 

369 Args: 

370 message: The message to log. 

371 **kwargs: Additional context. 

372 

373 """ 

374 if self._structured: 

375 self._logger.critical(message, **kwargs) 

376 else: 

377 self._logger.critical(self._format_message(message, **kwargs)) 

378 

379 def exception(self, message: str, **kwargs: object) -> None: 

380 """Log an exception with traceback. 

381 

382 Args: 

383 message: The message to log. 

384 **kwargs: Additional context. 

385 

386 """ 

387 if self._structured: 

388 self._logger.exception(message, **kwargs) 

389 else: 

390 self._logger.exception(self._format_message(message, **kwargs)) 

391 

392 

393def _configure_structlog(*, level: int | str | None = None) -> None: 

394 """Configure structlog with default processors and settings.""" 

395 if level is not None: 

396 if isinstance(level, str): 

397 level = getattr(_std_logging, level.upper(), _std_logging.INFO) 

398 _std_logging.basicConfig(level=level) 

399 

400 structlog.configure( 

401 processors=[ 

402 structlog.stdlib.filter_by_level, 

403 structlog.stdlib.add_logger_name, 

404 structlog.stdlib.add_log_level, 

405 structlog.processors.TimeStamper(fmt="iso"), 

406 correlation_id_processor, 

407 mask_sensitive_data_processor, 

408 structlog.processors.StackInfoRenderer(), 

409 structlog.processors.format_exc_info, 

410 structlog.processors.UnicodeDecoder(), 

411 structlog.processors.JSONRenderer(), 

412 ], 

413 wrapper_class=structlog.stdlib.BoundLogger, 

414 context_class=dict, 

415 logger_factory=structlog.stdlib.LoggerFactory(), 

416 ) 

417 

418 

419def _get_log_format(format_type: str) -> str: 

420 """Get the appropriate log format string.""" 

421 if format_type == "simple": 

422 return "%(levelname)s: %(message)s" 

423 elif format_type == "json": 

424 return JSON_FORMAT 

425 else: 

426 return DEFAULT_FORMAT 

427 

428 

429def setup_logging( 

430 level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO", 

431 *, 

432 format_type: Literal["simple", "detailed", "json"] = "detailed", 

433 log_file: str | None = None, 

434 use_structured: bool = False, 

435) -> None: 

436 """Configure the root logger. 

437 

438 Args: 

439 level: Log level to set. 

440 format_type: Output format type. 

441 log_file: Optional file to log to. 

442 use_structured: Use structlog if available. 

443 

444 """ 

445 # Configure structlog if available and requested 

446 if use_structured and HAS_STRUCTLOG: 

447 _configure_structlog() 

448 return 

449 

450 log_format = _get_log_format(format_type) 

451 

452 handlers: list[logging.Handler] = [logging.StreamHandler(sys.stdout)] 

453 

454 if log_file: 

455 handlers.append(logging.FileHandler(log_file, encoding="utf-8")) 

456 

457 logging.basicConfig( 

458 level=getattr(logging, level.upper()), 

459 format=log_format, 

460 handlers=handlers, 

461 force=True, 

462 ) 

463 

464 

465def get_logger( 

466 name: str = "stack", 

467 *, 

468 level: str = "INFO", 

469 use_structured: bool = False, 

470) -> StackLogger: 

471 """Get a configured logger instance. 

472 

473 Args: 

474 name: Logger name. 

475 level: Log level. 

476 use_structured: Use structlog if available. 

477 

478 Returns: 

479 Configured StackLogger instance. 

480 

481 Example: 

482 >>> logger = get_logger("my_module") 

483 >>> logger.bind(request_id="123").info("Processing request") 

484 

485 """ 

486 return StackLogger(name, level, use_structured=use_structured) 

487 

488 

489def log_operation( 

490 operation: str, 

491 *, 

492 logger: StackLogger | None = None, 

493 level: str = "INFO", 

494 expected_exceptions: tuple[type[Exception], ...] | type[Exception] = Exception, 

495) -> AbstractContextManager[StackLogger]: 

496 """Context manager for logging operations. 

497 

498 Args: 

499 operation: Name of the operation. 

500 logger: Logger to use (creates one if not provided). 

501 level: Log level for messages. 

502 expected_exceptions: Exceptions to catch and log as failures. 

503 

504 Yields: 

505 The logger instance. 

506 

507 Example: 

508 >>> with log_operation("setup") as logger: 

509 ... logger.info("Setting up environment") 

510 

511 """ 

512 

513 @contextmanager 

514 def _log_context() -> Iterator[StackLogger]: 

515 nonlocal logger 

516 if logger is None: 

517 logger = get_logger() 

518 

519 start_time = datetime.now(UTC) 

520 logger.bind(operation=operation) 

521 

522 log_method = getattr(logger, level.lower()) 

523 log_method(f"Starting: {operation}") 

524 

525 try: 

526 yield logger 

527 duration = (datetime.now(UTC) - start_time).total_seconds() 

528 log_method(f"Completed: {operation}", duration_seconds=duration) 

529 except expected_exceptions as e: 

530 duration = (datetime.now(UTC) - start_time).total_seconds() 

531 logger.exception( 

532 f"Failed: {operation}", 

533 duration_seconds=duration, 

534 error=str(e), 

535 ) 

536 raise 

537 finally: 

538 logger.unbind("operation") 

539 

540 return _log_context()