|
14 | 14 |
|
15 | 15 | import os |
16 | 16 |
|
17 | | -from newrelic.api.database_trace import (enable_datastore_instance_feature, |
18 | | - DatabaseTrace, register_database_client) |
| 17 | +from newrelic.api.database_trace import DatabaseTrace, register_database_client |
19 | 18 | from newrelic.api.function_trace import FunctionTrace |
20 | 19 | from newrelic.api.transaction import current_transaction |
21 | 20 | from newrelic.common.object_names import callable_name |
22 | 21 | from newrelic.common.object_wrapper import wrap_object |
| 22 | +from newrelic.hooks.database_dbapi2 import ConnectionFactory as DBAPI2ConnectionFactory |
| 23 | +from newrelic.hooks.database_dbapi2 import ConnectionWrapper as DBAPI2ConnectionWrapper |
23 | 24 |
|
24 | | -from newrelic.hooks.database_dbapi2 import (ConnectionWrapper as |
25 | | - DBAPI2ConnectionWrapper, ConnectionFactory as DBAPI2ConnectionFactory) |
26 | 25 |
|
27 | 26 | class ConnectionWrapper(DBAPI2ConnectionWrapper): |
28 | | - |
29 | 27 | def __enter__(self): |
30 | 28 | transaction = current_transaction() |
31 | 29 | name = callable_name(self.__wrapped__.__enter__) |
32 | 30 | with FunctionTrace(name, source=self.__wrapped__.__enter__): |
33 | | - cursor = self.__wrapped__.__enter__() |
| 31 | + cursor = self.__wrapped__.__enter__() |
34 | 32 |
|
35 | 33 | # The __enter__() method of original connection object returns |
36 | 34 | # a new cursor instance for use with 'as' assignment. We need |
37 | 35 | # to wrap that in a cursor wrapper otherwise we will not track |
38 | 36 | # any queries done via it. |
39 | 37 |
|
40 | | - return self.__cursor_wrapper__(cursor, self._nr_dbapi2_module, |
41 | | - self._nr_connect_params, None) |
| 38 | + return self.__cursor_wrapper__(cursor, self._nr_dbapi2_module, self._nr_connect_params, None) |
42 | 39 |
|
43 | 40 | def __exit__(self, exc, value, tb): |
44 | 41 | transaction = current_transaction() |
45 | 42 | name = callable_name(self.__wrapped__.__exit__) |
46 | 43 | with FunctionTrace(name, source=self.__wrapped__.__exit__): |
47 | 44 | if exc is None: |
48 | | - with DatabaseTrace('COMMIT', self._nr_dbapi2_module, self._nr_connect_params, source=self.__wrapped__.__exit__): |
| 45 | + with DatabaseTrace( |
| 46 | + "COMMIT", self._nr_dbapi2_module, self._nr_connect_params, source=self.__wrapped__.__exit__ |
| 47 | + ): |
49 | 48 | return self.__wrapped__.__exit__(exc, value, tb) |
50 | 49 | else: |
51 | | - with DatabaseTrace('ROLLBACK', self._nr_dbapi2_module, self._nr_connect_params, source=self.__wrapped__.__exit__): |
| 50 | + with DatabaseTrace( |
| 51 | + "ROLLBACK", self._nr_dbapi2_module, self._nr_connect_params, source=self.__wrapped__.__exit__ |
| 52 | + ): |
52 | 53 | return self.__wrapped__.__exit__(exc, value, tb) |
53 | 54 |
|
| 55 | + |
54 | 56 | class ConnectionFactory(DBAPI2ConnectionFactory): |
55 | 57 |
|
56 | 58 | __connection_wrapper__ = ConnectionWrapper |
57 | 59 |
|
| 60 | + |
58 | 61 | def instance_info(args, kwargs): |
59 | | - def _bind_params(host=None, user=None, passwd=None, db=None, port=None, |
60 | | - unix_socket=None, conv=None, connect_timeout=None, compress=None, |
61 | | - named_pipe=None, init_command=None, read_default_file=None, |
62 | | - read_default_group=None, *args, **kwargs): |
63 | | - return (host, port, db, unix_socket, |
64 | | - read_default_file, read_default_group) |
| 62 | + def _bind_params( |
| 63 | + host=None, |
| 64 | + user=None, |
| 65 | + passwd=None, |
| 66 | + db=None, |
| 67 | + port=None, |
| 68 | + unix_socket=None, |
| 69 | + conv=None, |
| 70 | + connect_timeout=None, |
| 71 | + compress=None, |
| 72 | + named_pipe=None, |
| 73 | + init_command=None, |
| 74 | + read_default_file=None, |
| 75 | + read_default_group=None, |
| 76 | + *args, |
| 77 | + **kwargs |
| 78 | + ): |
| 79 | + return (host, port, db, unix_socket, read_default_file, read_default_group) |
65 | 80 |
|
66 | 81 | params = _bind_params(*args, **kwargs) |
67 | 82 | host, port, db, unix_socket, read_default_file, read_default_group = params |
68 | 83 | explicit_host = host |
69 | 84 |
|
70 | 85 | port_path_or_id = None |
71 | 86 | if read_default_file or read_default_group: |
72 | | - host = host or 'default' |
73 | | - port_path_or_id = 'unknown' |
| 87 | + host = host or "default" |
| 88 | + port_path_or_id = "unknown" |
74 | 89 | elif not host: |
75 | | - host = 'localhost' |
| 90 | + host = "localhost" |
76 | 91 |
|
77 | | - if host == 'localhost': |
| 92 | + if host == "localhost": |
78 | 93 | # precedence: explicit -> cnf (if used) -> env -> 'default' |
79 | | - port_path_or_id = (unix_socket or |
80 | | - port_path_or_id or |
81 | | - os.getenv('MYSQL_UNIX_PORT', 'default')) |
| 94 | + port_path_or_id = unix_socket or port_path_or_id or os.getenv("MYSQL_UNIX_PORT", "default") |
82 | 95 | elif explicit_host: |
83 | 96 | # only reach here if host is explicitly passed in |
84 | 97 | port = port and str(port) |
85 | 98 | # precedence: explicit -> cnf (if used) -> env -> '3306' |
86 | | - port_path_or_id = (port or |
87 | | - port_path_or_id or |
88 | | - os.getenv('MYSQL_TCP_PORT', '3306')) |
| 99 | + port_path_or_id = port or port_path_or_id or os.getenv("MYSQL_TCP_PORT", "3306") |
89 | 100 |
|
90 | 101 | # There is no default database if omitted from the connect params |
91 | 102 | # In this case, we should report unknown |
92 | | - db = db or 'unknown' |
| 103 | + db = db or "unknown" |
93 | 104 |
|
94 | 105 | return (host, port_path_or_id, db) |
95 | 106 |
|
96 | | -def instrument_mysqldb(module): |
97 | | - register_database_client(module, database_product='MySQL', |
98 | | - quoting_style='single+double', explain_query='explain', |
99 | | - explain_stmts=('select',), instance_info=instance_info) |
100 | 107 |
|
101 | | - enable_datastore_instance_feature(module) |
| 108 | +def instrument_mysqldb(module): |
| 109 | + register_database_client( |
| 110 | + module, |
| 111 | + database_product="MySQL", |
| 112 | + quoting_style="single+double", |
| 113 | + explain_query="explain", |
| 114 | + explain_stmts=("select",), |
| 115 | + instance_info=instance_info, |
| 116 | + ) |
102 | 117 |
|
103 | | - wrap_object(module, 'connect', ConnectionFactory, (module,)) |
| 118 | + wrap_object(module, "connect", ConnectionFactory, (module,)) |
104 | 119 |
|
105 | 120 | # The connect() function is actually aliased with Connect() and |
106 | 121 | # Connection, the later actually being the Connection type object. |
107 | 122 | # Instrument Connect(), but don't instrument Connection in case that |
108 | 123 | # interferes with direct type usage. If people are using the |
109 | 124 | # Connection object directly, they should really be using connect(). |
110 | 125 |
|
111 | | - if hasattr(module, 'Connect'): |
112 | | - wrap_object(module, 'Connect', ConnectionFactory, (module,)) |
| 126 | + if hasattr(module, "Connect"): |
| 127 | + wrap_object(module, "Connect", ConnectionFactory, (module,)) |
0 commit comments