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.

Author: dhermes

gcloud/datastore/__init__.py

@@ -69,6 +69,50 @@
 _DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
 
 
+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, 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: Optional. A connection provided to be the default.
+
+    :rtype: string
+    :returns: The true / prefixed / un-aliased dataset ID.
+    """
+    if dataset_id.startswith('s~') or dataset_id.startswith('e~'):
+        return dataset_id
+
+    connection = connection or _implicit_environ.CONNECTION
+
+    # 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')
+
+    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 returned_pb.key.partition_id.dataset_id
+
+
 def set_default_dataset_id(dataset_id=None):
     """Set default dataset ID either explicitly or implicitly as fall-back.
 
@@ -91,6 +135,7 @@ def set_default_dataset_id(dataset_id=None):
         dataset_id = _implicit_environ.compute_engine_id()
 
     if dataset_id is not None:
+        dataset_id = _find_true_dataset_id(dataset_id)
         _implicit_environ.DATASET_ID = dataset_id
 
 
@@ -120,8 +165,9 @@ def set_defaults(dataset_id=None, connection=None):
     :type connection: :class:`gcloud.datastore.connection.Connection`
     :param connection: A connection provided to be the default.
     """
-    set_default_dataset_id(dataset_id=dataset_id)
+    # Set CONNECTION first in case _find_true_dataset_id needs it.
     set_default_connection(connection=connection)
+    set_default_dataset_id(dataset_id=dataset_id)
 
 
 def get_connection():

gcloud/datastore/batch.py

@@ -299,7 +299,6 @@ def _assign_entity_to_mutation(mutation_pb, entity, auto_id_entities):
     auto_id = entity.key.is_partial
 
     key_pb = entity.key.to_protobuf()
-    key_pb = helpers._prepare_key_for_request(key_pb)
 
     if auto_id:
         insert = mutation_pb.insert_auto_id.add()

gcloud/datastore/helpers.py

@@ -26,7 +26,6 @@
 import pytz
 import six
 
-from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 
@@ -280,33 +279,6 @@ def _set_protobuf_value(value_pb, val):
         setattr(value_pb, attr, val)
 
 
-def _prepare_key_for_request(key_pb):
-    """Add protobuf keys to a request object.
-
-    :type key_pb: :class:`gcloud.datastore._datastore_v1_pb2.Key`
-    :param key_pb: A key to be added to a request.
-
-    :rtype: :class:`gcloud.datastore._datastore_v1_pb2.Key`
-    :returns: A key which will be added to a request. It will be the
-              original if nothing needs to be changed.
-    """
-    if key_pb.partition_id.HasField('dataset_id'):
-        # We remove the dataset_id from the protobuf. This is because
-        # the backend fails a request if the key contains un-prefixed
-        # dataset ID. The backend fails because requests to
-        #     /datastore/.../datasets/foo/...
-        # and
-        #     /datastore/.../datasets/s~foo/...
-        # both go to the datastore given by 's~foo'. So if the key
-        # protobuf in the request body has dataset_id='foo', the
-        # backend will reject since 'foo' != 's~foo'.
-        new_key_pb = datastore_pb.Key()
-        new_key_pb.CopyFrom(key_pb)
-        new_key_pb.partition_id.ClearField('dataset_id')
-        key_pb = new_key_pb
-    return key_pb
-
-
 def _add_keys_to_request(request_field_pb, key_pbs):
     """Add protobuf keys to a request object.
 
@@ -317,5 +289,4 @@ def _add_keys_to_request(request_field_pb, key_pbs):
     :param key_pbs: The keys to add to a request.
     """
     for key_pb in key_pbs:
-        key_pb = _prepare_key_for_request(key_pb)
         request_field_pb.add().CopyFrom(key_pb)

gcloud/datastore/query.py

@@ -449,8 +449,7 @@ def _pb_from_query(query):
     composite_filter.operator = datastore_pb.CompositeFilter.AND
 
     if query.ancestor:
-        ancestor_pb = helpers._prepare_key_for_request(
-            query.ancestor.to_protobuf())
+        ancestor_pb = query.ancestor.to_protobuf()
 
         # Filter on __key__ HAS_ANCESTOR == ancestor.
         ancestor_filter = composite_filter.filter.add().property_filter
@@ -469,8 +468,7 @@ def _pb_from_query(query):
         # Set the value to filter on based on the type.
         if property_name == '__key__':
             key_pb = value.to_protobuf()
-            property_filter.value.key_value.CopyFrom(
-                helpers._prepare_key_for_request(key_pb))
+            property_filter.value.key_value.CopyFrom(key_pb)
         else:
             helpers._set_protobuf_value(property_filter.value, value)
 

gcloud/datastore/test___init__.py

@@ -45,22 +45,22 @@ def test_no_env_var_set(self):
 
     def test_set_from_env_var(self):
         from gcloud.datastore import _implicit_environ
-        IMPLICIT_DATASET_ID = 'IMPLICIT'
+        IMPLICIT_DATASET_ID = 's~IMPLICIT'
         with self._monkey(IMPLICIT_DATASET_ID):
             self._callFUT()
         self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
 
     def test_set_explicit_w_env_var_set(self):
         from gcloud.datastore import _implicit_environ
-        EXPLICIT_DATASET_ID = 'EXPLICIT'
+        EXPLICIT_DATASET_ID = 's~EXPLICIT'
         with self._monkey(None):
             self._callFUT(EXPLICIT_DATASET_ID)
         self.assertEqual(_implicit_environ.DATASET_ID, EXPLICIT_DATASET_ID)
 
     def test_set_explicit_no_env_var_set(self):
         from gcloud.datastore import _implicit_environ
-        IMPLICIT_DATASET_ID = 'IMPLICIT'
-        EXPLICIT_DATASET_ID = 'EXPLICIT'
+        IMPLICIT_DATASET_ID = 's~IMPLICIT'
+        EXPLICIT_DATASET_ID = 's~EXPLICIT'
         with self._monkey(IMPLICIT_DATASET_ID):
             self._callFUT(EXPLICIT_DATASET_ID)
         self.assertEqual(_implicit_environ.DATASET_ID, EXPLICIT_DATASET_ID)
@@ -73,7 +73,7 @@ def test_set_explicit_None_wo_env_var_set(self):
 
     def test_set_explicit_None_w_env_var_set(self):
         from gcloud.datastore import _implicit_environ
-        IMPLICIT_DATASET_ID = 'IMPLICIT'
+        IMPLICIT_DATASET_ID = 's~IMPLICIT'
         with self._monkey(IMPLICIT_DATASET_ID):
             self._callFUT(None)
         self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
@@ -82,7 +82,7 @@ def test_set_implicit_from_appengine(self):
         from gcloud._testing import _Monkey
         from gcloud.datastore import _implicit_environ
 
-        APP_ENGINE_ID = 'GAE'
+        APP_ENGINE_ID = 's~GAE'
         APP_IDENTITY = _AppIdentity(APP_ENGINE_ID)
 
         with self._monkey(None):
@@ -95,8 +95,8 @@ def test_set_implicit_both_env_and_appengine(self):
         from gcloud._testing import _Monkey
         from gcloud.datastore import _implicit_environ
 
-        IMPLICIT_DATASET_ID = 'IMPLICIT'
-        APP_IDENTITY = _AppIdentity('GAE')
+        IMPLICIT_DATASET_ID = 's~IMPLICIT'
+        APP_IDENTITY = _AppIdentity('s~GAE')
 
         with self._monkey(IMPLICIT_DATASET_ID):
             with _Monkey(_implicit_environ, app_identity=APP_IDENTITY):
@@ -108,7 +108,7 @@ def _implicit_compute_engine_helper(self, status):
         from gcloud._testing import _Monkey
         from gcloud.datastore import _implicit_environ
 
-        COMPUTE_ENGINE_ID = 'GCE'
+        COMPUTE_ENGINE_ID = 's~GCE'
         HTTPLIB2 = _Httplib2(COMPUTE_ENGINE_ID, status=status)
         if status == '200':
             EXPECTED_ID = COMPUTE_ENGINE_ID
@@ -148,9 +148,9 @@ def test_set_implicit_both_appengine_and_compute(self):
         from gcloud._testing import _Monkey
         from gcloud.datastore import _implicit_environ
 
-        APP_ENGINE_ID = 'GAE'
+        APP_ENGINE_ID = 's~GAE'
         APP_IDENTITY = _AppIdentity(APP_ENGINE_ID)
-        HTTPLIB2 = _Httplib2('GCE')
+        HTTPLIB2 = _Httplib2('s~GCE')
 
         with self._monkey(None):
             with _Monkey(_implicit_environ, app_identity=APP_IDENTITY,
@@ -164,9 +164,9 @@ def test_set_implicit_three_env_appengine_and_compute(self):
         from gcloud._testing import _Monkey
         from gcloud.datastore import _implicit_environ
 
-        IMPLICIT_DATASET_ID = 'IMPLICIT'
-        APP_IDENTITY = _AppIdentity('GAE')
-        HTTPLIB2 = _Httplib2('GCE')
+        IMPLICIT_DATASET_ID = 's~IMPLICIT'
+        APP_IDENTITY = _AppIdentity('s~GAE')
+        HTTPLIB2 = _Httplib2('s~GCE')
 
         with self._monkey(IMPLICIT_DATASET_ID):
             with _Monkey(_implicit_environ, app_identity=APP_IDENTITY,
@@ -177,6 +177,71 @@ def test_set_implicit_three_env_appengine_and_compute(self):
         self.assertEqual(len(HTTPLIB2._http_instances), 0)
 
 
+class Test__find_true_dataset_id(unittest2.TestCase):
+
+    def _callFUT(self, dataset_id, connection=None):
+        from gcloud.datastore import _find_true_dataset_id
+        return _find_true_dataset_id(dataset_id, connection=connection)
+
+    def test_prefixed(self):
+        PREFIXED = 's~DATASET'
+        result = self._callFUT(PREFIXED)
+        self.assertEqual(PREFIXED, result)
+
+    def test_unprefixed_no_connection(self):
+        from gcloud.datastore import _implicit_environ
+
+        UNPREFIXED = 'DATASET'
+        self.assertEqual(_implicit_environ.CONNECTION, None)
+        with self.assertRaises(AttributeError):
+            self._callFUT(UNPREFIXED)
+
+    def test_unprefixed_bogus_key_miss(self):
+        UNPREFIXED = 'DATASET'
+        PREFIX = 's~'
+        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, '__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=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, '__MissingLookupKind')
+        self.assertEqual(path_element[0].id, 1)
+        self.assertFalse(path_element[0].HasField('name'))
+
+        PREFIXED = PREFIX + UNPREFIXED
+        self.assertEqual(result, PREFIXED)
+
+
 class Test_set_default_connection(unittest2.TestCase):
 
     def setUp(self):
@@ -304,3 +369,32 @@ def Http(self, timeout=None):
         result = _Http(self)
         self._http_instances.append((timeout, result))
         return result
+
+
+class _Connection(object):
+
+    def __init__(self, prefix, from_missing=False):
+        self.prefix = prefix
+        self.from_missing = from_missing
+
+    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.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