feat(asyncpg): Add span streaming support#6215
Conversation
Migrates the asyncpg integration to use `record_sql_queries_supporting_streaming` and `StreamedSpan`, adding support for the `stream` trace lifecycle. Tests are parameterized to cover both static and streaming modes. Fixes PY-2305 Fixes #6003
|
bugbot run |
|
@sentry review |
Codecov Results 📊✅ 13 passed | Total: 13 | Pass Rate: 100% | Execution Time: 7.67s 📊 Comparison with Base Branch
✨ No test changes detected All tests are passing successfully. ❌ Patch coverage is 8.00%. Project has 15044 uncovered lines. Files with missing lines (1)
Coverage diff@@ Coverage Diff @@
## main #PR +/-##
==========================================
- Coverage 31.33% 31.30% -0.03%
==========================================
Files 190 190 —
Lines 21877 21898 +21
Branches 7360 7368 +8
==========================================
+ Hits 6853 6854 +1
- Misses 15024 15044 +20
- Partials 579 579 —Generated by Codecov Action |
![sentry[bot]](https://avatars.githubusercontent.com/in/12637?s=60&v=4){kind=small}
| integration = client.get_integration(AsyncPGIntegration) | ||
| if integration is not None and not integration._record_params: | ||
| params_list = None | ||
|
|
||
| param_style = "pyformat" if params_list else None | ||
|
|
||
| query = _normalize_query(query) | ||
| with record_sql_queries( | ||
| with record_sql_queries_supporting_streaming( | ||
| cursor=cursor, | ||
| query=query, | ||
| params_list=params_list, |
There was a problem hiding this comment.
Bug: The _wrap_connection_method and _wrap_cursor_creation wrappers are missing calls to add_query_source, preventing query source information from being captured for most asyncpg operations.
Severity: MEDIUM
Suggested Fix
Call add_query_source(span) within _wrap_connection_method and _wrap_cursor_creation before the span's context manager exits, similar to the implementation in _wrap_execute. This will ensure query source data is attached to the span.
Prompt for AI Agent
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.
Location: sentry_sdk/integrations/asyncpg.py#L112-L127
Potential issue: The `add_query_source` function is not called within the
`_wrap_connection_method` and `_wrap_cursor_creation` wrappers. As a result, when
`enable_db_query_source` is enabled, query source information like file path and line
number will be missing for most `asyncpg` operations, including `fetch`, `fetchrow`,
`execute`, `executemany`, and `prepare`. This is an incomplete implementation for the
new streaming mode, where `add_query_source` must be called before the span's context
manager exits to capture this data.
Did we get this right? 👍 / 👎 to inform future reviews.
| add_query_source(span) | ||
|
|
||
| with capture_internal_exceptions(): | ||
| add_query_source(span) | ||
| if not isinstance(span, StreamedSpan): |
There was a problem hiding this comment.
Bug: In _wrap_connect_addr, span attributes are set after the sampling decision is made, which can cause traces_sampler to make incorrect sampling choices for connection spans.
Severity: MEDIUM
Suggested Fix
Pass the span attributes directly to sentry_sdk.traces.start_span using the attributes keyword argument. This ensures the sampling context is fully populated before the traces_sampler is invoked, mirroring the pattern used elsewhere in the tracing utility functions.
Prompt for AI Agent
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.
Location: sentry_sdk/integrations/asyncpg.py#L94-L96
Potential issue: In `_wrap_connect_addr`, the span for a new database connection is
started before its attributes, such as `sentry.op` and `sentry.origin`, are set. The
sampling decision is made immediately when `sentry_sdk.traces.start_span` is called. If
a `traces_sampler` function is configured to make decisions based on these attributes
(e.g., to sample all DB operations), it will not have the necessary context. This can
lead to incorrect sampling decisions, where connection spans are dropped when they
should be kept.
Did we get this right? 👍 / 👎 to inform future reviews.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit fd3ee4e. Configure here.
| ) | ||
| res = await f(*args, **kwargs) | ||
| with sentry_sdk.traces.start_span(name="connect") as span: | ||
| span.set_attributes(span_attributes) |
There was a problem hiding this comment.
Connect attributes set too late
Medium Severity
Streamed asyncpg connect spans set sentry.op, sentry.origin, and DB attributes after span creation. This prevents traces_sampler and ignore_spans from correctly evaluating these spans based on their attributes.
Reviewed by Cursor Bugbot for commit fd3ee4e. Configure here.
|
|
||
| query = _normalize_query(query) | ||
| with record_sql_queries( | ||
| with record_sql_queries_supporting_streaming( |
There was a problem hiding this comment.
Streamed query metadata is dropped
Medium Severity
Streamed asyncpg query spans lose SQL metadata such as db.params, db.paramstyle, and db.executemany; record_sql_queries_supporting_streaming only attaches that data to breadcrumbs and legacy spans.
Reviewed by Cursor Bugbot for commit fd3ee4e. Configure here.
1 participant
Migrates the asyncpg integration to support the
streamtrace lifecycle.The integration now uses
record_sql_queries_supporting_streamingandStreamedSpanso that DB spans are emitted as soon as they complete rather than being buffered until transaction end.Tests are parameterized over both
span_streaming=Trueandspan_streaming=Falseto verify both paths.Fixes PY-2305, fixes #6003