asyncio.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. import sys
  2. import sentry_sdk
  3. from sentry_sdk.consts import OP
  4. from sentry_sdk.integrations import Integration, DidNotEnable
  5. from sentry_sdk.utils import event_from_exception, logger, reraise
  6. try:
  7. import asyncio
  8. from asyncio.tasks import Task
  9. except ImportError:
  10. raise DidNotEnable("asyncio not available")
  11. from typing import cast, TYPE_CHECKING
  12. if TYPE_CHECKING:
  13. from typing import Any
  14. from collections.abc import Coroutine
  15. from sentry_sdk._types import ExcInfo
  16. def get_name(coro):
  17. # type: (Any) -> str
  18. return (
  19. getattr(coro, "__qualname__", None)
  20. or getattr(coro, "__name__", None)
  21. or "coroutine without __name__"
  22. )
  23. def patch_asyncio():
  24. # type: () -> None
  25. orig_task_factory = None
  26. try:
  27. loop = asyncio.get_running_loop()
  28. orig_task_factory = loop.get_task_factory()
  29. def _sentry_task_factory(loop, coro, **kwargs):
  30. # type: (asyncio.AbstractEventLoop, Coroutine[Any, Any, Any], Any) -> asyncio.Future[Any]
  31. async def _task_with_sentry_span_creation():
  32. # type: () -> Any
  33. result = None
  34. with sentry_sdk.isolation_scope():
  35. with sentry_sdk.start_span(
  36. op=OP.FUNCTION,
  37. name=get_name(coro),
  38. origin=AsyncioIntegration.origin,
  39. ):
  40. try:
  41. result = await coro
  42. except StopAsyncIteration as e:
  43. raise e from None
  44. except Exception:
  45. reraise(*_capture_exception())
  46. return result
  47. task = None
  48. # Trying to use user set task factory (if there is one)
  49. if orig_task_factory:
  50. task = orig_task_factory(
  51. loop, _task_with_sentry_span_creation(), **kwargs
  52. )
  53. if task is None:
  54. # The default task factory in `asyncio` does not have its own function
  55. # but is just a couple of lines in `asyncio.base_events.create_task()`
  56. # Those lines are copied here.
  57. # WARNING:
  58. # If the default behavior of the task creation in asyncio changes,
  59. # this will break!
  60. task = Task(_task_with_sentry_span_creation(), loop=loop, **kwargs)
  61. if task._source_traceback: # type: ignore
  62. del task._source_traceback[-1] # type: ignore
  63. # Set the task name to include the original coroutine's name
  64. try:
  65. cast("asyncio.Task[Any]", task).set_name(
  66. f"{get_name(coro)} (Sentry-wrapped)"
  67. )
  68. except AttributeError:
  69. # set_name might not be available in all Python versions
  70. pass
  71. return task
  72. loop.set_task_factory(_sentry_task_factory) # type: ignore
  73. except RuntimeError:
  74. # When there is no running loop, we have nothing to patch.
  75. logger.warning(
  76. "There is no running asyncio loop so there is nothing Sentry can patch. "
  77. "Please make sure you call sentry_sdk.init() within a running "
  78. "asyncio loop for the AsyncioIntegration to work. "
  79. "See https://docs.sentry.io/platforms/python/integrations/asyncio/"
  80. )
  81. def _capture_exception():
  82. # type: () -> ExcInfo
  83. exc_info = sys.exc_info()
  84. client = sentry_sdk.get_client()
  85. integration = client.get_integration(AsyncioIntegration)
  86. if integration is not None:
  87. event, hint = event_from_exception(
  88. exc_info,
  89. client_options=client.options,
  90. mechanism={"type": "asyncio", "handled": False},
  91. )
  92. sentry_sdk.capture_event(event, hint=hint)
  93. return exc_info
  94. class AsyncioIntegration(Integration):
  95. identifier = "asyncio"
  96. origin = f"auto.function.{identifier}"
  97. @staticmethod
  98. def setup_once():
  99. # type: () -> None
  100. patch_asyncio()