utils.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. from sentry_sdk.consts import SPANDATA
  2. from sentry_sdk.integrations.redis.consts import (
  3. _COMMANDS_INCLUDING_SENSITIVE_DATA,
  4. _MAX_NUM_ARGS,
  5. _MAX_NUM_COMMANDS,
  6. _MULTI_KEY_COMMANDS,
  7. _SINGLE_KEY_COMMANDS,
  8. )
  9. from sentry_sdk.scope import should_send_default_pii
  10. from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
  11. from typing import TYPE_CHECKING
  12. if TYPE_CHECKING:
  13. from typing import Any, Optional, Sequence
  14. from sentry_sdk.tracing import Span
  15. def _get_safe_command(name, args):
  16. # type: (str, Sequence[Any]) -> str
  17. command_parts = [name]
  18. for i, arg in enumerate(args):
  19. if i > _MAX_NUM_ARGS:
  20. break
  21. name_low = name.lower()
  22. if name_low in _COMMANDS_INCLUDING_SENSITIVE_DATA:
  23. command_parts.append(SENSITIVE_DATA_SUBSTITUTE)
  24. continue
  25. arg_is_the_key = i == 0
  26. if arg_is_the_key:
  27. command_parts.append(repr(arg))
  28. else:
  29. if should_send_default_pii():
  30. command_parts.append(repr(arg))
  31. else:
  32. command_parts.append(SENSITIVE_DATA_SUBSTITUTE)
  33. command = " ".join(command_parts)
  34. return command
  35. def _safe_decode(key):
  36. # type: (Any) -> str
  37. if isinstance(key, bytes):
  38. try:
  39. return key.decode()
  40. except UnicodeDecodeError:
  41. return ""
  42. return str(key)
  43. def _key_as_string(key):
  44. # type: (Any) -> str
  45. if isinstance(key, (dict, list, tuple)):
  46. key = ", ".join(_safe_decode(x) for x in key)
  47. elif isinstance(key, bytes):
  48. key = _safe_decode(key)
  49. elif key is None:
  50. key = ""
  51. else:
  52. key = str(key)
  53. return key
  54. def _get_safe_key(method_name, args, kwargs):
  55. # type: (str, Optional[tuple[Any, ...]], Optional[dict[str, Any]]) -> Optional[tuple[str, ...]]
  56. """
  57. Gets the key (or keys) from the given method_name.
  58. The method_name could be a redis command or a django caching command
  59. """
  60. key = None
  61. if args is not None and method_name.lower() in _MULTI_KEY_COMMANDS:
  62. # for example redis "mget"
  63. key = tuple(args)
  64. elif args is not None and len(args) >= 1:
  65. # for example django "set_many/get_many" or redis "get"
  66. if isinstance(args[0], (dict, list, tuple)):
  67. key = tuple(args[0])
  68. else:
  69. key = (args[0],)
  70. elif kwargs is not None and "key" in kwargs:
  71. # this is a legacy case for older versions of Django
  72. if isinstance(kwargs["key"], (list, tuple)):
  73. if len(kwargs["key"]) > 0:
  74. key = tuple(kwargs["key"])
  75. else:
  76. if kwargs["key"] is not None:
  77. key = (kwargs["key"],)
  78. return key
  79. def _parse_rediscluster_command(command):
  80. # type: (Any) -> Sequence[Any]
  81. return command.args
  82. def _set_pipeline_data(
  83. span,
  84. is_cluster,
  85. get_command_args_fn,
  86. is_transaction,
  87. commands_seq,
  88. ):
  89. # type: (Span, bool, Any, bool, Sequence[Any]) -> None
  90. span.set_tag("redis.is_cluster", is_cluster)
  91. span.set_tag("redis.transaction", is_transaction)
  92. commands = []
  93. for i, arg in enumerate(commands_seq):
  94. if i >= _MAX_NUM_COMMANDS:
  95. break
  96. command = get_command_args_fn(arg)
  97. commands.append(_get_safe_command(command[0], command[1:]))
  98. span.set_data(
  99. "redis.commands",
  100. {
  101. "count": len(commands_seq),
  102. "first_ten": commands,
  103. },
  104. )
  105. def _set_client_data(span, is_cluster, name, *args):
  106. # type: (Span, bool, str, *Any) -> None
  107. span.set_tag("redis.is_cluster", is_cluster)
  108. if name:
  109. span.set_tag("redis.command", name)
  110. span.set_tag(SPANDATA.DB_OPERATION, name)
  111. if name and args:
  112. name_low = name.lower()
  113. if (name_low in _SINGLE_KEY_COMMANDS) or (
  114. name_low in _MULTI_KEY_COMMANDS and len(args) == 1
  115. ):
  116. span.set_tag("redis.key", args[0])