Add query source to DB spans (#2521)

Adding OTel compatible information to database spans that show the code location of the query.
Refs getsentry/team-sdks#40
---------

Co-authored-by: Ivana Kellyerova <[email protected]>

Author: antonpirker

sentry_sdk/consts.py

@@ -164,6 +164,30 @@ class SPANDATA:
     Example: 16456
     """
 
+    CODE_FILEPATH = "code.filepath"
+    """
+    The source code file name that identifies the code unit as uniquely as possible (preferably an absolute file path).
+    Example: "/app/myapplication/http/handler/server.py"
+    """
+
+    CODE_LINENO = "code.lineno"
+    """
+    The line number in `code.filepath` best representing the operation. It SHOULD point within the code unit named in `code.function`.
+    Example: 42
+    """
+
+    CODE_FUNCTION = "code.function"
+    """
+    The method or function name, or equivalent (usually rightmost part of the code unit's name).
+    Example: "server_request"
+    """
+
+    CODE_NAMESPACE = "code.namespace"
+    """
+    The "namespace" within which `code.function` is defined. Usually the qualified class or module name, such that `code.namespace` + some separator + `code.function` form a unique identifier for the code unit.
+    Example: "http.handler"
+    """
+
 
 class OP:
     CACHE_GET_ITEM = "cache.get_item"
@@ -264,6 +288,8 @@ def __init__(
         max_value_length=DEFAULT_MAX_VALUE_LENGTH,  # type: int
         enable_backpressure_handling=True,  # type: bool
         error_sampler=None,  # type: Optional[Callable[[Event, Hint], Union[float, bool]]]
+        enable_db_query_source=False,  # type: bool
+        db_query_source_threshold_ms=100,  # type: int
         spotlight=None,  # type: Optional[Union[bool, str]]
     ):
         # type: (...) -> None

sentry_sdk/tracing.py

@@ -479,6 +479,8 @@ def finish(self, hub=None, end_timestamp=None):
             self.timestamp = datetime_utcnow()
 
         maybe_create_breadcrumbs_from_span(hub, self)
+        add_additional_span_data(hub, self)
+
         return None
 
     def to_json(self):
@@ -998,6 +1000,7 @@ async def my_async_function():
 from sentry_sdk.tracing_utils import (
     Baggage,
     EnvironHeaders,
+    add_additional_span_data,
     extract_sentrytrace_data,
     has_tracing_enabled,
     maybe_create_breadcrumbs_from_span,

sentry_sdk/tracing_utils.py

@@ -1,14 +1,16 @@
-import re
 import contextlib
+import re
+import sys
 
 import sentry_sdk
-from sentry_sdk.consts import OP
+from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     Dsn,
     match_regex_list,
     to_string,
     is_sentry_url,
+    _is_external_source,
 )
 from sentry_sdk._compat import PY2, iteritems
 from sentry_sdk._types import TYPE_CHECKING
@@ -29,6 +31,8 @@
     from typing import Optional
     from typing import Union
 
+    from types import FrameType
+
 
 SENTRY_TRACE_REGEX = re.compile(
     "^[ \t]*"  # whitespace
@@ -162,6 +166,98 @@ def maybe_create_breadcrumbs_from_span(hub, span):
         )
 
 
+def add_query_source(hub, span):
+    # type: (sentry_sdk.Hub, sentry_sdk.tracing.Span) -> None
+    """
+    Adds OTel compatible source code information to the span
+    """
+    client = hub.client
+    if client is None:
+        return
+
+    if span.timestamp is None or span.start_timestamp is None:
+        return
+
+    should_add_query_source = client.options.get("enable_db_query_source", False)
+    if not should_add_query_source:
+        return
+
+    duration = span.timestamp - span.start_timestamp
+    threshold = client.options.get("db_query_source_threshold_ms", 0)
+    slow_query = duration.microseconds > threshold * 1000
+
+    if not slow_query:
+        return
+
+    project_root = client.options["project_root"]
+
+    # Find the correct frame
+    frame = sys._getframe()  # type: Union[FrameType, None]
+    while frame is not None:
+        try:
+            abs_path = frame.f_code.co_filename
+        except Exception:
+            abs_path = ""
+
+        try:
+            namespace = frame.f_globals.get("__name__")
+        except Exception:
+            namespace = None
+
+        is_sentry_sdk_frame = namespace is not None and namespace.startswith(
+            "sentry_sdk."
+        )
+        if (
+            abs_path.startswith(project_root)
+            and not _is_external_source(abs_path)
+            and not is_sentry_sdk_frame
+        ):
+            break
+        frame = frame.f_back
+    else:
+        frame = None
+
+    # Set the data
+    if frame is not None:
+        try:
+            lineno = frame.f_lineno
+        except Exception:
+            lineno = None
+        if lineno is not None:
+            span.set_data(SPANDATA.CODE_LINENO, frame.f_lineno)
+
+        try:
+            namespace = frame.f_globals.get("__name__")
+        except Exception:
+            namespace = None
+        if namespace is not None:
+            span.set_data(SPANDATA.CODE_NAMESPACE, namespace)
+
+        try:
+            filepath = frame.f_code.co_filename
+        except Exception:
+            filepath = None
+        if filepath is not None:
+            span.set_data(SPANDATA.CODE_FILEPATH, frame.f_code.co_filename)
+
+        try:
+            code_function = frame.f_code.co_name
+        except Exception:
+            code_function = None
+
+        if code_function is not None:
+            span.set_data(SPANDATA.CODE_FUNCTION, frame.f_code.co_name)
+
+
+def add_additional_span_data(hub, span):
+    # type: (sentry_sdk.Hub, sentry_sdk.tracing.Span) -> None
+    """
+    Adds additional data to the span
+    """
+    if span.op == OP.DB:
+        add_query_source(hub, span)
+
+
 def extract_sentrytrace_data(header):
     # type: (Optional[str]) -> Optional[Dict[str, Union[str, bool, None]]]
     """

tests/integrations/asyncpg/test_asyncpg.py

@@ -27,8 +27,9 @@
 
 from asyncpg import connect, Connection
 
-from sentry_sdk import capture_message
+from sentry_sdk import capture_message, start_transaction
 from sentry_sdk.integrations.asyncpg import AsyncPGIntegration
+from sentry_sdk.consts import SPANDATA
 
 
 PG_CONNECTION_URI = "postgresql://{}:{}@{}/{}".format(
@@ -460,3 +461,85 @@ async def test_connection_pool(sentry_init, capture_events) -> None:
             "type": "default",
         },
     ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("enable_db_query_source", [None, False])
+async def test_query_source_disabled(
+    sentry_init, capture_events, enable_db_query_source
+):
+    sentry_options = {
+        "integrations": [AsyncPGIntegration()],
+        "enable_tracing": True,
+    }
+    if enable_db_query_source is not None:
+        sentry_options["enable_db_query_source"] = enable_db_query_source
+        sentry_options["db_query_source_threshold_ms"] = 0
+
+    sentry_init(**sentry_options)
+
+    events = capture_events()
+
+    with start_transaction(name="test_transaction", sampled=True):
+        conn: Connection = await connect(PG_CONNECTION_URI)
+
+        await conn.execute(
+            "INSERT INTO users(name, password, dob) VALUES ('Alice', 'secret', '1990-12-25')",
+        )
+
+        await conn.close()
+
+    (event,) = events
+
+    span = event["spans"][-1]
+    assert span["description"].startswith("INSERT INTO")
+
+    data = span.get("data", {})
+
+    assert SPANDATA.CODE_LINENO not in data
+    assert SPANDATA.CODE_NAMESPACE not in data
+    assert SPANDATA.CODE_FILEPATH not in data
+    assert SPANDATA.CODE_FUNCTION not in data
+
+
+@pytest.mark.asyncio
+async def test_query_source(sentry_init, capture_events):
+    sentry_init(
+        integrations=[AsyncPGIntegration()],
+        enable_tracing=True,
+        enable_db_query_source=True,
+        db_query_source_threshold_ms=0,
+    )
+
+    events = capture_events()
+
+    with start_transaction(name="test_transaction", sampled=True):
+        conn: Connection = await connect(PG_CONNECTION_URI)
+
+        await conn.execute(
+            "INSERT INTO users(name, password, dob) VALUES ('Alice', 'secret', '1990-12-25')",
+        )
+
+        await conn.close()
+
+    (event,) = events
+
+    span = event["spans"][-1]
+    assert span["description"].startswith("INSERT INTO")
+
+    data = span.get("data", {})
+
+    assert SPANDATA.CODE_LINENO in data
+    assert SPANDATA.CODE_NAMESPACE in data
+    assert SPANDATA.CODE_FILEPATH in data
+    assert SPANDATA.CODE_FUNCTION in data
+
+    assert type(data.get(SPANDATA.CODE_LINENO)) == int
+    assert data.get(SPANDATA.CODE_LINENO) > 0
+    assert (
+        data.get(SPANDATA.CODE_NAMESPACE) == "tests.integrations.asyncpg.test_asyncpg"
+    )
+    assert data.get(SPANDATA.CODE_FILEPATH).endswith(
+        "tests/integrations/asyncpg/test_asyncpg.py"
+    )
+    assert data.get(SPANDATA.CODE_FUNCTION) == "test_query_source"

tests/integrations/django/myapp/urls.py

@@ -57,6 +57,7 @@ def path(path, *args, **kwargs):
     path("template-test2", views.template_test2, name="template_test2"),
     path("template-test3", views.template_test3, name="template_test3"),
     path("postgres-select", views.postgres_select, name="postgres_select"),
+    path("postgres-select-slow", views.postgres_select_orm, name="postgres_select_orm"),
     path(
         "permission-denied-exc",
         views.permission_denied_exc,

tests/integrations/django/myapp/views.py

@@ -193,6 +193,12 @@ def postgres_select(request, *args, **kwargs):
     return HttpResponse("ok")
 
 
+@csrf_exempt
+def postgres_select_orm(request, *args, **kwargs):
+    user = User.objects.using("postgres").all().first()
+    return HttpResponse("ok {}".format(user))
+
+
 @csrf_exempt
 def permission_denied_exc(*args, **kwargs):
     raise PermissionDenied("bye")