-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Expand file tree
/
Copy pathgrpc_error_interceptor.py
More file actions
48 lines (42 loc) · 1.85 KB
/
grpc_error_interceptor.py
File metadata and controls
48 lines (42 loc) · 1.85 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
import grpc
from feast.errors import FeastError
def exception_wrapper(behavior, request, context):
try:
return behavior(request, context)
except grpc.RpcError as e:
context.abort(e.code(), e.details())
except FeastError as e:
context.abort(
e.grpc_status_code(),
e.to_error_detail(),
)
class ErrorInterceptor(grpc.ServerInterceptor):
def intercept_service(self, continuation, handler_call_details):
handler = continuation(handler_call_details)
if handler is None:
return None
if handler.unary_unary:
return grpc.unary_unary_rpc_method_handler(
lambda req, ctx: exception_wrapper(handler.unary_unary, req, ctx),
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)
elif handler.unary_stream:
return grpc.unary_stream_rpc_method_handler(
lambda req, ctx: exception_wrapper(handler.unary_stream, req, ctx),
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)
elif handler.stream_unary:
return grpc.stream_unary_rpc_method_handler(
lambda req, ctx: exception_wrapper(handler.stream_unary, req, ctx),
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)
elif handler.stream_stream:
return grpc.stream_stream_rpc_method_handler(
lambda req, ctx: exception_wrapper(handler.stream_stream, req, ctx),
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)
return handler