fix(query-source): Fix query source duration check (#2675)

Author: sentrivana

sentry_sdk/_compat.py

@@ -1,6 +1,6 @@
 import sys
 import contextlib
-from datetime import datetime
+from datetime import datetime, timedelta
 from functools import wraps
 
 from sentry_sdk._types import TYPE_CHECKING
@@ -34,11 +34,19 @@
     binary_sequence_types = (bytearray, memoryview)
 
     def datetime_utcnow():
+        # type: () -> datetime
         return datetime.utcnow()
 
     def utc_from_timestamp(timestamp):
+        # type: (float) -> datetime
         return datetime.utcfromtimestamp(timestamp)
 
+    def duration_in_milliseconds(delta):
+        # type: (timedelta) -> float
+        seconds = delta.days * 24 * 60 * 60 + delta.seconds
+        milliseconds = seconds * 1000 + float(delta.microseconds) / 1000
+        return milliseconds
+
     def implements_str(cls):
         # type: (T) -> T
         cls.__unicode__ = cls.__str__
@@ -103,6 +111,10 @@ def utc_from_timestamp(timestamp):
         # type: (float) -> datetime
         return datetime.fromtimestamp(timestamp, timezone.utc)
 
+    def duration_in_milliseconds(delta):
+        # type: (timedelta) -> float
+        return delta / timedelta(milliseconds=1)
+
     def implements_str(x):
         # type: (T) -> T
         return x

sentry_sdk/integrations/gcp.py

@@ -7,7 +7,7 @@
 from sentry_sdk.consts import OP
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
-from sentry_sdk._compat import datetime_utcnow, reraise
+from sentry_sdk._compat import datetime_utcnow, duration_in_milliseconds, reraise
 from sentry_sdk.utils import (
     AnnotatedValue,
     capture_internal_exceptions,
@@ -158,7 +158,7 @@ def event_processor(event, hint):
         final_time = datetime_utcnow()
         time_diff = final_time - initial_time
 
-        execution_duration_in_millis = time_diff.microseconds / MILLIS_TO_SECONDS
+        execution_duration_in_millis = duration_in_milliseconds(time_diff)
 
         extra = event.setdefault("extra", {})
         extra["google cloud functions"] = {

sentry_sdk/tracing_utils.py

@@ -14,7 +14,7 @@
     _is_external_source,
     _module_in_list,
 )
-from sentry_sdk._compat import PY2, iteritems
+from sentry_sdk._compat import PY2, duration_in_milliseconds, iteritems
 from sentry_sdk._types import TYPE_CHECKING
 
 if PY2:
@@ -186,7 +186,7 @@ def add_query_source(hub, span):
 
     duration = span.timestamp - span.start_timestamp
     threshold = client.options.get("db_query_source_threshold_ms", 0)
-    slow_query = duration.microseconds > threshold * 1000
+    slow_query = duration_in_milliseconds(duration) > threshold
 
     if not slow_query:
         return

tests/integrations/asyncpg/test_asyncpg.py

@@ -31,6 +31,13 @@
 from sentry_sdk import capture_message, start_transaction
 from sentry_sdk.integrations.asyncpg import AsyncPGIntegration
 from sentry_sdk.consts import SPANDATA
+from sentry_sdk.tracing_utils import record_sql_queries
+from sentry_sdk._compat import contextmanager
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
 
 
 PG_CONNECTION_URI = "postgresql://{}:{}@{}/{}".format(
@@ -548,3 +555,110 @@ async def test_query_source(sentry_init, capture_events):
     assert is_relative_path
 
     assert data.get(SPANDATA.CODE_FUNCTION) == "test_query_source"
+
+
+@pytest.mark.asyncio
+async def test_no_query_source_if_duration_too_short(sentry_init, capture_events):
+    sentry_init(
+        integrations=[AsyncPGIntegration()],
+        enable_tracing=True,
+        enable_db_query_source=True,
+        db_query_source_threshold_ms=100,
+    )
+
+    events = capture_events()
+
+    with start_transaction(name="test_transaction", sampled=True):
+        conn: Connection = await connect(PG_CONNECTION_URI)
+
+        @contextmanager
+        def fake_record_sql_queries(*args, **kwargs):
+            with record_sql_queries(*args, **kwargs) as span:
+                pass
+            span.start_timestamp = datetime.datetime(2024, 1, 1, microsecond=0)
+            span.timestamp = datetime.datetime(2024, 1, 1, microsecond=99999)
+            yield span
+
+        with mock.patch(
+            "sentry_sdk.integrations.asyncpg.record_sql_queries",
+            fake_record_sql_queries,
+        ):
+            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_if_duration_over_threshold(sentry_init, capture_events):
+    sentry_init(
+        integrations=[AsyncPGIntegration()],
+        enable_tracing=True,
+        enable_db_query_source=True,
+        db_query_source_threshold_ms=100,
+    )
+
+    events = capture_events()
+
+    with start_transaction(name="test_transaction", sampled=True):
+        conn: Connection = await connect(PG_CONNECTION_URI)
+
+        @contextmanager
+        def fake_record_sql_queries(*args, **kwargs):
+            with record_sql_queries(*args, **kwargs) as span:
+                pass
+            span.start_timestamp = datetime.datetime(2024, 1, 1, microsecond=0)
+            span.timestamp = datetime.datetime(2024, 1, 1, microsecond=100001)
+            yield span
+
+        with mock.patch(
+            "sentry_sdk.integrations.asyncpg.record_sql_queries",
+            fake_record_sql_queries,
+        ):
+            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"
+    )
+
+    is_relative_path = data.get(SPANDATA.CODE_FILEPATH)[0] != os.sep
+    assert is_relative_path
+
+    assert (
+        data.get(SPANDATA.CODE_FUNCTION)
+        == "test_query_source_if_duration_over_threshold"
+    )

tests/integrations/django/test_db_query_data.py

@@ -2,6 +2,7 @@
 
 import os
 import pytest
+from datetime import datetime
 
 from django import VERSION as DJANGO_VERSION
 from django.db import connections
@@ -15,11 +16,17 @@
 
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.django import DjangoIntegration
+from sentry_sdk.tracing_utils import record_sql_queries
 
 from tests.conftest import unpack_werkzeug_response
 from tests.integrations.django.utils import pytest_mark_django_db_decorator
 from tests.integrations.django.myapp.wsgi import application
 
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
 
 @pytest.fixture
 def client():
@@ -228,3 +235,134 @@ def test_query_source_with_in_app_include(sentry_init, client, capture_events):
             break
     else:
         raise AssertionError("No db span found")
+
+
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_no_query_source_if_duration_too_short(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+        traces_sample_rate=1.0,
+        enable_db_query_source=True,
+        db_query_source_threshold_ms=100,
+    )
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    # trigger Django to open a new connection by marking the existing one as None.
+    connections["postgres"].connection = None
+
+    events = capture_events()
+
+    class fake_record_sql_queries:  # noqa: N801
+        def __init__(self, *args, **kwargs):
+            with record_sql_queries(*args, **kwargs) as span:
+                self.span = span
+
+            self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
+            self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
+
+        def __enter__(self):
+            return self.span
+
+        def __exit__(self, type, value, traceback):
+            pass
+
+    with mock.patch(
+        "sentry_sdk.integrations.django.record_sql_queries",
+        fake_record_sql_queries,
+    ):
+        _, status, _ = unpack_werkzeug_response(
+            client.get(reverse("postgres_select_orm"))
+        )
+
+    assert status == "200 OK"
+
+    (event,) = events
+    for span in event["spans"]:
+        if span.get("op") == "db" and "auth_user" in span.get("description"):
+            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
+
+            break
+    else:
+        raise AssertionError("No db span found")
+
+
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_query_source_if_duration_over_threshold(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+        traces_sample_rate=1.0,
+        enable_db_query_source=True,
+        db_query_source_threshold_ms=100,
+    )
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    # trigger Django to open a new connection by marking the existing one as None.
+    connections["postgres"].connection = None
+
+    events = capture_events()
+
+    class fake_record_sql_queries:  # noqa: N801
+        def __init__(self, *args, **kwargs):
+            with record_sql_queries(*args, **kwargs) as span:
+                self.span = span
+
+            self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
+            self.span.timestamp = datetime(2024, 1, 1, microsecond=101000)
+
+        def __enter__(self):
+            return self.span
+
+        def __exit__(self, type, value, traceback):
+            pass
+
+    with mock.patch(
+        "sentry_sdk.integrations.django.record_sql_queries",
+        fake_record_sql_queries,
+    ):
+        _, status, _ = unpack_werkzeug_response(
+            client.get(reverse("postgres_select_orm"))
+        )
+
+    assert status == "200 OK"
+
+    (event,) = events
+    for span in event["spans"]:
+        if span.get("op") == "db" and "auth_user" in span.get("description"):
+            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.django.myapp.views"
+            )
+            assert data.get(SPANDATA.CODE_FILEPATH).endswith(
+                "tests/integrations/django/myapp/views.py"
+            )
+
+            is_relative_path = data.get(SPANDATA.CODE_FILEPATH)[0] != os.sep
+            assert is_relative_path
+
+            assert data.get(SPANDATA.CODE_FUNCTION) == "postgres_select_orm"
+            break
+    else:
+        raise AssertionError("No db span found")