@@ -18,36 +18,53 @@ def setup_once():
18
18
# type: () -> None
19
19
import redis
20
20
21
- old_execute_command = redis .StrictRedis . execute_command
21
+ patch_redis_client ( redis .StrictRedis )
22
22
23
- def sentry_patched_execute_command (self , name , * args , ** kwargs ):
24
- # type: (redis.StrictRedis, str, *Any, **Any) -> Any
25
- hub = Hub .current
23
+ try :
24
+ import rb .clients # type: ignore
25
+ except ImportError :
26
+ pass
27
+ else :
28
+ patch_redis_client (rb .clients .FanoutClient )
29
+ patch_redis_client (rb .clients .MappingClient )
30
+ patch_redis_client (rb .clients .RoutingClient )
26
31
27
- if hub .get_integration (RedisIntegration ) is None :
28
- return old_execute_command (self , name , * args , ** kwargs )
29
32
30
- description = name
33
+ def patch_redis_client (cls ):
34
+ # type: (Any) -> None
35
+ """
36
+ This function can be used to instrument custom redis client classes or
37
+ subclasses.
38
+ """
31
39
32
- with capture_internal_exceptions ():
33
- description_parts = [name ]
34
- for i , arg in enumerate (args ):
35
- if i > 10 :
36
- break
40
+ old_execute_command = cls .execute_command
37
41
38
- description_parts .append (repr (arg ))
42
+ def sentry_patched_execute_command (self , name , * args , ** kwargs ):
43
+ # type: (Any, str, *Any, **Any) -> Any
44
+ hub = Hub .current
39
45
40
- description = " " .join (description_parts )
46
+ if hub .get_integration (RedisIntegration ) is None :
47
+ return old_execute_command (self , name , * args , ** kwargs )
41
48
42
- with hub .start_span (op = "redis" , description = description ) as span :
43
- if name :
44
- span .set_tag ("redis.command" , name )
49
+ description = name
45
50
46
- if name and args and name .lower () in ("get" , "set" , "setex" , "setnx" ):
47
- span .set_tag ("redis.key" , args [0 ])
51
+ with capture_internal_exceptions ():
52
+ description_parts = [name ]
53
+ for i , arg in enumerate (args ):
54
+ if i > 10 :
55
+ break
48
56
49
- return old_execute_command ( self , name , * args , ** kwargs )
57
+ description_parts . append ( repr ( arg ) )
50
58
51
- redis .StrictRedis .execute_command = ( # type: ignore
52
- sentry_patched_execute_command # type: ignore
53
- )
59
+ description = " " .join (description_parts )
60
+
61
+ with hub .start_span (op = "redis" , description = description ) as span :
62
+ if name :
63
+ span .set_tag ("redis.command" , name )
64
+
65
+ if name and args and name .lower () in ("get" , "set" , "setex" , "setnx" ):
66
+ span .set_tag ("redis.key" , args [0 ])
67
+
68
+ return old_execute_command (self , name , * args , ** kwargs )
69
+
70
+ cls .execute_command = sentry_patched_execute_command
0 commit comments