caching.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. import functools
  2. from typing import TYPE_CHECKING
  3. from sentry_sdk.integrations.redis.utils import _get_safe_key, _key_as_string
  4. from urllib3.util import parse_url as urlparse
  5. from django import VERSION as DJANGO_VERSION
  6. from django.core.cache import CacheHandler
  7. import sentry_sdk
  8. from sentry_sdk.consts import OP, SPANDATA
  9. from sentry_sdk.utils import (
  10. capture_internal_exceptions,
  11. ensure_integration_enabled,
  12. )
  13. if TYPE_CHECKING:
  14. from typing import Any
  15. from typing import Callable
  16. from typing import Optional
  17. METHODS_TO_INSTRUMENT = [
  18. "set",
  19. "set_many",
  20. "get",
  21. "get_many",
  22. ]
  23. def _get_span_description(method_name, args, kwargs):
  24. # type: (str, tuple[Any], dict[str, Any]) -> str
  25. return _key_as_string(_get_safe_key(method_name, args, kwargs))
  26. def _patch_cache_method(cache, method_name, address, port):
  27. # type: (CacheHandler, str, Optional[str], Optional[int]) -> None
  28. from sentry_sdk.integrations.django import DjangoIntegration
  29. original_method = getattr(cache, method_name)
  30. @ensure_integration_enabled(DjangoIntegration, original_method)
  31. def _instrument_call(
  32. cache, method_name, original_method, args, kwargs, address, port
  33. ):
  34. # type: (CacheHandler, str, Callable[..., Any], tuple[Any, ...], dict[str, Any], Optional[str], Optional[int]) -> Any
  35. is_set_operation = method_name.startswith("set")
  36. is_get_operation = not is_set_operation
  37. op = OP.CACHE_PUT if is_set_operation else OP.CACHE_GET
  38. description = _get_span_description(method_name, args, kwargs)
  39. with sentry_sdk.start_span(
  40. op=op,
  41. name=description,
  42. origin=DjangoIntegration.origin,
  43. ) as span:
  44. value = original_method(*args, **kwargs)
  45. with capture_internal_exceptions():
  46. if address is not None:
  47. span.set_data(SPANDATA.NETWORK_PEER_ADDRESS, address)
  48. if port is not None:
  49. span.set_data(SPANDATA.NETWORK_PEER_PORT, port)
  50. key = _get_safe_key(method_name, args, kwargs)
  51. if key is not None:
  52. span.set_data(SPANDATA.CACHE_KEY, key)
  53. item_size = None
  54. if is_get_operation:
  55. if value:
  56. item_size = len(str(value))
  57. span.set_data(SPANDATA.CACHE_HIT, True)
  58. else:
  59. span.set_data(SPANDATA.CACHE_HIT, False)
  60. else: # TODO: We don't handle `get_or_set` which we should
  61. arg_count = len(args)
  62. if arg_count >= 2:
  63. # 'set' command
  64. item_size = len(str(args[1]))
  65. elif arg_count == 1:
  66. # 'set_many' command
  67. item_size = len(str(args[0]))
  68. if item_size is not None:
  69. span.set_data(SPANDATA.CACHE_ITEM_SIZE, item_size)
  70. return value
  71. @functools.wraps(original_method)
  72. def sentry_method(*args, **kwargs):
  73. # type: (*Any, **Any) -> Any
  74. return _instrument_call(
  75. cache, method_name, original_method, args, kwargs, address, port
  76. )
  77. setattr(cache, method_name, sentry_method)
  78. def _patch_cache(cache, address=None, port=None):
  79. # type: (CacheHandler, Optional[str], Optional[int]) -> None
  80. if not hasattr(cache, "_sentry_patched"):
  81. for method_name in METHODS_TO_INSTRUMENT:
  82. _patch_cache_method(cache, method_name, address, port)
  83. cache._sentry_patched = True
  84. def _get_address_port(settings):
  85. # type: (dict[str, Any]) -> tuple[Optional[str], Optional[int]]
  86. location = settings.get("LOCATION")
  87. # TODO: location can also be an array of locations
  88. # see: https://docs.djangoproject.com/en/5.0/topics/cache/#redis
  89. # GitHub issue: https://github.com/getsentry/sentry-python/issues/3062
  90. if not isinstance(location, str):
  91. return None, None
  92. if "://" in location:
  93. parsed_url = urlparse(location)
  94. # remove the username and password from URL to not leak sensitive data.
  95. address = "{}://{}{}".format(
  96. parsed_url.scheme or "",
  97. parsed_url.hostname or "",
  98. parsed_url.path or "",
  99. )
  100. port = parsed_url.port
  101. else:
  102. address = location
  103. port = None
  104. return address, int(port) if port is not None else None
  105. def should_enable_cache_spans():
  106. # type: () -> bool
  107. from sentry_sdk.integrations.django import DjangoIntegration
  108. client = sentry_sdk.get_client()
  109. integration = client.get_integration(DjangoIntegration)
  110. from django.conf import settings
  111. return integration is not None and (
  112. (client.spotlight is not None and settings.DEBUG is True)
  113. or integration.cache_spans is True
  114. )
  115. def patch_caching():
  116. # type: () -> None
  117. if not hasattr(CacheHandler, "_sentry_patched"):
  118. if DJANGO_VERSION < (3, 2):
  119. original_get_item = CacheHandler.__getitem__
  120. @functools.wraps(original_get_item)
  121. def sentry_get_item(self, alias):
  122. # type: (CacheHandler, str) -> Any
  123. cache = original_get_item(self, alias)
  124. if should_enable_cache_spans():
  125. from django.conf import settings
  126. address, port = _get_address_port(
  127. settings.CACHES[alias or "default"]
  128. )
  129. _patch_cache(cache, address, port)
  130. return cache
  131. CacheHandler.__getitem__ = sentry_get_item
  132. CacheHandler._sentry_patched = True
  133. else:
  134. original_create_connection = CacheHandler.create_connection
  135. @functools.wraps(original_create_connection)
  136. def sentry_create_connection(self, alias):
  137. # type: (CacheHandler, str) -> Any
  138. cache = original_create_connection(self, alias)
  139. if should_enable_cache_spans():
  140. address, port = _get_address_port(self.settings[alias or "default"])
  141. _patch_cache(cache, address, port)
  142. return cache
  143. CacheHandler.create_connection = sentry_create_connection
  144. CacheHandler._sentry_patched = True