|
1 | | -import sys |
2 | 1 | import functools |
| 2 | +import sys |
3 | 3 |
|
4 | 4 | import sentry_sdk |
5 | 5 | from sentry_sdk.consts import OP |
6 | | -from sentry_sdk.integrations import Integration, DidNotEnable |
| 6 | +from sentry_sdk.integrations import DidNotEnable, Integration |
7 | 7 | from sentry_sdk.integrations._wsgi_common import nullcontext |
| 8 | +from sentry_sdk.traces import StreamedSpan |
| 9 | +from sentry_sdk.tracing import Span |
| 10 | +from sentry_sdk.tracing_utils import has_span_streaming_enabled |
| 11 | +from sentry_sdk.transport import AsyncHttpTransport |
8 | 12 | from sentry_sdk.utils import ( |
9 | 13 | event_from_exception, |
| 14 | + is_internal_task, |
10 | 15 | logger, |
11 | 16 | reraise, |
12 | | - is_internal_task, |
13 | 17 | ) |
14 | | -from sentry_sdk.transport import AsyncHttpTransport |
15 | 18 |
|
16 | 19 | try: |
17 | 20 | import asyncio |
18 | 21 | from asyncio.tasks import Task |
19 | 22 | except ImportError: |
20 | 23 | raise DidNotEnable("asyncio not available") |
21 | 24 |
|
22 | | -from typing import TYPE_CHECKING |
| 25 | +from typing import TYPE_CHECKING, Optional, Union |
23 | 26 |
|
24 | 27 | if TYPE_CHECKING: |
25 | | - from typing import Any, Callable, TypeVar |
26 | 28 | from collections.abc import Coroutine |
| 29 | + from typing import Any, Callable, TypeVar |
27 | 30 |
|
28 | 31 | from sentry_sdk._types import ExcInfo |
29 | 32 |
|
@@ -142,22 +145,31 @@ def _sentry_task_factory( |
142 | 145 | @_wrap_coroutine(coro) |
143 | 146 | async def _task_with_sentry_span_creation() -> "Any": |
144 | 147 | result = None |
145 | | - |
146 | | - integration = sentry_sdk.get_client().get_integration( |
147 | | - AsyncioIntegration |
148 | | - ) |
| 148 | + client = sentry_sdk.get_client() |
| 149 | + integration = client.get_integration(AsyncioIntegration) |
149 | 150 | task_spans = integration.task_spans if integration else False |
150 | 151 |
|
| 152 | + span_ctx: "Optional[Union[StreamedSpan, Span]]" = None |
| 153 | + is_span_streaming_enabled = has_span_streaming_enabled(client.options) |
| 154 | + |
151 | 155 | with sentry_sdk.isolation_scope(): |
152 | | - with ( |
153 | | - sentry_sdk.start_span( |
154 | | - op=OP.FUNCTION, |
155 | | - name=get_name(coro), |
156 | | - origin=AsyncioIntegration.origin, |
157 | | - ) |
158 | | - if task_spans |
159 | | - else nullcontext() |
160 | | - ): |
| 156 | + if task_spans: |
| 157 | + if is_span_streaming_enabled: |
| 158 | + span_ctx = sentry_sdk.traces.start_span( |
| 159 | + name=get_name(coro), |
| 160 | + attributes={ |
| 161 | + "sentry.op": OP.FUNCTION, |
| 162 | + "sentry.origin": AsyncioIntegration.origin, |
| 163 | + }, |
| 164 | + ) |
| 165 | + else: |
| 166 | + span_ctx = sentry_sdk.start_span( |
| 167 | + op=OP.FUNCTION, |
| 168 | + name=get_name(coro), |
| 169 | + origin=AsyncioIntegration.origin, |
| 170 | + ) |
| 171 | + |
| 172 | + with span_ctx if span_ctx else nullcontext(): |
161 | 173 | try: |
162 | 174 | result = await coro |
163 | 175 | except StopAsyncIteration as e: |
|
0 commit comments