Coverage for src/asgi_profiler/middleware.py: 95%

1"""Request capture.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6import threading 

7import time 

8import uuid 

9from typing import Any 

10 

11from starlette.types import ASGIApp, Message, Receive, Scope, Send 

12 

13from . import rules 

14from .config import ProfilerConfig 

15from .detectors import detect 

16from .instrument import current_origin, current_queries 

17from .models import Profile, Query 

18from .routing import request_path, route_path, route_pattern 

19from .storage import Storage 

20 

21logger = logging.getLogger("asgi_profiler") 

22 

23 

24class ProfilerMiddleware: 

25 """Pure ASGI, one `Profile` per HTTP request. 

26 

27 Pure ASGI rather than `BaseHTTPMiddleware` on purpose: the latter runs the 

28 downstream app in a separate task, which breaks contextvar propagation -- 

29 exactly what query attribution depends on. 

30 """ 

31 

32 def __init__( 

33 self, 

34 app: ASGIApp, 

35 storage: Storage, 

36 config: ProfilerConfig | None = None, 

37 ) -> None: 

38 self.app = app 

39 self.storage = storage 

40 self.config = config or ProfilerConfig() 

41 self._excludes = self.config.build_excludes() 

42 self._patterns = self.config.compiled_patterns() 

43 self._excludes_key = self._exclude_key() 

44 self._redacted = {h.lower() for h in self.config.redacted_headers} 

45 

46 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

47 if scope["type"] != "http": 

48 await self.app(scope, receive, send) 

49 return 

50 

51 original_root = scope.get("root_path", "") 

52 full_path = request_path(scope) 

53 # Match against the app-relative path: `mount_path`, `exclude_paths` 

54 # and the regexes are written against the application's own routing 

55 # table, not against whatever prefix it is deployed under. 

56 allowed = self._path_allows(route_path(scope)) 

57 if not allowed and not rules.any_includes(): 57 ↛ 62line 57 didn't jump to line 62 because the condition on line 57 was never true

58 # Nothing in this process can flip that decision later, so skip 

59 # before allocating a profile or capturing a single query. The 

60 # moment one `@profiler_include` exists anywhere, this shortcut 

61 # stops being sound and the decision moves below the response. 

62 await self.app(scope, receive, send) 

63 return 

64 

65 profile = Profile( 

66 id=uuid.uuid4().hex[:12], 

67 method=scope.get("method", "GET"), 

68 path=full_path, 

69 query_string=scope.get("query_string", b"").decode("latin-1"), 

70 client=self._client(scope), 

71 ) 

72 if self.config.capture_headers: 

73 profile.request_headers = self._headers(scope.get("headers", [])) 

74 

75 header_name = self.config.response_header 

76 stamp = header_name.lower().encode("latin-1") if header_name else None 

77 

78 queries: list[Query] = [] 

79 token = current_queries.set(queries) 

80 started = time.perf_counter() 

81 # Both halves of what a statement needs to place itself: when the 

82 # request began, and which thread the event loop is on. This runs on 

83 # the loop thread by definition -- an ASGI app is only ever called 

84 # from it -- so `get_ident()` here *is* the loop's thread. 

85 origin_token = current_origin.set((started, threading.get_ident())) 

86 

87 async def send_wrapper(message: Message) -> None: 

88 if message["type"] == "http.response.start": 

89 profile.status_code = message["status"] 

90 if stamp is not None: 

91 message = { 

92 **message, 

93 "headers": [ 

94 *message.get("headers", []), 

95 (stamp, profile.id.encode("latin-1")), 

96 ], 

97 } 

98 if self.config.capture_headers: 

99 profile.response_headers = self._headers(message.get("headers", [])) 

100 await send(message) 

101 

102 try: 

103 await self.app(scope, receive, send_wrapper) 

104 except Exception: 

105 # An unhandled exception is turned into a 500 by 

106 # ServerErrorMiddleware, which sits *outside* us -- so we never see 

107 # the http.response.start. Record it ourselves and re-raise. 

108 if not profile.status_code: 108 ↛ 110line 108 didn't jump to line 110 because the condition on line 108 was always true

109 profile.status_code = 500 

110 raise 

111 finally: 

112 profile.duration_ms = (time.perf_counter() - started) * 1000 

113 # A snapshot, not the live list: a fire-and-forget task inherits 

114 # this context and would otherwise keep appending after the 

115 # counters below are frozen, leaving the detail page showing 

116 # "0 queries" above a list of statements. 

117 profile.queries = list(queries) 

118 profile.route = route_pattern(scope, original_root) 

119 profile.finalise() 

120 # After the response has gone out, so the cost of concluding 

121 # things about a request is never paid by the person waiting for 

122 # it. `detect` swallows its own errors for the same reason 

123 # `_store` does. 

124 profile.problems = detect(profile, self.config.detector_settings()) 

125 current_queries.reset(token) 

126 current_origin.reset(origin_token) 

127 # Guard rather than an early `return`: a bare return inside 

128 # `finally` swallows whatever exception is in flight, so a handler 

129 # that raised would have been reported as a clean response. 

130 if self._keep(scope, allowed): 

131 self._store(profile) 

132 

133 def _store(self, profile: Profile) -> None: 

134 try: 

135 self.storage.add(profile) 

136 except Exception: # pragma: no cover - a wedged backend 

137 # The response has already gone out. A full disk or a locked 

138 # database must not turn every successful request into a logged 

139 # ASGI exception, and must never be able to take the application 

140 # down -- a profiler is not load-bearing. 

141 logger.warning("Could not record profile", exc_info=True) 

142 

143 # -- helpers --------------------------------------------------------- 

144 def _keep(self, scope: Scope, allowed: bool) -> bool: 

145 """The final verdict, once routing has named the endpoint. 

146 

147 A decorator on the endpoint is decisive in both directions -- it is the 

148 most specific statement of intent available, so it outranks every 

149 pattern. Only when there is no decorator does the path-based decision 

150 made before routing stand. 

151 """ 

152 decided = rules.endpoint_decision(scope) 

153 if decided is not None: 

154 return decided 

155 return allowed 

156 

157 def _exclude_key(self) -> tuple[Any, ...]: 

158 return ( 

159 self.config.mount_path, 

160 tuple(self.config.exclude_paths), 

161 self.config.include_regex, 

162 self.config.exclude_regex, 

163 ) 

164 

165 def _path_allows(self, path: str) -> bool: 

166 key = self._exclude_key() 

167 if key != self._excludes_key: 

168 self._excludes_key = key 

169 self._excludes = self.config.build_excludes() 

170 self._patterns = self.config.compiled_patterns() 

171 include_re, exclude_re = self._patterns 

172 return rules.path_allows( 

173 path, 

174 excludes=self._excludes, 

175 include_re=include_re, 

176 exclude_re=exclude_re, 

177 ) 

178 

179 def _headers(self, raw: Any) -> dict[str, str]: 

180 out: dict[str, str] = {} 

181 for key, value in raw: 

182 name = key.decode("latin-1") if isinstance(key, bytes) else str(key) 

183 name = name.lower() 

184 if name in self._redacted: 

185 out[name] = "<redacted>" 

186 continue 

187 out[name] = value.decode("latin-1") if isinstance(value, bytes) else str(value) 

188 return out 

189 

190 @staticmethod 

191 def _client(scope: Scope) -> str: 

192 client = scope.get("client") 

193 if not client: 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true

194 return "" 

195 host, port = client 

196 return f"{host}:{port}"