mirror of https://github.com/milvus-io/milvus.git
38 lines
1.1 KiB
Python
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)
|