Extract byte size from proto-plus messages

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
 

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),
         )
     )
 

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):