Add aiomysql instrumentation

.ci/.jenkins_exclude.yml

@@ -228,3 +228,10 @@ exclude:
  FRAMEWORK: aioredis-newest
  - PYTHON_VERSION: python-3.10-rc
  FRAMEWORK: httpx-0.12
+ # aiomysql
+ - PYTHON_VERSION: pypy-3
+ FRAMEWORK: aiomysql-newest
+ - PYTHON_VERSION: python-3.6
+ FRAMEWORK: aiomysql-newest
+ - PYTHON_VERSION: python-3.10-rc # getting "loop argument must agree with lock" error
+ FRAMEWORK: aiomysql-newest

.ci/.jenkins_framework.yml

@@ -46,3 +46,4 @@ FRAMEWORK:
  - httpx-newest
  - httplib2-newest
  - prometheus_client-newest
+ - aiomysql-newest

docs/supported-technologies.asciidoc

@@ -197,6 +197,20 @@ Collected trace data:
 
  * parametrized SQL query
 
+[float]
+[[automatic-instrumentation-db-aiomysql]]
+==== aiomysql
+
+Library: `aiomysql`
+
+Instrumented methods:
+
+ * `aiomysql.cursors.Cursor.execute`
+
+Collected trace data:
+
+ * parametrized SQL query
+
 [float]
 [[automatic-instrumentation-db-postgres]]
 ==== PostgreSQL

elasticapm/instrumentation/packages/asyncio/aiomysql.py

@@ -0,0 +1,66 @@
+# BSD 3-Clause License
+#
+# Copyright (c) 2019, Elasticsearch BV
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from elasticapm.contrib.asyncio.traces import async_capture_span
+from elasticapm.instrumentation.packages.asyncio.base import AsyncAbstractInstrumentedModule
+from elasticapm.instrumentation.packages.dbapi2 import extract_signature
+from elasticapm.utils.encoding import shorten
+
+
+class AioMySQLInstrumentation(AsyncAbstractInstrumentedModule):
+ name = "aiomysql"
+
+ instrument_list = [("aiomysql.cursors", "Cursor.execute")]
+
+ async def call(self, module, method, wrapped, instance, args, kwargs):
+ if method == "Cursor.execute":
+ query = args[0]
+ name = extract_signature(query)
+
+ # Truncate sql_string to 10000 characters to prevent large queries from
+ # causing an error to APM server.
+ query = shorten(query, string_length=10000)
+
+ context = {
+ "db": {"type": "sql", "statement": query},
+ "destination": {
+ "address": instance.connection.host,
+ "port": instance.connection.port,
+ "service": {"name": "mysql", "resource": "mysql", "type": "db"},
+ },
+ }
+ action = "query"
+ else:
+ raise AssertionError("call from uninstrumented method")
+
+ async with async_capture_span(
+ name, leaf=True, span_type="db", span_subtype="mysql", span_action=action, extra=context
+ ):
+ return await wrapped(*args, **kwargs)

elasticapm/instrumentation/packages/asyncio/aioredis.py

@@ -30,9 +30,9 @@
 
 from __future__ import absolute_import
 
+from elasticapm.contrib.asyncio.traces import async_capture_span
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
 from elasticapm.traces import execution_context
-from elasticapm.contrib.asyncio.traces import async_capture_span
 
 
 class RedisConnectionPoolInstrumentation(AbstractInstrumentedModule):

elasticapm/instrumentation/register.py

@@ -82,6 +82,7 @@
  "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionPoolInstrumentation",
  "elasticapm.instrumentation.packages.asyncio.aioredis.RedisPipelineInstrumentation",
  "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionInstrumentation",
+ "elasticapm.instrumentation.packages.asyncio.aiomysql.AioMySQLInstrumentation",
  ]
  )
 

setup.cfg

@@ -145,6 +145,7 @@ markers =
  pyodbc
  aiohttp
  aiopg
+ aiomysql
  asyncpg
  tornado
  starlette

tests/instrumentation/asyncio_tests/aiomysql_tests.py

@@ -0,0 +1,102 @@
+# BSD 3-Clause License
+#
+# Copyright (c) 2019, Elasticsearch BV
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import os
+
+import pytest
+
+from elasticapm.conf.constants import TRANSACTION
+from elasticapm.utils import default_ports
+
+aiomysql = pytest.importorskip("aiomysql")
+
+pytestmark = [pytest.mark.asyncio, pytest.mark.aiomysql]
+
+
+if "MYSQL_HOST" not in os.environ:
+ pytestmark.append(pytest.mark.skip("Skipping aiomysql tests, no MYSQL_HOST environment variable set"))
+
+
+@pytest.fixture(scope="function")
+async def aiomysql_connection(request, event_loop):
+ assert event_loop.is_running()
+ pool = await aiomysql.create_pool(
+ host=os.environ.get("MYSQL_HOST", "localhost"),
+ user=os.environ.get("MYSQL_USER", "eapm"),
+ password=os.environ.get("MYSQL_PASSWORD", ""),
+ db=os.environ.get("MYSQL_DATABASE", "eapm_tests"),
+ loop=event_loop,
+ )
+
+ try:
+ async with pool.acquire() as conn:
+ async with conn.cursor() as cursor:
+ await cursor.execute("CREATE TABLE `test` (`id` INT, `name` VARCHAR(5))")
+ await cursor.execute("INSERT INTO `test` (`id`, `name`) VALUES (1, 'one'), (2, 'two'), (3, 'three')")
+
+ yield conn
+
+ finally:
+ # Drop the testing table and close the connection pool after testcase
+ async with pool.acquire() as conn:
+ async with conn.cursor() as cursor:
+ await cursor.execute("DROP TABLE `test`")
+
+ pool.close()
+ await pool.wait_closed()
+
+
+@pytest.mark.integrationtest
+async def test_aiomysql_select(instrument, aiomysql_connection, elasticapm_client):
+ try:
+ elasticapm_client.begin_transaction("web.django")
+
+ async with aiomysql_connection.cursor() as cursor:
+ query = "SELECT * FROM test WHERE `name` LIKE 't%' ORDER BY id"
+ await cursor.execute(query)
+ assert await cursor.fetchall() == ((2, "two"), (3, "three"))
+
+ elasticapm_client.end_transaction(None, "test-transaction")
+ finally:
+ transactions = elasticapm_client.events[TRANSACTION]
+ spans = elasticapm_client.spans_for_transaction(transactions[0])
+ span = spans[0]
+ assert span["name"] == "SELECT FROM test"
+ assert span["type"] == "db"
+ assert span["subtype"] == "mysql"
+ assert span["action"] == "query"
+ assert "db" in span["context"]
+ assert span["context"]["db"]["type"] == "sql"
+ assert span["context"]["db"]["statement"] == query
+ assert span["context"]["destination"] == {
+ "address": os.environ.get("MYSQL_HOST", "localhost"),
+ "port": default_ports.get("mysql"),
+ "service": {"name": "mysql", "resource": "mysql", "type": "db"},
+ }

tests/requirements/reqs-aiomysql-newest.txt

@@ -0,0 +1,2 @@
+aiomysql
+-r reqs-base.txt

tests/scripts/envs/aiomysql.sh

@@ -0,0 +1,7 @@
+export PYTEST_MARKER="-m aiomysql"
+export DOCKER_DEPS="mysql"
+export MYSQL_HOST="mysql"
+export MYSQL_USER="eapm"
+export MYSQL_PASSWORD="Very(!)Secure"
+export WAIT_FOR_HOST="mysql"
+export WAIT_FOR_PORT=3306