googleapis · Mariatta · Feb 3, 2023 · Jan 11, 2023 · Jan 11, 2023
@@ -174,6 +174,7 @@ Running System Tests
 
 - You'll also need stored data in your dataset. To populate this data, run::
 
+   $ export SYSTEM_TESTS_DATABASE=system-tests-named-db
    $ python tests/system/utils/populate_datastore.py
 
 - If you make a mistake during development (i.e. a failing test that

@@ -34,9 +34,9 @@
 The main concepts with this API are:
 
 - :class:`~google.cloud.datastore.client.Client`
-  which represents a project (string) and namespace (string) bundled with
-  a connection and has convenience methods for constructing objects with that
-  project / namespace.
+  which represents a project (string), database (string), and namespace
+  (string) bundled with a connection and has convenience methods for
+  constructing objects with that project/database/namespace.
 
 - :class:`~google.cloud.datastore.entity.Entity`
   which represents a single entity in the datastore

@@ -22,6 +22,7 @@
 from google.cloud.datastore_v1.types import query as query_pb2
 from google.cloud.datastore import helpers
 from google.cloud.datastore.query import _pb_from_query
+from google.cloud.datastore.constants import DEFAULT_DATABASE
 
 
 _NOT_FINISHED = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED
@@ -123,6 +124,18 @@ def project(self):
         """
         return self._nested_query._project or self._client.project
 
+    @property
+    def database(self):
+        """Get the database for this AggregationQuery.
+        :rtype: str
+        :returns: The database for the query.
+        """
+        if self._nested_query._database or (
+            self._nested_query._database == DEFAULT_DATABASE
+        ):
+            return self._nested_query._database
+        return self._client.database
+
     @property
     def namespace(self):
         """The nested query's namespace
@@ -376,6 +389,7 @@ def _next_page(self):
 
         partition_id = entity_pb2.PartitionId(
             project_id=self._aggregation_query.project,
+            database_id=self._aggregation_query.database,
             namespace_id=self._aggregation_query.namespace,
         )
 
@@ -390,6 +404,7 @@ def _next_page(self):
         response_pb = self.client._datastore_api.run_aggregation_query(
             request={
                 "project_id": self._aggregation_query.project,
+                "database_id": self._aggregation_query.database,
                 "partition_id": partition_id,
                 "read_options": read_options,
                 "aggregation_query": query_pb,
@@ -409,6 +424,7 @@ def _next_page(self):
             response_pb = self.client._datastore_api.run_aggregation_query(
                 request={
                     "project_id": self._aggregation_query.project,
+                    "database_id": self._aggregation_query.database,
                     "partition_id": partition_id,
                     "read_options": read_options,
                     "aggregation_query": query_pb,

@@ -23,6 +23,7 @@
 
 from google.cloud.datastore import helpers
 from google.cloud.datastore_v1.types import datastore as _datastore_pb2
+from google.cloud.datastore.constants import DEFAULT_DATABASE
 
 
 class Batch(object):
@@ -122,6 +123,15 @@ def project(self):
         """
         return self._client.project
 
+    @property
+    def database(self):
+        """Getter for database in which the batch will run.
+
+        :rtype: :class:`str`
+        :returns: The database in which the batch will run.
+        """
+        return self._client.database
+
     @property
     def namespace(self):
         """Getter for namespace in which the batch will run.
@@ -218,6 +228,12 @@ def put(self, entity):
         if self.project != entity.key.project:
             raise ValueError("Key must be from same project as batch")
 
+        entity_key_database = entity.key.database
+        if entity_key_database is None:
+            entity_key_database = DEFAULT_DATABASE
+        if self.database != entity_key_database:
+            raise ValueError("Key must be from same database as batch")
+
         if entity.key.is_partial:
             entity_pb = self._add_partial_key_entity_pb()
             self._partial_key_entities.append(entity)
@@ -245,6 +261,12 @@ def delete(self, key):
         if self.project != key.project:
             raise ValueError("Key must be from same project as batch")
 
+        key_db = key.database
+        if key_db is None:
+            key_db = DEFAULT_DATABASE
+        if self.database != key_db:
+            raise ValueError("Key must be from same database as batch")
+
         key_pb = key.to_protobuf()
         self._add_delete_key_pb()._pb.CopyFrom(key_pb._pb)
 
@@ -284,6 +306,7 @@ def _commit(self, retry, timeout):
         commit_response_pb = self._client._datastore_api.commit(
             request={
                 "project_id": self.project,
+                "database_id": self.database,
                 "mode": mode,
                 "transaction": self._id,
                 "mutations": self._mutations,

@@ -25,6 +25,7 @@
 from google.cloud.datastore import helpers
 from google.cloud.datastore._http import HTTPDatastoreAPI
 from google.cloud.datastore.batch import Batch
+from google.cloud.datastore.constants import DEFAULT_DATABASE
 from google.cloud.datastore.entity import Entity
 from google.cloud.datastore.key import Key
 from google.cloud.datastore.query import Query
@@ -126,6 +127,7 @@ def _extended_lookup(
     retry=None,
     timeout=None,
     read_time=None,
+    database=DEFAULT_DATABASE,
 ):
     """Repeat lookup until all keys found (unless stop requested).
 
@@ -179,6 +181,10 @@ def _extended_lookup(
         ``eventual==True`` or ``transaction_id``.
         This feature is in private preview.
 
+    :type database: str
+    :param database:
+        (Optional) Database from which to fetch data. Defaults to the (default) database.
+
     :rtype: list of :class:`.entity_pb2.Entity`
     :returns: The requested entities.
     :raises: :class:`ValueError` if missing / deferred are not null or
@@ -201,6 +207,7 @@ def _extended_lookup(
         lookup_response = datastore_api.lookup(
             request={
                 "project_id": project,
+                "database_id": database,
                 "keys": key_pbs,
                 "read_options": read_options,
             },
@@ -276,6 +283,9 @@ class Client(ClientWithProject):
                       environment variable.
                       This parameter should be considered private, and could
                       change in the future.
+
+    :type database: str
+    :param database: (Optional) database to pass to proxied API methods.
     """
 
     SCOPE = ("https://www.googleapis.com/auth/datastore",)
@@ -290,6 +300,8 @@ def __init__(
         client_options=None,
         _http=None,
         _use_grpc=None,
+        *,
+        database=DEFAULT_DATABASE,
     ):
         emulator_host = os.getenv(DATASTORE_EMULATOR_HOST)
 
@@ -306,6 +318,7 @@ def __init__(
             client_options=client_options,
             _http=_http,
         )
+        self.database = database
         self.namespace = namespace
         self._client_info = client_info
         self._client_options = client_options
@@ -549,6 +562,7 @@ def get_multi(
         entity_pbs = _extended_lookup(
             datastore_api=self._datastore_api,
             project=self.project,
+            database=self.database,
             key_pbs=[key.to_protobuf() for key in keys],
             eventual=eventual,
             missing=missing,
@@ -740,7 +754,11 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
 
         response_pb = self._datastore_api.allocate_ids(
-            request={"project_id": incomplete_key.project, "keys": incomplete_key_pbs},
+            request={
+                "project_id": incomplete_key.project,
+                "database_id": incomplete_key.database,
+                "keys": incomplete_key_pbs,
+            },
             **kwargs,
         )
         allocated_ids = [
@@ -753,11 +771,14 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
     def key(self, *path_args, **kwargs):
         """Proxy to :class:`google.cloud.datastore.key.Key`.
 
-        Passes our ``project``.
+        Passes our ``project`` and our ``database``.
         """
         if "project" in kwargs:
             raise TypeError("Cannot pass project")
         kwargs["project"] = self.project
+        if "database" in kwargs:
+            raise TypeError("Cannot pass database")
+        kwargs["database"] = self.database
         if "namespace" not in kwargs:
             kwargs["namespace"] = self.namespace
         return Key(*path_args, **kwargs)
@@ -780,7 +801,7 @@ def transaction(self, **kwargs):
     def query(self, **kwargs):
         """Proxy to :class:`google.cloud.datastore.query.Query`.
 
-        Passes our ``project``.
+        Passes our ``project``  and our ``database``.
 
         Using query to search a datastore:
 
@@ -834,7 +855,10 @@ def do_something_with(entity):
             raise TypeError("Cannot pass client")
         if "project" in kwargs:
             raise TypeError("Cannot pass project")
+        if "database" in kwargs:
+            raise TypeError("Cannot pass database")
         kwargs["project"] = self.project
+        kwargs["database"] = self.database
         if "namespace" not in kwargs:
             kwargs["namespace"] = self.namespace
         return Query(self, **kwargs)
@@ -963,18 +987,26 @@ def reserve_ids_sequential(self, complete_key, num_ids, retry=None, timeout=None
         key_class = type(complete_key)
         namespace = complete_key._namespace
         project = complete_key._project
+        database = complete_key._database
         flat_path = list(complete_key._flat_path[:-1])
         start_id = complete_key._flat_path[-1]
 
         key_pbs = []
         for id in range(start_id, start_id + num_ids):
             path = flat_path + [id]
-            key = key_class(*path, project=project, namespace=namespace)
+            key = key_class(
+                *path, project=project, database=database, namespace=namespace
+            )
             key_pbs.append(key.to_protobuf())
 
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
         self._datastore_api.reserve_ids(
-            request={"project_id": complete_key.project, "keys": key_pbs}, **kwargs
+            request={
+                "project_id": complete_key.project,
+                "database_id": complete_key.database,
+                "keys": key_pbs,
+            },
+            **kwargs,
         )
         return None
 
@@ -1021,7 +1053,12 @@ def reserve_ids_multi(self, complete_keys, retry=None, timeout=None):
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
         key_pbs = [key.to_protobuf() for key in complete_keys]
         self._datastore_api.reserve_ids(
-            request={"project_id": complete_keys[0].project, "keys": key_pbs}, **kwargs
+            request={
+                "project_id": complete_keys[0].project,
+                "database_id": complete_keys[0].database,
+                "keys": key_pbs,
+            },
+            **kwargs,
         )
 
         return None
@@ -0,0 +1,4 @@
+"""Constants for Datastore."""
+
+DEFAULT_DATABASE = ""
+"""Datastore default database."""
@@ -29,6 +29,7 @@
 from google.cloud.datastore_v1.types import entity as entity_pb2
 from google.cloud.datastore.entity import Entity
 from google.cloud.datastore.key import Key
+from google.cloud.datastore.constants import DEFAULT_DATABASE
 from google.protobuf import timestamp_pb2
 
 
@@ -300,11 +301,14 @@ def key_from_protobuf(pb):
     project = None
     if pb.partition_id.project_id:  # Simple field (string)
         project = pb.partition_id.project_id
+    database = DEFAULT_DATABASE
+    if pb.partition_id.database_id:  # Simple field (string)
+        database = pb.partition_id.database_id
     namespace = None
     if pb.partition_id.namespace_id:  # Simple field (string)
         namespace = pb.partition_id.namespace_id
 
-    return Key(*path_args, namespace=namespace, project=project)
+    return Key(*path_args, namespace=namespace, project=project, database=database)
 
 
 def _pb_attr_value(val):