1+ # Copyright 2010 New Relic, Inc.
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ from newrelic .api .datastore_trace import DatastoreTrace
16+ from newrelic .api .time_trace import current_trace
17+ from newrelic .api .transaction import current_transaction
18+ from newrelic .common .object_wrapper import wrap_function_wrapper
19+ from newrelic .hooks .datastore_redis import (
20+ _redis_client_methods ,
21+ _redis_multipart_commands ,
22+ _redis_operation_re ,
23+ )
24+
25+
26+ def _conn_attrs_to_dict (connection ):
27+ host = getattr (connection , "host" , None )
28+ port = getattr (connection , "port" , None )
29+ if not host and not port and hasattr (connection , "_address" ):
30+ host , port = connection ._address
31+ return {
32+ "host" : host ,
33+ "port" : port ,
34+ "path" : getattr (connection , "path" , None ),
35+ "db" : getattr (connection , "db" , getattr (connection , "_db" , None )),
36+ }
37+
38+
39+ def _instance_info (kwargs ):
40+ host = kwargs .get ("host" ) or "localhost"
41+ port_path_or_id = str (kwargs .get ("port" ) or kwargs .get ("path" , 6379 ))
42+ db = str (kwargs .get ("db" ) or 0 )
43+
44+ return (host , port_path_or_id , db )
45+
46+
47+ def _wrap_AioRedis_method_wrapper (module , instance_class_name , operation ):
48+ async def _nr_wrapper_AioRedis_method_ (wrapped , instance , args , kwargs ):
49+ transaction = current_transaction ()
50+ if transaction is None :
51+ return await wrapped (* args , ** kwargs )
52+
53+ with DatastoreTrace (product = "Redis" , target = None , operation = operation ):
54+ return await wrapped (* args , ** kwargs )
55+
56+ name = "%s.%s" % (instance_class_name , operation )
57+ wrap_function_wrapper (module , name , _nr_wrapper_AioRedis_method_ )
58+
59+
60+ async def wrap_Connection_send_command (wrapped , instance , args , kwargs ):
61+ transaction = current_transaction ()
62+ if not transaction :
63+ return await wrapped (* args , ** kwargs )
64+
65+ host , port_path_or_id , db = (None , None , None )
66+
67+ try :
68+ dt = transaction .settings .datastore_tracer
69+ if dt .instance_reporting .enabled or dt .database_name_reporting .enabled :
70+ conn_kwargs = _conn_attrs_to_dict (instance )
71+ host , port_path_or_id , db = _instance_info (conn_kwargs )
72+ except Exception :
73+ pass
74+
75+ # Older Redis clients would when sending multi part commands pass
76+ # them in as separate arguments to send_command(). Need to therefore
77+ # detect those and grab the next argument from the set of arguments.
78+
79+ operation = args [0 ].strip ().lower ()
80+
81+ # If it's not a multi part command, there's no need to trace it, so
82+ # we can return early.
83+
84+ if operation .split ()[0 ] not in _redis_multipart_commands : # Set the datastore info on the DatastoreTrace containing this function call.
85+ trace = current_trace ()
86+
87+ # Find DatastoreTrace no matter how many other traces are inbetween
88+ while trace is not None and not isinstance (trace , DatastoreTrace ):
89+ trace = getattr (trace , "parent" , None )
90+
91+ if trace is not None :
92+ trace .host = host
93+ trace .port_path_or_id = port_path_or_id
94+ trace .database_name = db
95+
96+ return await wrapped (* args , ** kwargs )
97+
98+ # Convert multi args to single arg string
99+
100+ if operation in _redis_multipart_commands and len (args ) > 1 :
101+ operation = "%s %s" % (operation , args [1 ].strip ().lower ())
102+
103+ operation = _redis_operation_re .sub ("_" , operation )
104+
105+ with DatastoreTrace (
106+ product = "Redis" , target = None , operation = operation , host = host , port_path_or_id = port_path_or_id , database_name = db
107+ ):
108+ return await wrapped (* args , ** kwargs )
109+
110+
111+ def instrument_aioredis_client (module ):
112+ # StrictRedis is just an alias of Redis, no need to wrap it as well.
113+ if hasattr (module , "Redis" ):
114+ class_ = getattr (module , "Redis" )
115+ for operation in _redis_client_methods :
116+ if hasattr (class_ , operation ):
117+ _wrap_AioRedis_method_wrapper (module , "Redis" , operation )
118+
119+
120+ def instrument_aioredis_connection (module ):
121+ if hasattr (module , "Connection" ):
122+ if hasattr (module .Connection , "send_command" ):
123+ wrap_function_wrapper (module , "Connection.send_command" , wrap_Connection_send_command )
124+
125+ if hasattr (module , "RedisConnection" ):
126+ if hasattr (module .RedisConnection , "execute" ):
127+ wrap_function_wrapper (module , "RedisConnection.execute" , wrap_Connection_send_command )
0 commit comments