feat: named db support (#882)

CONTRIBUTING.rst

@@ -146,6 +146,7 @@ Running System Tests
 
 - To run system tests for a given package, you can execute::
 
+   $ export SYSTEM_TESTS_DATABASE=system-tests-named-db
    $ nox -e system
 
   .. note::
@@ -188,6 +189,7 @@ Running System Tests
 
    # Create the indexes
    $ gcloud datastore indexes create tests/system/index.yaml
+   $ gcloud alpha datastore indexes create --database=$SYSTEM_TESTS_DATABASE tests/system/index.yaml
 
 
 *************

google/cloud/ndb/_datastore_api.py

@@ -19,6 +19,7 @@
 import logging
 
 from google.api_core import exceptions as core_exceptions
+from google.api_core import gapic_v1
 from google.cloud.datastore import helpers
 from google.cloud.datastore_v1.types import datastore as datastore_pb2
 from google.cloud.datastore_v1.types import entity as entity_pb2
@@ -56,7 +57,7 @@ def stub():
     return context.client.stub
 
 
-def make_call(rpc_name, request, retries=None, timeout=None):
+def make_call(rpc_name, request, retries=None, timeout=None, metadata=()):
     """Make a call to the Datastore API.
 
     Args:
@@ -68,6 +69,8 @@ def make_call(rpc_name, request, retries=None, timeout=None):
             If :data:`0` is passed, the call is attempted only once.
         timeout (float): Timeout, in seconds, to pass to gRPC call. If
             :data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
 
     Returns:
         tasklets.Future: Future for the eventual response for the API call.
@@ -85,7 +88,7 @@ def make_call(rpc_name, request, retries=None, timeout=None):
     def rpc_call():
         context = context_module.get_toplevel_context()
 
-        call = method.future(request, timeout=timeout)
+        call = method.future(request, timeout=timeout, metadata=metadata)
         rpc = _remote.RemoteCall(call, rpc_name)
         utils.logging_debug(log, rpc)
         utils.logging_debug(log, "timeout={}", timeout)
@@ -282,7 +285,7 @@ def lookup_callback(self, rpc):
                 future.set_result(entity)
 
 
-def _datastore_lookup(keys, read_options, retries=None, timeout=None):
+def _datastore_lookup(keys, read_options, retries=None, timeout=None, metadata=()):
     """Issue a Lookup call to Datastore using gRPC.
 
     Args:
@@ -295,18 +298,24 @@ def _datastore_lookup(keys, read_options, retries=None, timeout=None):
             If :data:`0` is passed, the call is attempted only once.
         timeout (float): Timeout, in seconds, to pass to gRPC call. If
             :data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
 
     Returns:
         tasklets.Future: Future object for eventual result of lookup.
     """
     client = context_module.get_context().client
     request = datastore_pb2.LookupRequest(
         project_id=client.project,
+        database_id=client.database,
         keys=[key for key in keys],
         read_options=read_options,
     )
+    metadata = _add_routing_info(metadata, request)
 
-    return make_call("lookup", request, retries=retries, timeout=timeout)
+    return make_call(
+        "lookup", request, retries=retries, timeout=timeout, metadata=metadata
+    )
 
 
 def get_read_options(options, default_read_consistency=None):
@@ -843,7 +852,7 @@ def _complete(key_pb):
     return False
 
 
-def _datastore_commit(mutations, transaction, retries=None, timeout=None):
+def _datastore_commit(mutations, transaction, retries=None, timeout=None, metadata=()):
     """Call Commit on Datastore.
 
     Args:
@@ -857,6 +866,8 @@ def _datastore_commit(mutations, transaction, retries=None, timeout=None):
             If :data:`0` is passed, the call is attempted only once.
         timeout (float): Timeout, in seconds, to pass to gRPC call. If
             :data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
 
     Returns:
         tasklets.Tasklet: A future for
@@ -870,12 +881,16 @@ def _datastore_commit(mutations, transaction, retries=None, timeout=None):
     client = context_module.get_context().client
     request = datastore_pb2.CommitRequest(
         project_id=client.project,
+        database_id=client.database,
         mode=mode,
         mutations=mutations,
         transaction=transaction,
     )
+    metadata = _add_routing_info(metadata, request)
 
-    return make_call("commit", request, retries=retries, timeout=timeout)
+    return make_call(
+        "commit", request, retries=retries, timeout=timeout, metadata=metadata
+    )
 
 
 def allocate(keys, options):
@@ -973,7 +988,7 @@ def allocate_ids_callback(self, rpc):
             future.set_result(key)
 
 
-def _datastore_allocate_ids(keys, retries=None, timeout=None):
+def _datastore_allocate_ids(keys, retries=None, timeout=None, metadata=()):
     """Calls ``AllocateIds`` on Datastore.
 
     Args:
@@ -984,15 +999,22 @@ def _datastore_allocate_ids(keys, retries=None, timeout=None):
             If :data:`0` is passed, the call is attempted only once.
         timeout (float): Timeout, in seconds, to pass to gRPC call. If
             :data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
 
     Returns:
         tasklets.Future: A future for
             :class:`google.cloud.datastore_v1.datastore_pb2.AllocateIdsResponse`
     """
     client = context_module.get_context().client
-    request = datastore_pb2.AllocateIdsRequest(project_id=client.project, keys=keys)
+    request = datastore_pb2.AllocateIdsRequest(
+        project_id=client.project, database_id=client.database, keys=keys
+    )
+    metadata = _add_routing_info(metadata, request)
 
-    return make_call("allocate_ids", request, retries=retries, timeout=timeout)
+    return make_call(
+        "allocate_ids", request, retries=retries, timeout=timeout, metadata=metadata
+    )
 
 
 @tasklets.tasklet
@@ -1018,7 +1040,7 @@ def begin_transaction(read_only, retries=None, timeout=None):
     raise tasklets.Return(response.transaction)
 
 
-def _datastore_begin_transaction(read_only, retries=None, timeout=None):
+def _datastore_begin_transaction(read_only, retries=None, timeout=None, metadata=()):
     """Calls ``BeginTransaction`` on Datastore.
 
     Args:
@@ -1029,6 +1051,8 @@ def _datastore_begin_transaction(read_only, retries=None, timeout=None):
             If :data:`0` is passed, the call is attempted only once.
         timeout (float): Timeout, in seconds, to pass to gRPC call. If
             :data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
 
     Returns:
         tasklets.Tasklet: A future for
@@ -1045,10 +1069,19 @@ def _datastore_begin_transaction(read_only, retries=None, timeout=None):
         )
 
     request = datastore_pb2.BeginTransactionRequest(
-        project_id=client.project, transaction_options=options
+        project_id=client.project,
+        database_id=client.database,
+        transaction_options=options,
+    )
+    metadata = _add_routing_info(metadata, request)
+
+    return make_call(
+        "begin_transaction",
+        request,
+        retries=retries,
+        timeout=timeout,
+        metadata=metadata,
     )
-
-    return make_call("begin_transaction", request, retries=retries, timeout=timeout)
 
 
 @tasklets.tasklet
@@ -1069,7 +1102,7 @@ def rollback(transaction, retries=None, timeout=None):
     yield _datastore_rollback(transaction, retries=retries, timeout=timeout)
 
 
-def _datastore_rollback(transaction, retries=None, timeout=None):
+def _datastore_rollback(transaction, retries=None, timeout=None, metadata=()):
     """Calls Rollback in Datastore.
 
     Args:
@@ -1079,14 +1112,50 @@ def _datastore_rollback(transaction, retries=None, timeout=None):
             If :data:`0` is passed, the call is attempted only once.
         timeout (float): Timeout, in seconds, to pass to gRPC call. If
             :data:`None` is passed, will use :data:`_DEFAULT_TIMEOUT`.
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
 
     Returns:
         tasklets.Tasklet: Future for
             :class:`google.cloud.datastore_v1.datastore_pb2.RollbackResponse`
     """
     client = context_module.get_context().client
     request = datastore_pb2.RollbackRequest(
-        project_id=client.project, transaction=transaction
+        project_id=client.project,
+        database_id=client.database,
+        transaction=transaction,
     )
+    metadata = _add_routing_info(metadata, request)
+
+    return make_call(
+        "rollback", request, retries=retries, timeout=timeout, metadata=metadata
+    )
+
+
+def _add_routing_info(metadata, request):
+    """Adds routing header info to the given metadata.
+
+    Args:
+        metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata. Not modified.
+        request (Any): An appropriate request object for the call, eg,
+            `entity_pb2.LookupRequest` for calling ``Lookup``.
+
+    Returns:
+        Sequence[Tuple[str, str]]: Sequence with routing info added,
+            if it is included in the request.
+    """
+    header_params = {}
+
+    if request.project_id:
+        header_params["project_id"] = request.project_id
+
+    if request.database_id:
+        header_params["database_id"] = request.database_id
+
+    if header_params:
+        return tuple(metadata) + (
+            gapic_v1.routing_header.to_grpc_metadata(header_params),
+        )
 
-    return make_call("rollback", request, retries=retries, timeout=timeout)
+    return tuple(metadata)

google/cloud/ndb/_datastore_query.py

@@ -1010,17 +1010,22 @@ def _datastore_run_query(query):
     """
     query_pb = _query_to_protobuf(query)
     partition_id = entity_pb2.PartitionId(
-        project_id=query.project, namespace_id=query.namespace
+        project_id=query.project,
+        database_id=query.database,
+        namespace_id=query.namespace,
     )
     read_options = _datastore_api.get_read_options(query)
     request = datastore_pb2.RunQueryRequest(
         project_id=query.project,
+        database_id=query.database,
         partition_id=partition_id,
         query=query_pb,
         read_options=read_options,
     )
+    metadata = _datastore_api._add_routing_info((), request)
+
     response = yield _datastore_api.make_call(
-        "run_query", request, timeout=query.timeout
+        "run_query", request, timeout=query.timeout, metadata=metadata
     )
     utils.logging_debug(log, response)
     raise tasklets.Return(response)

google/cloud/ndb/_gql.py

@@ -98,8 +98,7 @@ def __init__(self, query_string, _app=None, _auth_domain=None, namespace=None):
 
         Args:
             query_string (str): properly formatted GQL query string.
-            namespace (str): the namespace to use for this query.
-
+            namespace (str): The namespace to use for this query. Defaults to the client's value.
         Raises:
             exceptions.BadQueryError: if the query is not parsable.
         """
@@ -853,7 +852,10 @@ def _key_function(values):
         context = context_module.get_context()
         client = context.client
         return key.Key(
-            *values, namespace=context.get_namespace(), project=client.project
+            *values,
+            project=client.project,
+            database=client.database,
+            namespace=context.get_namespace(),
         )
     _raise_cast_error(
         "Key requires even number of operands or single string, {}".format(values)

google/cloud/ndb/client.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""A client for NDB which manages credentials, project, namespace."""
+"""A client for NDB which manages credentials, project, namespace, and database."""
 
 import contextlib
 import grpc
@@ -92,17 +92,25 @@ class Client(google_client.ClientWithProject):
         client_options (Optional[:class:`~google.api_core.client_options.ClientOptions` or :class:`dict`])
             Client options used to set user options on the client.
             API Endpoint should be set through client_options.
+        database (Optional[str]): Database to access. Defaults to the (default) database.
     """
 
     SCOPE = ("https://www.googleapis.com/auth/datastore",)
     """The scopes required for authenticating as a Cloud Datastore consumer."""
 
     def __init__(
-        self, project=None, namespace=None, credentials=None, client_options=None
+        self,
+        project=None,
+        namespace=None,
+        credentials=None,
+        client_options=None,
+        database=None,
     ):
         self.namespace = namespace
+        self.host = os.environ.get(environment_vars.GCD_HOST, DATASTORE_API_HOST)
         self.client_info = _CLIENT_INFO
         self._client_options = client_options
+        self.database = database
 
         # Use insecure connection when using Datastore Emulator, otherwise
         # use secure connection