Coverage for src/asgi_profiler/__init__.py: 99%

1"""asgi-profiler -- request and SQL profiling for ASGI applications. 

2 

3 from fastapi import FastAPI 

4 from asgi_profiler import install 

5 

6 app = FastAPI() 

7 install(app) # viewer at /profiler 

8 

9Records every request the app handles along with the SQL each one ran, and 

10serves a browsable UI. Works with Starlette, FastAPI, SQLAlchemy and 

11SQLModel, sync and async. 

12""" 

13 

14from __future__ import annotations 

15 

16import contextlib 

17import dataclasses 

18from importlib.metadata import PackageNotFoundError 

19from importlib.metadata import version as _installed_version 

20from typing import Any 

21 

22from .config import ProfilerConfig 

23from .detectors import TYPES as DETECTOR_TYPES 

24from .detectors import Settings as DetectorSettings 

25from .detectors import detect 

26from .fingerprint import normalise_sql, sql_hash 

27from .instrument import install as install_sql_hooks 

28from .instrument import uninstall as uninstall_sql_hooks 

29from .middleware import ProfilerMiddleware 

30from .models import ( 

31 PathSummary, 

32 Problem, 

33 Profile, 

34 Query, 

35 QueryGroup, 

36 StatementSummary, 

37 group_queries, 

38) 

39from .routing import route_pattern 

40from .rules import profiler_exclude, profiler_include 

41from .storage import ( 

42 BaseStorage, 

43 Filters, 

44 MemoryStorage, 

45 Page, 

46 SQLiteStorage, 

47 Storage, 

48 aggregate_statements, 

49 summarise, 

50) 

51from .viewer import build_viewer 

52 

53__all__ = [ 

54 "BaseStorage", 

55 "DETECTOR_TYPES", 

56 "DetectorSettings", 

57 "Filters", 

58 "MemoryStorage", 

59 "Page", 

60 "PathSummary", 

61 "Problem", 

62 "Profile", 

63 "Profiler", 

64 "ProfilerConfig", 

65 "ProfilerMiddleware", 

66 "Query", 

67 "QueryGroup", 

68 "SQLiteStorage", 

69 "StatementSummary", 

70 "Storage", 

71 "aggregate_statements", 

72 "build_viewer", 

73 "detect", 

74 "group_queries", 

75 "install", 

76 "install_sql_hooks", 

77 "normalise_sql", 

78 "profiler_exclude", 

79 "profiler_include", 

80 "route_pattern", 

81 "sql_hash", 

82 "summarise", 

83 "uninstall_sql_hooks", 

84] 

85 

86try: 

87 #: Read from the installed distribution rather than repeated here. A 

88 #: hand-maintained literal drifts from `pyproject.toml` silently, and the 

89 #: release only compares the tag against the packaged version -- so a 

90 #: wheel can ship saying 0.1.0 in its metadata and 0.0.1 in its code. 

91 __version__ = _installed_version("asgi-profiler") 

92except PackageNotFoundError: # pragma: no cover - running from a source tree 

93 __version__ = "0.0.0.dev0" 

94 

95#: Marker set on an app that already has the profiler, so a second 

96#: `install()` is refused rather than silently doubling every count. 

97#: 

98#: An attribute on the app rather than a set of `id()`s: CPython reuses the 

99#: address of a collected object, so an id-keyed set reports a false positive 

100#: as soon as an earlier app is garbage collected. 

101_MARKER = "_asgi_profiler_installed" 

102 

103 

104class Profiler: 

105 """Handle returned by :func:`install`, for programmatic access.""" 

106 

107 def __init__(self, storage: Storage, config: ProfilerConfig) -> None: 

108 self.storage = storage 

109 self.config = config 

110 

111 @property 

112 def profiles(self) -> list[Profile]: 

113 """Every retained profile. 

114 

115 Convenient, but it hydrates the whole history -- on a large 

116 :class:`SQLiteStorage` that is every statement of every request. 

117 Prefer :meth:`slowest` or :meth:`search` when the history is big. 

118 """ 

119 return self.storage.list() 

120 

121 def search(self, filters: Filters | None = None, **kwargs: Any) -> Page: 

122 """One page of profiles, filtered and ordered by the backend.""" 

123 return self.storage.search(filters or Filters(), **kwargs) 

124 

125 def slowest(self, count: int = 10) -> list[Profile]: 

126 return self.storage.search(Filters(order="slowest"), page=1, size=count).items 

127 

128 def summary(self) -> list[PathSummary]: 

129 return self.storage.summarise() 

130 

131 def statements(self, limit: int = 100) -> list[StatementSummary]: 

132 return self.storage.statements(limit) 

133 

134 def exclude(self, *paths: str) -> None: 

135 """Stop recording requests under `paths`, after `install()`. 

136 

137 For integrations that only learn what to exclude later -- mounting the 

138 viewer inside an admin, say, where the admin's own URL is not known 

139 until it is registered. Matched on segment boundaries like 

140 `ProfilerConfig.exclude_paths`, which this appends to. 

141 """ 

142 wanted = [p for p in paths if p and p not in self.config.exclude_paths] 

143 if wanted: 143 ↛ exitline 143 didn't return from function 'exclude' because the condition on line 143 was always true

144 self.config.exclude_paths = (*self.config.exclude_paths, *wanted) 

145 

146 def clear(self) -> None: 

147 self.storage.clear() 

148 

149 def close(self) -> None: 

150 """Release the storage backend. Safe to call on any backend.""" 

151 closer = getattr(self.storage, "close", None) 

152 if callable(closer): 

153 closer() 

154 

155 

156def install( 

157 app: Any, 

158 config: ProfilerConfig | None = None, 

159 storage: Storage | None = None, 

160 **options: Any, 

161) -> Profiler: 

162 """Add the middleware and mount the viewer. 

163 

164 Args: 

165 app: a Starlette or FastAPI application. Call this before the app 

166 starts serving -- middleware cannot be added afterwards. Pass 

167 `mount_path=None` to record without mounting the viewer. 

168 config: a :class:`ProfilerConfig`. Keyword options are merged into a 

169 *copy*, so a shared config object is never rewritten under you. 

170 storage: a custom store; defaults to :class:`MemoryStorage`. Pass 

171 :class:`SQLiteStorage` to keep history across restarts and share 

172 it between workers. 

173 **options: shorthand for `ProfilerConfig` fields, e.g. 

174 ``install(app, mount_path="/_perf", capture_stacks=False)``. 

175 

176 Returns: 

177 A :class:`Profiler` handle wrapping the storage. 

178 

179 Raises: 

180 RuntimeError: if this app already has the profiler installed. 

181 TypeError: on an unknown option, or an `authorize` that is not usable. 

182 """ 

183 config = dataclasses.replace(config) if config is not None else ProfilerConfig() 

184 for key, value in options.items(): 

185 if not hasattr(config, key): 

186 raise TypeError(f"Unknown profiler option: {key!r}") 

187 setattr(config, key, value) 

188 

189 if getattr(app, _MARKER, False): 

190 raise RuntimeError( 

191 "The profiler is already installed on this app. Installing twice " 

192 "adds a second middleware and a second mount, so every request is " 

193 "recorded twice into two separate stores." 

194 ) 

195 

196 if storage is None: 

197 # `is None`, never `or`: an empty MemoryStorage is falsy, so `or` threw 

198 # away the caller's store and quietly recorded into a different one. 

199 storage = MemoryStorage(max_requests=config.max_requests) 

200 

201 install_sql_hooks( 

202 capture_stacks=config.capture_stacks, 

203 stack_depth=config.stack_depth, 

204 ) 

205 app.add_middleware(ProfilerMiddleware, storage=storage, config=config) 

206 if config.mount_path is not None: 

207 app.mount( 

208 config.mount_path, 

209 build_viewer(storage, config, prefix=config.mount_path), 

210 name="profiler", 

211 ) 

212 with contextlib.suppress(AttributeError): # an app defining __slots__ 

213 setattr(app, _MARKER, True) 

214 

215 return Profiler(storage, config)