Skip to content

Commit d1fc138

Browse files
ericapisaniericapisani
authored
feat(asyncpg): Add query source to execute, executemany, and cursor methods (#6242)
1 parent c416afe commit d1fc138

2 files changed

Lines changed: 295 additions & 55 deletions

File tree

‎sentry_sdk/integrations/asyncpg.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ def setup_once() -> None:
5454
asyncpg.Connection.prepare = _wrap_connection_method(asyncpg.Connection.prepare)
5555

5656
BaseCursor._bind_exec = _wrap_cursor_method(BaseCursor._bind_exec)
57-
BaseCursor._bind = _wrap_cursor_method(BaseCursor._bind)
5857
BaseCursor._exec = _wrap_cursor_method(BaseCursor._exec)
5958

6059
asyncpg.connect_utils._connect_addr = _wrap_connect_addr(
@@ -146,8 +145,17 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
146145
params_list = args[2] if len(args) > 2 else None
147146
with _record(None, query, params_list, executemany=executemany) as span:
148147
_set_db_data(span, args[0])
148+
149149
res = await f(*args, **kwargs)
150150

151+
if isinstance(span, StreamedSpan):
152+
with capture_internal_exceptions():
153+
add_query_source(span)
154+
155+
if not isinstance(span, StreamedSpan):
156+
with capture_internal_exceptions():
157+
add_query_source(span)
158+
151159
return res
152160

153161
return _inner
@@ -174,6 +182,14 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
174182
_set_db_data(span, cursor._connection)
175183
res = await f(*args, **kwargs)
176184

185+
if isinstance(span, StreamedSpan):
186+
with capture_internal_exceptions():
187+
add_query_source(span)
188+
189+
if not isinstance(span, StreamedSpan):
190+
with capture_internal_exceptions():
191+
add_query_source(span)
192+
177193
return res
178194

179195
return _inner

0 commit comments

Comments
 (0)