Add asyncpg instrumentation in Elastic APM (#889)

* Add a new file

* Add AsyncPGInstrumentation to register.py

* Fix a typo in register.py

* Use correct function for registering

* Update the docs + add a template for the tests

* Update variable names in the tests

* Add asyncpg.sh for the tests

* Remove unused docstrings

* Improve tests

* Remove callproc test

* PR suggestions

* FIx tests: use URI instead of explicit parameters

* FIx tests: implement further PR suggestions

* FIx tests: implement more PR suggestions

* Use Connection.execute and Connection.executemany instead of Connection._do_execute

* Remove unnecessary requirements

* Fix a type in the tests

* Use yield and proper awaiting in the tests

* Use yield and proper awaiting in the tests + fix a typo

* Fix tests

* Additional fixes to tests

* fix typo

Co-authored-by: Colton Myers <colton.myers@gmail.com>
Co-authored-by: Benjamin Wohlwend <bw@piquadrat.ch>

.ci/.jenkins_exclude.yml

@@ -267,6 +267,17 @@ exclude:
     FRAMEWORK: aiopg-newest
   - PYTHON_VERSION: python-3.6
     FRAMEWORK: aiopg-newest
+  # asyncpg
+  - PYTHON_VERSION: python-2.7
+    FRAMEWORK: asyncpg-newest
+  - PYTHON_VERSION: pypy-2
+    FRAMEWORK: asyncpg-newest
+  - PYTHON_VERSION: pypy-3
+    FRAMEWORK: asyncpg-newest
+  - PYTHON_VERSION: python-3.5
+    FRAMEWORK: asyncpg-newest
+  - PYTHON_VERSION: python-3.6
+    FRAMEWORK: asyncpg-newest
   # psutil
   - PYTHON_VERSION: pypy-2  #currently fails on pypy2 (https://github.com/giampaolo/psutil/issues/1659)
     FRAMEWORK: psutil-newest

.ci/.jenkins_framework.yml

@@ -36,6 +36,7 @@ FRAMEWORK:
   - mysqlclient-newest
   - aiohttp-newest
   - aiopg-newest
+  - asyncpg-newest
   - tornado-newest
   - starlette-newest
   - pymemcache-newest

.ci/.jenkins_framework_full.yml

@@ -65,6 +65,7 @@ FRAMEWORK:
   - aiohttp-4.0
   - aiohttp-newest
   - aiopg-newest
+  - asyncpg-newest
   - tornado-newest
   - starlette-0.13
   - starlette-newest

docs/supported-technologies.asciidoc

@@ -212,6 +212,22 @@ Instrumented methods:
  * `aiopg.cursor.Cursor.execute`
  * `aiopg.cursor.Cursor.callproc`
 
+Collected trace data:
+
+ * parametrized SQL query
+
+[float]
+[[automatic-instrumentation-db-asyncg]]
+==== asyncpg
+
+Library: `asyncpg` (`>=0.20`)
+
+Instrumented methods:
+
+ * `asyncpg.connection.Connection.execute`
+ * `asyncpg.connection.Connection.executemany`
+
+
 Collected trace data:
 
  * parametrized SQL query

elasticapm/instrumentation/packages/asyncio/asyncpg.py

@@ -0,0 +1,59 @@
+#  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
+
+
+class AsyncPGInstrumentation(AsyncAbstractInstrumentedModule):
+    """
+    Implement asyncpg instrumentation with two methods Connection.execute
+    and Connection.executemany since Connection._do_execute is not called
+    given a prepared query is passed to a connection. As in:
+    https://github.com/MagicStack/asyncpg/blob/master/asyncpg/connection.py#L294-L297
+    """
+
+    name = "asyncpg"
+
+    instrument_list = [
+        ("asyncpg.connection", "Connection.execute"),
+        ("asyncpg.connection", "Connection.executemany"),
+    ]
+
+    async def call(self, module, method, wrapped, instance, args, kwargs):
+        query = args[0] if len(args) else kwargs["query"]
+        name = extract_signature(query)
+        context = {"db": {"type": "sql", "statement": query}}
+        action = "query"
+        async with async_capture_span(
+            name, leaf=True, span_type="db", span_subtype="postgres", span_action=action, extra=context
+        ):
+            return await wrapped(*args, **kwargs)

elasticapm/instrumentation/register.py

@@ -75,6 +75,7 @@
             "elasticapm.instrumentation.packages.asyncio.elasticsearch.ElasticSearchAsyncConnection",
             "elasticapm.instrumentation.packages.asyncio.elasticsearch.AsyncElasticsearchInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aiopg.AioPGInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.asyncpg.AsyncPGInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoRequestExecuteInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoHandleRequestExceptionInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoRenderInstrumentation",

setup.cfg

@@ -149,6 +149,7 @@ markers =
     pymssql
     aiohttp
     aiopg
+    asyncpg
     tornado
     starlette
     graphene

tests/instrumentation/asyncio/asyncpg_tests.py

@@ -0,0 +1,109 @@
+#  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 import constants
+
+asyncpg = pytest.importorskip("asyncpg")  # isort:skip
+pytestmark = [pytest.mark.asyncpg, pytest.mark.asyncio]
+
+if "POSTGRES_DB" not in os.environ:
+    pytestmark.append(
+        pytest.mark.skip(
+            "Skipping asyncpg tests, no POSTGRES_DB environment variable set"
+        )
+    )
+
+
+def dsn():
+    return "postgres://{user}:{password}@{host}:{port}/{database}".format(
+        **{
+            "database": os.environ.get("POSTGRES_DB", "elasticapm_test"),
+            "user": os.environ.get("POSTGRES_USER", "postgres"),
+            "password": os.environ.get("POSTGRES_PASSWORD", "postgres"),
+            "host": os.environ.get("POSTGRES_HOST", "localhost"),
+            "port": os.environ.get("POSTGRES_PORT", "5432"),
+        }
+    )
+
+
+@pytest.fixture()
+async def connection(request):
+    conn = await asyncpg.connect(dsn())
+
+    await conn.execute(
+        "BEGIN;"
+        "CREATE TABLE test(id int, name VARCHAR(5) NOT NULL);"
+        "INSERT INTO test VALUES (1, 'one'), (2, 'two'), (3, 'three');"
+    )
+    yield conn
+
+    await conn.execute("ROLLBACK")
+    await conn.close()
+
+
+async def test_execute_with_sleep(instrument, connection, elasticapm_client):
+    elasticapm_client.begin_transaction("test")
+    await connection.execute("SELECT pg_sleep(0.1);")
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[constants.TRANSACTION][0]
+    spans = elasticapm_client.spans_for_transaction(transaction)
+
+    assert len(spans) == 1
+    span = spans[0]
+    assert 100 < span["duration"] < 110
+    assert transaction["id"] == span["transaction_id"]
+    assert span["type"] == "db"
+    assert span["subtype"] == "postgres"
+    assert span["action"] == "query"
+    assert span["sync"] == False
+    assert span["name"] == "SELECT FROM"
+
+
+async def test_executemany(instrument, connection, elasticapm_client):
+    elasticapm_client.begin_transaction("test")
+    await connection.executemany(
+        "INSERT INTO test VALUES ($1, $2);", [(1, "uno"), (2, "due")]
+    )
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[constants.TRANSACTION][0]
+    spans = elasticapm_client.spans_for_transaction(transaction)
+
+    assert len(spans) == 1
+    span = spans[0]
+    assert transaction["id"] == span["transaction_id"]
+    assert span["subtype"] == "postgres"
+    assert span["action"] == "query"
+    assert span["sync"] == False
+    assert span["name"] == "INSERT INTO test"

tests/requirements/reqs-asyncpg-newest.txt

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

tests/scripts/envs/asyncpg.sh

@@ -0,0 +1,8 @@
+export PYTEST_MARKER="-m asyncpg"
+export DOCKER_DEPS="postgres"
+export POSTGRES_HOST="postgres"
+export POSTGRES_USER="postgres"
+export POSTGRES_PASSWORD="postgres"
+export POSTGRES_DB="elasticapm_test"
+export POSTGRES_HOST="postgres"
+export POSTGRES_PORT="5432"