Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: optimized protobuf access for performance #155

Merged
merged 4 commits into from Apr 20, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
43 changes: 24 additions & 19 deletions google/cloud/datastore/helpers.py
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
craiglabenz marked this conversation as resolved.
Show resolved Hide resolved
else:
proto_pb = pb
pb = pb._pb
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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 = [
Expand All @@ -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":
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/test_helpers.py
Expand Up @@ -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
Expand Down