Instrumentation for asyncpg

docs-requirements.txt

@@ -4,6 +4,7 @@ sphinx-autodoc-typehints~=1.10.2
 
 # Required by ext packages
 asgiref~=3.0
+asyncpg>=0.12.0
 ddtrace>=0.34.0
 aiohttp~= 3.0
 Deprecated>=1.2.6

docs/ext/asyncpg/asyncpg.rst

@@ -0,0 +1,10 @@
+opentelemetry.ext.asyncpg package
+=================================
+
+Module contents
+---------------
+
+.. automodule:: opentelemetry.ext.asyncpg
+    :members:
+    :undoc-members:
+    :show-inheritance:

ext/opentelemetry-ext-asyncpg/CHANGELOG.md

@@ -0,0 +1,5 @@
+# Changelog
+
+## Unreleased
+
+- Initial Release

ext/opentelemetry-ext-asyncpg/README.rst

@@ -0,0 +1,23 @@
+OpenTelemetry asyncpg Integration
+==================================
+
+|pypi|
+
+.. |pypi| image:: https://badge.fury.io/py/opentelemetry-ext-asyncpg.svg
+   :target: https://pypi.org/project/opentelemetry-ext-asyncpg/
+
+This library allows tracing PostgreSQL queries made by the
+`asyncpg <https://magicstack.github.io/asyncpg/current/>`_ library.
+
+Installation
+------------
+
+::
+
+     pip install opentelemetry-ext-asyncpg
+
+References
+----------
+
+* `OpenTelemetry asyncpg Integration <https://opentelemetry-python.readthedocs.io/en/latest/ext/asyncpg/asyncpg.html>`_
+* `OpenTelemetry Project <https://opentelemetry.io/>`_

ext/opentelemetry-ext-asyncpg/setup.cfg

@@ -0,0 +1,50 @@
+# Copyright The OpenTelemetry Authors
+#
+# 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.
+#
+[metadata]
+name = opentelemetry-ext-asyncpg
+description = AsyncPG for OpenTelemetry
+long_description = file: README.rst
+long_description_content_type = text/x-rst
+author = OpenTelemetry Authors
+author_email = cncf-opentelemetry-contributors@lists.cncf.io
+url = https://github.com/open-telemetry/opentelemetry-python/ext/opentelemetry-ext-asyncpg
+platforms = any
+license = Apache-2.0
+classifiers =
+    Development Status :: 4 - Beta
+    Intended Audience :: Developers
+    License :: OSI Approved :: Apache Software License
+    Programming Language :: Python
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.5
+    Programming Language :: Python :: 3.6
+    Programming Language :: Python :: 3.7
+    Programming Language :: Python :: 3.8
+
+[options]
+python_requires = >=3.5
+package_dir=
+    =src
+packages=find_namespace:
+install_requires =
+    opentelemetry-api == 0.9.dev0
+    asyncpg >= 0.12.0
+
+[options.extras_require]
+test =
+    opentelemetry-ext-testutil
+
+[options.packages.find]
+where = src

ext/opentelemetry-ext-asyncpg/setup.py

@@ -0,0 +1,26 @@
+# Copyright The OpenTelemetry Authors
+#
+# 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 os
+
+import setuptools
+
+BASE_DIR = os.path.dirname(__file__)
+VERSION_FILENAME = os.path.join(
+    BASE_DIR, "src", "opentelemetry", "ext", "asyncpg", "version.py"
+)
+PACKAGE_INFO = {}
+with open(VERSION_FILENAME) as f:
+    exec(f.read(), PACKAGE_INFO)
+
+setuptools.setup(version=PACKAGE_INFO["__version__"])

ext/opentelemetry-ext-asyncpg/src/opentelemetry/ext/asyncpg/__init__.py

@@ -0,0 +1,146 @@
+# Copyright The OpenTelemetry Authors
+#
+# 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.
+
+"""
+This library allows tracing PostgreSQL queries made by the
+`asyncpg <https://magicstack.github.io/asyncpg/current/>`_ library.
+
+Usage
+-----
+
+.. code-block:: python
+
+    import asyncpg
+    import opentelemetry.ext.asyncpg
+
+    # You can optionally pass a custom TracerProvider to AsyncPGInstrumentor.instrument()
+    opentelemetry.ext.asyncpg.AsyncPGInstrumentor().instrument()
+    conn = await asyncpg.connect(user='user', password='password',
+                                 database='database', host='127.0.0.1')
+    values = await conn.fetch('''SELECT 42;''')
+
+API
+---
+"""
+
+import functools
+
+from asyncpg import Connection, exceptions
+
+from opentelemetry import trace
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCanonicalCode
+
+_APPLIED = "_opentelemetry_ext_asyncpg_applied"
+
+
+def _exception_to_canonical_code(exc: Exception) -> StatusCanonicalCode:
+    if isinstance(exc, (exceptions.InterfaceError,),):
+        return StatusCanonicalCode.INVALID_ARGUMENT
+    if isinstance(exc, exceptions.IdleInTransactionSessionTimeoutError):
+        return StatusCanonicalCode.DEADLINE_EXCEEDED
+    return StatusCanonicalCode.UNKNOWN
+
+
+def _hydrate_span_from_args(span, *args, **__):
+    span.set_attribute("db.type", "sql")
+
+    if len(args) <= 0:
+        return span
+
+    connection = args[0]
+    if connection is not None:
+        params = getattr(connection, "_params", None)
+        if params is not None:
+            database_name = getattr(params, "database", None)
+            if database_name is not None:
+                span.set_attribute("db.instance", database_name)
+
+            database_user = getattr(params, "user", None)
+            if database_user is not None:
+                span.set_attribute("db.user", database_user)
+
+    if len(args) > 1 and args[1] is not None:
+        span.set_attribute("db.statement", args[1])
+
+    if len(args) > 2 and args[2] is not None and len(args[2]) > 0:
+        span.set_attribute("db.statement.parameters", args[2])
+
+    return span
+
+
+def _execute(wrapped, tracer_provider):
+    tracer = trace.get_tracer(__name__, "0.8", tracer_provider)
+
+    @functools.wraps(wrapped)
+    async def _method(*args, **kwargs):
+
+        exception = None
+
+        with tracer.start_as_current_span(
+            "postgresql", kind=SpanKind.CLIENT
+        ) as span:
+
+            span = _hydrate_span_from_args(span, *args, **kwargs)
+
+            try:
+                result = await wrapped(*args, **kwargs)
+            except Exception as exc:  # pylint: disable=W0703
+                exception = exc
+
+            if exception is not None:
+                span.set_status(
+                    Status(_exception_to_canonical_code(exception))
+                )
+            else:
+                span.set_status(Status(StatusCanonicalCode.OK))
+
+        if exception is not None:
+            raise exception.with_traceback(exception.__traceback__)
+
+        return result
+
+    setattr(_method, _APPLIED, True)
+    return _method
+
+
+class AsyncPGInstrumentor(BaseInstrumentor):
+    def instrument(self, **kwargs):
+        self._instrument(**kwargs)
+
+    def uninstrument(self, **kwargs):
+        self._uninstrument(**kwargs)
+
+    @staticmethod
+    def _instrument(**kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+
+        for method in ["_execute", "_executemany"]:
+            _original = getattr(Connection, method, None)
+            if hasattr(_original, _APPLIED) is False:
+                setattr(
+                    Connection, method, _execute(_original, tracer_provider)
+                )
+
+    @staticmethod
+    def _uninstrument(**__):
+        for method in ["_execute", "_executemany"]:
+            _connection_method = getattr(Connection, method, None)
+            if _connection_method is not None and getattr(
+                _connection_method, _APPLIED, False
+            ):
+                original = getattr(_connection_method, "__wrapped__", None)
+                if original is not None:
+                    setattr(Connection, method, original)

ext/opentelemetry-ext-asyncpg/src/opentelemetry/ext/asyncpg/version.py

@@ -0,0 +1,15 @@
+# Copyright The OpenTelemetry Authors
+#
+# 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.
+
+__version__ = "0.9.dev0"