forked from getsentry/sentry-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathserver.py
More file actions
134 lines (113 loc) · 5.58 KB
/
Copy pathserver.py
File metadata and controls
134 lines (113 loc) · 5.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
from typing import TYPE_CHECKING
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
from sentry_sdk.traces import SegmentSource
from sentry_sdk.tracing import TransactionSource
from sentry_sdk.tracing_utils import has_span_streaming_enabled
from sentry_sdk.utils import event_from_exception
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from typing import Any, Optional
try:
import grpc
from grpc import HandlerCallDetails, RpcMethodHandler
from grpc.aio import AbortError, ServicerContext
except ImportError:
raise DidNotEnable("grpcio is not installed")
class ServerInterceptor(grpc.aio.ServerInterceptor): # type: ignore
def __init__(
self: "ServerInterceptor",
find_name: "Callable[[ServicerContext], str] | None" = None,
) -> None:
self._custom_find_name = find_name
super().__init__()
async def intercept_service(
self: "ServerInterceptor",
continuation: "Callable[[HandlerCallDetails], Awaitable[RpcMethodHandler]]",
handler_call_details: "HandlerCallDetails",
) -> "Optional[Awaitable[RpcMethodHandler]]":
handler = await continuation(handler_call_details)
if handler is None:
return None
method_name = handler_call_details.method
custom_find_name = self._custom_find_name
if not handler.request_streaming and not handler.response_streaming:
handler_factory = grpc.unary_unary_rpc_method_handler
async def wrapped(request: "Any", context: "ServicerContext") -> "Any":
with sentry_sdk.isolation_scope():
name = (
custom_find_name(context) if custom_find_name else method_name
)
if not name:
return await handler(request, context)
span_streaming = has_span_streaming_enabled(
sentry_sdk.get_client().options
)
if span_streaming:
# What if the headers are empty?
sentry_sdk.traces.continue_trace(
dict(context.invocation_metadata())
)
with sentry_sdk.traces.start_span(
name=name,
attributes={
"sentry.op": OP.GRPC_SERVER,
"sentry.span.source": SegmentSource.CUSTOM.value,
"sentry.origin": SPAN_ORIGIN,
},
parent_span=None,
):
try:
return await handler.unary_unary(request, context)
except AbortError:
raise
except Exception as exc:
event, hint = event_from_exception(
exc,
mechanism={"type": "grpc", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
raise
else:
# What if the headers are empty?
transaction = sentry_sdk.continue_trace(
dict(context.invocation_metadata()),
op=OP.GRPC_SERVER,
name=name,
source=TransactionSource.CUSTOM,
origin=SPAN_ORIGIN,
)
with sentry_sdk.start_transaction(transaction=transaction):
try:
return await handler.unary_unary(request, context)
except AbortError:
raise
except Exception as exc:
event, hint = event_from_exception(
exc,
mechanism={"type": "grpc", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
raise
elif not handler.request_streaming and handler.response_streaming:
handler_factory = grpc.unary_stream_rpc_method_handler
async def wrapped(request: "Any", context: "ServicerContext") -> "Any": # type: ignore
async for r in handler.unary_stream(request, context):
yield r
elif handler.request_streaming and not handler.response_streaming:
handler_factory = grpc.stream_unary_rpc_method_handler
async def wrapped(request: "Any", context: "ServicerContext") -> "Any":
response = handler.stream_unary(request, context)
return await response
elif handler.request_streaming and handler.response_streaming:
handler_factory = grpc.stream_stream_rpc_method_handler
async def wrapped(request: "Any", context: "ServicerContext") -> "Any": # type: ignore
async for r in handler.stream_stream(request, context):
yield r
return handler_factory(
wrapped,
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)