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

1""" 

2Observability Context Module. 

3 

4Provides context variables and context managers for tracing and observability, 

5such as the correlation ID. 

6""" 

7 

8from collections.abc import Iterator 

9from contextlib import contextmanager 

10from contextvars import ContextVar 

11 

12__all__ = [ 

13 "correlation_id_var", 

14 "correlation_scope", 

15 "get_correlation_id", 

16 "set_correlation_id", 

17] 

18 

19# Context variables for observability 

20correlation_id_var: ContextVar[str | None] = ContextVar( 

21 "correlation_id", 

22 default=None, 

23) 

24 

25 

26def get_correlation_id() -> str | None: 

27 """Get the current correlation ID. 

28 

29 Returns: 

30 The correlation ID if set, otherwise None. 

31 

32 """ 

33 return correlation_id_var.get() 

34 

35 

36def set_correlation_id(correlation_id: str | None) -> None: 

37 """Set the correlation ID for the current context. 

38 

39 Args: 

40 correlation_id: The correlation ID string, or None to clear. 

41 

42 """ 

43 correlation_id_var.set(correlation_id) 

44 

45 

46@contextmanager 

47def correlation_scope(correlation_id: str | None) -> Iterator[None]: 

48 """Context manager to set the correlation ID and restore it after. 

49 

50 Args: 

51 correlation_id: The correlation ID to set for the duration of the scope. 

52 

53 Yields: 

54 None 

55 

56 """ 

57 token = correlation_id_var.set(correlation_id) 

58 try: 

59 yield 

60 finally: 

61 correlation_id_var.reset(token)