diff --git a/contrib/opencensus-ext-grpc/CHANGELOG.md b/contrib/opencensus-ext-grpc/CHANGELOG.md index 6ff58e2d5..d06dd0866 100644 --- a/contrib/opencensus-ext-grpc/CHANGELOG.md +++ b/contrib/opencensus-ext-grpc/CHANGELOG.md @@ -2,6 +2,9 @@ ## Unreleased +- Extract byte size from proto-plus messages + ([#976](https://github.com/census-instrumentation/opencensus-python/pull/976)) + ## 0.7.1 Released 2019-08-05 diff --git a/contrib/opencensus-ext-grpc/opencensus/ext/grpc/utils.py b/contrib/opencensus-ext-grpc/opencensus/ext/grpc/utils.py index 0cc213bb3..7a36c3d84 100644 --- a/contrib/opencensus-ext-grpc/opencensus/ext/grpc/utils.py +++ b/contrib/opencensus-ext-grpc/opencensus/ext/grpc/utils.py @@ -6,6 +6,17 @@ from opencensus.trace import execution_context, time_event +def extract_byte_size(proto_message): + """Gets the byte size from a google.protobuf or proto-plus message""" + if hasattr(proto_message, "ByteSize"): + # google.protobuf message + return proto_message.ByteSize() + if hasattr(type(proto_message), "pb"): + # proto-plus message + return type(proto_message).pb(proto_message).ByteSize() + return None + + def add_message_event(proto_message, span, message_event_type, message_id=1): """Adds a MessageEvent to the span based off of the given protobuf message @@ -15,7 +26,7 @@ def add_message_event(proto_message, span, message_event_type, message_id=1): datetime.utcnow(), message_id, type=message_event_type, - uncompressed_size_bytes=proto_message.ByteSize() + uncompressed_size_bytes=extract_byte_size(proto_message), ) ) diff --git a/contrib/opencensus-ext-grpc/tests/test_server_interceptor.py b/contrib/opencensus-ext-grpc/tests/test_server_interceptor.py index 680de1002..850c036fa 100644 --- a/contrib/opencensus-ext-grpc/tests/test_server_interceptor.py +++ b/contrib/opencensus-ext-grpc/tests/test_server_interceptor.py @@ -167,6 +167,20 @@ def test__wrap_rpc_behavior_none(self): new_handler = server_interceptor._wrap_rpc_behavior(None, lambda: None) self.assertEqual(new_handler, None) + def test_extract_byte_size(self): + # should work with a google.protobuf message + google_protobuf_mock = mock.Mock() + google_protobuf_mock.ByteSize.return_value = 5 + self.assertEqual(grpc_utils.extract_byte_size(google_protobuf_mock), 5) + + # should work with a proto-plus style message + protoplus_protobuf_mock = mock.Mock(spec=[]) + type(protoplus_protobuf_mock).pb = mock.Mock() + type(protoplus_protobuf_mock).pb.return_value.ByteSize.return_value = 5 + self.assertEqual( + grpc_utils.extract_byte_size(protoplus_protobuf_mock), 5 + ) + class MockTracer(object): def __init__(self, *args, **kwargs):