feat(dependencies)!: Upgrade to google-cloud-datastore 2.x (#841)

* feat(dependencies)!: Upgrade to google-cloud-datastore >= 2.7.2 This is the first version of google-cloud-datastore that does not support Python 3.6, which aligns it with our supported versions. Also upgrade some other minimum dependencies to align. BREAKING CHANGE: If you are using both the NDB and the base google-cloud-datastore classes in your codebase, you will have to update your use of those google-cloud-datastore classes. See https://github.com/googleapis/python-datastore/blob/main/UPGRADING.md for recommendations for upgrading to google-cloud-datastore 2.x. tweak deps * fix: Update module imports * fix: Fix enum namespaces * fix: Update datastore stub creation * fix: Update API capitalization/casing The new Datastore GRPC transport has different method naming conventions than the old stub did. * fix: Correct access to SerializeToString, CopyFrom, and MergeFromString * test: Fix tests

Author: rwhogg

google/cloud/ndb/_cache.py

@@ -738,4 +738,4 @@ def global_cache_key(key):
  Returns:
  bytes: The cache key.
  """
- return _PREFIX + key.to_protobuf().SerializeToString()
+ return _PREFIX + key.to_protobuf()._pb.SerializeToString()

google/cloud/ndb/_datastore_api.py

@@ -20,8 +20,8 @@
 
 from google.api_core import exceptions as core_exceptions
 from google.cloud.datastore import helpers
-from google.cloud.datastore_v1.proto import datastore_pb2
-from google.cloud.datastore_v1.proto import entity_pb2
+from google.cloud.datastore_v1.types import datastore as datastore_pb2
+from google.cloud.datastore_v1.types import entity as entity_pb2
 
 from google.cloud.ndb import context as context_module
 from google.cloud.ndb import _batch
@@ -33,9 +33,9 @@
 from google.cloud.ndb import tasklets
 from google.cloud.ndb import utils
 
-EVENTUAL = datastore_pb2.ReadOptions.EVENTUAL
+EVENTUAL = datastore_pb2.ReadOptions.ReadConsistency.EVENTUAL
 EVENTUAL_CONSISTENCY = EVENTUAL # Legacy NDB
-STRONG = datastore_pb2.ReadOptions.STRONG
+STRONG = datastore_pb2.ReadOptions.ReadConsistency.STRONG
 
 _DEFAULT_TIMEOUT = None
 _NOT_FOUND = object()
@@ -144,7 +144,7 @@ def lookup(key, options):
  if not key_locked:
  if result:
  entity_pb = entity_pb2.Entity()
- entity_pb.MergeFromString(result)
+ entity_pb._pb.MergeFromString(result)
 
  elif use_datastore:
  lock = yield _cache.global_lock_for_read(cache_key, result)
@@ -165,7 +165,7 @@ def lookup(key, options):
  if use_global_cache and not key_locked:
  if entity_pb is not _NOT_FOUND:
  expires = context._global_cache_timeout(key, options)
- serialized = entity_pb.SerializeToString()
+ serialized = entity_pb._pb.SerializeToString()
  yield _cache.global_compare_and_swap(
  cache_key, serialized, expires=expires
  )
@@ -211,7 +211,7 @@ def add(self, key):
  Returns:
  tasklets.Future: A future for the eventual result.
  """
- todo_key = key.to_protobuf().SerializeToString()
+ todo_key = key.to_protobuf()._pb.SerializeToString()
  future = tasklets.Future(info="Lookup({})".format(key))
  self.todo.setdefault(todo_key, []).append(future)
  return future
@@ -221,7 +221,7 @@ def idle_callback(self):
  keys = []
  for todo_key in self.todo.keys():
  key_pb = entity_pb2.Key()
- key_pb.ParseFromString(todo_key)
+ key_pb._pb.ParseFromString(todo_key)
  keys.append(key_pb)
 
  read_options = get_read_options(self.options)
@@ -264,20 +264,20 @@ def lookup_callback(self, rpc):
  if results.deferred:
  next_batch = _batch.get_batch(type(self), self.options)
  for key in results.deferred:
- todo_key = key.SerializeToString()
+ todo_key = key._pb.SerializeToString()
  next_batch.todo.setdefault(todo_key, []).extend(self.todo[todo_key])
 
  # For all missing keys, set result to _NOT_FOUND and let callers decide
  # how to handle
  for result in results.missing:
- todo_key = result.entity.key.SerializeToString()
+ todo_key = result.entity.key._pb.SerializeToString()
  for future in self.todo[todo_key]:
  future.set_result(_NOT_FOUND)
 
  # For all found entities, set the result on their corresponding futures
  for result in results.found:
  entity = result.entity
- todo_key = entity.key.SerializeToString()
+ todo_key = entity.key._pb.SerializeToString()
  for future in self.todo[todo_key]:
  future.set_result(entity)
 
@@ -306,7 +306,7 @@ def _datastore_lookup(keys, read_options, retries=None, timeout=None):
  read_options=read_options,
  )
 
- return make_call("Lookup", request, retries=retries, timeout=timeout)
+ return make_call("lookup", request, retries=retries, timeout=timeout)
 
 
 def get_read_options(options, default_read_consistency=None):
@@ -375,7 +375,7 @@ def put(entity, options):
  lock = yield _cache.global_lock_for_write(cache_key)
  else:
  expires = context._global_cache_timeout(entity.key, options)
- cache_value = entity_pb.SerializeToString()
+ cache_value = entity_pb._pb.SerializeToString()
  yield _cache.global_set(cache_key, cache_value, expires=expires)
 
  if use_datastore:
@@ -725,7 +725,7 @@ def allocate_ids_callback(self, rpc, mutations, futures):
  # Update mutations with complete keys
  response = rpc.result()
  for mutation, key, future in zip(mutations, response.keys, futures):
- mutation.upsert.key.CopyFrom(key)
+ mutation.upsert.key._pb.CopyFrom(key._pb)
  future.set_result(key)
 
  @tasklets.tasklet
@@ -863,9 +863,9 @@ def _datastore_commit(mutations, transaction, retries=None, timeout=None):
  :class:`google.cloud.datastore_v1.datastore_pb2.CommitResponse`
  """
  if transaction is None:
- mode = datastore_pb2.CommitRequest.NON_TRANSACTIONAL
+ mode = datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL
  else:
- mode = datastore_pb2.CommitRequest.TRANSACTIONAL
+ mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL
 
  client = context_module.get_context().client
  request = datastore_pb2.CommitRequest(
@@ -875,7 +875,7 @@ def _datastore_commit(mutations, transaction, retries=None, timeout=None):
  transaction=transaction,
  )
 
- return make_call("Commit", request, retries=retries, timeout=timeout)
+ return make_call("commit", request, retries=retries, timeout=timeout)
 
 
 def allocate(keys, options):
@@ -992,7 +992,7 @@ def _datastore_allocate_ids(keys, retries=None, timeout=None):
  client = context_module.get_context().client
  request = datastore_pb2.AllocateIdsRequest(project_id=client.project, keys=keys)
 
- return make_call("AllocateIds", request, retries=retries, timeout=timeout)
+ return make_call("allocate_ids", request, retries=retries, timeout=timeout)
 
 
 @tasklets.tasklet
@@ -1048,7 +1048,7 @@ def _datastore_begin_transaction(read_only, retries=None, timeout=None):
  project_id=client.project, transaction_options=options
  )
 
- return make_call("BeginTransaction", request, retries=retries, timeout=timeout)
+ return make_call("begin_transaction", request, retries=retries, timeout=timeout)
 
 
 @tasklets.tasklet
@@ -1089,4 +1089,4 @@ def _datastore_rollback(transaction, retries=None, timeout=None):
  project_id=client.project, transaction=transaction
  )
 
- return make_call("Rollback", request, retries=retries, timeout=timeout)
+ return make_call("rollback", request, retries=retries, timeout=timeout)

google/cloud/ndb/_datastore_query.py

@@ -22,9 +22,9 @@
 
 from google.cloud import environment_vars
 
-from google.cloud.datastore_v1.proto import datastore_pb2
-from google.cloud.datastore_v1.proto import entity_pb2
-from google.cloud.datastore_v1.proto import query_pb2
+from google.cloud.datastore_v1.types import datastore as datastore_pb2
+from google.cloud.datastore_v1.types import entity as entity_pb2
+from google.cloud.datastore_v1.types import query as query_pb2
 from google.cloud.datastore import helpers, Key
 
 from google.cloud.ndb import context as context_module
@@ -38,24 +38,24 @@
 log = logging.getLogger(__name__)
 
 MoreResultsType = query_pb2.QueryResultBatch.MoreResultsType
-NO_MORE_RESULTS = MoreResultsType.Value("NO_MORE_RESULTS")
-NOT_FINISHED = MoreResultsType.Value("NOT_FINISHED")
-MORE_RESULTS_AFTER_LIMIT = MoreResultsType.Value("MORE_RESULTS_AFTER_LIMIT")
+NO_MORE_RESULTS = MoreResultsType.NO_MORE_RESULTS
+NOT_FINISHED = MoreResultsType.NOT_FINISHED
+MORE_RESULTS_AFTER_LIMIT = MoreResultsType.MORE_RESULTS_AFTER_LIMIT
 
 ResultType = query_pb2.EntityResult.ResultType
-RESULT_TYPE_FULL = ResultType.Value("FULL")
-RESULT_TYPE_KEY_ONLY = ResultType.Value("KEY_ONLY")
-RESULT_TYPE_PROJECTION = ResultType.Value("PROJECTION")
+RESULT_TYPE_FULL = ResultType.FULL
+RESULT_TYPE_KEY_ONLY = ResultType.KEY_ONLY
+RESULT_TYPE_PROJECTION = ResultType.PROJECTION
 
-DOWN = query_pb2.PropertyOrder.DESCENDING
-UP = query_pb2.PropertyOrder.ASCENDING
+DOWN = query_pb2.PropertyOrder.Direction.DESCENDING
+UP = query_pb2.PropertyOrder.Direction.ASCENDING
 
 FILTER_OPERATORS = {
- "=": query_pb2.PropertyFilter.EQUAL,
- "<": query_pb2.PropertyFilter.LESS_THAN,
- "<=": query_pb2.PropertyFilter.LESS_THAN_OR_EQUAL,
- ">": query_pb2.PropertyFilter.GREATER_THAN,
- ">=": query_pb2.PropertyFilter.GREATER_THAN_OR_EQUAL,
+ "=": query_pb2.PropertyFilter.Operator.EQUAL,
+ "<": query_pb2.PropertyFilter.Operator.LESS_THAN,
+ "<=": query_pb2.PropertyFilter.Operator.LESS_THAN_OR_EQUAL,
+ ">": query_pb2.PropertyFilter.Operator.GREATER_THAN,
+ ">=": query_pb2.PropertyFilter.Operator.GREATER_THAN_OR_EQUAL,
 }
 
 _KEY_NOT_IN_CACHE = object()
@@ -77,7 +77,7 @@ def make_filter(name, op, value):
  property=query_pb2.PropertyReference(name=name),
  op=FILTER_OPERATORS[op],
  )
- helpers._set_protobuf_value(filter_pb.value, value)
+ helpers._set_protobuf_value(filter_pb.value._pb, value)
  return filter_pb
 
 
@@ -92,7 +92,7 @@ def make_composite_and_filter(filter_pbs):
  query_pb2.CompositeFilter: The new composite filter.
  """
  return query_pb2.CompositeFilter(
- op=query_pb2.CompositeFilter.AND,
+ op=query_pb2.CompositeFilter.Operator.AND,
  filters=[_filter_pb(filter_pb) for filter_pb in filter_pbs],
  )
 
@@ -683,7 +683,7 @@ def has_next_async(self):
  next_result = result_sets[0].next()
 
  # Check to see if it's a duplicate
- hash_key = next_result.result_pb.entity.key.SerializeToString()
+ hash_key = next_result.result_pb.entity.key._pb.SerializeToString()
  if hash_key in self._seen_keys:
  continue
 
@@ -811,9 +811,9 @@ def _compare(self, other):
  ).flat_path
  else:
  this_value_pb = self.result_pb.entity.properties[order.name]
- this_value = helpers._get_value_from_value_pb(this_value_pb)
+ this_value = helpers._get_value_from_value_pb(this_value_pb._pb)
  other_value_pb = other.result_pb.entity.properties[order.name]
- other_value = helpers._get_value_from_value_pb(other_value_pb)
+ other_value = helpers._get_value_from_value_pb(other_value_pb._pb)
 
  # Compare key paths if ordering by key property
  if isinstance(this_value, Key):
@@ -935,19 +935,19 @@ def _query_to_protobuf(query):
  ancestor_pb = query.ancestor._key.to_protobuf()
  ancestor_filter_pb = query_pb2.PropertyFilter(
  property=query_pb2.PropertyReference(name="__key__"),
- op=query_pb2.PropertyFilter.HAS_ANCESTOR,
+ op=query_pb2.PropertyFilter.Operator.HAS_ANCESTOR,
  )
- ancestor_filter_pb.value.key_value.CopyFrom(ancestor_pb)
+ ancestor_filter_pb.value.key_value._pb.CopyFrom(ancestor_pb._pb)
 
  if filter_pb is None:
  filter_pb = ancestor_filter_pb
 
  elif isinstance(filter_pb, query_pb2.CompositeFilter):
- filter_pb.filters.add(property_filter=ancestor_filter_pb)
+ filter_pb.filters._pb.add(property_filter=ancestor_filter_pb._pb)
 
  else:
  filter_pb = query_pb2.CompositeFilter(
- op=query_pb2.CompositeFilter.AND,
+ op=query_pb2.CompositeFilter.Operator.AND,
  filters=[
  _filter_pb(filter_pb),
  _filter_pb(ancestor_filter_pb),
@@ -969,7 +969,7 @@ def _query_to_protobuf(query):
  query_pb.offset = query.offset
 
  if query.limit:
- query_pb.limit.value = query.limit
+ query_pb._pb.limit.value = query.limit
 
  return query_pb
 
@@ -1016,7 +1016,7 @@ def _datastore_run_query(query):
  read_options=read_options,
  )
  response = yield _datastore_api.make_call(
- "RunQuery", request, timeout=query.timeout
+ "run_query", request, timeout=query.timeout
  )
  utils.logging_debug(log, response)
  raise tasklets.Return(response)

google/cloud/ndb/client.py

@@ -19,12 +19,13 @@
 import os
 import requests
 
-from google.api_core import client_info
+from google.api_core.gapic_v1 import client_info
 from google.cloud import environment_vars
 from google.cloud import _helpers
 from google.cloud import client as google_client
-from google.cloud.datastore_v1.gapic import datastore_client
-from google.cloud.datastore_v1.proto import datastore_pb2_grpc
+from google.cloud.datastore_v1.services.datastore.transports import (
+ grpc as datastore_grpc,
+)
 
 from google.cloud.ndb import __version__
 from google.cloud.ndb import context as context_module
@@ -35,7 +36,7 @@
  user_agent="google-cloud-ndb/{}".format(__version__)
 )
 
-DATASTORE_API_HOST = datastore_client.DatastoreClient.SERVICE_ADDRESS.rsplit(":", 1)[0]
+DATASTORE_API_HOST = "datastore.googleapis.com"
 
 
 def _get_gcd_project():
@@ -114,14 +115,17 @@ def __init__(self, project=None, namespace=None, credentials=None):
 
  if emulator:
  channel = grpc.insecure_channel(self.host)
-
  else:
- user_agent = _CLIENT_INFO.to_user_agent()
+ user_agent = self.client_info.to_user_agent()
  channel = _helpers.make_secure_channel(
  self._credentials, user_agent, self.host
  )
-
- self.stub = datastore_pb2_grpc.DatastoreStub(channel)
+ self.stub = datastore_grpc.DatastoreGrpcTransport(
+ host=self.host,
+ credentials=credentials,
+ client_info=self.client_info,
+ channel=channel,
+ )
 
  @contextlib.contextmanager
  def context(

google/cloud/ndb/key.py

@@ -204,7 +204,7 @@ class Key(object):
  >>> reference
  app: "example"
  path {
- Element {
+ element {
  type: "Kind"
  id: 1337
  }
@@ -681,13 +681,13 @@ def reference(self):
  >>> key = ndb.Key("Trampoline", 88, project="xy", namespace="zt")
  >>> key.reference()
  app: "xy"
+ name_space: "zt"
  path {
- Element {
+ element {
  type: "Trampoline"
  id: 88
  }
  }
- name_space: "zt"
  <BLANKLINE>
  """
  if self._reference is None: