Skip to content

Commit f02c204

Browse files
sentrivanasentrivana
authored
Fix asyncpg breadcrumbs (#3685)
1 parent 7422a4c commit f02c204

File tree

2 files changed

+45
-54
lines changed

2 files changed

+45
-54
lines changed

‎sentry_sdk/integrations/asyncpg.py

Lines changed: 39 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import annotations
22
import contextlib
3-
from typing import Any, TypeVar, Callable, Awaitable, Iterator
3+
from typing import Any, TypeVar, Callable, Awaitable, Iterator, Optional
44

55
import sentry_sdk
66
from sentry_sdk.consts import OP, SPANDATA
@@ -21,6 +21,7 @@
2121
except ImportError:
2222
raise DidNotEnable("asyncpg not installed.")
2323

24+
2425
# asyncpg.__version__ is a string containing the semantic version in the form of "<major>.<minor>.<patch>"
2526
asyncpg_version = parse_version(asyncpg.__version__)
2627

@@ -123,10 +124,13 @@ def _wrap_connection_method(
123124
async def _inner(*args: Any, **kwargs: Any) -> T:
124125
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
125126
return await f(*args, **kwargs)
127+
126128
query = args[1]
127129
params_list = args[2] if len(args) > 2 else None
130+
128131
with _record(None, query, params_list, executemany=executemany) as span:
129-
_set_db_data(span, args[0])
132+
data = _get_db_data(conn=args[0])
133+
_set_on_span(span, data)
130134
res = await f(*args, **kwargs)
131135

132136
return res
@@ -146,7 +150,8 @@ def _inner(*args: Any, **kwargs: Any) -> T: # noqa: N807
146150
params_list,
147151
executemany=False,
148152
) as span:
149-
_set_db_data(span, args[0])
153+
data = _get_db_data(conn=args[0])
154+
_set_on_span(span, data)
150155
res = f(*args, **kwargs)
151156
span.set_attribute("db.cursor", _serialize_span_attribute(res))
152157

@@ -160,41 +165,19 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
160165
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
161166
return await f(*args, **kwargs)
162167

163-
user = kwargs["params"].user
164-
database = kwargs["params"].database
165-
166168
with sentry_sdk.start_span(
167169
op=OP.DB,
168170
name="connect",
169171
origin=AsyncPGIntegration.origin,
170172
) as span:
171-
span.set_attribute(SPANDATA.DB_SYSTEM, "postgresql")
172-
addr = kwargs.get("addr")
173-
if addr:
174-
try:
175-
span.set_attribute(SPANDATA.SERVER_ADDRESS, addr[0])
176-
span.set_attribute(SPANDATA.SERVER_PORT, addr[1])
177-
except IndexError:
178-
pass
179-
180-
span.set_attribute(SPANDATA.DB_NAME, database)
181-
span.set_attribute(SPANDATA.DB_USER, user)
173+
data = _get_db_data(
174+
addr=kwargs.get("addr"),
175+
database=kwargs["params"].database,
176+
user=kwargs["params"].user,
177+
)
178+
_set_on_span(span, data)
182179

183180
with capture_internal_exceptions():
184-
data = {}
185-
for attr in (
186-
"db.cursor",
187-
"db.params",
188-
"db.paramstyle",
189-
SPANDATA.DB_NAME,
190-
SPANDATA.DB_SYSTEM,
191-
SPANDATA.DB_USER,
192-
SPANDATA.SERVER_ADDRESS,
193-
SPANDATA.SERVER_PORT,
194-
):
195-
if span.get_attribute(attr):
196-
data[attr] = span.get_attribute(attr)
197-
198181
sentry_sdk.add_breadcrumb(
199182
message="connect", category="query", data=data
200183
)
@@ -206,21 +189,37 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
206189
return _inner
207190

208191

209-
def _set_db_data(span: Span, conn: Any) -> None:
210-
span.set_attribute(SPANDATA.DB_SYSTEM, "postgresql")
192+
def _get_db_data(
193+
conn: Any = None,
194+
addr: Optional[tuple[str]] = None,
195+
database: Optional[str] = None,
196+
user: Optional[str] = None,
197+
) -> dict[str, str]:
198+
if conn is not None:
199+
addr = conn._addr
200+
database = conn._params.database
201+
user = conn._params.user
202+
203+
data = {
204+
SPANDATA.DB_SYSTEM: "postgresql",
205+
}
211206

212-
addr = conn._addr
213207
if addr:
214208
try:
215-
span.set_attribute(SPANDATA.SERVER_ADDRESS, addr[0])
216-
span.set_attribute(SPANDATA.SERVER_PORT, addr[1])
209+
data[SPANDATA.SERVER_ADDRESS] = addr[0]
210+
data[SPANDATA.SERVER_PORT] = addr[1]
217211
except IndexError:
218212
pass
219213

220-
database = conn._params.database
221214
if database:
222-
span.set_attribute(SPANDATA.DB_NAME, database)
215+
data[SPANDATA.DB_NAME] = database
223216

224-
user = conn._params.user
225217
if user:
226-
span.set_attribute(SPANDATA.DB_USER, user)
218+
data[SPANDATA.DB_USER] = user
219+
220+
return data
221+
222+
223+
def _set_on_span(span: Span, data: dict[str, Any]):
224+
for key, value in data.items():
225+
span.set_attribute(key, value)

‎tests/integrations/asyncpg/test_asyncpg.py

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,6 @@ async def _clean_pg():
7474
async def test_connect(sentry_init, capture_events) -> None:
7575
sentry_init(
7676
integrations=[AsyncPGIntegration()],
77-
traces_sample_rate=1.0,
7877
_experiments={"record_sql_params": True},
7978
)
8079
events = capture_events()
@@ -97,7 +96,6 @@ async def test_connect(sentry_init, capture_events) -> None:
9796
async def test_execute(sentry_init, capture_events) -> None:
9897
sentry_init(
9998
integrations=[AsyncPGIntegration()],
100-
traces_sample_rate=1.0,
10199
_experiments={"record_sql_params": True},
102100
)
103101
events = capture_events()
@@ -163,7 +161,6 @@ async def test_execute(sentry_init, capture_events) -> None:
163161
async def test_execute_many(sentry_init, capture_events) -> None:
164162
sentry_init(
165163
integrations=[AsyncPGIntegration()],
166-
traces_sample_rate=1.0,
167164
_experiments={"record_sql_params": True},
168165
)
169166
events = capture_events()
@@ -202,7 +199,6 @@ async def test_execute_many(sentry_init, capture_events) -> None:
202199
async def test_record_params(sentry_init, capture_events) -> None:
203200
sentry_init(
204201
integrations=[AsyncPGIntegration(record_params=True)],
205-
traces_sample_rate=1.0,
206202
_experiments={"record_sql_params": True},
207203
)
208204
events = capture_events()
@@ -243,7 +239,6 @@ async def test_record_params(sentry_init, capture_events) -> None:
243239
async def test_cursor(sentry_init, capture_events) -> None:
244240
sentry_init(
245241
integrations=[AsyncPGIntegration()],
246-
traces_sample_rate=1.0,
247242
_experiments={"record_sql_params": True},
248243
)
249244
events = capture_events()
@@ -308,7 +303,6 @@ async def test_cursor(sentry_init, capture_events) -> None:
308303
async def test_cursor_manual(sentry_init, capture_events) -> None:
309304
sentry_init(
310305
integrations=[AsyncPGIntegration()],
311-
traces_sample_rate=1.0,
312306
_experiments={"record_sql_params": True},
313307
)
314308
events = capture_events()
@@ -375,7 +369,6 @@ async def test_cursor_manual(sentry_init, capture_events) -> None:
375369
async def test_prepared_stmt(sentry_init, capture_events) -> None:
376370
sentry_init(
377371
integrations=[AsyncPGIntegration()],
378-
traces_sample_rate=1.0,
379372
_experiments={"record_sql_params": True},
380373
)
381374
events = capture_events()
@@ -425,7 +418,6 @@ async def test_prepared_stmt(sentry_init, capture_events) -> None:
425418
async def test_connection_pool(sentry_init, capture_events) -> None:
426419
sentry_init(
427420
integrations=[AsyncPGIntegration()],
428-
traces_sample_rate=1.0,
429421
_experiments={"record_sql_params": True},
430422
)
431423
events = capture_events()
@@ -497,7 +489,7 @@ async def test_connection_pool(sentry_init, capture_events) -> None:
497489
async def test_query_source_disabled(sentry_init, capture_events):
498490
sentry_options = {
499491
"integrations": [AsyncPGIntegration()],
500-
"enable_tracing": True,
492+
"traces_sample_rate": 1.0,
501493
"enable_db_query_source": False,
502494
"db_query_source_threshold_ms": 0,
503495
}
@@ -535,7 +527,7 @@ async def test_query_source_enabled(
535527
):
536528
sentry_options = {
537529
"integrations": [AsyncPGIntegration()],
538-
"enable_tracing": True,
530+
"traces_sample_rate": 1.0,
539531
"db_query_source_threshold_ms": 0,
540532
}
541533
if enable_db_query_source is not None:
@@ -571,7 +563,7 @@ async def test_query_source_enabled(
571563
async def test_query_source(sentry_init, capture_events):
572564
sentry_init(
573565
integrations=[AsyncPGIntegration()],
574-
enable_tracing=True,
566+
traces_sample_rate=1.0,
575567
enable_db_query_source=True,
576568
db_query_source_threshold_ms=0,
577569
)
@@ -621,7 +613,7 @@ async def test_query_source_with_module_in_search_path(sentry_init, capture_even
621613
"""
622614
sentry_init(
623615
integrations=[AsyncPGIntegration()],
624-
enable_tracing=True,
616+
traces_sample_rate=1.0,
625617
enable_db_query_source=True,
626618
db_query_source_threshold_ms=0,
627619
)
@@ -667,7 +659,7 @@ async def test_query_source_with_module_in_search_path(sentry_init, capture_even
667659
async def test_no_query_source_if_duration_too_short(sentry_init, capture_events):
668660
sentry_init(
669661
integrations=[AsyncPGIntegration()],
670-
enable_tracing=True,
662+
traces_sample_rate=1.0,
671663
enable_db_query_source=True,
672664
db_query_source_threshold_ms=100,
673665
)
@@ -714,7 +706,7 @@ def fake_record_sql_queries(*args, **kwargs):
714706
async def test_query_source_if_duration_over_threshold(sentry_init, capture_events):
715707
sentry_init(
716708
integrations=[AsyncPGIntegration()],
717-
enable_tracing=True,
709+
traces_sample_rate=1.0,
718710
enable_db_query_source=True,
719711
db_query_source_threshold_ms=100,
720712
)

0 commit comments

Comments
 (0)