quart.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237
  1. import asyncio
  2. import inspect
  3. from functools import wraps
  4. import sentry_sdk
  5. from sentry_sdk.integrations import DidNotEnable, Integration
  6. from sentry_sdk.integrations._wsgi_common import _filter_headers
  7. from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
  8. from sentry_sdk.scope import should_send_default_pii
  9. from sentry_sdk.tracing import SOURCE_FOR_STYLE
  10. from sentry_sdk.utils import (
  11. capture_internal_exceptions,
  12. ensure_integration_enabled,
  13. event_from_exception,
  14. )
  15. from typing import TYPE_CHECKING
  16. if TYPE_CHECKING:
  17. from typing import Any
  18. from typing import Union
  19. from sentry_sdk._types import Event, EventProcessor
  20. try:
  21. import quart_auth # type: ignore
  22. except ImportError:
  23. quart_auth = None
  24. try:
  25. from quart import ( # type: ignore
  26. has_request_context,
  27. has_websocket_context,
  28. Request,
  29. Quart,
  30. request,
  31. websocket,
  32. )
  33. from quart.signals import ( # type: ignore
  34. got_background_exception,
  35. got_request_exception,
  36. got_websocket_exception,
  37. request_started,
  38. websocket_started,
  39. )
  40. except ImportError:
  41. raise DidNotEnable("Quart is not installed")
  42. else:
  43. # Quart 0.19 is based on Flask and hence no longer has a Scaffold
  44. try:
  45. from quart.scaffold import Scaffold # type: ignore
  46. except ImportError:
  47. from flask.sansio.scaffold import Scaffold # type: ignore
  48. TRANSACTION_STYLE_VALUES = ("endpoint", "url")
  49. class QuartIntegration(Integration):
  50. identifier = "quart"
  51. origin = f"auto.http.{identifier}"
  52. transaction_style = ""
  53. def __init__(self, transaction_style="endpoint"):
  54. # type: (str) -> None
  55. if transaction_style not in TRANSACTION_STYLE_VALUES:
  56. raise ValueError(
  57. "Invalid value for transaction_style: %s (must be in %s)"
  58. % (transaction_style, TRANSACTION_STYLE_VALUES)
  59. )
  60. self.transaction_style = transaction_style
  61. @staticmethod
  62. def setup_once():
  63. # type: () -> None
  64. request_started.connect(_request_websocket_started)
  65. websocket_started.connect(_request_websocket_started)
  66. got_background_exception.connect(_capture_exception)
  67. got_request_exception.connect(_capture_exception)
  68. got_websocket_exception.connect(_capture_exception)
  69. patch_asgi_app()
  70. patch_scaffold_route()
  71. def patch_asgi_app():
  72. # type: () -> None
  73. old_app = Quart.__call__
  74. async def sentry_patched_asgi_app(self, scope, receive, send):
  75. # type: (Any, Any, Any, Any) -> Any
  76. if sentry_sdk.get_client().get_integration(QuartIntegration) is None:
  77. return await old_app(self, scope, receive, send)
  78. middleware = SentryAsgiMiddleware(
  79. lambda *a, **kw: old_app(self, *a, **kw),
  80. span_origin=QuartIntegration.origin,
  81. asgi_version=3,
  82. )
  83. return await middleware(scope, receive, send)
  84. Quart.__call__ = sentry_patched_asgi_app
  85. def patch_scaffold_route():
  86. # type: () -> None
  87. old_route = Scaffold.route
  88. def _sentry_route(*args, **kwargs):
  89. # type: (*Any, **Any) -> Any
  90. old_decorator = old_route(*args, **kwargs)
  91. def decorator(old_func):
  92. # type: (Any) -> Any
  93. if inspect.isfunction(old_func) and not asyncio.iscoroutinefunction(
  94. old_func
  95. ):
  96. @wraps(old_func)
  97. @ensure_integration_enabled(QuartIntegration, old_func)
  98. def _sentry_func(*args, **kwargs):
  99. # type: (*Any, **Any) -> Any
  100. current_scope = sentry_sdk.get_current_scope()
  101. if current_scope.transaction is not None:
  102. current_scope.transaction.update_active_thread()
  103. sentry_scope = sentry_sdk.get_isolation_scope()
  104. if sentry_scope.profile is not None:
  105. sentry_scope.profile.update_active_thread_id()
  106. return old_func(*args, **kwargs)
  107. return old_decorator(_sentry_func)
  108. return old_decorator(old_func)
  109. return decorator
  110. Scaffold.route = _sentry_route
  111. def _set_transaction_name_and_source(scope, transaction_style, request):
  112. # type: (sentry_sdk.Scope, str, Request) -> None
  113. try:
  114. name_for_style = {
  115. "url": request.url_rule.rule,
  116. "endpoint": request.url_rule.endpoint,
  117. }
  118. scope.set_transaction_name(
  119. name_for_style[transaction_style],
  120. source=SOURCE_FOR_STYLE[transaction_style],
  121. )
  122. except Exception:
  123. pass
  124. async def _request_websocket_started(app, **kwargs):
  125. # type: (Quart, **Any) -> None
  126. integration = sentry_sdk.get_client().get_integration(QuartIntegration)
  127. if integration is None:
  128. return
  129. if has_request_context():
  130. request_websocket = request._get_current_object()
  131. if has_websocket_context():
  132. request_websocket = websocket._get_current_object()
  133. # Set the transaction name here, but rely on ASGI middleware
  134. # to actually start the transaction
  135. _set_transaction_name_and_source(
  136. sentry_sdk.get_current_scope(), integration.transaction_style, request_websocket
  137. )
  138. scope = sentry_sdk.get_isolation_scope()
  139. evt_processor = _make_request_event_processor(app, request_websocket, integration)
  140. scope.add_event_processor(evt_processor)
  141. def _make_request_event_processor(app, request, integration):
  142. # type: (Quart, Request, QuartIntegration) -> EventProcessor
  143. def inner(event, hint):
  144. # type: (Event, dict[str, Any]) -> Event
  145. # if the request is gone we are fine not logging the data from
  146. # it. This might happen if the processor is pushed away to
  147. # another thread.
  148. if request is None:
  149. return event
  150. with capture_internal_exceptions():
  151. # TODO: Figure out what to do with request body. Methods on request
  152. # are async, but event processors are not.
  153. request_info = event.setdefault("request", {})
  154. request_info["url"] = request.url
  155. request_info["query_string"] = request.query_string
  156. request_info["method"] = request.method
  157. request_info["headers"] = _filter_headers(dict(request.headers))
  158. if should_send_default_pii():
  159. request_info["env"] = {"REMOTE_ADDR": request.access_route[0]}
  160. _add_user_to_event(event)
  161. return event
  162. return inner
  163. async def _capture_exception(sender, exception, **kwargs):
  164. # type: (Quart, Union[ValueError, BaseException], **Any) -> None
  165. integration = sentry_sdk.get_client().get_integration(QuartIntegration)
  166. if integration is None:
  167. return
  168. event, hint = event_from_exception(
  169. exception,
  170. client_options=sentry_sdk.get_client().options,
  171. mechanism={"type": "quart", "handled": False},
  172. )
  173. sentry_sdk.capture_event(event, hint=hint)
  174. def _add_user_to_event(event):
  175. # type: (Event) -> None
  176. if quart_auth is None:
  177. return
  178. user = quart_auth.current_user
  179. if user is None:
  180. return
  181. with capture_internal_exceptions():
  182. user_info = event.setdefault("user", {})
  183. user_info["id"] = quart_auth.current_user._auth_id