Repository: elastic/apm-agent-python

Issue:
[Feature] Add instrumentation for aiomysql
Ref elastic/apm-agent-python#740 

https://github.com/aio-libs/aiomysql

Another bridge driver, it depends on `PyMySQL` so I'm not sure if we need additional instrumentation or not.


Discussion:

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/asyncio/aiomysql.py b/elasticapm/instrumentation/packages/asyncio/aiomysql.py
new file mode 100644
--- /dev/null
+++ b/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)
diff --git a/elasticapm/instrumentation/packages/asyncio/aioredis.py b/elasticapm/instrumentation/packages/asyncio/aioredis.py
--- a/elasticapm/instrumentation/packages/asyncio/aioredis.py
+++ b/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):
diff --git a/elasticapm/instrumentation/register.py b/elasticapm/instrumentation/register.py
--- a/elasticapm/instrumentation/register.py
+++ b/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",
         ]
     )
 


Test patch:
diff --git a/tests/instrumentation/asyncio_tests/aiomysql_tests.py b/tests/instrumentation/asyncio_tests/aiomysql_tests.py
new file mode 100644
--- /dev/null
+++ b/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"},
+        }
diff --git a/tests/requirements/reqs-aiomysql-newest.txt b/tests/requirements/reqs-aiomysql-newest.txt
new file mode 100644
--- /dev/null
+++ b/tests/requirements/reqs-aiomysql-newest.txt
@@ -0,0 +1,2 @@
+aiomysql
+-r reqs-base.txt
diff --git a/tests/scripts/envs/aiomysql.sh b/tests/scripts/envs/aiomysql.sh
new file mode 100644
--- /dev/null
+++ b/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