fastapi.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. import asyncio
  2. from copy import deepcopy
  3. from functools import wraps
  4. import sentry_sdk
  5. from sentry_sdk.integrations import DidNotEnable
  6. from sentry_sdk.scope import should_send_default_pii
  7. from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource
  8. from sentry_sdk.utils import transaction_from_function
  9. from typing import TYPE_CHECKING
  10. if TYPE_CHECKING:
  11. from typing import Any, Callable, Dict
  12. from sentry_sdk._types import Event
  13. try:
  14. from sentry_sdk.integrations.starlette import (
  15. StarletteIntegration,
  16. StarletteRequestExtractor,
  17. )
  18. except DidNotEnable:
  19. raise DidNotEnable("Starlette is not installed")
  20. try:
  21. import fastapi # type: ignore
  22. except ImportError:
  23. raise DidNotEnable("FastAPI is not installed")
  24. _DEFAULT_TRANSACTION_NAME = "generic FastAPI request"
  25. class FastApiIntegration(StarletteIntegration):
  26. identifier = "fastapi"
  27. @staticmethod
  28. def setup_once():
  29. # type: () -> None
  30. patch_get_request_handler()
  31. def _set_transaction_name_and_source(scope, transaction_style, request):
  32. # type: (sentry_sdk.Scope, str, Any) -> None
  33. name = ""
  34. if transaction_style == "endpoint":
  35. endpoint = request.scope.get("endpoint")
  36. if endpoint:
  37. name = transaction_from_function(endpoint) or ""
  38. elif transaction_style == "url":
  39. route = request.scope.get("route")
  40. if route:
  41. path = getattr(route, "path", None)
  42. if path is not None:
  43. name = path
  44. if not name:
  45. name = _DEFAULT_TRANSACTION_NAME
  46. source = TransactionSource.ROUTE
  47. else:
  48. source = SOURCE_FOR_STYLE[transaction_style]
  49. scope.set_transaction_name(name, source=source)
  50. def patch_get_request_handler():
  51. # type: () -> None
  52. old_get_request_handler = fastapi.routing.get_request_handler
  53. def _sentry_get_request_handler(*args, **kwargs):
  54. # type: (*Any, **Any) -> Any
  55. dependant = kwargs.get("dependant")
  56. if (
  57. dependant
  58. and dependant.call is not None
  59. and not asyncio.iscoroutinefunction(dependant.call)
  60. ):
  61. old_call = dependant.call
  62. @wraps(old_call)
  63. def _sentry_call(*args, **kwargs):
  64. # type: (*Any, **Any) -> Any
  65. current_scope = sentry_sdk.get_current_scope()
  66. if current_scope.transaction is not None:
  67. current_scope.transaction.update_active_thread()
  68. sentry_scope = sentry_sdk.get_isolation_scope()
  69. if sentry_scope.profile is not None:
  70. sentry_scope.profile.update_active_thread_id()
  71. return old_call(*args, **kwargs)
  72. dependant.call = _sentry_call
  73. old_app = old_get_request_handler(*args, **kwargs)
  74. async def _sentry_app(*args, **kwargs):
  75. # type: (*Any, **Any) -> Any
  76. integration = sentry_sdk.get_client().get_integration(FastApiIntegration)
  77. if integration is None:
  78. return await old_app(*args, **kwargs)
  79. request = args[0]
  80. _set_transaction_name_and_source(
  81. sentry_sdk.get_current_scope(), integration.transaction_style, request
  82. )
  83. sentry_scope = sentry_sdk.get_isolation_scope()
  84. extractor = StarletteRequestExtractor(request)
  85. info = await extractor.extract_request_info()
  86. def _make_request_event_processor(req, integration):
  87. # type: (Any, Any) -> Callable[[Event, Dict[str, Any]], Event]
  88. def event_processor(event, hint):
  89. # type: (Event, Dict[str, Any]) -> Event
  90. # Extract information from request
  91. request_info = event.get("request", {})
  92. if info:
  93. if "cookies" in info and should_send_default_pii():
  94. request_info["cookies"] = info["cookies"]
  95. if "data" in info:
  96. request_info["data"] = info["data"]
  97. event["request"] = deepcopy(request_info)
  98. return event
  99. return event_processor
  100. sentry_scope._name = FastApiIntegration.identifier
  101. sentry_scope.add_event_processor(
  102. _make_request_event_processor(request, integration)
  103. )
  104. return await old_app(*args, **kwargs)
  105. return _sentry_app
  106. fastapi.routing.get_request_handler = _sentry_get_request_handler