Update Protobuf and GCP dependencies in Beam Python SDK

sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py

@@ -477,7 +477,8 @@ def process(self, element):
       client_element = self.element_to_client_batch_item(element)
       self._batch_elements.append(client_element)
       self.add_to_batch(client_element)
-      self._batch_bytes_size += self._batch.mutations[-1].ByteSize()
+      self._batch_bytes_size += util.extract_byte_size(
+          self._batch.mutations[-1])
 
       if (len(self._batch.mutations) >= self._target_batch_size or
           self._batch_bytes_size > util.WRITE_BATCH_MAX_BYTES_SIZE):

sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py

@@ -66,9 +66,9 @@ def __init__(self, entity=None, key=None):
 
   def ByteSize(self):
     if self.entity is not None:
-      return helpers.entity_to_protobuf(self.entity).ByteSize()
+      return util.extract_byte_size(helpers.entity_to_protobuf(self.entity))
     else:
-      return self.key.to_protobuf().ByteSize()
+      return util.extract_byte_size(self.key.to_protobuf())
 
 
 class FakeBatch(object):

sdks/python/apache_beam/io/gcp/datastore/v1new/util.py

@@ -137,3 +137,22 @@ def report_latency(self, now, latency_ms, num_mutations):
       num_mutations: int, number of mutations contained in the RPC.
     """
     self._commit_time_per_entity_ms.add(now, latency_ms / num_mutations)
+
+
+def extract_byte_size(proto_message):
+  """
+    Gets the byte size from a google.protobuf or proto-plus message
+
+    google-cloud-datastore moved from using protobuf to using
+    proto-plus messages.
+    protobuf object has attribute ByteSize() but proto.Message() objects
+    don't. Workaround:
+    https://github.com/googleapis/proto-plus-python/issues/163
+    """
+  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