Skip to content

Commit 9a0c133

Browse files
sl0thentr0pysl0thentr0py
authored
fix(sqlalchemy): Use context instead of connection in sqlalchemy integration (getsentry#1388)
* Revert "fix(sqlalchemy): Change context manager type to avoid race in threads (getsentry#1368)" This reverts commit de0bc50. This caused a regression (getsentry#1385) since the span finishes immediately in __enter__ and so all db spans have wrong time durations. * Use context instead of conn in sqlalchemy hooks
1 parent 4703bc3 commit 9a0c133

File tree

3 files changed

+62
-67
lines changed

3 files changed

+62
-67
lines changed

‎sentry_sdk/integrations/django/__init__.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from sentry_sdk.hub import Hub, _should_send_default_pii
1010
from sentry_sdk.scope import add_global_event_processor
1111
from sentry_sdk.serializer import add_global_repr_processor
12-
from sentry_sdk.tracing_utils import RecordSqlQueries
12+
from sentry_sdk.tracing_utils import record_sql_queries
1313
from sentry_sdk.utils import (
1414
HAS_REAL_CONTEXTVARS,
1515
CONTEXTVARS_ERROR_MESSAGE,
@@ -538,7 +538,7 @@ def execute(self, sql, params=None):
538538
if hub.get_integration(DjangoIntegration) is None:
539539
return real_execute(self, sql, params)
540540

541-
with RecordSqlQueries(
541+
with record_sql_queries(
542542
hub, self.cursor, sql, params, paramstyle="format", executemany=False
543543
):
544544
return real_execute(self, sql, params)
@@ -549,7 +549,7 @@ def executemany(self, sql, param_list):
549549
if hub.get_integration(DjangoIntegration) is None:
550550
return real_executemany(self, sql, param_list)
551551

552-
with RecordSqlQueries(
552+
with record_sql_queries(
553553
hub, self.cursor, sql, param_list, paramstyle="format", executemany=True
554554
):
555555
return real_executemany(self, sql, param_list)

‎sentry_sdk/integrations/sqlalchemy.py

+15-12
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from sentry_sdk._types import MYPY
44
from sentry_sdk.hub import Hub
55
from sentry_sdk.integrations import Integration, DidNotEnable
6-
from sentry_sdk.tracing_utils import RecordSqlQueries
6+
from sentry_sdk.tracing_utils import record_sql_queries
77

88
try:
99
from sqlalchemy.engine import Engine # type: ignore
@@ -50,37 +50,40 @@ def _before_cursor_execute(
5050
if hub.get_integration(SqlalchemyIntegration) is None:
5151
return
5252

53-
ctx_mgr = RecordSqlQueries(
53+
ctx_mgr = record_sql_queries(
5454
hub,
5555
cursor,
5656
statement,
5757
parameters,
5858
paramstyle=context and context.dialect and context.dialect.paramstyle or None,
5959
executemany=executemany,
6060
)
61-
conn._sentry_sql_span_manager = ctx_mgr
61+
context._sentry_sql_span_manager = ctx_mgr
6262

6363
span = ctx_mgr.__enter__()
6464

6565
if span is not None:
66-
conn._sentry_sql_span = span
66+
context._sentry_sql_span = span
6767

6868

69-
def _after_cursor_execute(conn, cursor, statement, *args):
70-
# type: (Any, Any, Any, *Any) -> None
69+
def _after_cursor_execute(conn, cursor, statement, parameters, context, *args):
70+
# type: (Any, Any, Any, Any, Any, *Any) -> None
7171
ctx_mgr = getattr(
72-
conn, "_sentry_sql_span_manager", None
72+
context, "_sentry_sql_span_manager", None
7373
) # type: ContextManager[Any]
7474

7575
if ctx_mgr is not None:
76-
conn._sentry_sql_span_manager = None
76+
context._sentry_sql_span_manager = None
7777
ctx_mgr.__exit__(None, None, None)
7878

7979

8080
def _handle_error(context, *args):
8181
# type: (Any, *Any) -> None
82-
conn = context.connection
83-
span = getattr(conn, "_sentry_sql_span", None) # type: Optional[Span]
82+
execution_context = context.execution_context
83+
if execution_context is None:
84+
return
85+
86+
span = getattr(execution_context, "_sentry_sql_span", None) # type: Optional[Span]
8487

8588
if span is not None:
8689
span.set_status("internal_error")
@@ -89,9 +92,9 @@ def _handle_error(context, *args):
8992
# from SQLAlchemy codebase it does seem like any error coming into this
9093
# handler is going to be fatal.
9194
ctx_mgr = getattr(
92-
conn, "_sentry_sql_span_manager", None
95+
execution_context, "_sentry_sql_span_manager", None
9396
) # type: ContextManager[Any]
9497

9598
if ctx_mgr is not None:
96-
conn._sentry_sql_span_manager = None
99+
execution_context._sentry_sql_span_manager = None
97100
ctx_mgr.__exit__(None, None, None)

‎sentry_sdk/tracing_utils.py

+44-52
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import re
2+
import contextlib
23
import json
34
import math
45

@@ -105,58 +106,6 @@ def __iter__(self):
105106
yield k[len(self.prefix) :]
106107

107108

108-
class RecordSqlQueries:
109-
def __init__(
110-
self,
111-
hub, # type: sentry_sdk.Hub
112-
cursor, # type: Any
113-
query, # type: Any
114-
params_list, # type: Any
115-
paramstyle, # type: Optional[str]
116-
executemany, # type: bool
117-
):
118-
# type: (...) -> None
119-
# TODO: Bring back capturing of params by default
120-
self._hub = hub
121-
if self._hub.client and self._hub.client.options["_experiments"].get(
122-
"record_sql_params", False
123-
):
124-
if not params_list or params_list == [None]:
125-
params_list = None
126-
127-
if paramstyle == "pyformat":
128-
paramstyle = "format"
129-
else:
130-
params_list = None
131-
paramstyle = None
132-
133-
self._query = _format_sql(cursor, query)
134-
135-
self._data = {}
136-
if params_list is not None:
137-
self._data["db.params"] = params_list
138-
if paramstyle is not None:
139-
self._data["db.paramstyle"] = paramstyle
140-
if executemany:
141-
self._data["db.executemany"] = True
142-
143-
def __enter__(self):
144-
# type: () -> Span
145-
with capture_internal_exceptions():
146-
self._hub.add_breadcrumb(
147-
message=self._query, category="query", data=self._data
148-
)
149-
150-
with self._hub.start_span(op="db", description=self._query) as span:
151-
for k, v in self._data.items():
152-
span.set_data(k, v)
153-
return span
154-
155-
def __exit__(self, exc_type, exc_val, exc_tb):
156-
# type: (Any, Any, Any) -> None
157-
pass
158-
159-
160109
def has_tracing_enabled(options):
161110
# type: (Dict[str, Any]) -> bool
162111
"""
@@ -201,6 +150,49 @@ def is_valid_sample_rate(rate):
201150
return True
202151

203152

153+
@contextlib.contextmanager
154+
def record_sql_queries(
155+
hub, # type: sentry_sdk.Hub
156+
cursor, # type: Any
157+
query, # type: Any
158+
params_list, # type: Any
159+
paramstyle, # type: Optional[str]
160+
executemany, # type: bool
161+
):
162+
# type: (...) -> Generator[Span, None, None]
163+
164+
# TODO: Bring back capturing of params by default
165+
if hub.client and hub.client.options["_experiments"].get(
166+
"record_sql_params", False
167+
):
168+
if not params_list or params_list == [None]:
169+
params_list = None
170+
171+
if paramstyle == "pyformat":
172+
paramstyle = "format"
173+
else:
174+
params_list = None
175+
paramstyle = None
176+
177+
query = _format_sql(cursor, query)
178+
179+
data = {}
180+
if params_list is not None:
181+
data["db.params"] = params_list
182+
if paramstyle is not None:
183+
data["db.paramstyle"] = paramstyle
184+
if executemany:
185+
data["db.executemany"] = True
186+
187+
with capture_internal_exceptions():
188+
hub.add_breadcrumb(message=query, category="query", data=data)
189+
190+
with hub.start_span(op="db", description=query) as span:
191+
for k, v in data.items():
192+
span.set_data(k, v)
193+
yield span
194+
195+
204196
def maybe_create_breadcrumbs_from_span(hub, span):
205197
# type: (sentry_sdk.Hub, Span) -> None
206198
if span.op == "redis":

0 commit comments

Comments
 (0)