milvus/shards/mishards/grpc_utils/__init__.py

38 lines
1.1 KiB
Python

from grpc_opentracing import SpanDecorator
from milvus.grpc_gen import status_pb2
class GrpcSpanDecorator(SpanDecorator):
def __call__(self, span, rpc_info):
status = None
if not rpc_info.response:
return
if isinstance(rpc_info.response, status_pb2.Status):
status = rpc_info.response
else:
try:
status = rpc_info.response.status
except Exception as e:
status = status_pb2.Status(error_code=status_pb2.UNEXPECTED_ERROR,
reason='Should not happen')
if status.error_code == 0:
return
error_log = {'event': 'error',
'request': rpc_info.request,
'response': rpc_info.response
}
span.set_tag('error', True)
span.log_kv(error_log)
def mark_grpc_method(func):
setattr(func, 'grpc_method', True)
return func
def is_grpc_method(func):
if not func:
return False
return getattr(func, 'grpc_method', False)