diff --git a/gcloud/datastore/helpers.py b/gcloud/datastore/helpers.py index bb6f3fb0422e..47673d468869 100644 --- a/gcloud/datastore/helpers.py +++ b/gcloud/datastore/helpers.py @@ -33,6 +33,46 @@ INT_VALUE_CHECKER = Int64ValueChecker() +def find_true_dataset_id(dataset_id, connection): + """Find the true (unaliased) dataset ID. + + If the given ID already has a 's~' or 'e~' prefix, does nothing. + Otherwise, looks up a bogus Key('__MissingLookupKind', 1) and reads the + true prefixed dataset ID from the response (either from found or from + missing). + + For some context, see: + github.com/GoogleCloudPlatform/gcloud-python/pull/528 + github.com/GoogleCloudPlatform/google-cloud-datastore/issues/59 + + :type dataset_id: string + :param dataset_id: The dataset ID to un-alias / prefix. + + :type connection: :class:`gcloud.datastore.connection.Connection` + :param connection: A connection provided to connection to the dataset. + + :rtype: string + :returns: The true / prefixed / un-aliased dataset ID. + """ + if dataset_id.startswith('s~') or dataset_id.startswith('e~'): + return dataset_id + + # Create the bogus Key protobuf to be looked up and remove + # the dataset ID so the backend won't complain. + bogus_key_pb = Key('__MissingLookupKind', 1, + dataset_id=dataset_id).to_protobuf() + bogus_key_pb.partition_id.ClearField('dataset_id') + + found_pbs, missing_pbs, _ = connection.lookup(dataset_id, [bogus_key_pb]) + # By not passing in `deferred`, lookup will continue until + # all results are `found` or `missing`. + all_pbs = missing_pbs + found_pbs + # We only asked for one, so should only receive one. + returned_pb, = all_pbs + + return returned_pb.key.partition_id.dataset_id + + def entity_from_protobuf(pb): """Factory method for creating an entity based on a protobuf. diff --git a/gcloud/datastore/test_helpers.py b/gcloud/datastore/test_helpers.py index 61289adced4b..29914c919eb0 100644 --- a/gcloud/datastore/test_helpers.py +++ b/gcloud/datastore/test_helpers.py @@ -574,3 +574,96 @@ def test_prepare_dataset_id_unset(self): key = datastore_pb.Key() new_key = self._callFUT(key) self.assertTrue(new_key is key) + + +class Test_find_true_dataset_id(unittest2.TestCase): + + def setUp(self): + from gcloud.datastore._testing import _setup_defaults + _setup_defaults(self) + + def tearDown(self): + from gcloud.datastore._testing import _tear_down_defaults + _tear_down_defaults(self) + + def _callFUT(self, dataset_id, connection): + from gcloud.datastore.helpers import find_true_dataset_id + return find_true_dataset_id(dataset_id, connection) + + def test_prefixed(self): + PREFIXED = 's~DATASET' + result = self._callFUT(PREFIXED, object()) + self.assertEqual(PREFIXED, result) + + def test_unprefixed_bogus_key_miss(self): + UNPREFIXED = 'DATASET' + PREFIX = 's~' + CONNECTION = _Connection(PREFIX, from_missing=False) + result = self._callFUT(UNPREFIXED, CONNECTION) + + self.assertEqual(CONNECTION._called_dataset_id, UNPREFIXED) + + self.assertEqual(len(CONNECTION._lookup_result), 1) + + # Make sure just one. + called_key_pb, = CONNECTION._called_key_pbs + path_element = called_key_pb.path_element + self.assertEqual(len(path_element), 1) + self.assertEqual(path_element[0].kind, '__MissingLookupKind') + self.assertEqual(path_element[0].id, 1) + self.assertFalse(path_element[0].HasField('name')) + + PREFIXED = PREFIX + UNPREFIXED + self.assertEqual(result, PREFIXED) + + def test_unprefixed_bogus_key_hit(self): + UNPREFIXED = 'DATASET' + PREFIX = 'e~' + CONNECTION = _Connection(PREFIX, from_missing=True) + result = self._callFUT(UNPREFIXED, CONNECTION) + + self.assertEqual(CONNECTION._called_dataset_id, UNPREFIXED) + self.assertEqual(CONNECTION._lookup_result, []) + + # Make sure just one. + called_key_pb, = CONNECTION._called_key_pbs + path_element = called_key_pb.path_element + self.assertEqual(len(path_element), 1) + self.assertEqual(path_element[0].kind, '__MissingLookupKind') + self.assertEqual(path_element[0].id, 1) + self.assertFalse(path_element[0].HasField('name')) + + PREFIXED = PREFIX + UNPREFIXED + self.assertEqual(result, PREFIXED) + + +class _Connection(object): + + _called_dataset_id = _called_key_pbs = _lookup_result = None + + def __init__(self, prefix, from_missing=False): + self.prefix = prefix + self.from_missing = from_missing + + def lookup(self, dataset_id, key_pbs): + from gcloud.datastore import _datastore_v1_pb2 as datastore_pb + + # Store the arguments called with. + self._called_dataset_id = dataset_id + self._called_key_pbs = key_pbs + + key_pb, = key_pbs + + response = datastore_pb.Entity() + response.key.CopyFrom(key_pb) + response.key.partition_id.dataset_id = self.prefix + dataset_id + + missing = [] + deferred = [] + if self.from_missing: + missing[:] = [response] + self._lookup_result = [] + else: + self._lookup_result = [response] + + return self._lookup_result, missing, deferred