10000 feat: named database support (#439) · googleapis/python-datastore@abf0060 · GitHub
[go: up one dir, main page]
Skip to content

Commit abf0060

Browse files
Mariattarwhogggcf-owl-bot[bot]vishwarajanandmeredithslota
authored
feat: named database support (#439)
* feat: named database support (#398) * feat: Add named database support * test: Use named db in system tests * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Handle the case when client doesn't have database property * fix: add custom routing headers * Fixing tests for easier merge * fixing code coverage * addressing pr comments * feat: Multi db test parametrization (#436) * Feat: Parametrize the tests for multidb support Remove "database" argument from Query and AggregationQuery constructors. Use the "database" from the client instead. Once set in the client, the "database" will be used throughout and cannot be re-set. Parametrize the tests where-ever clients are used. Use the `system-tests-named-db` in the system test. * Add test case for when parent database name != child database name * Update owlbot, removing the named db parameter * Reverted test fixes * fixing tests * fix code coverage * pr suggestion * address pr comments --------- Co-authored-by: Vishwaraj Anand <vishwaraj.anand00@gmail.com> --------- Co-authored-by: Bob "Wombat" Hogg <wombat@rwhogg.site> Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com> Co-authored-by: Vishwaraj Anand <vishwaraj.anand00@gmail.com> Co-authored-by: meredithslota <meredithslota@google.com>
1 parent a12971c commit abf0060

29 files changed

+1842
-921
lines changed

‎google/cloud/datastore/__init__.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@
3434
The main concepts with this API are:
3535
3636
- :class:`~google.cloud.datastore.client.Client`
37-
which represents a project (string) and namespace (string) bundled with
38-
a connection and has convenience methods for constructing objects with that
39-
project / namespace.
37+
which represents a project (string), database (string), and namespace
38+
(string) bundled with a connection and has convenience methods for
39+
constructing objects with that project/database/namespace.
4040
4141
- :class:`~google.cloud.datastore.entity.Entity`
4242
which represents a single entity in the datastore

‎google/cloud/datastore/_http.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ def _request(
5959
data,
6060
base_url,
6161
client_info,
62+
database,
6263
retry=None,
6364
timeout=None,
6465
):
@@ -84,6 +85,9 @@ def _request(
8485
:type client_info: :class:`google.api_core.client_info.ClientInfo`
8586
:param client_info: used to generate user agent.
8687
88+
:type database: str
89+
:param database: The database to make the request for.
90+
8791
:type retry: :class:`google.api_core.retry.Retry`
8892
:param retry: (Optional) retry policy for the request
8993
@@ -101,6 +105,7 @@ def _request(
101105
"User-Agent": user_agent,
102106
connection_module.CLIENT_INFO_HEADER: user_agent,
103107
}
108+
_update_headers(headers, project, database)
104109
api_url = build_api_url(project, method, base_url)
105110

106111
requester = http.request
@@ -136,6 +141,7 @@ def _rpc(
136141
client_info,
137142
request_pb,
138143
response_pb_cls,
144+
database,
139145
retry=None,
140146
timeout=None,
141147
):
@@ -165,6 +171,9 @@ def _rpc(
165171
:param response_pb_cls: The class used to unmarshall the response
166172
protobuf.
167173
174+
:type database: str
175+
:param database: The database to make the request for.
176+
168177
:type retry: :class:`google.api_core.retry.Retry`
169178
:param retry: (Optional) retry policy for the request
170179
@@ -177,7 +186,7 @@ def _rpc(
177186
req_data = request_pb._pb.SerializeToString()
178187
kwargs = _make_retry_timeout_kwargs(retry, timeout)
179188
response = _request(
180-
http, project, method, req_data, base_url, client_info, **kwargs
189+
http, project, method, req_data, base_url, client_info, database, **kwargs
181190
)
182191
return response_pb_cls.deserialize(response)
183192

@@ -236,6 +245,7 @@ def lookup(self, request, retry=None, timeout=None):
236245
"""
237246
request_pb = _make_request_pb(request, _datastore_pb2.LookupRequest)
238247
project_id = request_pb.project_id
248+
database_id = request_pb.database_id
239249

240250
return _rpc(
241251
self.client._http,
@@ -245,6 +255,7 @@ def lookup(self, request, retry=None, timeout=None):
245255
self.client._client_info,
246256
request_pb,
247257
_datastore_pb2.LookupResponse,
258+
database_id,
248259
retry=retry,
249260
timeout=timeout,
250261
)
@@ -267,6 +278,7 @@ def run_query(self, request, retry=None, timeout=None):
267278
"""
268279
request_pb = _make_request_pb(request, _datastore_pb2.RunQueryRequest)
269280
project_id = request_pb.project_id
281+
database_id = request_pb.database_id
270282

271283
return _rpc(
272284
self.client._http,
@@ -276,6 +288,7 @@ def run_query(self, request, retry=None, timeout=None):
276288
self.client._client_info,
277289
request_pb,
278290
_datastore_pb2.RunQueryResponse,
291+
database_id,
279292
retry=retry,
280293
timeout=timeout,
281294
)
@@ -300,6 +313,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None):
300313
request, _datastore_pb2.RunAggregationQueryRequest
301314
)
302315
project_id = request_pb.project_id
316+
database_id = request_pb.database_id
303317

304318
return _rpc(
305319
self.client._http,
@@ -309,6 +323,7 @@ def run_aggregation_query(self, request, retry=None, timeout=None):
309323
self.client._client_info,
310324
request_pb,
311325
_datastore_pb2.RunAggregationQueryResponse,
326+
database_id,
312327
retry=retry,
313328
timeout=timeout,
314329
)
@@ -331,6 +346,7 @@ def begin_transaction(self, request, retry=None, timeout=None):
331346
"""
332347
request_pb = _make_request_pb(request, _datastore_pb2.BeginTransactionRequest)
333348
project_id = request_pb.project_id
349+
database_id = request_pb.database_id
334350

335351
return _rpc(
336352
self.client._http,
@@ -340,6 +356,7 @@ def begin_transaction(self, request, retry=None, timeout=None):
340356
self.client._client_info,
341357
request_pb,
342358
_datastore_pb2.BeginTransactionResponse,
359+
database_id,
343360
retry=retry,
344361
timeout=timeout,
345362
)
@@ -362,6 +379,7 @@ def commit(self, request, retry=None, timeout=None):
362379
"""
363380
request_pb = _make_request_pb(request, _datastore_pb2.CommitRequest)
364381
project_id = request_pb.project_id
382+
database_id = request_pb.database_id
365383

366384
return _rpc(
367385
self.client._http,
@@ -371,6 +389,7 @@ def commit(self, request, retry=None, timeout=None):
371389
self.client._client_info,
372390
request_pb,
373391
_datastore_pb2.CommitResponse,
392+
database_id,
374393
retry=retry,
375394
timeout=timeout,
376395
)
@@ -393,6 +412,7 @@ def rollback(self, request, retry=None, timeout=None):
393412
"""
394413
request_pb = _make_request_pb(request, _datastore_pb2.RollbackRequest)
395414
project_id = request_pb.project_id
415+
database_id = request_pb.database_id
396416

397417
return _rpc(
398418
self.client._http,
@@ -402,6 +422,7 @@ def rollback(self, request, retry=None, timeout=None):
402422
self.client._client_info,
403423
request_pb,
404424
_datastore_pb2.RollbackResponse,
425+
database_id,
405426
retry=retry,
406427
timeout=timeout,
407428
)
@@ -424,6 +445,7 @@ def allocate_ids(self, request, retry=None, timeout=None):
424445
"""
425446
request_pb = _make_request_pb(request, _datastore_pb2.AllocateIdsRequest)
426447
project_id = request_pb.project_id
448+
database_id = request_pb.database_id
427449

428450
return _rpc(
429451
self.client._http,
@@ -433,6 +455,7 @@ def allocate_ids(self, request, retry=None, timeout=None):
433455
self.client._client_info,
434456
request_pb,
435457
_datastore_pb2.AllocateIdsResponse,
458+
database_id,
436459
retry=retry,
437460
timeout=timeout,
438461
)
@@ -455,6 +478,7 @@ def reserve_ids(self, request, retry=None, timeout=None):
455478
"""
456479
request_pb = _make_request_pb(request, _datastore_pb2.ReserveIdsRequest)
457480
project_id = request_pb.project_id
481+
database_id = request_pb.database_id
458482

459483
return _rpc(
460484
self.client._http,
@@ -464,6 +488,18 @@ def reserve_ids(self, request, retry=None, timeout=None):
464488
self.client._client_info,
465489
request_pb,
466490
_datastore_pb2.ReserveIdsResponse,
491+
database_id,
467492
retry=retry,
468493
timeout=timeout,
469494
)
495+
496+
497+
def _update_headers(headers, project_id, database_id=None):
498+
"""Update the request headers.
499+
Pass the project id, or optionally the database_id if provided.
500+
"""
501+
headers["x-goog-request-params"] = f"project_id={project_id}"
502+
if database_id:
503+
headers[
504+
"x-goog-request-params"
505+
] = f"project_id={project_id}&database_id={database_id}"

‎google/cloud/datastore/aggregation.py

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,7 @@ def _next_page(self):
376376

377377
partition_id = entity_pb2.PartitionId(
378378
project_id=self._aggregation_query.project,
379+
database_id=self.client.database,
379380
namespace_id=self._aggregation_query.namespace,
380381
)
381382

@@ -386,14 +387,15 @@ def _next_page(self):
386387

387388
if self._timeout is not None:
388389
kwargs["timeout"] = self._timeout
389-
390+
request = {
391+
"project_id": self._aggregation_query.project,
392+
"partition_id": partition_id,
393+
"read_options": read_options,
394+
"aggregation_query": query_pb,
395+
}
396+
helpers.set_database_id_to_request(request, self.client.database)
390397
response_pb = self.client._datastore_api.run_aggregation_query(
391-
request={
392-
"project_id": self._aggregation_query.project,
393-
"partition_id": partition_id,
394-
"read_options": read_options,
395-
"aggregation_query": query_pb,
396-
},
398+
request=request,
397399
**kwargs,
398400
)
399401

@@ -406,13 +408,15 @@ def _next_page(self):
406408
query_pb = query_pb2.AggregationQuery()
407409
query_pb._pb.CopyFrom(old_query_pb._pb) # copy for testability
408410

411+
request = {
412+
"project_id": self._aggregation_query.project,
413+
"partition_id": partition_id,
414+
"read_options": read_options,
415+
"aggregation_query": query_pb,
416+
}
417+
helpers.set_database_id_to_request(request, self.client.database)
409418
response_pb = self.client._datastore_api.run_aggregation_query(
410-
request={
411-
"project_id": self._aggregation_query.project,
412-
"partition_id": partition_id,
413-
"read_options": read_options,
414-
"aggregation_query": query_pb,
415-
},
419+
request=request,
416420
**kwargs,
417421
)
418422

‎google/cloud/datastore/batch.py

Lines changed: 25 additions & 6 deletions
< 5898 tr class="diff-line-row">
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,15 @@ def project(self):
122122
"""
123123
return self._client.project
124124

125+
@property
126+
def database(self):
127+
"""Getter for database in which the batch will run.
128+
129+
:rtype: :class:`str`
130+
:returns: The database in which the batch will run.
131+
"""
132+
return self._client.database
133+
125134
@property
126135
def namespace(self):
127136
"""Getter for namespace in which the batch will run.
@@ -218,6 +227,9 @@ def put(self, entity):
218227
if self.project != entity.key.project:
219228
raise ValueError("Key must be from same project as batch")
220229

230+
if self.database != entity.key.database:
231+
raise ValueError("Key must be from same database as batch")
232+
221233
if entity.key.is_partial:
222234
entity_pb = self._add_partial_key_entity_pb()
223235
self._partial_key_entities.append(entity)
@@ -245,6 +257,9 @@ def delete(self, key):
245257
if self.project != key.project:
246258
raise ValueError("Key must be from same project as batch")
247259

260+
if self.database != key.database:
261+
raise ValueError("Key must be from same database as batch")
262+
248263
key_pb = key.to_protobuf()
249264
self._add_delete_key_pb()._pb.CopyFrom(key_pb._pb)
250265

@@ -281,13 +296,17 @@ def _commit(self, retry, timeout):
281296
if timeout is not None:
282297
kwargs["timeout"] = timeout
283298

299+
request = {
300+
"project_id": self.project,
301+
"mode": mode,
302+
"transaction": self._id,
303+
"mutations": self._mutations,
304+
}
305+
306+
helpers.set_database_id_to_request(request, self._client.database)
307+
284308
commit_response_pb = self._client._datastore_api.commit(
285-
request={
286-
"project_id": self.project,
287-
"mode": mode,
288-
"transaction": self._id,
289-
"mutations": self._mutations,
290-
},
309+
request=request,
291310
**kwargs,
292311
)
293312

0 commit comments

Comments
 (0)
0