Skip to content

Commit 1a9225c

Browse files
sentrivanasentrivana
authored
fix(query-source): Fix query source duration check (#2675)
1 parent e864eab commit 1a9225c

File tree

7 files changed

+436
-5
lines changed

7 files changed

+436
-5
lines changed

‎sentry_sdk/_compat.py

+13-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import sys
22
import contextlib
3-
from datetime import datetime
3+
from datetime import datetime, timedelta
44
from functools import wraps
55

66
from sentry_sdk._types import TYPE_CHECKING
@@ -34,11 +34,19 @@
3434
binary_sequence_types = (bytearray, memoryview)
3535

3636
def datetime_utcnow():
37+
# type: () -> datetime
3738
return datetime.utcnow()
3839

3940
def utc_from_timestamp(timestamp):
41+
# type: (float) -> datetime
4042
return datetime.utcfromtimestamp(timestamp)
4143

44+
def duration_in_milliseconds(delta):
45+
# type: (timedelta) -> float
46+
seconds = delta.days * 24 * 60 * 60 + delta.seconds
47+
milliseconds = seconds * 1000 + float(delta.microseconds) / 1000
48+
return milliseconds
49+
4250
def implements_str(cls):
4351
# type: (T) -> T
4452
cls.__unicode__ = cls.__str__
@@ -103,6 +111,10 @@ def utc_from_timestamp(timestamp):
103111
# type: (float) -> datetime
104112
return datetime.fromtimestamp(timestamp, timezone.utc)
105113

114+
def duration_in_milliseconds(delta):
115+
# type: (timedelta) -> float
116+
return delta / timedelta(milliseconds=1)
117+
106118
def implements_str(x):
107119
# type: (T) -> T
108120
return x

‎sentry_sdk/integrations/gcp.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from sentry_sdk.consts import OP
88
from sentry_sdk.hub import Hub, _should_send_default_pii
99
from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
10-
from sentry_sdk._compat import datetime_utcnow, reraise
10+
from sentry_sdk._compat import datetime_utcnow, duration_in_milliseconds, reraise
1111
from sentry_sdk.utils import (
1212
AnnotatedValue,
1313
capture_internal_exceptions,
@@ -158,7 +158,7 @@ def event_processor(event, hint):
158158
final_time = datetime_utcnow()
159159
time_diff = final_time - initial_time
160160

161-
execution_duration_in_millis = time_diff.microseconds / MILLIS_TO_SECONDS
161+
execution_duration_in_millis = duration_in_milliseconds(time_diff)
162162

163163
extra = event.setdefault("extra", {})
164164
extra["google cloud functions"] = {

‎sentry_sdk/tracing_utils.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
_is_external_source,
1515
_module_in_list,
1616
)
17-
from sentry_sdk._compat import PY2, iteritems
17+
from sentry_sdk._compat import PY2, duration_in_milliseconds, iteritems
1818
from sentry_sdk._types import TYPE_CHECKING
1919

2020
if PY2:
@@ -186,7 +186,7 @@ def add_query_source(hub, span):
186186

187187
duration = span.timestamp - span.start_timestamp
188188
threshold = client.options.get("db_query_source_threshold_ms", 0)
189-
slow_query = duration.microseconds > threshold * 1000
189+
slow_query = duration_in_milliseconds(duration) > threshold
190190

191191
if not slow_query:
192192
return

‎tests/integrations/asyncpg/test_asyncpg.py

+114
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,13 @@
3131
from sentry_sdk import capture_message, start_transaction
3232
from sentry_sdk.integrations.asyncpg import AsyncPGIntegration
3333
from sentry_sdk.consts import SPANDATA
34+
from sentry_sdk.tracing_utils import record_sql_queries
35+
from sentry_sdk._compat import contextmanager
36+
37+
try:
38+
from unittest import mock
39+
except ImportError:
40+
import mock
3441

3542

3643
PG_CONNECTION_URI = "postgresql://{}:{}@{}/{}".format(
@@ -548,3 +555,110 @@ async def test_query_source(sentry_init, capture_events):
548555
assert is_relative_path
549556

550557
assert data.get(SPANDATA.CODE_FUNCTION) == "test_query_source"
558+
559+
560+
@pytest.mark.asyncio
561+
async def test_no_query_source_if_duration_too_short(sentry_init, capture_events):
562+
sentry_init(
563+
integrations=[AsyncPGIntegration()],
564+
enable_tracing=True,
565+
enable_db_query_source=True,
566+
db_query_source_threshold_ms=100,
567+
)
568+
569+
events = capture_events()
570+
571+
with start_transaction(name="test_transaction", sampled=True):
572+
conn: Connection = await connect(PG_CONNECTION_URI)
573+
574+
@contextmanager
575+
def fake_record_sql_queries(*args, **kwargs):
576+
with record_sql_queries(*args, **kwargs) as span:
577+
pass
578+
span.start_timestamp = datetime.datetime(2024, 1, 1, microsecond=0)
579+
span.timestamp = datetime.datetime(2024, 1, 1, microsecond=99999)
580+
yield span
581+
582+
with mock.patch(
583+
"sentry_sdk.integrations.asyncpg.record_sql_queries",
584+
fake_record_sql_queries,
585+
):
586+
await conn.execute(
587+
"INSERT INTO users(name, password, dob) VALUES ('Alice', 'secret', '1990-12-25')",
588+
)
589+
590+
await conn.close()
591+
592+
(event,) = events
593+
594+
span = event["spans"][-1]
595+
assert span["description"].startswith("INSERT INTO")
596+
597+
data = span.get("data", {})
598+
599+
assert SPANDATA.CODE_LINENO not in data
600+
assert SPANDATA.CODE_NAMESPACE not in data
601+
assert SPANDATA.CODE_FILEPATH not in data
602+
assert SPANDATA.CODE_FUNCTION not in data
603+
604+
605+
@pytest.mark.asyncio
606+
async def test_query_source_if_duration_over_threshold(sentry_init, capture_events):
607+
sentry_init(
608+
integrations=[AsyncPGIntegration()],
609+
enable_tracing=True,
610+
enable_db_query_source=True,
611+
db_query_source_threshold_ms=100,
612+
)
613+
614+
events = capture_events()
615+
616+
with start_transaction(name="test_transaction", sampled=True):
617+
conn: Connection = await connect(PG_CONNECTION_URI)
618+
619+
@contextmanager
620+
def fake_record_sql_queries(*args, **kwargs):
621+
with record_sql_queries(*args, **kwargs) as span:
622+
pass
623+
span.start_timestamp = datetime.datetime(2024, 1, 1, microsecond=0)
624+
span.timestamp = datetime.datetime(2024, 1, 1, microsecond=100001)
625+
yield span
626+
627+
with mock.patch(
628+
"sentry_sdk.integrations.asyncpg.record_sql_queries",
629+
fake_record_sql_queries,
630+
):
631+
await conn.execute(
632+
"INSERT INTO users(name, password, dob) VALUES ('Alice', 'secret', '1990-12-25')",
633+
)
634+
635+
await conn.close()
636+
637+
(event,) = events
638+
639+
span = event["spans"][-1]
640+
assert span["description"].startswith("INSERT INTO")
641+
642+
data = span.get("data", {})
643+
644+
assert SPANDATA.CODE_LINENO in data
645+
assert SPANDATA.CODE_NAMESPACE in data
646+
assert SPANDATA.CODE_FILEPATH in data
647+
assert SPANDATA.CODE_FUNCTION in data
648+
649+
assert type(data.get(SPANDATA.CODE_LINENO)) == int
650+
assert data.get(SPANDATA.CODE_LINENO) > 0
651+
assert (
652+
data.get(SPANDATA.CODE_NAMESPACE) == "tests.integrations.asyncpg.test_asyncpg"
653+
)
654+
assert data.get(SPANDATA.CODE_FILEPATH).endswith(
655+
"tests/integrations/asyncpg/test_asyncpg.py"
656+
)
657+
658+
is_relative_path = data.get(SPANDATA.CODE_FILEPATH)[0] != os.sep
659+
assert is_relative_path
660+
661+
assert (
662+
data.get(SPANDATA.CODE_FUNCTION)
663+
== "test_query_source_if_duration_over_threshold"
664+
)

‎tests/integrations/django/test_db_query_data.py

+138
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import os
44
import pytest
5+
from datetime import datetime
56

67
from django import VERSION as DJANGO_VERSION
78
from django.db import connections
@@ -15,11 +16,17 @@
1516

1617
from sentry_sdk.consts import SPANDATA
1718
from sentry_sdk.integrations.django import DjangoIntegration
19+
from sentry_sdk.tracing_utils import record_sql_queries
1820

1921
from tests.conftest import unpack_werkzeug_response
2022
from tests.integrations.django.utils import pytest_mark_django_db_decorator
2123
from tests.integrations.django.myapp.wsgi import application
2224

25+
try:
26+
from unittest import mock
27+
except ImportError:
28+
import mock
29+
2330

2431
@pytest.fixture
2532
def client():
@@ -228,3 +235,134 @@ def test_query_source_with_in_app_include(sentry_init, client, capture_events):
228235
break
229236
else:
230237
raise AssertionError("No db span found")
238+
239+
240+
@pytest.mark.forked
241+
@pytest_mark_django_db_decorator(transaction=True)
242+
def test_no_query_source_if_duration_too_short(sentry_init, client, capture_events):
243+
sentry_init(
244+
integrations=[DjangoIntegration()],
245+
send_default_pii=True,
246+
traces_sample_rate=1.0,
247+
enable_db_query_source=True,
248+
db_query_source_threshold_ms=100,
249+
)
250+
251+
if "postgres" not in connections:
252+
pytest.skip("postgres tests disabled")
253+
254+
# trigger Django to open a new connection by marking the existing one as None.
255+
connections["postgres"].connection = None
256+
257+
events = capture_events()
258+
259+
class fake_record_sql_queries: # noqa: N801
260+
def __init__(self, *args, **kwargs):
261+
with record_sql_queries(*args, **kwargs) as span:
262+
self.span = span
263+
264+
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
265+
self.span.timestamp = datetime(2024, 1, 1, microsecond=99999)
266+
267+
def __enter__(self):
268+
return self.span
269+
270+
def __exit__(self, type, value, traceback):
271+
pass
272+
273+
with mock.patch(
274+
"sentry_sdk.integrations.django.record_sql_queries",
275+
fake_record_sql_queries,
276+
):
277+
_, status, _ = unpack_werkzeug_response(
278+
client.get(reverse("postgres_select_orm"))
279+
)
280+
281+
assert status == "200 OK"
282+
283+
(event,) = events
284+
for span in event["spans"]:
285+
if span.get("op") == "db" and "auth_user" in span.get("description"):
286+
data = span.get("data", {})
287+
288+
assert SPANDATA.CODE_LINENO not in data
289+
assert SPANDATA.CODE_NAMESPACE not in data
290+
assert SPANDATA.CODE_FILEPATH not in data
291+
assert SPANDATA.CODE_FUNCTION not in data
292+
293+
break
294+
else:
295+
raise AssertionError("No db span found")
296+
297+
298+
@pytest.mark.forked
299+
@pytest_mark_django_db_decorator(transaction=True)
300+
def test_query_source_if_duration_over_threshold(sentry_init, client, capture_events):
301+
sentry_init(
302+
integrations=[DjangoIntegration()],
303+
send_default_pii=True,
304+
traces_sample_rate=1.0,
305+
enable_db_query_source=True,
306+
db_query_source_threshold_ms=100,
307+
)
308+
309+
if "postgres" not in connections:
310+
pytest.skip("postgres tests disabled")
311+
312+
# trigger Django to open a new connection by marking the existing one as None.
313+
connections["postgres"].connection = None
314+
315+
events = capture_events()
316+
317+
class fake_record_sql_queries: # noqa: N801
318+
def __init__(self, *args, **kwargs):
319+
with record_sql_queries(*args, **kwargs) as span:
320+
self.span = span
321+
322+
self.span.start_timestamp = datetime(2024, 1, 1, microsecond=0)
323+
self.span.timestamp = datetime(2024, 1, 1, microsecond=101000)
324+
325+
def __enter__(self):
326+
return self.span
327+
328+
def __exit__(self, type, value, traceback):
329+
pass
330+
331+
with mock.patch(
332+
"sentry_sdk.integrations.django.record_sql_queries",
333+
fake_record_sql_queries,
334+
):
335+
_, status, _ = unpack_werkzeug_response(
336+
client.get(reverse("postgres_select_orm"))
337+
)
338+
339+
assert status == "200 OK"
340+
341+
(event,) = events
342+
for span in event["spans"]:
343+
if span.get("op") == "db" and "auth_user" in span.get("description"):
344+
data = span.get("data", {})
345+
346+
assert SPANDATA.CODE_LINENO in data
347+
assert SPANDATA.CODE_NAMESPACE in data
348+
assert SPANDATA.CODE_FILEPATH in data
349+
assert SPANDATA.CODE_FUNCTION in data
350+
351+
assert type(data.get(SPANDATA.CODE_LINENO)) == int
352+
assert data.get(SPANDATA.CODE_LINENO) > 0
353+
354+
assert (
355+
data.get(SPANDATA.CODE_NAMESPACE)
356+
== "tests.integrations.django.myapp.views"
357+
)
358+
assert data.get(SPANDATA.CODE_FILEPATH).endswith(
359+
"tests/integrations/django/myapp/views.py"
360+
)
361+
362+
is_relative_path = data.get(SPANDATA.CODE_FILEPATH)[0] != os.sep
363+
assert is_relative_path
364+
365+
assert data.get(SPANDATA.CODE_FUNCTION) == "postgres_select_orm"
366+
break
367+
else:
368+
raise AssertionError("No db span found")

0 commit comments

Comments
 (0)