Skip to content

Commit a7f0b75

Browse files
HiveTraumtoumorokoshi
HiveTraum
and
toumorokoshi
authored
Instrumentation for asyncpg (#814)
Co-authored-by: Yusuke Tsutsumi <[email protected]>
1 parent 9d74918 commit a7f0b75

File tree

12 files changed

+555
-0
lines changed

12 files changed

+555
-0
lines changed

‎docs-requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ sphinx-autodoc-typehints~=1.10.2
44

55
# Required by ext packages
66
asgiref~=3.0
7+
asyncpg>=0.12.0
78
ddtrace>=0.34.0
89
aiohttp~= 3.0
910
Deprecated>=1.2.6

‎docs/ext/asyncpg/asyncpg.rst‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
opentelemetry.ext.asyncpg package
2+
=================================
3+
4+
Module contents
5+
---------------
6+
7+
.. automodule:: opentelemetry.ext.asyncpg
8+
:members:
9+
:undoc-members:
10+
:show-inheritance:

‎ext/opentelemetry-ext-asyncpg/CHANGELOG.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# Changelog
2+
3+
## Unreleased
4+
5+
- Initial Release ([#814](https://github.com/open-telemetry/opentelemetry-python/pull/814))

‎ext/opentelemetry-ext-asyncpg/README.rst‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
OpenTelemetry asyncpg Integration
2+
=================================
3+
4+
|pypi|
5+
6+
.. |pypi| image:: https://badge.fury.io/py/opentelemetry-ext-asyncpg.svg
7+
:target: https://pypi.org/project/opentelemetry-ext-asyncpg/
8+
9+
This library allows tracing PostgreSQL queries made by the
10+
`asyncpg <https://magicstack.github.io/asyncpg/current/>`_ library.
11+
12+
Installation
13+
------------
14+
15+
::
16+
17+
pip install opentelemetry-ext-asyncpg
18+
19+
References
20+
----------
21+
22+
* `OpenTelemetry asyncpg Integration <https://opentelemetry-python.readthedocs.io/en/latest/ext/asyncpg/asyncpg.html>`_
23+
* `OpenTelemetry Project <https://opentelemetry.io/>`_

‎ext/opentelemetry-ext-asyncpg/setup.cfg‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
[metadata]
16+
name = opentelemetry-ext-asyncpg
17+
description = OpenTelemetry instrumentation for AsyncPG
18+
long_description = file: README.rst
19+
long_description_content_type = text/x-rst
20+
author = OpenTelemetry Authors
21+
author_email = [email protected]
22+
url = https://github.com/open-telemetry/opentelemetry-python/ext/opentelemetry-ext-asyncpg
23+
platforms = any
24+
license = Apache-2.0
25+
classifiers =
26+
Development Status :: 4 - Beta
27+
Intended Audience :: Developers
28+
License :: OSI Approved :: Apache Software License
29+
Programming Language :: Python
30+
Programming Language :: Python :: 3
31+
Programming Language :: Python :: 3.5
32+
Programming Language :: Python :: 3.6
33+
Programming Language :: Python :: 3.7
34+
Programming Language :: Python :: 3.8
35+
36+
[options]
37+
python_requires = >=3.5
38+
package_dir=
39+
=src
40+
packages=find_namespace:
41+
install_requires =
42+
opentelemetry-api == 0.10.dev0
43+
opentelemetry-instrumentation == 0.10.dev0
44+
asyncpg >= 0.12.0
45+
46+
[options.extras_require]
47+
test =
48+
opentelemetry-test == 0.10.dev0
49+
50+
[options.packages.find]
51+
where = src
52+
53+
[options.entry_points]
54+
opentelemetry_instrumentor =
55+
asyncpg = opentelemetry.ext.asyncpg:AsyncPGInstrumentor

‎ext/opentelemetry-ext-asyncpg/setup.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
import os
15+
16+
import setuptools
17+
18+
BASE_DIR = os.path.dirname(__file__)
19+
VERSION_FILENAME = os.path.join(
20+
BASE_DIR, "src", "opentelemetry", "ext", "asyncpg", "version.py"
21+
)
22+
PACKAGE_INFO = {}
23+
with open(VERSION_FILENAME) as f:
24+
exec(f.read(), PACKAGE_INFO)
25+
26+
setuptools.setup(version=PACKAGE_INFO["__version__"])

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

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
This library allows tracing PostgreSQL queries made by the
17+
`asyncpg <https://magicstack.github.io/asyncpg/current/>`_ library.
18+
19+
Usage
20+
-----
21+
22+
.. code-block:: python
23+
24+
import asyncpg
25+
from opentelemetry.ext.asyncpg import AsyncPGInstrumentor
26+
27+
# You can optionally pass a custom TracerProvider to AsyncPGInstrumentor.instrument()
28+
AsyncPGInstrumentor().instrument()
29+
conn = await asyncpg.connect(user='user', password='password',
30+
database='database', host='127.0.0.1')
31+
values = await conn.fetch('''SELECT 42;''')
32+
33+
API
34+
---
35+
"""
36+
37+
import asyncpg
38+
import wrapt
39+
from asyncpg import exceptions
40+
41+
from opentelemetry import trace
42+
from opentelemetry.ext.asyncpg.version import __version__
43+
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
44+
from opentelemetry.instrumentation.utils import unwrap
45+
from opentelemetry.trace import SpanKind
46+
from opentelemetry.trace.status import Status, StatusCanonicalCode
47+
48+
_APPLIED = "_opentelemetry_tracer"
49+
50+
51+
def _exception_to_canonical_code(exc: Exception) -> StatusCanonicalCode:
52+
if isinstance(
53+
exc, (exceptions.InterfaceError, exceptions.SyntaxOrAccessError),
54+
):
55+
return StatusCanonicalCode.INVALID_ARGUMENT
56+
if isinstance(exc, exceptions.IdleInTransactionSessionTimeoutError):
57+
return StatusCanonicalCode.DEADLINE_EXCEEDED
58+
return StatusCanonicalCode.UNKNOWN
59+
60+
61+
def _hydrate_span_from_args(connection, query, parameters) -> dict:
62+
span_attributes = {"db.type": "sql"}
63+
64+
params = getattr(connection, "_params", None)
65+
span_attributes["db.instance"] = getattr(params, "database", None)
66+
span_attributes["db.user"] = getattr(params, "user", None)
67+
68+
if query is not None:
69+
span_attributes["db.statement"] = query
70+
71+
if parameters is not None and len(parameters) > 0:
72+
span_attributes["db.statement.parameters"] = str(parameters)
73+
74+
return span_attributes
75+
76+
77+
async def _do_execute(func, instance, args, kwargs):
78+
span_attributes = _hydrate_span_from_args(instance, args[0], args[1:])
79+
tracer = getattr(asyncpg, _APPLIED)
80+
81+
exception = None
82+
83+
with tracer.start_as_current_span(
84+
"postgresql", kind=SpanKind.CLIENT
85+
) as span:
86+
87+
for attribute, value in span_attributes.items():
88+
span.set_attribute(attribute, value)
89+
90+
try:
91+
result = await func(*args, **kwargs)
92+
except Exception as exc: # pylint: disable=W0703
93+
exception = exc
94+
raise
95+
finally:
96+
if exception is not None:
97+
span.set_status(
98+
Status(_exception_to_canonical_code(exception))
99+
)
100+
else:
101+
span.set_status(Status(StatusCanonicalCode.OK))
102+
103+
return result
104+
105+
106+
class AsyncPGInstrumentor(BaseInstrumentor):
107+
def _instrument(self, **kwargs):
108+
tracer_provider = kwargs.get(
109+
"tracer_provider", trace.get_tracer_provider()
110+
)
111+
setattr(
112+
asyncpg,
113+
_APPLIED,
114+
tracer_provider.get_tracer("asyncpg", __version__),
115+
)
116+
for method in [
117+
"Connection.execute",
118+
"Connection.executemany",
119+
"Connection.fetch",
120+
"Connection.fetchval",
121+
"Connection.fetchrow",
122+
]:
123+
wrapt.wrap_function_wrapper(
124+
"asyncpg.connection", method, _do_execute
125+
)
126+
127+
def _uninstrument(self, **__):
128+
delattr(asyncpg, _APPLIED)
129+
for method in [
130+
"execute",
131+
"executemany",
132+
"fetch",
133+
"fetchval",
134+
"fetchrow",
135+
]:
136+
unwrap(asyncpg.Connection, method)

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
__version__ = "0.10.dev0"

‎ext/opentelemetry-ext-asyncpg/tests/__init__.py‎

Whitespace-only changes.

‎ext/opentelemetry-ext-asyncpg/tests/test_asyncpg_wrapper.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import asyncpg
2+
from asyncpg import Connection
3+
4+
from opentelemetry.ext.asyncpg import AsyncPGInstrumentor
5+
from opentelemetry.test.test_base import TestBase
6+
7+
8+
class TestAsyncPGInstrumentation(TestBase):
9+
def test_instrumentation_flags(self):
10+
AsyncPGInstrumentor().instrument()
11+
self.assertTrue(hasattr(asyncpg, "_opentelemetry_tracer"))
12+
AsyncPGInstrumentor().uninstrument()
13+
self.assertFalse(hasattr(asyncpg, "_opentelemetry_tracer"))
14+
15+
def test_duplicated_instrumentation(self):
16+
AsyncPGInstrumentor().instrument()
17+
AsyncPGInstrumentor().instrument()
18+
AsyncPGInstrumentor().instrument()
19+
AsyncPGInstrumentor().uninstrument()
20+
for method_name in ["execute", "fetch"]:
21+
method = getattr(Connection, method_name, None)
22+
self.assertFalse(
23+
hasattr(method, "_opentelemetry_ext_asyncpg_applied")
24+
)
25+
26+
def test_duplicated_uninstrumentation(self):
27+
AsyncPGInstrumentor().instrument()
28+
AsyncPGInstrumentor().uninstrument()
29+
AsyncPGInstrumentor().uninstrument()
30+
AsyncPGInstrumentor().uninstrument()
31+
for method_name in ["execute", "fetch"]:
32+
method = getattr(Connection, method_name, None)
33+
self.assertFalse(
34+
hasattr(method, "_opentelemetry_ext_asyncpg_applied")
35+
)

0 commit comments

Comments
 (0)