Coverage for polar/logfire.py: 86%

1import os 2 ctx1ba

2from collections.abc import Callable, Sequence 2 ctx1ba

3from typing import TYPE_CHECKING, Any, Literal 2 ctx1ba

4 

5import httpx 2 ctx1ba

6import logfire 2 ctx1ba

7from fastapi import FastAPI 2 ctx1ba

8from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor 2 ctx1ba

9from opentelemetry.sdk.trace.sampling import ( 2 ctx1ba

10 ALWAYS_OFF, 

11 ALWAYS_ON, 

12 ParentBased, 

13 Sampler, 

14 SamplingResult, 

15) 

16 

17if TYPE_CHECKING: 17 ↛ 18line 17 didn't jump to line 18 because the condition on line 17 was never true2 ctx1ba

18 from opentelemetry.context import Context 

19 from opentelemetry.trace import Link, SpanKind 

20 from opentelemetry.trace.span import TraceState 

21 from opentelemetry.util.types import Attributes 

22 

23from polar.config import settings 2 ctx1ba

24from polar.kit.db.postgres import Engine 2 ctx1ba

25 

26Matcher = Callable[[str, "Attributes | None"], bool] 2 ctx1ba

27 

28 

29class IgnoreSampler(Sampler): 2 ctx1ba

30 def __init__(self, matchers: Sequence[Matcher]) -> None: 2 ctx1ba

31 super().__init__() 1 ctx1a

32 self.matchers = matchers 1 ctx1a

33 

34 def should_sample( 2 ctx1ba

35 self, 

36 parent_context: "Context | None", 

37 trace_id: int, 

38 name: str, 

39 kind: "SpanKind | None" = None, 

40 attributes: "Attributes | None" = None, 

41 links: Sequence["Link"] | None = None, 

42 trace_state: "TraceState | None" = None, 

43 ) -> SamplingResult: 

44 sampler = ALWAYS_ON 2 ctx1dc

45 

46 for matcher in self.matchers: 2 ctx1dc

47 if matcher(name, attributes): 2 ctx1dc

48 sampler = ALWAYS_OFF (empty)

49 break (empty)

50 

51 return sampler.should_sample( 2 ctx1dc

52 parent_context, 

53 trace_id, 

54 name, 

55 kind, 

56 attributes, 

57 links, 

58 trace_state, 

59 ) 

60 

61 def get_description(self) -> str: 2 ctx1ba

62 return "IgnoreSampler" 

63 

64 

65def _healthz_matcher(name: str, attributes: "Attributes | None") -> bool: 2 ctx1ba

66 return attributes is not None and attributes.get("http.route") == "/healthz" 2 ctx1dc

67 

68 

69def _worker_health_matcher(name: str, attributes: "Attributes | None") -> bool: 2 ctx1ba

70 lower_name = name.lower() 2 ctx1dc

71 return lower_name.startswith("recording health:") or lower_name.startswith( 2 ctx1dc

72 "health check successful" 

73 ) 

74 

75 

76def _scrubbing_callback(match: logfire.ScrubMatch) -> Any | None: 2 ctx1ba

77 # Don't scrub auth subject in log messages 

78 if match.path == ("attributes", "subject"): 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true1 ctx1d

79 return match.value 

80 return None 1 ctx1d

81 

82 

83def configure_logfire(service_name: Literal["server", "worker"]) -> None: 2 ctx1ba

84 logfire.configure( 2 ctx1ac

85 send_to_logfire="if-token-present", 

86 token=settings.LOGFIRE_TOKEN, 

87 service_name=service_name, 

88 service_version=os.environ.get("RELEASE_VERSION", "development"), 

89 console=False, 

90 sampling=logfire.SamplingOptions( 

91 head=ParentBased(IgnoreSampler((_healthz_matcher, _worker_health_matcher))), 

92 ), 

93 scrubbing=logfire.ScrubbingOptions(callback=_scrubbing_callback), 

94 ) 

95 

96 

97def instrument_httpx(client: httpx.AsyncClient | httpx.Client | None = None) -> None: 2 ctx1ba

98 if client: 3 ctx1bac

99 HTTPXClientInstrumentor().instrument_client(client) 2 ctx1ba

100 else: 

101 HTTPXClientInstrumentor().instrument() (empty)

102 

103 

104def instrument_fastapi(app: FastAPI) -> None: 2 ctx1ba

105 logfire.instrument_fastapi(app, capture_headers=True) (empty)

106 

107 

108def instrument_sqlalchemy(engines: Sequence[Engine]) -> None: 2 ctx1ba

109 logfire.instrument_sqlalchemy(engines=engines) (empty)

110 

111 

112__all__ = [ 2 ctx1ba

113 "configure_logfire", 

114 "instrument_fastapi", 

115 "instrument_sqlalchemy", 

116]