fix: optimized protobuf access for performance (#155)

craiglabenz · web-flow · commit 5b67daa3b2da · 2021-04-20T16:57:02.000-04:00
More efficiently uses proto-plus wrappers, as well as inner protobuf attribute access, to greatly reduce the performance costs seen in version 2.0.0 (which stemmed from the introduction of proto-plus). The size of the performance improvement scales with the number of attributes on each Entity, but in general, speeds once again closely approximate those from 1.15. Fixes #145 Fixes #150
diff --git a/google/cloud/datastore/helpers.py b/google/cloud/datastore/helpers.py
@@ -49,16 +49,21 @@ def _get_meaning(value_pb, is_list=False):
     """
     meaning = None
     if is_list:
+
+        values = (
+            value_pb._pb.array_value.values
+            if hasattr(value_pb, "_pb")
+            else value_pb.array_value.values
+        )
+
         # An empty list will have no values, hence no shared meaning
         # set among them.
-        if len(value_pb.array_value.values) == 0:
+        if len(values) == 0:
             return None
 
         # We check among all the meanings, some of which may be None,
         # the rest which may be enum/int values.
-        all_meanings = [
-            _get_meaning(sub_value_pb) for sub_value_pb in value_pb.array_value.values
-        ]
+        all_meanings = [_get_meaning(sub_value_pb) for sub_value_pb in values]
         unique_meanings = set(all_meanings)
         if len(unique_meanings) == 1:
             # If there is a unique meaning, we preserve it.
@@ -119,11 +124,8 @@ def entity_from_protobuf(pb):
     :rtype: :class:`google.cloud.datastore.entity.Entity`
     :returns: The entity derived from the protobuf.
     """
-
-    if not getattr(pb, "_pb", False):
-        # Coerce raw pb type into proto-plus pythonic type.
-        proto_pb = entity_pb2.Entity(pb)
-        pb = pb
+    if not isinstance(pb, entity_pb2.Entity):
+        proto_pb = entity_pb2.Entity.wrap(pb)
     else:
         proto_pb = pb
         pb = pb._pb
@@ -152,7 +154,7 @@ def entity_from_protobuf(pb):
         if is_list and len(value) > 0:
             exclude_values = set(
                 value_pb.exclude_from_indexes
-                for value_pb in value_pb.array_value.values
+                for value_pb in value_pb._pb.array_value.values
             )
             if len(exclude_values) != 1:
                 raise ValueError(
@@ -402,33 +404,36 @@ def _get_value_from_value_pb(value):
     """
     if not getattr(value, "_pb", False):
         # Coerce raw pb type into proto-plus pythonic type.
-        value = entity_pb2.Value(value)
+        value = entity_pb2.Value.wrap(value)
 
     value_type = value._pb.WhichOneof("value_type")
 
     if value_type == "timestamp_value":
+        # Do not access `._pb` here, as that returns a Timestamp proto,
+        # but this should return a Pythonic `DatetimeWithNanoseconds` value,
+        # which is found at `value.timestamp_value`
         result = value.timestamp_value
 
     elif value_type == "key_value":
-        result = key_from_protobuf(value.key_value)
+        result = key_from_protobuf(value._pb.key_value)
 
     elif value_type == "boolean_value":
-        result = value.boolean_value
+        result = value._pb.boolean_value
 
     elif value_type == "double_value":
-        result = value.double_value
+        result = value._pb.double_value
 
     elif value_type == "integer_value":
-        result = value.integer_value
+        result = value._pb.integer_value
 
     elif value_type == "string_value":
-        result = value.string_value
+        result = value._pb.string_value
 
     elif value_type == "blob_value":
-        result = value.blob_value
+        result = value._pb.blob_value
 
     elif value_type == "entity_value":
-        result = entity_from_protobuf(value.entity_value)
+        result = entity_from_protobuf(value._pb.entity_value)
 
     elif value_type == "array_value":
         result = [
@@ -437,7 +442,7 @@ def _get_value_from_value_pb(value):
 
     elif value_type == "geo_point_value":
         result = GeoPoint(
-            value.geo_point_value.latitude, value.geo_point_value.longitude,
+            value._pb.geo_point_value.latitude, value._pb.geo_point_value.longitude,
         )
 
     elif value_type == "null_value":
diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py
@@ -140,6 +140,15 @@ def test_entity_no_key(self):
         self.assertIsNone(entity.key)
         self.assertEqual(dict(entity), {})
 
+    def test_pb2_entity_no_key(self):
+        from google.cloud.datastore_v1.types import entity as entity_pb2
+
+        entity_pb = entity_pb2.Entity()
+        entity = self._call_fut(entity_pb)
+
+        self.assertIsNone(entity.key)
+        self.assertEqual(dict(entity), {})
+
     def test_entity_with_meaning(self):
         from google.cloud.datastore_v1.types import entity as entity_pb2
         from google.cloud.datastore.helpers import _new_value_pb
