10000 feat: Add SQLAlchemy instrumentation (#462) · etherscan-io/sentry-python@211f7f4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 211f7f4

Browse files
untitakeruntitaker
authored
feat: Add SQLAlchemy instrumentation (getsentry#462)
* ref: Do not add PII to SQL queries * feat: SQLAlchemy instrumentation * fix: Add information about SQL dialect used * ref: Apply feedback and remove dead code
1 parent ac8d5b8 commit 211f7f4

File tree

11 files changed

+238
-266
lines changed

11 files changed

+238
-266
lines changed

‎sentry_sdk/hub.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -446,10 +446,10 @@ def span(
446446
try:
447447
yield span
448448
except Exception:
449-
span.set_tag("error", True)
449+
span.set_failure()
450450
raise
451451
else:
452-
span.set_tag("error", False)
452+
span.set_success()
453453
finally:
454454
try:
455455
span.finish()

‎sentry_sdk/integrations/_sql_common.py

Lines changed: 0 additions & 81 deletions
This file was deleted.

‎sentry_sdk/integrations/django/__init__.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@
4646
from sentry_sdk.integrations.logging import ignore_logger
4747
from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
4848
from sentry_sdk.integrations._wsgi_common import RequestExtractor
49-
from sentry_sdk.integrations._sql_common import format_sql
5049
from sentry_sdk.integrations.django.transactions import LEGACY_RESOLVER
5150
from sentry_sdk.integrations.django.templates import get_template_frame_from_exception
5251

@@ -391,7 +390,7 @@ def execute(self, sql, params=None):
391390
return real_execute(self, sql, params)
392391

393392
with record_sql_queries(
394-
hub, [format_sql(sql, params, self.cursor)], label="Django: "
393+
hub, self.cursor, sql, params, paramstyle="format", executemany=False
395394
):
396395
return real_execute(self, sql, params)
397396

@@ -401,9 +400,7 @@ def executemany(self, sql, param_list):
401400
return real_executemany(self, sql, param_list)
402401

403402
with record_sql_queries(
404-
hub,
405-
[format_sql(sql, params, self.cursor) for params in param_list],
406-
label="Django: ",
403+
hub, self.cursor, sql, param_list, paramstyle="format", executemany=True
407404
):
408405
return real_executemany(self, sql, param_list)
409406

‎sentry_sdk/integrations/sqlalchemy.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
from __future__ import absolute_import
2+
3+
from sentry_sdk._types import MYPY
4+
from sentry_sdk.hub import Hub
5+
from sentry_sdk.integrations import Integration
6+
from sentry_sdk.tracing import record_sql_queries
7+
8+
from sqlalchemy.engine import Engine # type: ignore
9+
from sqlalchemy.event import listen # type: ignore
10+
11+
if MYPY:
12+
from typing import Any
13+
from typing import ContextManager
14+
from typing import Optional
15+
16+
from sentry_sdk.tracing import Span
17+
18+
19+
class SqlalchemyIntegration(Integration):
20+
identifier = "sqlalchemy"
21+
22+
@staticmethod
23+
def setup_once():
24+
# type: () -> None
25+
26+
listen(Engine, "before_cursor_execute", _before_cursor_execute)
27+
listen(Engine, "after_cursor_execute", _after_cursor_execute)
28+
listen(Engine, "dbapi_error", _dbapi_error)
29+
30+
31+
def _before_cursor_execute(
32+
conn, cursor, statement, parameters, context, executemany, *args
33+
):
34+
# type: (Any, Any, Any, Any, Any, bool, *Any) -> None
35+
hub = Hub.current
36+
if hub.get_integration(SqlalchemyIntegration) is None:
37+
return
38+
39+
ctx_mgr = record_sql_queries(
40+
hub,
41+
cursor,
42+
statement,
43+
parameters,
44+
paramstyle=context and context.dialect and context.dialect.paramstyle or None,
45+
executemany=executemany,
46+
)
47+
conn._sentry_sql_span_manager = ctx_mgr
48+
49+
span = ctx_mgr.__enter__()
50+
51+
if span is not None:
52+
span.set_success() # might be overwritten later
53+
conn._sentry_sql_span = span
54+
55+
56+
def _after_cursor_execute(conn, cursor, statement, *args):
57+
# type: (Any, Any, Any, *Any) -> None
58+
ctx_mgr = getattr(conn, "_sentry_sql_span_manager", None) # type: ContextManager
59+
60+
if ctx_mgr is not None:
61+
conn._sentry_sql_span_manager = None
62+
ctx_mgr.__exit__(None, None, None)
63+
64+
65+
def _dbapi_error(conn, *args):
66+
# type: (Any, *Any) -> None
67+
span = getattr(conn, "_sentry_sql_span", None) # type: Optional[Span]
68+
69+
if span is not None:
70+
span.set_failure()

‎sentry_sdk/tracing.py

Lines changed: 58 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from datetime import datetime
66

7-
from sentry_sdk.utils import capture_internal_exceptions, concat_strings
7+
from sentry_sdk.utils import capture_internal_exceptions
88
from sentry_sdk._compat import PY2
99
from sentry_sdk._types import MYPY
1010

@@ -16,11 +16,14 @@
1616
if MYPY:
1717
import typing
1818

19+
from typing import Generator
1920
from typing import Optional
2021
from typing import Any
2122
from typing import Dict
2223
from typing import List
2324

25+
from sentry_sdk import Hub
26+
2427
_traceparent_header_format_re = re.compile(
2528
"^[ \t]*" # whitespace
2629
"([0-9a-f]{32})?" # trace_id
@@ -188,6 +191,12 @@ def set_tag(self, key, value):
188191
def set_data(self, key, value):
189192
self._data[key] = value
190193

194+
def set_failure(self):
195+
self.set_tag("error", True)
196+
197+
def set_success(self):
198+
self.set_tag("error", False)
199+
191200
def finish(self):
192201
self.timestamp = datetime.now()
193202
if self._finished_spans is not None:
@@ -218,25 +227,55 @@ def get_trace_context(self):
218227
}
219228

220229

230+
def _format_sql(cursor, sql):
231+
# type: (Any, str) -> Optional[str]
232+
233+
real_sql = None
234+
235+
# If we're using psycopg2, it could be that we're
236+
# looking at a query that uses Composed objects. Use psycopg2's mogrify
237+
# function to format the query. We lose per-parameter trimming but gain
238+
# accuracy in formatting.
239+
try:
240+
if hasattr(cursor, "mogrify"):
241+
real_sql = cursor.mogrify(sql)
242+
if isinstance(real_sql, bytes):
243+
real_sql = real_sql.decode(cursor.connection.encoding)
244+
except Exception:
245+
real_sql = None
246+
247+
return real_sql or str(sql)
248+
249+
221250
@contextlib.contextmanager
222-
def record_sql_queries(hub, queries, label=""):
223-
if not queries:
224-
yield None
225-
else:
226-
description = None
227-
with capture_internal_exceptions():
228-
strings = [label]
229-
for query in queries:
230-
hub.add_breadcrumb(message=query, category="query")
231-
strings.append(query)
232-
233-
description = concat_strings(strings)
234-
235-
if description is None:
236-
yield None
237-
else:
238-
with hub.span(op="db", description=description) as span:
239-
yield span
251+
def record_sql_queries(
252+
hub, # type: Hub
253+
cursor, # type: Any
254+
query, # type: Any
255+
params_list, # type: Any
256+
paramstyle, # type: Optional[str]
257+
executemany, # type: bool
258+
):
259+
# type: (...) -> Generator[Optional[Span], None, None]
260+
if not params_list or params_list == [None]:
261+
params_list = None
262+
263+
if paramstyle == "pyformat":
264+
paramstyle = "format"
265+
266+
query = _format_sql(cursor, query)
267+
268+
data = {"db.params": params_list, "db.paramstyle": paramstyle}
269+
if executemany:
270+
data["db.executemany"] = True
271+
272+
with capture_internal_exceptions():
273+
hub.add_breadcrumb(message=query, category="query", data=data)
274+
275+
with hub.span(op="db", description=query) as span:
276+
for k, v in data.items():
277+
span.set_data(k, v)
278+
yield span
240279

241280

242281
@contextlib.contextmanager

0 commit comments

Comments
 (0)
0