fix(sqlalchemy): Use context instead of connection in sqlalchemy integration (#1388)

* Revert "fix(sqlalchemy): Change context manager type to avoid race in threads (#1368)"

This reverts commit de0bc50.
This caused a regression (#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

Author: sl0thentr0py

sentry_sdk/integrations/django/__init__.py

@@ -9,7 +9,7 @@
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.scope import add_global_event_processor
 from sentry_sdk.serializer import add_global_repr_processor
-from sentry_sdk.tracing_utils import RecordSqlQueries
+from sentry_sdk.tracing_utils import record_sql_queries
 from sentry_sdk.utils import (
     HAS_REAL_CONTEXTVARS,
     CONTEXTVARS_ERROR_MESSAGE,
@@ -538,7 +538,7 @@ def execute(self, sql, params=None):
         if hub.get_integration(DjangoIntegration) is None:
             return real_execute(self, sql, params)
 
-        with RecordSqlQueries(
+        with record_sql_queries(
             hub, self.cursor, sql, params, paramstyle="format", executemany=False
         ):
             return real_execute(self, sql, params)
@@ -549,7 +549,7 @@ def executemany(self, sql, param_list):
         if hub.get_integration(DjangoIntegration) is None:
             return real_executemany(self, sql, param_list)
 
-        with RecordSqlQueries(
+        with record_sql_queries(
             hub, self.cursor, sql, param_list, paramstyle="format", executemany=True
         ):
             return real_executemany(self, sql, param_list)

sentry_sdk/integrations/sqlalchemy.py

@@ -3,7 +3,7 @@
 from sentry_sdk._types import MYPY
 from sentry_sdk.hub import Hub
 from sentry_sdk.integrations import Integration, DidNotEnable
-from sentry_sdk.tracing_utils import RecordSqlQueries
+from sentry_sdk.tracing_utils import record_sql_queries
 
 try:
     from sqlalchemy.engine import Engine  # type: ignore
@@ -50,37 +50,40 @@ def _before_cursor_execute(
     if hub.get_integration(SqlalchemyIntegration) is None:
         return
 
-    ctx_mgr = RecordSqlQueries(
+    ctx_mgr = record_sql_queries(
         hub,
         cursor,
         statement,
         parameters,
         paramstyle=context and context.dialect and context.dialect.paramstyle or None,
         executemany=executemany,
     )
-    conn._sentry_sql_span_manager = ctx_mgr
+    context._sentry_sql_span_manager = ctx_mgr
 
     span = ctx_mgr.__enter__()
 
     if span is not None:
-        conn._sentry_sql_span = span
+        context._sentry_sql_span = span
 
 
-def _after_cursor_execute(conn, cursor, statement, *args):
-    # type: (Any, Any, Any, *Any) -> None
+def _after_cursor_execute(conn, cursor, statement, parameters, context, *args):
+    # type: (Any, Any, Any, Any, Any, *Any) -> None
     ctx_mgr = getattr(
-        conn, "_sentry_sql_span_manager", None
+        context, "_sentry_sql_span_manager", None
     )  # type: ContextManager[Any]
 
     if ctx_mgr is not None:
-        conn._sentry_sql_span_manager = None
+        context._sentry_sql_span_manager = None
         ctx_mgr.__exit__(None, None, None)
 
 
 def _handle_error(context, *args):
     # type: (Any, *Any) -> None
-    conn = context.connection
-    span = getattr(conn, "_sentry_sql_span", None)  # type: Optional[Span]
+    execution_context = context.execution_context
+    if execution_context is None:
+        return
+
+    span = getattr(execution_context, "_sentry_sql_span", None)  # type: Optional[Span]
 
     if span is not None:
         span.set_status("internal_error")
@@ -89,9 +92,9 @@ def _handle_error(context, *args):
     # from SQLAlchemy codebase it does seem like any error coming into this
     # handler is going to be fatal.
     ctx_mgr = getattr(
-        conn, "_sentry_sql_span_manager", None
+        execution_context, "_sentry_sql_span_manager", None
     )  # type: ContextManager[Any]
 
     if ctx_mgr is not None:
-        conn._sentry_sql_span_manager = None
+        execution_context._sentry_sql_span_manager = None
         ctx_mgr.__exit__(None, None, None)

sentry_sdk/tracing_utils.py

@@ -1,4 +1,5 @@
 import re
+import contextlib
 import json
 import math
 
@@ -105,58 +106,6 @@ def __iter__(self):
             yield k[len(self.prefix) :]
 
 
-class RecordSqlQueries:
-    def __init__(
-        self,
-        hub,  # type: sentry_sdk.Hub
-        cursor,  # type: Any
-        query,  # type: Any
-        params_list,  # type:  Any
-        paramstyle,  # type: Optional[str]
-        executemany,  # type: bool
-    ):
-        # type: (...) -> None
-        # TODO: Bring back capturing of params by default
-        self._hub = hub
-        if self._hub.client and self._hub.client.options["_experiments"].get(
-            "record_sql_params", False
-        ):
-            if not params_list or params_list == [None]:
-                params_list = None
-
-            if paramstyle == "pyformat":
-                paramstyle = "format"
-        else:
-            params_list = None
-            paramstyle = None
-
-        self._query = _format_sql(cursor, query)
-
-        self._data = {}
-        if params_list is not None:
-            self._data["db.params"] = params_list
-        if paramstyle is not None:
-            self._data["db.paramstyle"] = paramstyle
-        if executemany:
-            self._data["db.executemany"] = True
-
-    def __enter__(self):
-        # type: () -> Span
-        with capture_internal_exceptions():
-            self._hub.add_breadcrumb(
-                message=self._query, category="query", data=self._data
-            )
-
-        with self._hub.start_span(op="db", description=self._query) as span:
-            for k, v in self._data.items():
-                span.set_data(k, v)
-            return span
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        # type: (Any, Any, Any) -> None
-        pass
-
-
 def has_tracing_enabled(options):
     # type: (Dict[str, Any]) -> bool
     """
@@ -201,6 +150,49 @@ def is_valid_sample_rate(rate):
     return True
 
 
+@contextlib.contextmanager
+def record_sql_queries(
+    hub,  # type: sentry_sdk.Hub
+    cursor,  # type: Any
+    query,  # type: Any
+    params_list,  # type:  Any
+    paramstyle,  # type: Optional[str]
+    executemany,  # type: bool
+):
+    # type: (...) -> Generator[Span, None, None]
+
+    # TODO: Bring back capturing of params by default
+    if hub.client and hub.client.options["_experiments"].get(
+        "record_sql_params", False
+    ):
+        if not params_list or params_list == [None]:
+            params_list = None
+
+        if paramstyle == "pyformat":
+            paramstyle = "format"
+    else:
+        params_list = None
+        paramstyle = None
+
+    query = _format_sql(cursor, query)
+
+    data = {}
+    if params_list is not None:
+        data["db.params"] = params_list
+    if paramstyle is not None:
+        data["db.paramstyle"] = paramstyle
+    if executemany:
+        data["db.executemany"] = True
+
+    with capture_internal_exceptions():
+        hub.add_breadcrumb(message=query, category="query", data=data)
+
+    with hub.start_span(op="db", description=query) as span:
+        for k, v in data.items():
+            span.set_data(k, v)
+        yield span
+
+
 def maybe_create_breadcrumbs_from_span(hub, span):
     # type: (sentry_sdk.Hub, Span) -> None
     if span.op == "redis":