From ab4c10308efd9605d305a42763c49d1c8dbda7c5 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Fri, 9 Jan 2015 11:26:44 -0800
Subject: [PATCH] Using lookup() to determine true dataset instead of
 allocate_ids().

---
 gcloud/datastore/__init__.py      | 24 ++++++++++----
 gcloud/datastore/test___init__.py | 54 +++++++++++++++++++++++++------
 2 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
index b8973da666d60..3a97fa977cb61 100644
--- a/gcloud/datastore/__init__.py
+++ b/gcloud/datastore/__init__.py
@@ -72,8 +72,13 @@ def _find_true_dataset_id(dataset_id, connection=None):
     """Find the true (unaliased) dataset ID.
 
     If the given ID already has a 's~' or 'e~' prefix, does nothing.
-    Otherwise, allocates a single key of bogus kind '__Foo' and reads
-    the true prefixed dataset ID from the response.
+    Otherwise, looks up a bogus Key('__Foo', 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.
@@ -89,14 +94,21 @@ def _find_true_dataset_id(dataset_id, connection=None):
 
     connection = connection or _implicit_environ.CONNECTION
 
-    # Create the partial Key protobuf to be allocated and remove
+    # Create the bogus Key protobuf to be looked up and remove
     # the dataset ID so the backend won't complain.
-    bogus_key_pb = Key('__Foo', dataset_id=dataset_id).to_protobuf()
+    bogus_key_pb = Key('__Foo', 1, dataset_id=dataset_id).to_protobuf()
     bogus_key_pb.partition_id.ClearField('dataset_id')
 
-    allocated_pb, = connection.allocate_ids(dataset_id, [bogus_key_pb])
+    missing_pbs = []
+    found_pbs = connection.lookup(dataset_id, [bogus_key_pb],
+                                  missing=missing_pbs)
+    # 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 allocated_pb.partition_id.dataset_id
+    return returned_pb.key.partition_id.dataset_id
 
 
 def set_default_dataset_id(dataset_id=None):
diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
index 6e4d1ca440958..29b24abd954e9 100644
--- a/gcloud/datastore/test___init__.py
+++ b/gcloud/datastore/test___init__.py
@@ -98,19 +98,46 @@ def test_unprefixed_no_connection(self):
         with self.assertRaises(AttributeError):
             self._callFUT(UNPREFIXED)
 
-    def test_unprefixed_explicit_connection(self):
+    def test_unprefixed_bogus_key_miss(self):
         UNPREFIXED = 'DATASET'
         PREFIX = 's~'
-        CONNECTION = _Connection(PREFIX)
+        CONNECTION = _Connection(PREFIX, from_missing=False)
         result = self._callFUT(UNPREFIXED, connection=CONNECTION)
 
         self.assertEqual(CONNECTION._called_dataset_id, UNPREFIXED)
+
+        self.assertEqual(len(CONNECTION._lookup_result), 1)
+        self.assertEqual(CONNECTION._called_missing, [])
+
+        # 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, '__Foo')
+        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=CONNECTION)
+
+        self.assertEqual(CONNECTION._called_dataset_id, UNPREFIXED)
+
+        self.assertEqual(CONNECTION._lookup_result, [])
+        # Though missing=[] was called, it is copied in place.
+        self.assertEqual(len(CONNECTION._called_missing), 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, '__Foo')
-        self.assertFalse(path_element[0].HasField('id'))
+        self.assertEqual(path_element[0].id, 1)
         self.assertFalse(path_element[0].HasField('name'))
 
         PREFIXED = PREFIX + UNPREFIXED
@@ -207,19 +234,28 @@ def test_it(self):
 
 class _Connection(object):
 
-    def __init__(self, prefix):
+    def __init__(self, prefix, from_missing=False):
         self.prefix = prefix
+        self.from_missing = from_missing
 
-    def allocate_ids(self, dataset_id, key_pbs):
+    def lookup(self, dataset_id, key_pbs, missing=None):
         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
+        self._called_missing = missing
 
         key_pb, = key_pbs
 
-        response = datastore_pb.Key()
-        response.CopyFrom(key_pb)
-        response.partition_id.dataset_id = self.prefix + dataset_id
-        return [response]
+        response = datastore_pb.Entity()
+        response.key.CopyFrom(key_pb)
+        response.key.partition_id.dataset_id = self.prefix + dataset_id
+
+        if self.from_missing:
+            missing[:] = [response]
+            self._lookup_result = []
+        else:
+            self._lookup_result = [response]
+
+        return self._lookup_result