Skip to content

Commit 227c52a

Browse files
ericapisaniericapisani
authored
Merge branch 'py-2305-migrate-asyncpg' into py-2305-add-query-source-to-methods-that-were-missing-it
2 parents 060924e + 6bb9b6c commit 227c52a

17 files changed

Lines changed: 1407 additions & 499 deletions

File tree

‎sentry_sdk/consts.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -899,6 +899,12 @@ class SPANDATA:
899899
Example: "com.example.ExampleService/exampleMethod"
900900
"""
901901

902+
RPC_RESPONSE_STATUS_CODE = "rpc.response.status_code"
903+
"""
904+
Status code of the RPC returned by the RPC server or generated by the client.
905+
Example: "DEADLINE_EXCEEDED"
906+
"""
907+
902908
SERVER_ADDRESS = "server.address"
903909
"""
904910
Name of the database host.

‎sentry_sdk/integrations/aiohttp.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
request_body_within_bounds,
1717
)
1818
from sentry_sdk.integrations.logging import ignore_logger
19-
from sentry_sdk.scope import should_send_default_pii
19+
from sentry_sdk.scope import should_send_default_pii, Scope
2020
from sentry_sdk.sessions import track_session
2121
from sentry_sdk.traces import (
2222
SOURCE_FOR_STYLE as SEGMENT_SOURCE_FOR_STYLE,
@@ -139,9 +139,7 @@ async def sentry_app_handle(
139139
span_ctx: "ContextManager[Union[Span, StreamedSpan]]"
140140
if is_span_streaming_enabled:
141141
sentry_sdk.traces.continue_trace(headers)
142-
sentry_sdk.get_current_scope().set_custom_sampling_context(
143-
{"aiohttp_request": request}
144-
)
142+
Scope.set_custom_sampling_context({"aiohttp_request": request})
145143

146144
header_attributes: "dict[str, Any]" = {}
147145
for header, header_value in _filter_headers(

‎sentry_sdk/integrations/asgi.py‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
capture_internal_exceptions,
4747
qualname_from_function,
4848
)
49+
from sentry_sdk.scope import Scope
4950

5051
from typing import TYPE_CHECKING
5152

@@ -257,9 +258,7 @@ async def _run_app(
257258
):
258259
sentry_sdk.traces.continue_trace(_get_headers(scope))
259260

260-
sentry_scope.set_custom_sampling_context(
261-
{"asgi_scope": scope}
262-
)
261+
Scope.set_custom_sampling_context({"asgi_scope": scope})
263262

264263
attributes["sentry.op"] = f"{ty}.server"
265264
segment = sentry_sdk.traces.start_span(
@@ -270,9 +269,7 @@ async def _run_app(
270269
else:
271270
sentry_sdk.traces.new_trace()
272271

273-
sentry_scope.set_custom_sampling_context(
274-
{"asgi_scope": scope}
275-
)
272+
Scope.set_custom_sampling_context({"asgi_scope": scope})
276273

277274
attributes["sentry.op"] = OP.HTTP_SERVER
278275
segment = sentry_sdk.traces.start_span(

‎sentry_sdk/integrations/celery/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
)
1515
from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
1616
from sentry_sdk.integrations.logging import ignore_logger
17-
from sentry_sdk.scope import should_send_default_pii
17+
from sentry_sdk.scope import should_send_default_pii, Scope
1818
from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
1919
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource
2020
from sentry_sdk.tracing_utils import Baggage, has_span_streaming_enabled
@@ -361,7 +361,7 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any":
361361
headers = args[3].get("headers") or {}
362362
if span_streaming:
363363
sentry_sdk.traces.continue_trace(headers)
364-
scope.set_custom_sampling_context(custom_sampling_context)
364+
Scope.set_custom_sampling_context(custom_sampling_context)
365365
span = sentry_sdk.traces.start_span(
366366
name=task_name,
367367
parent_span=None, # make this a segment

‎sentry_sdk/integrations/grpc/aio/client.py‎

Lines changed: 82 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from typing import Callable, Union, AsyncIterable, Any
22

33
import sentry_sdk
4-
from sentry_sdk.consts import OP
4+
from sentry_sdk.consts import OP, SPANDATA
55
from sentry_sdk.integrations import DidNotEnable
66
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
7+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
78

89
try:
910
from grpc.aio import (
@@ -49,23 +50,47 @@ async def intercept_unary_unary(
4950
) -> "Union[UnaryUnaryCall, Message]":
5051
method = client_call_details.method
5152

52-
with sentry_sdk.start_span(
53-
op=OP.GRPC_CLIENT,
54-
name="unary unary call to %s" % method.decode(),
55-
origin=SPAN_ORIGIN,
56-
) as span:
57-
span.set_data("type", "unary unary")
58-
span.set_data("method", method)
59-
60-
client_call_details = self._update_client_call_details_metadata_from_scope(
61-
client_call_details
62-
)
63-
64-
response = await continuation(client_call_details, request)
65-
status_code = await response.code()
66-
span.set_data("code", status_code.name)
67-
68-
return response
53+
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
54+
if span_streaming:
55+
with sentry_sdk.traces.start_span(
56+
name="unary unary call to %s" % method.decode(),
57+
attributes={
58+
"sentry.op": OP.GRPC_CLIENT,
59+
"sentry.origin": SPAN_ORIGIN,
60+
SPANDATA.RPC_METHOD: method.decode(),
61+
},
62+
) as span:
63+
client_call_details = (
64+
self._update_client_call_details_metadata_from_scope(
65+
client_call_details
66+
)
67+
)
68+
69+
response = await continuation(client_call_details, request)
70+
status_code = await response.code()
71+
span.set_attribute(SPANDATA.RPC_RESPONSE_STATUS_CODE, status_code.name)
72+
73+
return response
74+
else:
75+
with sentry_sdk.start_span(
76+
op=OP.GRPC_CLIENT,
77+
name="unary unary call to %s" % method.decode(),
78+
origin=SPAN_ORIGIN,
79+
) as span:
80+
span.set_data("type", "unary unary")
81+
span.set_data("method", method)
82+
83+
client_call_details = (
84+
self._update_client_call_details_metadata_from_scope(
85+
client_call_details
86+
)
87+
)
88+
89+
response = await continuation(client_call_details, request)
90+
status_code = await response.code()
91+
span.set_data("code", status_code.name)
92+
93+
return response
6994

7095

7196
class SentryUnaryStreamClientInterceptor(
@@ -80,20 +105,42 @@ async def intercept_unary_stream(
80105
) -> "Union[AsyncIterable[Any], UnaryStreamCall]":
81106
method = client_call_details.method
82107

83-
with sentry_sdk.start_span(
84-
op=OP.GRPC_CLIENT,
85-
name="unary stream call to %s" % method.decode(),
86-
origin=SPAN_ORIGIN,
87-
) as span:
88-
span.set_data("type", "unary stream")
89-
span.set_data("method", method)
90-
91-
client_call_details = self._update_client_call_details_metadata_from_scope(
92-
client_call_details
93-
)
94-
95-
response = await continuation(client_call_details, request)
96-
# status_code = await response.code()
97-
# span.set_data("code", status_code)
98-
99-
return response
108+
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
109+
if span_streaming:
110+
with sentry_sdk.traces.start_span(
111+
name="unary stream call to %s" % method.decode(),
112+
attributes={
113+
"sentry.op": OP.GRPC_CLIENT,
114+
"sentry.origin": SPAN_ORIGIN,
115+
SPANDATA.RPC_METHOD: method.decode(),
116+
},
117+
) as span:
118+
client_call_details = (
119+
self._update_client_call_details_metadata_from_scope(
120+
client_call_details
121+
)
122+
)
123+
124+
response = await continuation(client_call_details, request)
125+
126+
return response
127+
else:
128+
with sentry_sdk.start_span(
129+
op=OP.GRPC_CLIENT,
130+
name="unary stream call to %s" % method.decode(),
131+
origin=SPAN_ORIGIN,
132+
) as span:
133+
span.set_data("type", "unary stream")
134+
span.set_data("method", method)
135+
136+
client_call_details = (
137+
self._update_client_call_details_metadata_from_scope(
138+
client_call_details
139+
)
140+
)
141+
142+
response = await continuation(client_call_details, request)
143+
# status_code = await response.code()
144+
# span.set_data("code", status_code)
145+
146+
return response

‎sentry_sdk/integrations/grpc/aio/server.py‎

Lines changed: 51 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
55
from sentry_sdk.tracing import TransactionSource
66
from sentry_sdk.utils import event_from_exception
7+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
78

89
from typing import TYPE_CHECKING
910

@@ -52,27 +53,57 @@ async def wrapped(request: "Any", context: "ServicerContext") -> "Any":
5253
if not name:
5354
return await handler(request, context)
5455

55-
# What if the headers are empty?
56-
transaction = sentry_sdk.continue_trace(
57-
dict(context.invocation_metadata()),
58-
op=OP.GRPC_SERVER,
59-
name=name,
60-
source=TransactionSource.CUSTOM,
61-
origin=SPAN_ORIGIN,
56+
span_streaming = has_span_streaming_enabled(
57+
sentry_sdk.get_client().options
6258
)
63-
64-
with sentry_sdk.start_transaction(transaction=transaction):
65-
try:
66-
return await handler.unary_unary(request, context)
67-
except AbortError:
68-
raise
69-
except Exception as exc:
70-
event, hint = event_from_exception(
71-
exc,
72-
mechanism={"type": "grpc", "handled": False},
73-
)
74-
sentry_sdk.capture_event(event, hint=hint)
75-
raise
59+
if span_streaming:
60+
# What if the headers are empty?
61+
sentry_sdk.traces.continue_trace(
62+
dict(context.invocation_metadata())
63+
)
64+
65+
with sentry_sdk.traces.start_span(
66+
name=name,
67+
attributes={
68+
"sentry.op": OP.GRPC_SERVER,
69+
"sentry.span.source": TransactionSource.CUSTOM.value,
70+
"sentry.origin": SPAN_ORIGIN,
71+
},
72+
parent_span=None,
73+
):
74+
try:
75+
return await handler.unary_unary(request, context)
76+
except AbortError:
77+
raise
78+
except Exception as exc:
79+
event, hint = event_from_exception(
80+
exc,
81+
mechanism={"type": "grpc", "handled": False},
82+
)
83+
sentry_sdk.capture_event(event, hint=hint)
84+
raise
85+
else:
86+
# What if the headers are empty?
87+
transaction = sentry_sdk.continue_trace(
88+
dict(context.invocation_metadata()),
89+
op=OP.GRPC_SERVER,
90+
name=name,
91+
source=TransactionSource.CUSTOM,
92+
origin=SPAN_ORIGIN,
93+
)
94+
95+
with sentry_sdk.start_transaction(transaction=transaction):
96+
try:
97+
return await handler.unary_unary(request, context)
98+
except AbortError:
99+
raise
100+
except Exception as exc:
101+
event, hint = event_from_exception(
102+
exc,
103+
mechanism={"type": "grpc", "handled": False},
104+
)
105+
sentry_sdk.capture_event(event, hint=hint)
106+
raise
76107

77108
elif not handler.request_streaming and handler.response_streaming:
78109
handler_factory = grpc.unary_stream_rpc_method_handler

0 commit comments

Comments
 (0)