From 01d0354f4d8b3a34af0967f06ec29cf9552813bb Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Tue, 31 Mar 2015 18:08:52 -0700
Subject: [PATCH] Adding one-time RPC to find unaliased / true dataset ID.

---
 gcloud/datastore/helpers.py      | 40 ++++++++++++++
 gcloud/datastore/test_helpers.py | 93 ++++++++++++++++++++++++++++++++
 2 files changed, 133 insertions(+)

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