Skip to content

Commit e50b270

Browse files
dhermesdhermes
committed
Adding one-time RPC to find unaliased / true dataset ID.
Also removing unnecessary functions that are no longer needed since there is no need to muck with the dataset ID.
1 parent 0525530 commit e50b270

File tree

11 files changed

+124
-117
lines changed

11 files changed

+124
-117
lines changed

‎gcloud/datastore/__init__.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
from gcloud.datastore.api import allocate_ids
5858
from gcloud.datastore.api import get
5959
from gcloud.datastore.connection import Connection
60+
from gcloud.datastore.key import Key
6061

6162

6263
SCOPE = ('https://www.googleapis.com/auth/datastore',
@@ -66,6 +67,37 @@
6667
_DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
6768

6869

70+
def _find_true_dataset_id(dataset_id, connection=None):
71+
"""Find the true (unaliased) dataset ID.
72+
73+
If the given ID already has a 's~' or 'e~' prefix, does nothing.
74+
Otherwise, allocates a single key of bogus kind '__Foo' and reads
75+
the true prefixed dataset ID from the response.
76+
77+
:type dataset_id: string
78+
:param dataset_id: The dataset ID to un-alias / prefix.
79+
80+
:type connection: :class:`gcloud.datastore.connection.Connection`
81+
:param connection: Optional. A connection provided to be the default.
82+
83+
:rtype: string
84+
:returns: The true / prefixed / un-aliased dataset ID.
85+
"""
86+
if dataset_id.startswith('s~') or dataset_id.startswith('e~'):
87+
return dataset_id
88+
89+
connection = connection or _implicit_environ.CONNECTION
90+
91+
# Create the partial Key protobuf to be allocated and remove
92+
# the dataset ID so the backend won't complain.
93+
bogus_key_pb = Key('__Foo', dataset_id=dataset_id).to_protobuf()
94+
bogus_key_pb.partition_id.ClearField('dataset_id')
95+
96+
allocated_pb, = connection.allocate_ids(dataset_id, [bogus_key_pb])
97+
98+
return allocated_pb.partition_id.dataset_id
99+
100+
69101
def set_default_dataset_id(dataset_id=None):
70102
"""Set default dataset ID either explicitly or implicitly as fall-back.
71103
@@ -82,6 +114,7 @@ def set_default_dataset_id(dataset_id=None):
82114
dataset_id = os.getenv(_DATASET_ENV_VAR_NAME)
83115

84116
if dataset_id is not None:
117+
dataset_id = _find_true_dataset_id(dataset_id)
85118
_implicit_environ.DATASET_ID = dataset_id
86119

87120

‎gcloud/datastore/connection.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -458,8 +458,6 @@ def save_entity(self, dataset_id, key_pb, properties,
458458
in_batch = False
459459
mutation = self.mutation()
460460

461-
key_pb = helpers._prepare_key_for_request(key_pb)
462-
463461
# If the Key is complete, we should upsert
464462
# instead of using insert_auto_id.
465463
path = key_pb.path_element[-1]

‎gcloud/datastore/helpers.py‎

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import pytz
2626
import six
2727

28-
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
2928
from gcloud.datastore.entity import Entity
3029
from gcloud.datastore.key import Key
3130

@@ -261,33 +260,6 @@ def _set_protobuf_value(value_pb, val):
261260
setattr(value_pb, attr, val)
262261

263262

264-
def _prepare_key_for_request(key_pb):
265-
"""Add protobuf keys to a request object.
266-
267-
:type key_pb: :class:`gcloud.datastore.datastore_v1_pb2.Key`
268-
:param key_pb: A key to be added to a request.
269-
270-
:rtype: :class:`gcloud.datastore.datastore_v1_pb2.Key`
271-
:returns: A key which will be added to a request. It will be the
272-
original if nothing needs to be changed.
273-
"""
274-
if key_pb.partition_id.HasField('dataset_id'):
275-
# We remove the dataset_id from the protobuf. This is because
276-
# the backend fails a request if the key contains un-prefixed
277-
# dataset ID. The backend fails because requests to
278-
# /datastore/.../datasets/foo/...
279-
# and
280-
# /datastore/.../datasets/s~foo/...
281-
# both go to the datastore given by 's~foo'. So if the key
282-
# protobuf in the request body has dataset_id='foo', the
283-
# backend will reject since 'foo' != 's~foo'.
284-
new_key_pb = datastore_pb.Key()
285-
new_key_pb.CopyFrom(key_pb)
286-
new_key_pb.partition_id.ClearField('dataset_id')
287-
key_pb = new_key_pb
288-
return key_pb
289-
290-
291263
def _add_keys_to_request(request_field_pb, key_pbs):
292264
"""Add protobuf keys to a request object.
293265
@@ -298,5 +270,4 @@ def _add_keys_to_request(request_field_pb, key_pbs):
298270
:param key_pbs: The keys to add to a request.
299271
"""
300272
for key_pb in key_pbs:
301-
key_pb = _prepare_key_for_request(key_pb)
302273
request_field_pb.add().CopyFrom(key_pb)

‎gcloud/datastore/query.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -445,8 +445,7 @@ def _pb_from_query(query):
445445
composite_filter.operator = datastore_pb.CompositeFilter.AND
446446

447447
if query.ancestor:
448-
ancestor_pb = helpers._prepare_key_for_request(
449-
query.ancestor.to_protobuf())
448+
ancestor_pb = query.ancestor.to_protobuf()
450449

451450
# Filter on __key__ HAS_ANCESTOR == ancestor.
452451
ancestor_filter = composite_filter.filter.add().property_filter
@@ -465,8 +464,7 @@ def _pb_from_query(query):
465464
# Set the value to filter on based on the type.
466465
if property_name == '__key__':
467466
key_pb = value.to_protobuf()
468-
property_filter.value.key_value.CopyFrom(
469-
helpers._prepare_key_for_request(key_pb))
467+
property_filter.value.key_value.CopyFrom(key_pb)
470468
else:
471469
helpers._set_protobuf_value(property_filter.value, value)
472470

‎gcloud/datastore/test___init__.py‎

Lines changed: 63 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,22 +45,22 @@ def test_no_env_var_set(self):
4545

4646
def test_set_from_env_var(self):
4747
from gcloud.datastore import _implicit_environ
48-
IMPLICIT_DATASET_ID = 'IMPLICIT'
48+
IMPLICIT_DATASET_ID = 's~IMPLICIT'
4949
with self._monkey(IMPLICIT_DATASET_ID):
5050
self._callFUT()
5151
self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
5252

5353
def test_set_explicit_w_env_var_set(self):
5454
from gcloud.datastore import _implicit_environ
55-
EXPLICIT_DATASET_ID = 'EXPLICIT'
55+
EXPLICIT_DATASET_ID = 's~EXPLICIT'
5656
with self._monkey(None):
5757
self._callFUT(EXPLICIT_DATASET_ID)
5858
self.assertEqual(_implicit_environ.DATASET_ID, EXPLICIT_DATASET_ID)
5959

6060
def test_set_explicit_no_env_var_set(self):
6161
from gcloud.datastore import _implicit_environ
62-
IMPLICIT_DATASET_ID = 'IMPLICIT'
63-
EXPLICIT_DATASET_ID = 'EXPLICIT'
62+
IMPLICIT_DATASET_ID = 's~IMPLICIT'
63+
EXPLICIT_DATASET_ID = 's~EXPLICIT'
6464
with self._monkey(IMPLICIT_DATASET_ID):
6565
self._callFUT(EXPLICIT_DATASET_ID)
6666
self.assertEqual(_implicit_environ.DATASET_ID, EXPLICIT_DATASET_ID)
@@ -73,12 +73,50 @@ def test_set_explicit_None_wo_env_var_set(self):
7373

7474
def test_set_explicit_None_w_env_var_set(self):
7575
from gcloud.datastore import _implicit_environ
76-
IMPLICIT_DATASET_ID = 'IMPLICIT'
76+
IMPLICIT_DATASET_ID = 's~IMPLICIT'
7777
with self._monkey(IMPLICIT_DATASET_ID):
7878
self._callFUT(None)
7979
self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
8080

8181

82+
class Test__find_true_dataset_id(unittest2.TestCase):
83+
84+
def _callFUT(self, dataset_id, connection=None):
85+
from gcloud.datastore import _find_true_dataset_id
86+
return _find_true_dataset_id(dataset_id, connection=connection)
87+
88+
def test_prefixed(self):
89+
PREFIXED = 's~DATASET'
90+
result = self._callFUT(PREFIXED)
91+
self.assertEqual(PREFIXED, result)
92+
93+
def test_unprefixed_no_connection(self):
94+
from gcloud.datastore import _implicit_environ
95+
96+
UNPREFIXED = 'DATASET'
97+
self.assertEqual(_implicit_environ.CONNECTION, None)
98+
with self.assertRaises(AttributeError):
99+
self._callFUT(UNPREFIXED)
100+
101+
def test_unprefixed_explicit_connection(self):
102+
UNPREFIXED = 'DATASET'
103+
PREFIX = 's~'
104+
CONNECTION = _Connection(PREFIX)
105+
result = self._callFUT(UNPREFIXED, connection=CONNECTION)
106+
107+
self.assertEqual(CONNECTION._called_dataset_id, UNPREFIXED)
108+
# Make sure just one.
109+
called_key_pb, = CONNECTION._called_key_pbs
110+
path_element = called_key_pb.path_element
111+
self.assertEqual(len(path_element), 1)
112+
self.assertEqual(path_element[0].kind, '__Foo')
113+
self.assertFalse(path_element[0].HasField('id'))
114+
self.assertFalse(path_element[0].HasField('name'))
115+
116+
PREFIXED = PREFIX + UNPREFIXED
117+
self.assertEqual(result, PREFIXED)
118+
119+
82120
class Test_set_default_connection(unittest2.TestCase):
83121

84122
def setUp(self):
@@ -134,3 +172,23 @@ def test_it(self):
134172
self.assertTrue(isinstance(found, Connection))
135173
self.assertTrue(found._credentials is client._signed)
136174
self.assertTrue(client._get_app_default_called)
175+
176+
177+
class _Connection(object):
178+
179+
def __init__(self, prefix):
180+
self.prefix = prefix
181+
182+
def allocate_ids(self, dataset_id, key_pbs):
183+
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
184+
185+
# Store the arguments called with.
186+
self._called_dataset_id = dataset_id
187+
self._called_key_pbs = key_pbs
188+
189+
key_pb, = key_pbs
190+
191+
response = datastore_pb.Key()
192+
response.CopyFrom(key_pb)
193+
response.partition_id.dataset_id = self.prefix + dataset_id
194+
return [response]

‎gcloud/datastore/test_connection.py‎

Lines changed: 17 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ def test_lookup_single_key_empty_response(self):
237237
request.ParseFromString(cw['body'])
238238
keys = list(request.key)
239239
self.assertEqual(len(keys), 1)
240-
_compare_key_pb_after_request(self, key_pb, keys[0])
240+
self.assertEqual(key_pb, keys[0])
241241

242242
def test_lookup_single_key_empty_response_w_eventual(self):
243243
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
@@ -263,7 +263,7 @@ def test_lookup_single_key_empty_response_w_eventual(self):
263263
request.ParseFromString(cw['body'])
264264
keys = list(request.key)
265265
self.assertEqual(len(keys), 1)
266-
_compare_key_pb_after_request(self, key_pb, keys[0])
266+
self.assertEqual(key_pb, keys[0])
267267
self.assertEqual(request.read_options.read_consistency,
268268
datastore_pb.ReadOptions.EVENTUAL)
269269
self.assertEqual(request.read_options.transaction, '')
@@ -303,7 +303,7 @@ def test_lookup_single_key_empty_response_w_transaction(self):
303303
request.ParseFromString(cw['body'])
304304
keys = list(request.key)
305305
self.assertEqual(len(keys), 1)
306-
_compare_key_pb_after_request(self, key_pb, keys[0])
306+
self.assertEqual(key_pb, keys[0])
307307
self.assertEqual(request.read_options.transaction, TRANSACTION)
308308

309309
def test_lookup_single_key_nonempty_response(self):
@@ -335,7 +335,7 @@ def test_lookup_single_key_nonempty_response(self):
335335
request.ParseFromString(cw['body'])
336336
keys = list(request.key)
337337
self.assertEqual(len(keys), 1)
338-
_compare_key_pb_after_request(self, key_pb, keys[0])
338+
self.assertEqual(key_pb, keys[0])
339339

340340
def test_lookup_multiple_keys_empty_response(self):
341341
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
@@ -362,8 +362,8 @@ def test_lookup_multiple_keys_empty_response(self):
362362
request.ParseFromString(cw['body'])
363363
keys = list(request.key)
364364
self.assertEqual(len(keys), 2)
365-
_compare_key_pb_after_request(self, key_pb1, keys[0])
366-
_compare_key_pb_after_request(self, key_pb2, keys[1])
365+
self.assertEqual(key_pb1, keys[0])
366+
self.assertEqual(key_pb2, keys[1])
367367

368368
def test_lookup_multiple_keys_w_missing(self):
369369
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
@@ -398,8 +398,8 @@ def test_lookup_multiple_keys_w_missing(self):
398398
request.ParseFromString(cw['body'])
399399
keys = list(request.key)
400400
self.assertEqual(len(keys), 2)
401-
_compare_key_pb_after_request(self, key_pb1, keys[0])
402-
_compare_key_pb_after_request(self, key_pb2, keys[1])
401+
self.assertEqual(key_pb1, keys[0])
402+
self.assertEqual(key_pb2, keys[1])
403403

404404
def test_lookup_multiple_keys_w_missing_non_empty(self):
405405
DATASET_ID = 'DATASET'
@@ -446,8 +446,8 @@ def test_lookup_multiple_keys_w_deferred(self):
446446
request.ParseFromString(cw['body'])
447447
keys = list(request.key)
448448
self.assertEqual(len(keys), 2)
449-
_compare_key_pb_after_request(self, key_pb1, keys[0])
450-
_compare_key_pb_after_request(self, key_pb2, keys[1])
449+
self.assertEqual(key_pb1, keys[0])
450+
self.assertEqual(key_pb2, keys[1])
451451

452452
def test_lookup_multiple_keys_w_deferred_non_empty(self):
453453
DATASET_ID = 'DATASET'
@@ -502,8 +502,8 @@ def test_lookup_multiple_keys_w_deferred_from_backend_but_not_passed(self):
502502
request.ParseFromString(cw[0]['body'])
503503
keys = list(request.key)
504504
self.assertEqual(len(keys), 2)
505-
_compare_key_pb_after_request(self, key_pb1, keys[0])
506-
_compare_key_pb_after_request(self, key_pb2, keys[1])
505+
self.assertEqual(key_pb1, keys[0])
506+
self.assertEqual(key_pb2, keys[1])
507507

508508
self._verifyProtobufCall(cw[1], URI, conn)
509509
request.ParseFromString(cw[1]['body'])
@@ -904,7 +904,7 @@ def test_allocate_ids_non_empty(self):
904904
request.ParseFromString(cw['body'])
905905
self.assertEqual(len(request.key), len(before_key_pbs))
906906
for key_before, key_after in zip(before_key_pbs, request.key):
907-
_compare_key_pb_after_request(self, key_before, key_after)
907+
self.assertEqual(key_before, key_after)
908908

909909
def test_save_entity_wo_transaction_w_upsert(self):
910910
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
@@ -935,7 +935,7 @@ def test_save_entity_wo_transaction_w_upsert(self):
935935
upserts = list(mutation.upsert)
936936
self.assertEqual(len(upserts), 1)
937937
upsert = upserts[0]
938-
_compare_key_pb_after_request(self, key_pb, upsert.key)
938+
self.assertEqual(key_pb, upsert.key)
939939
props = list(upsert.property)
940940
self.assertEqual(len(props), 1)
941941
self.assertEqual(props[0].name, 'foo')
@@ -976,7 +976,7 @@ def test_save_entity_w_exclude_from_indexes(self):
976976
upserts = list(mutation.upsert)
977977
self.assertEqual(len(upserts), 1)
978978
upsert = upserts[0]
979-
_compare_key_pb_after_request(self, key_pb, upsert.key)
979+
self.assertEqual(key_pb, upsert.key)
980980
props = sorted(upsert.property,
981981
key=operator.attrgetter('name'),
982982
reverse=True)
@@ -1025,7 +1025,7 @@ def test_save_entity_wo_transaction_w_auto_id(self):
10251025
mutation = request.mutation
10261026
inserts = list(mutation.insert_auto_id)
10271027
insert = inserts[0]
1028-
_compare_key_pb_after_request(self, key_pb, insert.key)
1028+
self.assertEqual(key_pb, insert.key)
10291029
props = list(insert.property)
10301030
self.assertEqual(len(props), 1)
10311031
self.assertEqual(props[0].name, 'foo')
@@ -1124,7 +1124,7 @@ def test_delete_entities_wo_transaction(self):
11241124
deletes = list(mutation.delete)
11251125
self.assertEqual(len(deletes), 1)
11261126
delete = deletes[0]
1127-
_compare_key_pb_after_request(self, key_pb, delete)
1127+
self.assertEqual(key_pb, delete)
11281128
self.assertEqual(request.mode, rq_class.NON_TRANSACTIONAL)
11291129

11301130
def test_delete_entities_w_transaction(self):
@@ -1196,16 +1196,6 @@ def id(self):
11961196
return self._id
11971197

11981198

1199-
def _compare_key_pb_after_request(test, key_before, key_after):
1200-
test.assertFalse(key_after.partition_id.HasField('dataset_id'))
1201-
test.assertEqual(key_before.partition_id.namespace,
1202-
key_after.partition_id.namespace)
1203-
test.assertEqual(len(key_before.path_element),
1204-
len(key_after.path_element))
1205-
for elt1, elt2 in zip(key_before.path_element, key_after.path_element):
1206-
test.assertEqual(elt1, elt2)
1207-
1208-
12091199
class _Connection(object):
12101200
_called_with = None
12111201
_missing = _deferred = ()

0 commit comments

Comments
 (0)