feat: named database support

google/cloud/datastore/__init__.py

@@ -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

google/cloud/datastore/_http.py

@@ -59,6 +59,7 @@ def _request(
     data,
     base_url,
     client_info,
+    database=None,
     retry=None,
     timeout=None,
 ):
@@ -84,6 +85,9 @@ def _request(
     :type client_info: :class:`google.api_core.client_info.ClientInfo`
     :param client_info: used to generate user agent.
 
+    :type database: str
+    :param database: (Optional) The database to make the request for.
+
     :type retry: :class:`google.api_core.retry.Retry`
     :param retry: (Optional) retry policy for the request
 
@@ -101,6 +105,7 @@ def _request(
         "User-Agent": user_agent,
         connection_module.CLIENT_INFO_HEADER: user_agent,
     }
+    _update_headers(headers, project, database)
     api_url = build_api_url(project, method, base_url)
 
     requester = http.request
@@ -136,6 +141,7 @@ def _rpc(
     client_info,
     request_pb,
     response_pb_cls,
+    database=None,
     retry=None,
     timeout=None,
 ):
@@ -165,6 +171,9 @@ def _rpc(
     :param response_pb_cls: The class used to unmarshall the response
                             protobuf.
 
+    :type database: str
+    :param database: (Optional) The database to make the request for.
+
     :type retry: :class:`google.api_core.retry.Retry`
     :param retry: (Optional) retry policy for the request
 
@@ -177,7 +186,7 @@ def _rpc(
     req_data = request_pb._pb.SerializeToString()
     kwargs = _make_retry_timeout_kwargs(retry, timeout)
     response = _request(
-        http, project, method, req_data, base_url, client_info, **kwargs
+        http, project, method, req_data, base_url, client_info, database, **kwargs
     )
     return response_pb_cls.deserialize(response)
 
@@ -236,6 +245,7 @@ def lookup(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.LookupRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -245,6 +255,7 @@ def lookup(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.LookupResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -267,6 +278,7 @@ def run_query(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.RunQueryRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -276,6 +288,7 @@ def run_query(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.RunQueryResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -300,6 +313,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None):
             request, _datastore_pb2.RunAggregationQueryRequest
         )
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -309,6 +323,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.RunAggregationQueryResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -331,6 +346,7 @@ def begin_transaction(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.BeginTransactionRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -340,6 +356,7 @@ def begin_transaction(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.BeginTransactionResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -362,6 +379,7 @@ def commit(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.CommitRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -371,6 +389,7 @@ def commit(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.CommitResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -393,6 +412,7 @@ def rollback(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.RollbackRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -402,6 +422,7 @@ def rollback(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.RollbackResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -424,6 +445,7 @@ def allocate_ids(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.AllocateIdsRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -433,6 +455,7 @@ def allocate_ids(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.AllocateIdsResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
@@ -455,6 +478,7 @@ def reserve_ids(self, request, retry=None, timeout=None):
         """
         request_pb = _make_request_pb(request, _datastore_pb2.ReserveIdsRequest)
         project_id = request_pb.project_id
+        database_id = request_pb.database_id
 
         return _rpc(
             self.client._http,
@@ -464,6 +488,18 @@ def reserve_ids(self, request, retry=None, timeout=None):
             self.client._client_info,
             request_pb,
             _datastore_pb2.ReserveIdsResponse,
+            database_id,
             retry=retry,
             timeout=timeout,
         )
+
+
+def _update_headers(headers, project_id, database_id=None):
+    """Update the request headers.
+    Pass the project id, or optionally the database_id if provided.
+    """
+    headers["x-goog-request-params"] = f"project_id={project_id}"
+    if database_id:
+        headers[
+            "x-goog-request-params"
+        ] = f"project_id={project_id}&database_id={database_id}"

google/cloud/datastore/aggregation.py

@@ -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,

google/cloud/datastore/batch.py

@@ -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,