Skip to content

Commit 2a3a3dc

Browse files
hmstepanekhmstepanek
authored
Merge pull request #1265 from newrelic/feature-aiomysql
aiomysql Instrumentation
2 parents 75d6a4d + da57d59 commit 2a3a3dc

File tree

5 files changed

+309
-1
lines changed

5 files changed

+309
-1
lines changed

‎newrelic/config.py

+3
Original file line numberDiff line numberDiff line change
@@ -3166,6 +3166,9 @@ def _process_module_builtin_defaults():
31663166
_process_module_definition("MySQLdb", "newrelic.hooks.database_mysqldb", "instrument_mysqldb")
31673167
_process_module_definition("pymysql", "newrelic.hooks.database_pymysql", "instrument_pymysql")
31683168

3169+
_process_module_definition("aiomysql", "newrelic.hooks.database_aiomysql", "instrument_aiomysql")
3170+
_process_module_definition("aiomysql.pool", "newrelic.hooks.database_aiomysql", "instrument_aiomysql_pool")
3171+
31693172
_process_module_definition("pyodbc", "newrelic.hooks.database_pyodbc", "instrument_pyodbc")
31703173

31713174
_process_module_definition("pymssql", "newrelic.hooks.database_pymssql", "instrument_pymssql")

‎newrelic/hooks/database_aiomysql.py

+106
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
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+
import sys
16+
17+
from newrelic.api.database_trace import register_database_client
18+
from newrelic.api.function_trace import FunctionTrace
19+
from newrelic.common.object_names import callable_name
20+
from newrelic.common.object_wrapper import (
21+
ObjectProxy,
22+
wrap_function_wrapper,
23+
wrap_object,
24+
)
25+
from newrelic.hooks.database_dbapi2_async import (
26+
AsyncConnectionFactory as DBAPI2AsyncConnectionFactory,
27+
)
28+
from newrelic.hooks.database_dbapi2_async import (
29+
AsyncConnectionWrapper as DBAPI2AsyncConnectionWrapper,
30+
)
31+
from newrelic.hooks.database_dbapi2_async import (
32+
AsyncCursorWrapper as DBAPI2AsyncCursorWrapper,
33+
)
34+
35+
36+
class AsyncCursorContextManagerWrapper(ObjectProxy):
37+
38+
__cursor_wrapper__ = DBAPI2AsyncCursorWrapper
39+
40+
def __init__(self, context_manager, dbapi2_module, connect_params, cursor_args):
41+
super().__init__(context_manager)
42+
self._nr_dbapi2_module = dbapi2_module
43+
self._nr_connect_params = connect_params
44+
self._nr_cursor_args = cursor_args
45+
46+
async def __aenter__(self):
47+
cursor = await self.__wrapped__.__aenter__()
48+
return self.__cursor_wrapper__(cursor, self._nr_dbapi2_module, self._nr_connect_params, self._nr_cursor_args)
49+
50+
async def __aexit__(self, exc, val, tb):
51+
return await self.__wrapped__.__aexit__(exc, val, tb)
52+
53+
54+
class AsyncConnectionWrapper(DBAPI2AsyncConnectionWrapper):
55+
56+
__cursor_wrapper__ = AsyncCursorContextManagerWrapper
57+
58+
59+
class AsyncConnectionFactory(DBAPI2AsyncConnectionFactory):
60+
61+
__connection_wrapper__ = AsyncConnectionWrapper
62+
63+
64+
def wrap_pool__acquire(dbapi2_module):
65+
async def _wrap_pool__acquire(wrapped, instance, args, kwargs):
66+
rollup = ["Datastore/all", f"Datastore/{dbapi2_module._nr_database_product}/all"]
67+
68+
with FunctionTrace(name=callable_name(wrapped), terminal=True, rollup=rollup, source=wrapped):
69+
connection = await wrapped(*args, **kwargs)
70+
connection_kwargs = getattr(instance, "_conn_kwargs", {})
71+
return AsyncConnectionWrapper(connection, dbapi2_module, (((), connection_kwargs)))
72+
73+
return _wrap_pool__acquire
74+
75+
76+
def instance_info(args, kwargs):
77+
def _bind_params(host=None, user=None, password=None, db=None, port=None, *args, **kwargs):
78+
return host, port, db
79+
80+
host, port, db = _bind_params(*args, **kwargs)
81+
82+
return (host, port, db)
83+
84+
85+
def instrument_aiomysql(module):
86+
register_database_client(
87+
module,
88+
database_product="MySQL",
89+
quoting_style="single+double",
90+
explain_query="explain",
91+
explain_stmts=("select",),
92+
instance_info=instance_info,
93+
)
94+
95+
# Only instrument the connect method directly, don't instrument
96+
# Connection. This follows the DBAPI2 spec and what was done for
97+
# PyMySQL which this library is based on.
98+
99+
wrap_object(module, "connect", AsyncConnectionFactory, (module,))
100+
101+
102+
def instrument_aiomysql_pool(module):
103+
dbapi2_module = sys.modules["aiomysql"]
104+
if hasattr(module, "Pool"):
105+
if hasattr(module.Pool, "_acquire"):
106+
wrap_function_wrapper(module, "Pool._acquire", wrap_pool__acquire(dbapi2_module))

‎tests/datastore_aiomysql/conftest.py

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
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+
16+
from testing_support.fixture.event_loop import ( # noqa: F401; pylint: disable=W0611
17+
event_loop as loop,
18+
)
19+
from testing_support.fixtures import ( # noqa: F401; pylint: disable=W0611
20+
collector_agent_registration_fixture,
21+
collector_available_fixture,
22+
)
23+
24+
_default_settings = {
25+
"package_reporting.enabled": False, # Turn off package reporting for testing as it causes slow downs.
26+
"transaction_tracer.explain_threshold": 0.0,
27+
"transaction_tracer.transaction_threshold": 0.0,
28+
"transaction_tracer.stack_trace_threshold": 0.0,
29+
"debug.log_data_collector_payloads": True,
30+
"debug.record_transaction_failure": True,
31+
"debug.log_explain_plan_queries": True,
32+
}
33+
34+
collector_agent_registration = collector_agent_registration_fixture(
35+
app_name="Python Agent Test (datastore_aiomysql)",
36+
default_settings=_default_settings,
37+
linked_applications=["Python Agent Test (datastore)"],
38+
)

‎tests/datastore_aiomysql/test_database.py

+157
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
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+
import aiomysql
16+
from testing_support.db_settings import mysql_settings
17+
from testing_support.util import instance_hostname
18+
from testing_support.validators.validate_database_trace_inputs import (
19+
validate_database_trace_inputs,
20+
)
21+
from testing_support.validators.validate_transaction_metrics import (
22+
validate_transaction_metrics,
23+
)
24+
25+
from newrelic.api.background_task import background_task
26+
27+
DB_SETTINGS = mysql_settings()[0]
28+
TABLE_NAME = f"datastore_aiomysql_{DB_SETTINGS['namespace']}"
29+
PROCEDURE_NAME = f"hello_{DB_SETTINGS['namespace']}"
30+
31+
HOST = instance_hostname(DB_SETTINGS["host"])
32+
PORT = DB_SETTINGS["port"]
33+
34+
35+
async def execute_db_calls_with_cursor(cursor):
36+
await cursor.execute(f"""drop table if exists {TABLE_NAME}""")
37+
38+
await cursor.execute(f"create table {TABLE_NAME} (a integer, b real, c text)")
39+
40+
await cursor.executemany(
41+
f"insert into {TABLE_NAME} values (%s, %s, %s)",
42+
[(1, 1.0, "1.0"), (2, 2.2, "2.2"), (3, 3.3, "3.3")],
43+
)
44+
45+
await cursor.execute(f"""select * from {TABLE_NAME}""")
46+
47+
async for _ in cursor:
48+
pass
49+
50+
await cursor.execute(f"update {TABLE_NAME} set a=%s, b=%s, c=%s where a=%s", (4, 4.0, "4.0", 1))
51+
52+
await cursor.execute(f"""delete from {TABLE_NAME} where a=2""")
53+
await cursor.execute(f"""drop procedure if exists {PROCEDURE_NAME}""")
54+
await cursor.execute(
55+
f"""CREATE PROCEDURE {PROCEDURE_NAME}()
56+
BEGIN
57+
SELECT 'Hello World!';
58+
END"""
59+
)
60+
61+
await cursor.callproc(PROCEDURE_NAME)
62+
63+
64+
SCOPED_METRICS = [
65+
(f"Datastore/statement/MySQL/{TABLE_NAME}/select", 1),
66+
(f"Datastore/statement/MySQL/{TABLE_NAME}/insert", 1),
67+
(f"Datastore/statement/MySQL/{TABLE_NAME}/update", 1),
68+
(f"Datastore/statement/MySQL/{TABLE_NAME}/delete", 1),
69+
("Datastore/operation/MySQL/drop", 2),
70+
("Datastore/operation/MySQL/create", 2),
71+
(f"Datastore/statement/MySQL/{PROCEDURE_NAME}/call", 1),
72+
("Datastore/operation/MySQL/commit", 2),
73+
("Datastore/operation/MySQL/rollback", 1),
74+
]
75+
76+
ROLLUP_METRICS = [
77+
("Datastore/all", 13),
78+
("Datastore/allOther", 13),
79+
("Datastore/MySQL/all", 13),
80+
("Datastore/MySQL/allOther", 13),
81+
(f"Datastore/statement/MySQL/{TABLE_NAME}/select", 1),
82+
(f"Datastore/statement/MySQL/{TABLE_NAME}/insert", 1),
83+
(f"Datastore/statement/MySQL/{TABLE_NAME}/update", 1),
84+
(f"Datastore/statement/MySQL/{TABLE_NAME}/delete", 1),
85+
("Datastore/operation/MySQL/select", 1),
86+
("Datastore/operation/MySQL/insert", 1),
87+
("Datastore/operation/MySQL/update", 1),
88+
("Datastore/operation/MySQL/delete", 1),
89+
(f"Datastore/statement/MySQL/{PROCEDURE_NAME}/call", 1),
90+
("Datastore/operation/MySQL/call", 1),
91+
("Datastore/operation/MySQL/drop", 2),
92+
("Datastore/operation/MySQL/create", 2),
93+
("Datastore/operation/MySQL/commit", 2),
94+
("Datastore/operation/MySQL/rollback", 1),
95+
(f"Datastore/instance/MySQL/{HOST}/{PORT}", 12),
96+
]
97+
98+
99+
@validate_transaction_metrics(
100+
"test_database:test_execute_via_connection",
101+
scoped_metrics=list(SCOPED_METRICS) + [("Function/aiomysql.connection:connect", 1)],
102+
rollup_metrics=list(ROLLUP_METRICS) + [("Function/aiomysql.connection:connect", 1)],
103+
background_task=True,
104+
)
105+
@validate_database_trace_inputs(sql_parameters_type=tuple)
106+
@background_task()
107+
def test_execute_via_connection(loop):
108+
async def _test():
109+
connection = await aiomysql.connect(
110+
db=DB_SETTINGS["name"],
111+
user=DB_SETTINGS["user"],
112+
password=DB_SETTINGS["password"],
113+
host=DB_SETTINGS["host"],
114+
port=DB_SETTINGS["port"],
115+
)
116+
117+
async with connection:
118+
async with connection.cursor() as cursor:
119+
await execute_db_calls_with_cursor(cursor)
120+
121+
await connection.commit()
122+
await connection.rollback()
123+
await connection.commit()
124+
125+
loop.run_until_complete(_test())
126+
127+
128+
@validate_transaction_metrics(
129+
"test_database:test_execute_via_pool",
130+
scoped_metrics=list(SCOPED_METRICS) + [("Function/aiomysql.pool:Pool._acquire", 1)],
131+
rollup_metrics=list(ROLLUP_METRICS) + [("Function/aiomysql.pool:Pool._acquire", 1)],
132+
background_task=True,
133+
)
134+
@validate_database_trace_inputs(sql_parameters_type=tuple)
135+
@background_task()
136+
def test_execute_via_pool(loop):
137+
async def _test():
138+
pool = await aiomysql.create_pool(
139+
db=DB_SETTINGS["name"],
140+
user=DB_SETTINGS["user"],
141+
password=DB_SETTINGS["password"],
142+
host=DB_SETTINGS["host"],
143+
port=DB_SETTINGS["port"],
144+
loop=loop,
145+
)
146+
async with pool.acquire() as connection:
147+
async with connection.cursor() as cursor:
148+
await execute_db_calls_with_cursor(cursor)
149+
150+
await connection.commit()
151+
await connection.rollback()
152+
await connection.commit()
153+
154+
pool.close()
155+
await pool.wait_closed()
156+
157+
loop.run_until_complete(_test())

‎tox.ini

+5-1
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ envlist =
6161
mongodb8-datastore_motor-{py37,py38,py39,py310,py311,py312,py313}-motorlatest,
6262
mongodb3-datastore_pymongo-{py37,py38,py39,py310,py311,py312}-pymongo03,
6363
mongodb8-datastore_pymongo-{py37,py38,py39,py310,py311,py312,py313,pypy310}-pymongo04,
64+
mysql-datastore_aiomysql-{py37,py38,py39,py310,py311,py312,py313,pypy310},
6465
mssql-datastore_pymssql-{py37,py38,py39,py310,py311,py312,py313},
6566
mysql-datastore_mysql-mysqllatest-{py37,py38,py39,py310,py311,py312,py313},
6667
mysql-datastore_pymysql-{py37,py38,py39,py310,py311,py312,py313,pypy310},
@@ -251,6 +252,8 @@ deps =
251252
cross_agent: requests
252253
datastore_asyncpg: asyncpg
253254
datastore_aiomcache: aiomcache
255+
datastore_aiomysql: aiomysql
256+
datastore_aiomysql: cryptography
254257
datastore_bmemcached: python-binary-memcached
255258
datastore_elasticsearch: requests
256259
datastore_elasticsearch-elasticsearch07: elasticsearch<8.0
@@ -466,9 +469,10 @@ changedir =
466469
component_tastypie: tests/component_tastypie
467470
coroutines_asyncio: tests/coroutines_asyncio
468471
cross_agent: tests/cross_agent
472+
datastore_aiomcache: tests/datastore_aiomcache
473+
datastore_aiomysql: tests/datastore_aiomysql
469474
datastore_asyncpg: tests/datastore_asyncpg
470475
datastore_bmemcached: tests/datastore_bmemcached
471-
datastore_aiomcache: tests/datastore_aiomcache
472476
datastore_elasticsearch: tests/datastore_elasticsearch
473477
datastore_firestore: tests/datastore_firestore
474478
datastore_memcache: tests/datastore_memcache

0 commit comments

Comments
 (0)