diff --git a/newrelic/config.py b/newrelic/config.py index b4368fe70..9896d6a5e 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -3166,6 +3166,9 @@ def _process_module_builtin_defaults(): _process_module_definition("MySQLdb", "newrelic.hooks.database_mysqldb", "instrument_mysqldb") _process_module_definition("pymysql", "newrelic.hooks.database_pymysql", "instrument_pymysql") + _process_module_definition("aiomysql", "newrelic.hooks.database_aiomysql", "instrument_aiomysql") + _process_module_definition("aiomysql.pool", "newrelic.hooks.database_aiomysql", "instrument_aiomysql_pool") + _process_module_definition("pyodbc", "newrelic.hooks.database_pyodbc", "instrument_pyodbc") _process_module_definition("pymssql", "newrelic.hooks.database_pymssql", "instrument_pymssql") diff --git a/newrelic/hooks/database_aiomysql.py b/newrelic/hooks/database_aiomysql.py new file mode 100644 index 000000000..847c961d8 --- /dev/null +++ b/newrelic/hooks/database_aiomysql.py @@ -0,0 +1,106 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys + +from newrelic.api.database_trace import register_database_client +from newrelic.api.function_trace import FunctionTrace +from newrelic.common.object_names import callable_name +from newrelic.common.object_wrapper import ( + ObjectProxy, + wrap_function_wrapper, + wrap_object, +) +from newrelic.hooks.database_dbapi2_async import ( + AsyncConnectionFactory as DBAPI2AsyncConnectionFactory, +) +from newrelic.hooks.database_dbapi2_async import ( + AsyncConnectionWrapper as DBAPI2AsyncConnectionWrapper, +) +from newrelic.hooks.database_dbapi2_async import ( + AsyncCursorWrapper as DBAPI2AsyncCursorWrapper, +) + + +class AsyncCursorContextManagerWrapper(ObjectProxy): + + __cursor_wrapper__ = DBAPI2AsyncCursorWrapper + + def __init__(self, context_manager, dbapi2_module, connect_params, cursor_args): + super().__init__(context_manager) + self._nr_dbapi2_module = dbapi2_module + self._nr_connect_params = connect_params + self._nr_cursor_args = cursor_args + + async def __aenter__(self): + cursor = await self.__wrapped__.__aenter__() + return self.__cursor_wrapper__(cursor, self._nr_dbapi2_module, self._nr_connect_params, self._nr_cursor_args) + + async def __aexit__(self, exc, val, tb): + return await self.__wrapped__.__aexit__(exc, val, tb) + + +class AsyncConnectionWrapper(DBAPI2AsyncConnectionWrapper): + + __cursor_wrapper__ = AsyncCursorContextManagerWrapper + + +class AsyncConnectionFactory(DBAPI2AsyncConnectionFactory): + + __connection_wrapper__ = AsyncConnectionWrapper + + +def wrap_pool__acquire(dbapi2_module): + async def _wrap_pool__acquire(wrapped, instance, args, kwargs): + rollup = ["Datastore/all", f"Datastore/{dbapi2_module._nr_database_product}/all"] + + with FunctionTrace(name=callable_name(wrapped), terminal=True, rollup=rollup, source=wrapped): + connection = await wrapped(*args, **kwargs) + connection_kwargs = getattr(instance, "_conn_kwargs", {}) + return AsyncConnectionWrapper(connection, dbapi2_module, (((), connection_kwargs))) + + return _wrap_pool__acquire + + +def instance_info(args, kwargs): + def _bind_params(host=None, user=None, password=None, db=None, port=None, *args, **kwargs): + return host, port, db + + host, port, db = _bind_params(*args, **kwargs) + + return (host, port, db) + + +def instrument_aiomysql(module): + register_database_client( + module, + database_product="MySQL", + quoting_style="single+double", + explain_query="explain", + explain_stmts=("select",), + instance_info=instance_info, + ) + + # Only instrument the connect method directly, don't instrument + # Connection. This follows the DBAPI2 spec and what was done for + # PyMySQL which this library is based on. + + wrap_object(module, "connect", AsyncConnectionFactory, (module,)) + + +def instrument_aiomysql_pool(module): + dbapi2_module = sys.modules["aiomysql"] + if hasattr(module, "Pool"): + if hasattr(module.Pool, "_acquire"): + wrap_function_wrapper(module, "Pool._acquire", wrap_pool__acquire(dbapi2_module)) diff --git a/tests/datastore_aiomysql/conftest.py b/tests/datastore_aiomysql/conftest.py new file mode 100644 index 000000000..8c8f3c5c7 --- /dev/null +++ b/tests/datastore_aiomysql/conftest.py @@ -0,0 +1,38 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from testing_support.fixture.event_loop import ( # noqa: F401; pylint: disable=W0611 + event_loop as loop, +) +from testing_support.fixtures import ( # noqa: F401; pylint: disable=W0611 + collector_agent_registration_fixture, + collector_available_fixture, +) + +_default_settings = { + "package_reporting.enabled": False, # Turn off package reporting for testing as it causes slow downs. + "transaction_tracer.explain_threshold": 0.0, + "transaction_tracer.transaction_threshold": 0.0, + "transaction_tracer.stack_trace_threshold": 0.0, + "debug.log_data_collector_payloads": True, + "debug.record_transaction_failure": True, + "debug.log_explain_plan_queries": True, +} + +collector_agent_registration = collector_agent_registration_fixture( + app_name="Python Agent Test (datastore_aiomysql)", + default_settings=_default_settings, + linked_applications=["Python Agent Test (datastore)"], +) diff --git a/tests/datastore_aiomysql/test_database.py b/tests/datastore_aiomysql/test_database.py new file mode 100644 index 000000000..b05359ac2 --- /dev/null +++ b/tests/datastore_aiomysql/test_database.py @@ -0,0 +1,157 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import aiomysql +from testing_support.db_settings import mysql_settings +from testing_support.util import instance_hostname +from testing_support.validators.validate_database_trace_inputs import ( + validate_database_trace_inputs, +) +from testing_support.validators.validate_transaction_metrics import ( + validate_transaction_metrics, +) + +from newrelic.api.background_task import background_task + +DB_SETTINGS = mysql_settings()[0] +TABLE_NAME = f"datastore_aiomysql_{DB_SETTINGS['namespace']}" +PROCEDURE_NAME = f"hello_{DB_SETTINGS['namespace']}" + +HOST = instance_hostname(DB_SETTINGS["host"]) +PORT = DB_SETTINGS["port"] + + +async def execute_db_calls_with_cursor(cursor): + await cursor.execute(f"""drop table if exists {TABLE_NAME}""") + + await cursor.execute(f"create table {TABLE_NAME} (a integer, b real, c text)") + + await cursor.executemany( + f"insert into {TABLE_NAME} values (%s, %s, %s)", + [(1, 1.0, "1.0"), (2, 2.2, "2.2"), (3, 3.3, "3.3")], + ) + + await cursor.execute(f"""select * from {TABLE_NAME}""") + + async for _ in cursor: + pass + + await cursor.execute(f"update {TABLE_NAME} set a=%s, b=%s, c=%s where a=%s", (4, 4.0, "4.0", 1)) + + await cursor.execute(f"""delete from {TABLE_NAME} where a=2""") + await cursor.execute(f"""drop procedure if exists {PROCEDURE_NAME}""") + await cursor.execute( + f"""CREATE PROCEDURE {PROCEDURE_NAME}() + BEGIN + SELECT 'Hello World!'; + END""" + ) + + await cursor.callproc(PROCEDURE_NAME) + + +SCOPED_METRICS = [ + (f"Datastore/statement/MySQL/{TABLE_NAME}/select", 1), + (f"Datastore/statement/MySQL/{TABLE_NAME}/insert", 1), + (f"Datastore/statement/MySQL/{TABLE_NAME}/update", 1), + (f"Datastore/statement/MySQL/{TABLE_NAME}/delete", 1), + ("Datastore/operation/MySQL/drop", 2), + ("Datastore/operation/MySQL/create", 2), + (f"Datastore/statement/MySQL/{PROCEDURE_NAME}/call", 1), + ("Datastore/operation/MySQL/commit", 2), + ("Datastore/operation/MySQL/rollback", 1), +] + +ROLLUP_METRICS = [ + ("Datastore/all", 13), + ("Datastore/allOther", 13), + ("Datastore/MySQL/all", 13), + ("Datastore/MySQL/allOther", 13), + (f"Datastore/statement/MySQL/{TABLE_NAME}/select", 1), + (f"Datastore/statement/MySQL/{TABLE_NAME}/insert", 1), + (f"Datastore/statement/MySQL/{TABLE_NAME}/update", 1), + (f"Datastore/statement/MySQL/{TABLE_NAME}/delete", 1), + ("Datastore/operation/MySQL/select", 1), + ("Datastore/operation/MySQL/insert", 1), + ("Datastore/operation/MySQL/update", 1), + ("Datastore/operation/MySQL/delete", 1), + (f"Datastore/statement/MySQL/{PROCEDURE_NAME}/call", 1), + ("Datastore/operation/MySQL/call", 1), + ("Datastore/operation/MySQL/drop", 2), + ("Datastore/operation/MySQL/create", 2), + ("Datastore/operation/MySQL/commit", 2), + ("Datastore/operation/MySQL/rollback", 1), + (f"Datastore/instance/MySQL/{HOST}/{PORT}", 12), +] + + +@validate_transaction_metrics( + "test_database:test_execute_via_connection", + scoped_metrics=list(SCOPED_METRICS) + [("Function/aiomysql.connection:connect", 1)], + rollup_metrics=list(ROLLUP_METRICS) + [("Function/aiomysql.connection:connect", 1)], + background_task=True, +) +@validate_database_trace_inputs(sql_parameters_type=tuple) +@background_task() +def test_execute_via_connection(loop): + async def _test(): + connection = await aiomysql.connect( + db=DB_SETTINGS["name"], + user=DB_SETTINGS["user"], + password=DB_SETTINGS["password"], + host=DB_SETTINGS["host"], + port=DB_SETTINGS["port"], + ) + + async with connection: + async with connection.cursor() as cursor: + await execute_db_calls_with_cursor(cursor) + + await connection.commit() + await connection.rollback() + await connection.commit() + + loop.run_until_complete(_test()) + + +@validate_transaction_metrics( + "test_database:test_execute_via_pool", + scoped_metrics=list(SCOPED_METRICS) + [("Function/aiomysql.pool:Pool._acquire", 1)], + rollup_metrics=list(ROLLUP_METRICS) + [("Function/aiomysql.pool:Pool._acquire", 1)], + background_task=True, +) +@validate_database_trace_inputs(sql_parameters_type=tuple) +@background_task() +def test_execute_via_pool(loop): + async def _test(): + pool = await aiomysql.create_pool( + db=DB_SETTINGS["name"], + user=DB_SETTINGS["user"], + password=DB_SETTINGS["password"], + host=DB_SETTINGS["host"], + port=DB_SETTINGS["port"], + loop=loop, + ) + async with pool.acquire() as connection: + async with connection.cursor() as cursor: + await execute_db_calls_with_cursor(cursor) + + await connection.commit() + await connection.rollback() + await connection.commit() + + pool.close() + await pool.wait_closed() + + loop.run_until_complete(_test()) diff --git a/tox.ini b/tox.ini index 6a5965ecb..81df5d404 100644 --- a/tox.ini +++ b/tox.ini @@ -61,6 +61,7 @@ envlist = mongodb8-datastore_motor-{py37,py38,py39,py310,py311,py312,py313}-motorlatest, mongodb3-datastore_pymongo-{py37,py38,py39,py310,py311,py312}-pymongo03, mongodb8-datastore_pymongo-{py37,py38,py39,py310,py311,py312,py313,pypy310}-pymongo04, + mysql-datastore_aiomysql-{py37,py38,py39,py310,py311,py312,py313,pypy310}, mssql-datastore_pymssql-{py37,py38,py39,py310,py311,py312,py313}, mysql-datastore_mysql-mysqllatest-{py37,py38,py39,py310,py311,py312,py313}, mysql-datastore_pymysql-{py37,py38,py39,py310,py311,py312,py313,pypy310}, @@ -251,6 +252,8 @@ deps = cross_agent: requests datastore_asyncpg: asyncpg datastore_aiomcache: aiomcache + datastore_aiomysql: aiomysql + datastore_aiomysql: cryptography datastore_bmemcached: python-binary-memcached datastore_elasticsearch: requests datastore_elasticsearch-elasticsearch07: elasticsearch<8.0 @@ -466,9 +469,10 @@ changedir = component_tastypie: tests/component_tastypie coroutines_asyncio: tests/coroutines_asyncio cross_agent: tests/cross_agent + datastore_aiomcache: tests/datastore_aiomcache + datastore_aiomysql: tests/datastore_aiomysql datastore_asyncpg: tests/datastore_asyncpg datastore_bmemcached: tests/datastore_bmemcached - datastore_aiomcache: tests/datastore_aiomcache datastore_elasticsearch: tests/datastore_elasticsearch datastore_firestore: tests/datastore_firestore datastore_memcache: tests/datastore_memcache