Skip to content

Commit 994a45b

Browse files
Zhenayantonpirker
Zhenay
and
antonpirker
authored
Redis: Add support for redis.asyncio (#1933)
 --------- Co-authored-by: Anton Pirker <anton.pirker@sentry.io>
1 parent b89fa8d commit 994a45b

File tree

8 files changed

+310
-83
lines changed

8 files changed

+310
-83
lines changed

‎.github/workflows/test-integration-redis.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ jobs:
3131
strategy:
3232
fail-fast: false
3333
matrix:
34-
python-version: ["3.7","3.8","3.9"]
34+
python-version: ["3.7","3.8","3.9","3.10","3.11"]
3535
# python3.6 reached EOL and is no longer being supported on
3636
# new versions of hosted runners on Github Actions
3737
# ubuntu-20.04 is the last version that supported python3.6

sentry_sdk/integrations/redis.py renamed to sentry_sdk/integrations/redis/__init__.py

Lines changed: 127 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
if TYPE_CHECKING:
1616
from typing import Any, Sequence
17+
from sentry_sdk.tracing import Span
1718

1819
_SINGLE_KEY_COMMANDS = frozenset(
1920
["decr", "decrby", "get", "incr", "incrby", "pttl", "set", "setex", "setnx", "ttl"]
@@ -25,10 +26,64 @@
2526
]
2627

2728
_MAX_NUM_ARGS = 10 # Trim argument lists to this many values
29+
_MAX_NUM_COMMANDS = 10 # Trim command lists to this many values
2830

2931
_DEFAULT_MAX_DATA_SIZE = 1024
3032

3133

34+
def _get_safe_command(name, args):
35+
# type: (str, Sequence[Any]) -> str
36+
command_parts = [name]
37+
38+
for i, arg in enumerate(args):
39+
if i > _MAX_NUM_ARGS:
40+
break
41+
42+
name_low = name.lower()
43+
44+
if name_low in _COMMANDS_INCLUDING_SENSITIVE_DATA:
45+
command_parts.append(SENSITIVE_DATA_SUBSTITUTE)
46+
continue
47+
48+
arg_is_the_key = i == 0
49+
if arg_is_the_key:
50+
command_parts.append(repr(arg))
51+
52+
else:
53+
if _should_send_default_pii():
54+
command_parts.append(repr(arg))
55+
else:
56+
command_parts.append(SENSITIVE_DATA_SUBSTITUTE)
57+
58+
command = " ".join(command_parts)
59+
return command
60+
61+
62+
def _set_pipeline_data(
63+
span, is_cluster, get_command_args_fn, is_transaction, command_stack
64+
):
65+
# type: (Span, bool, Any, bool, Sequence[Any]) -> None
66+
span.set_tag("redis.is_cluster", is_cluster)
67+
transaction = is_transaction if not is_cluster else False
68+
span.set_tag("redis.transaction", transaction)
69+
70+
commands = []
71+
for i, arg in enumerate(command_stack):
72+
if i >= _MAX_NUM_COMMANDS:
73+
break
74+
75+
command = get_command_args_fn(arg)
76+
commands.append(_get_safe_command(command[0], command[1:]))
77+
78+
span.set_data(
79+
"redis.commands",
80+
{
81+
"count": len(command_stack),
82+
"first_ten": commands,
83+
},
84+
)
85+
86+
3287
def patch_redis_pipeline(pipeline_cls, is_cluster, get_command_args_fn):
3388
# type: (Any, bool, Any) -> None
3489
old_execute = pipeline_cls.execute
@@ -44,24 +99,12 @@ def sentry_patched_execute(self, *args, **kwargs):
4499
op=OP.DB_REDIS, description="redis.pipeline.execute"
45100
) as span:
46101
with capture_internal_exceptions():
47-
span.set_tag("redis.is_cluster", is_cluster)
48-
transaction = self.transaction if not is_cluster else False
49-
span.set_tag("redis.transaction", transaction)
50-
51-
commands = []
52-
for i, arg in enumerate(self.command_stack):
53-
if i > _MAX_NUM_ARGS:
54-
break
55-
command_args = []
56-
for j, command_arg in enumerate(get_command_args_fn(arg)):
57-
if j > 0:
58-
command_arg = repr(command_arg)
59-
command_args.append(command_arg)
60-
commands.append(" ".join(command_args))
61-
62-
span.set_data(
63-
"redis.commands",
64-
{"count": len(self.command_stack), "first_ten": commands},
102+
_set_pipeline_data(
103+
span,
104+
is_cluster,
105+
get_command_args_fn,
106+
self.transaction,
107+
self.command_stack,
65108
)
66109
span.set_data(SPANDATA.DB_SYSTEM, "redis")
67110

@@ -80,6 +123,43 @@ def _parse_rediscluster_command(command):
80123
return command.args
81124

82125

126+
def _patch_redis(StrictRedis, client): # noqa: N803
127+
# type: (Any, Any) -> None
128+
patch_redis_client(StrictRedis, is_cluster=False)
129+
patch_redis_pipeline(client.Pipeline, False, _get_redis_command_args)
130+
try:
131+
strict_pipeline = client.StrictPipeline
132+
except AttributeError:
133+
pass
134+
else:
135+
patch_redis_pipeline(strict_pipeline, False, _get_redis_command_args)
136+
137+
try:
138+
import redis.asyncio
139+
except ImportError:
140+
pass
141+
else:
142+
from sentry_sdk.integrations.redis.asyncio import (
143+
patch_redis_async_client,
144+
patch_redis_async_pipeline,
145+
)
146+
147+
patch_redis_async_client(redis.asyncio.client.StrictRedis)
148+
patch_redis_async_pipeline(redis.asyncio.client.Pipeline)
149+
150+
151+
def _patch_rb():
152+
# type: () -> None
153+
try:
154+
import rb.clients # type: ignore
155+
except ImportError:
156+
pass
157+
else:
158+
patch_redis_client(rb.clients.FanoutClient, is_cluster=False)
159+
patch_redis_client(rb.clients.MappingClient, is_cluster=False)
160+
patch_redis_client(rb.clients.RoutingClient, is_cluster=False)
161+
162+
83163
def _patch_rediscluster():
84164
# type: () -> None
85165
try:
@@ -119,30 +199,40 @@ def setup_once():
119199
except ImportError:
120200
raise DidNotEnable("Redis client not installed")
121201

122-
patch_redis_client(StrictRedis, is_cluster=False)
123-
patch_redis_pipeline(client.Pipeline, False, _get_redis_command_args)
124-
try:
125-
strict_pipeline = client.StrictPipeline # type: ignore
126-
except AttributeError:
127-
pass
128-
else:
129-
patch_redis_pipeline(strict_pipeline, False, _get_redis_command_args)
130-
131-
try:
132-
import rb.clients # type: ignore
133-
except ImportError:
134-
pass
135-
else:
136-
patch_redis_client(rb.clients.FanoutClient, is_cluster=False)
137-
patch_redis_client(rb.clients.MappingClient, is_cluster=False)
138-
patch_redis_client(rb.clients.RoutingClient, is_cluster=False)
202+
_patch_redis(StrictRedis, client)
203+
_patch_rb()
139204

140205
try:
141206
_patch_rediscluster()
142207
except Exception:
143208
logger.exception("Error occurred while patching `rediscluster` library")
144209

145210

211+
def _get_span_description(name, *args):
212+
# type: (str, *Any) -> str
213+
description = name
214+
215+
with capture_internal_exceptions():
216+
description = _get_safe_command(name, args)
217+
218+
return description
219+
220+
221+
def _set_client_data(span, is_cluster, name, *args):
222+
# type: (Span, bool, str, *Any) -> None
223+
span.set_tag("redis.is_cluster", is_cluster)
224+
if name:
225+
span.set_tag("redis.command", name)
226+
span.set_tag(SPANDATA.DB_OPERATION, name)
227+
228+
if name and args:
229+
name_low = name.lower()
230+
if (name_low in _SINGLE_KEY_COMMANDS) or (
231+
name_low in _MULTI_KEY_COMMANDS and len(args) == 1
232+
):
233+
span.set_tag("redis.key", args[0])
234+
235+
146236
def patch_redis_client(cls, is_cluster):
147237
# type: (Any, bool) -> None
148238
"""
@@ -159,31 +249,7 @@ def sentry_patched_execute_command(self, name, *args, **kwargs):
159249
if integration is None:
160250
return old_execute_command(self, name, *args, **kwargs)
161251

162-
description = name
163-
164-
with capture_internal_exceptions():
165-
description_parts = [name]
166-
for i, arg in enumerate(args):
167-
if i > _MAX_NUM_ARGS:
168-
break
169-
170-
name_low = name.lower()
171-
172-
if name_low in _COMMANDS_INCLUDING_SENSITIVE_DATA:
173-
description_parts.append(SENSITIVE_DATA_SUBSTITUTE)
174-
continue
175-
176-
arg_is_the_key = i == 0
177-
if arg_is_the_key:
178-
description_parts.append(repr(arg))
179-
180-
else:
181-
if _should_send_default_pii():
182-
description_parts.append(repr(arg))
183-
else:
184-
description_parts.append(SENSITIVE_DATA_SUBSTITUTE)
185-
186-
description = " ".join(description_parts)
252+
description = _get_span_description(name, *args)
187253

188254
data_should_be_truncated = (
189255
integration.max_data_size and len(description) > integration.max_data_size
@@ -192,18 +258,7 @@ def sentry_patched_execute_command(self, name, *args, **kwargs):
192258
description = description[: integration.max_data_size - len("...")] + "..."
193259

194260
with hub.start_span(op=OP.DB_REDIS, description=description) as span:
195-
span.set_tag("redis.is_cluster", is_cluster)
196-
197-
if name:
198-
span.set_tag("redis.command", name)
199-
span.set_tag(SPANDATA.DB_OPERATION, name)
200-
201-
if name and args:
202-
name_low = name.lower()
203-
if (name_low in _SINGLE_KEY_COMMANDS) or (
204-
name_low in _MULTI_KEY_COMMANDS and len(args) == 1
205-
):
206-
span.set_tag("redis.key", args[0])
261+
_set_client_data(span, is_cluster, name, *args)
207262

208263
return old_execute_command(self, name, *args, **kwargs)
209264

‎sentry_sdk/integrations/redis/asyncio.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
from __future__ import absolute_import
2+
3+
from sentry_sdk import Hub
4+
from sentry_sdk.consts import OP
5+
from sentry_sdk.utils import capture_internal_exceptions
6+
from sentry_sdk.integrations.redis import (
7+
RedisIntegration,
8+
_get_redis_command_args,
9+
_get_span_description,
10+
_set_client_data,
11+
_set_pipeline_data,
12+
)
13+
14+
15+
from sentry_sdk._types import MYPY
16+
17+
if MYPY:
18+
from typing import Any
19+
20+
21+
def patch_redis_async_pipeline(pipeline_cls):
22+
# type: (Any) -> None
23+
old_execute = pipeline_cls.execute
24+
25+
async def _sentry_execute(self, *args, **kwargs):
26+
# type: (Any, *Any, **Any) -> Any
27+
hub = Hub.current
28+
29+
if hub.get_integration(RedisIntegration) is None:
30+
return await old_execute(self, *args, **kwargs)
31+
32+
with hub.start_span(
33+
op=OP.DB_REDIS, description="redis.pipeline.execute"
34+
) as span:
35+
with capture_internal_exceptions():
36+
_set_pipeline_data(
37+
span,
38+
False,
39+
_get_redis_command_args,
40+
self.is_transaction,
41+
self.command_stack,
42+
)
43+
44+
return await old_execute(self, *args, **kwargs)
45+
46+
pipeline_cls.execute = _sentry_execute
47+
48+
49+
def patch_redis_async_client(cls):
50+
# type: (Any) -> None
51+
old_execute_command = cls.execute_command
52+
53+
async def _sentry_execute_command(self, name, *args, **kwargs):
54+
# type: (Any, str, *Any, **Any) -> Any
55+
hub = Hub.current
56+
57+
if hub.get_integration(RedisIntegration) is None:
58+
return await old_execute_command(self, name, *args, **kwargs)
59+
60+
description = _get_span_description(name, *args)
61+
62+
with hub.start_span(op=OP.DB_REDIS, description=description) as span:
63+
_set_client_data(span, False, name, *args)
64+
65+
return await old_execute_command(self, name, *args, **kwargs)
66+
67+
cls.execute_command = _sentry_execute_command

‎tests/integrations/redis/asyncio/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
import pytest
2+
3+
pytest.importorskip("fakeredis.aioredis")

0 commit comments

Comments
 (0)