Skip to content

Commit 4d68edb

Browse files
committed
Tortoise ORM instrumentation
1 parent 50e7b1b commit 4d68edb

File tree

6 files changed

+420
-0
lines changed

6 files changed

+420
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
OpenTelemetry Tortoise ORM Instrumentation
2+
==========================================
3+
4+
|pypi|
5+
6+
.. |pypi| image:: https://badge.fury.io/py/opentelemetry-instrumentation-tortoiseorm.svg
7+
:target: https://pypi.org/project/opentelemetry-instrumentation-tortoiseorm/
8+
9+
This library allows tracing queries made by tortoise ORM backends, mysql, postgres and sqlite.
10+
11+
Installation
12+
------------
13+
14+
::
15+
16+
pip install opentelemetry-instrumentation-tortoiseorm
17+
18+
References
19+
----------
20+
21+
* `OpenTelemetry Project <https://opentelemetry.io/>`_
22+
* `OpenTelemetry Python Examples <https://github.com/open-telemetry/opentelemetry-python/tree/main/docs/examples>`_
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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-instrumentation-tortoiseorm
17+
description = OpenTelemetry instrumentation for Tortoise ORM
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-contrib/tree/main/instrumentation/opentelemetry-instrumentation-tortoiseorm
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.6
32+
Programming Language :: Python :: 3.7
33+
Programming Language :: Python :: 3.8
34+
35+
[options]
36+
python_requires = >=3.6
37+
package_dir=
38+
=src
39+
packages=find_namespace:
40+
install_requires =
41+
opentelemetry-api ~= 1.3
42+
opentelemetry-semantic-conventions == 0.24b0
43+
opentelemetry-instrumentation == 0.24b0
44+
45+
[options.extras_require]
46+
test =
47+
opentelemetry-test == 0.24b0
48+
49+
[options.packages.find]
50+
where = src
51+
52+
[options.entry_points]
53+
opentelemetry_instrumentor =
54+
tortoise = opentelemetry.instrumentation.tortoiseorm:TortoiseORMInstrumentor
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
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+
# DO NOT EDIT. THIS FILE WAS AUTOGENERATED FROM templates/instrumentation_setup.py.txt.
17+
# RUN `python scripts/generate_setup.py` TO REGENERATE.
18+
19+
20+
import distutils.cmd
21+
import json
22+
import os
23+
from configparser import ConfigParser
24+
25+
import setuptools
26+
27+
config = ConfigParser()
28+
config.read("setup.cfg")
29+
30+
# We provide extras_require parameter to setuptools.setup later which
31+
# overwrites the extra_require section from setup.cfg. To support extra_require
32+
# secion in setup.cfg, we load it here and merge it with the extra_require param.
33+
extras_require = {}
34+
if "options.extras_require" in config:
35+
for key, value in config["options.extras_require"].items():
36+
extras_require[key] = [v for v in value.split("\n") if v.strip()]
37+
38+
BASE_DIR = os.path.dirname(__file__)
39+
PACKAGE_INFO = {}
40+
41+
VERSION_FILENAME = os.path.join(
42+
BASE_DIR,
43+
"src",
44+
"opentelemetry",
45+
"instrumentation",
46+
"tortoiseorm",
47+
"version.py",
48+
)
49+
with open(VERSION_FILENAME) as f:
50+
exec(f.read(), PACKAGE_INFO)
51+
52+
PACKAGE_FILENAME = os.path.join(
53+
BASE_DIR,
54+
"src",
55+
"opentelemetry",
56+
"instrumentation",
57+
"tortoiseorm",
58+
"package.py",
59+
)
60+
with open(PACKAGE_FILENAME) as f:
61+
exec(f.read(), PACKAGE_INFO)
62+
63+
# Mark any instruments/runtime dependencies as test dependencies as well.
64+
extras_require["instruments"] = PACKAGE_INFO["_instruments"]
65+
test_deps = extras_require.get("test", [])
66+
for dep in extras_require["instruments"]:
67+
test_deps.append(dep)
68+
69+
extras_require["test"] = test_deps
70+
71+
72+
class JSONMetadataCommand(distutils.cmd.Command):
73+
74+
description = (
75+
"print out package metadata as JSON. This is used by OpenTelemetry dev scripts to ",
76+
"auto-generate code in other places",
77+
)
78+
user_options = []
79+
80+
def initialize_options(self):
81+
pass
82+
83+
def finalize_options(self):
84+
pass
85+
86+
def run(self):
87+
metadata = {
88+
"name": config["metadata"]["name"],
89+
"version": PACKAGE_INFO["__version__"],
90+
"instruments": PACKAGE_INFO["_instruments"],
91+
}
92+
print(json.dumps(metadata))
93+
94+
95+
setuptools.setup(
96+
cmdclass={"meta": JSONMetadataCommand},
97+
version=PACKAGE_INFO["__version__"],
98+
extras_require=extras_require,
99+
)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
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+
Instrument `tortoise-orm`_ to report SQL queries.
17+
Usage
18+
-----
19+
.. code:: python
20+
from opentelemetry.instrumentation.tortoiseorm import TortoiseORMInstrumentor
21+
from tortoise.contrib.fastapi import register_tortoise
22+
23+
register_tortoise(
24+
app,
25+
db_url=settings.db_url,
26+
modules={"models": ["example_app.db_models"]},
27+
generate_schemas=True,
28+
add_exception_handlers=True,
29+
)
30+
31+
TortoiseORMInstrumentor().instrument(tracer_provider=tracer)
32+
API
33+
---
34+
"""
35+
from typing import Collection
36+
37+
try:
38+
import tortoise.backends.asyncpg.client
39+
40+
TORTOISE_POSTGRES_SUPPORT = True
41+
except ModuleNotFoundError:
42+
TORTOISE_POSTGRES_SUPPORT = False
43+
44+
try:
45+
import tortoise.backends.mysql.client
46+
47+
TORTOISE_MYSQL_SUPPORT = True
48+
except ModuleNotFoundError:
49+
TORTOISE_MYSQL_SUPPORT = False
50+
51+
try:
52+
import tortoise.backends.sqlite.client
53+
54+
TORTOISE_SQLITE_SUPPORT = True
55+
except ModuleNotFoundError:
56+
TORTOISE_SQLITE_SUPPORT = False
57+
58+
import wrapt
59+
60+
from opentelemetry import trace
61+
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
62+
from opentelemetry.instrumentation.tortoiseorm.package import _instruments
63+
from opentelemetry.instrumentation.tortoiseorm.version import __version__
64+
from opentelemetry.instrumentation.utils import unwrap
65+
from opentelemetry.semconv.trace import DbSystemValues, SpanAttributes
66+
from opentelemetry.trace import SpanKind
67+
from opentelemetry.trace.status import Status, StatusCode
68+
69+
70+
def _hydrate_span_from_args(connection, query, parameters) -> dict:
71+
"""Get network and database attributes from connection."""
72+
span_attributes = {}
73+
capabilities = getattr(connection, "capabilities", None)
74+
if capabilities:
75+
if capabilities.dialect == "sqlite":
76+
span_attributes[SpanAttributes.DB_SYSTEM] = DbSystemValues.SQLITE.value
77+
elif capabilities.dialect == "postgres":
78+
span_attributes[SpanAttributes.DB_SYSTEM] = DbSystemValues.POSTGRESQL.value
79+
elif capabilities.dialect == "mysql":
80+
span_attributes[SpanAttributes.DB_SYSTEM] = DbSystemValues.MYSQL.value
81+
dbname = getattr(connection, "filename", None)
82+
if dbname:
83+
span_attributes[SpanAttributes.DB_NAME] = dbname
84+
dbname = getattr(connection, "database", None)
85+
if dbname:
86+
span_attributes[SpanAttributes.DB_NAME] = dbname
87+
if query is not None:
88+
span_attributes[SpanAttributes.DB_STATEMENT] = query
89+
user = getattr(connection, "user", None)
90+
if user:
91+
span_attributes[SpanAttributes.DB_USER] = user
92+
host = getattr(connection, "host", None)
93+
if host:
94+
span_attributes[SpanAttributes.NET_PEER_NAME] = host
95+
port = getattr(connection, "port", None)
96+
if port:
97+
span_attributes[SpanAttributes.NET_PEER_PORT] = port
98+
99+
if parameters is not None and len(parameters) > 0:
100+
span_attributes["db.statement.parameters"] = str(parameters)
101+
102+
return span_attributes
103+
104+
105+
class TortoiseORMInstrumentor(BaseInstrumentor):
106+
"""An instrumentor for Tortoise-ORM
107+
See `BaseInstrumentor`
108+
"""
109+
110+
def instrumentation_dependencies(self) -> Collection[str]:
111+
return _instruments
112+
113+
def _instrument(self, **kwargs):
114+
"""Instruments Tortoise ORM backend methods.
115+
Args:
116+
**kwargs: Optional arguments
117+
``tracer_provider``: a TracerProvider, defaults to global
118+
Returns:
119+
None
120+
"""
121+
tracer_provider = kwargs.get("tracer_provider")
122+
self._tracer = trace.get_tracer(__name__, __version__, tracer_provider)
123+
if TORTOISE_SQLITE_SUPPORT:
124+
funcs = [
125+
"SqliteClient.execute_many",
126+
"SqliteClient.execute_query",
127+
"SqliteClient.execute_insert",
128+
"SqliteClient.execute_query_dict",
129+
"SqliteClient.execute_script",
130+
]
131+
for f in funcs:
132+
wrapt.wrap_function_wrapper(
133+
"tortoise.backends.sqlite.client",
134+
f,
135+
self._do_execute,
136+
)
137+
138+
if TORTOISE_POSTGRES_SUPPORT:
139+
funcs = [
140+
"AsyncpgDBClient.execute_many",
141+
"AsyncpgDBClient.execute_query",
142+
"AsyncpgDBClient.execute_insert",
143+
"AsyncpgDBClient.execute_query_dict",
144+
"AsyncpgDBClient.execute_script",
145+
]
146+
for f in funcs:
147+
wrapt.wrap_function_wrapper(
148+
"tortoise.backends.asyncpg.client",
149+
f,
150+
self._do_execute,
151+
)
152+
153+
if TORTOISE_MYSQL_SUPPORT:
154+
funcs = [
155+
"MySQLClient.execute_many",
156+
"MySQLClient.execute_query",
157+
"MySQLClient.execute_insert",
158+
"MySQLClient.execute_query_dict",
159+
"MySQLClient.execute_script",
160+
]
161+
for f in funcs:
162+
wrapt.wrap_function_wrapper(
163+
"tortoise.backends.mysql.client",
164+
f,
165+
self._do_execute,
166+
)
167+
168+
def _uninstrument(self, **kwargs):
169+
if TORTOISE_SQLITE_SUPPORT:
170+
unwrap(tortoise.backends.sqlite.client.SqliteClient, "execute_query")
171+
unwrap(tortoise.backends.sqlite.client.SqliteClient, "execute_many")
172+
unwrap(tortoise.backends.sqlite.client.SqliteClient, "execute_insert")
173+
unwrap(tortoise.backends.sqlite.client.SqliteClient, "execute_query_dict")
174+
unwrap(tortoise.backends.sqlite.client.SqliteClient, "execute_script")
175+
if TORTOISE_MYSQL_SUPPORT:
176+
unwrap(tortoise.backends.mysql.client.MySQLClient, "execute_query")
177+
unwrap(tortoise.backends.mysql.client.MySQLClient, "execute_many")
178+
unwrap(tortoise.backends.mysql.client.MySQLClient, "execute_insert")
179+
unwrap(tortoise.backends.mysql.client.MySQLClient, "execute_query_dict")
180+
unwrap(tortoise.backends.mysql.client.MySQLClient, "execute_script")
181+
if self.TORTOISE_POSTGRES_SUPPORT:
182+
unwrap(tortoise.backends.asyncpg.client.AsyncpgDBClient, "execute_query")
183+
unwrap(tortoise.backends.asyncpg.client.AsyncpgDBClient, "execute_many")
184+
unwrap(tortoise.backends.asyncpg.client.AsyncpgDBClient, "execute_insert")
185+
unwrap(
186+
tortoise.backends.asyncpg.client.AsyncpgDBClient, "execute_query_dict"
187+
)
188+
unwrap(tortoise.backends.asyncpg.client.AsyncpgDBClient, "execute_script")
189+
190+
async def _do_execute(self, func, instance, args, kwargs):
191+
192+
exception = None
193+
name = args[0]
194+
195+
with self._tracer.start_as_current_span(name, kind=SpanKind.CLIENT) as span:
196+
if span.is_recording():
197+
span_attributes = _hydrate_span_from_args(
198+
instance,
199+
args[0],
200+
args[1:],
201+
)
202+
for attribute, value in span_attributes.items():
203+
span.set_attribute(attribute, value)
204+
205+
try:
206+
result = await func(*args, **kwargs)
207+
except Exception as exc: # pylint: disable=W0703
208+
exception = exc
209+
raise
210+
finally:
211+
if span.is_recording() and exception is not None:
212+
span.set_status(Status(StatusCode.ERROR))
213+
214+
return result

0 commit comments

Comments
 (0)