Merge pull request #897 from dhermes/fix-701-2nd-attempt

Implementing get_multi in datastore.

Author: dhermes

gcloud/datastore/__init__.py

@@ -56,6 +56,7 @@
 from gcloud.datastore.api import allocate_ids
 from gcloud.datastore.api import delete
 from gcloud.datastore.api import get
+from gcloud.datastore.api import get_multi
 from gcloud.datastore.api import put
 from gcloud.datastore.batch import Batch
 from gcloud.datastore.connection import SCOPE

gcloud/datastore/api.py

@@ -166,7 +166,8 @@ def _extended_lookup(connection, dataset_id, key_pbs,
     return results
 
 
-def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
+def get_multi(keys, missing=None, deferred=None,
+              connection=None, dataset_id=None):
     """Retrieves entities, along with their attributes.
 
     :type keys: list of :class:`gcloud.datastore.key.Key`
@@ -234,6 +235,45 @@ def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
     return entities
 
 
+def get(key, missing=None, deferred=None, connection=None, dataset_id=None):
+    """Retrieves entity from a single key (if it exists).
+
+    .. note::
+
+       This is just a thin wrapper over :func:`gcloud.datastore.get_multi`.
+       The backend API does not make a distinction between a single key or
+       multiple keys in a lookup request.
+
+    :type key: :class:`gcloud.datastore.key.Key`
+    :param key: The key to be retrieved from the datastore.
+
+    :type missing: an empty list or None.
+    :param missing: If a list is passed, the key-only entities returned
+                    by the backend as "missing" will be copied into it.
+                    Use only as a keyword param.
+
+    :type deferred: an empty list or None.
+    :param deferred: If a list is passed, the keys returned
+                     by the backend as "deferred" will be copied into it.
+                     Use only as a keyword param.
+
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: Optional. The connection used to connect to datastore.
+                       If not passed, inferred from the environment.
+
+    :type dataset_id: :class:`gcloud.datastore.connection.Connection`
+    :param dataset_id: Optional. The dataset ID used to connect to datastore.
+                       If not passed, inferred from the environment.
+
+    :rtype: :class:`gcloud.datastore.entity.Entity` or ``NoneType``
+    :returns: The requested entity if it exists.
+    """
+    entities = get_multi([key], missing=missing, deferred=deferred,
+                         connection=connection, dataset_id=dataset_id)
+    if entities:
+        return entities[0]
+
+
 def put(entities, connection=None, dataset_id=None):
     """Save the entities in the Cloud Datastore.
 

gcloud/datastore/connection.py

@@ -153,7 +153,7 @@ def lookup(self, dataset_id, key_pbs,
 
         >>> from gcloud import datastore
         >>> key = datastore.Key('MyKind', 1234, dataset_id='dataset-id')
-        >>> datastore.get([key])
+        >>> datastore.get(key)
         [<Entity object>]
 
         Using the ``connection`` class directly:

gcloud/datastore/dataset.py

@@ -15,6 +15,7 @@
 
 from gcloud.datastore.api import delete
 from gcloud.datastore.api import get
+from gcloud.datastore.api import get_multi
 from gcloud.datastore.api import put
 from gcloud.datastore.batch import Batch
 from gcloud.datastore.key import Key
@@ -38,14 +39,23 @@ def __init__(self, dataset_id, connection=None):
         self.dataset_id = dataset_id
         self.connection = connection
 
-    def get(self, keys, missing=None, deferred=None):
+    def get(self, key, missing=None, deferred=None):
         """Proxy to :func:`gcloud.datastore.api.get`.
 
         Passes our ``dataset_id``.
         """
-        return get(keys, missing=missing, deferred=deferred,
+        return get(key, missing=missing, deferred=deferred,
                    connection=self.connection, dataset_id=self.dataset_id)
 
+    def get_multi(self, keys, missing=None, deferred=None):
+        """Proxy to :func:`gcloud.datastore.api.get_multi`.
+
+        Passes our ``dataset_id``.
+        """
+        return get_multi(keys, missing=missing, deferred=deferred,
+                         connection=self.connection,
+                         dataset_id=self.dataset_id)
+
     def put(self, entities):
         """Proxy to :func:`gcloud.datastore.api.put`.
 

gcloud/datastore/demo/demo.py

@@ -33,13 +33,13 @@
 datastore.put([toy])
 
 # If we look it up by its key, we should find it...
-print(datastore.get([toy.key]))
+print(datastore.get(toy.key))
 
 # And we should be able to delete it...
 datastore.delete([toy.key])
 
 # Since we deleted it, if we do another lookup it shouldn't be there again:
-print(datastore.get([toy.key]))
+print(datastore.get(toy.key))
 
 # Now let's try a more advanced query.
 # First, let's create some entities.
@@ -104,7 +104,7 @@
     xact.rollback()
 
 # Let's check if the entity was actually created:
-created = datastore.get([key])
+created = datastore.get(key)
 print('yes' if created else 'no')
 
 # Remember, a key won't be complete until the transaction is commited.

gcloud/datastore/test_api.py

@@ -15,6 +15,22 @@
 import unittest2
 
 
+def _make_entity_pb(dataset_id, kind, integer_id, name=None, str_val=None):
+    from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
+
+    entity_pb = datastore_pb.Entity()
+    entity_pb.key.partition_id.dataset_id = dataset_id
+    path_element = entity_pb.key.path_element.add()
+    path_element.kind = kind
+    path_element.id = integer_id
+    if name is not None and str_val is not None:
+        prop = entity_pb.property.add()
+        prop.name = name
+        prop.value.string_value = str_val
+
+    return entity_pb
+
+
 class Test__require_dataset_id(unittest2.TestCase):
 
     _MARKER = object()
@@ -158,7 +174,7 @@ def test_implicit_set_passed_explicitly(self):
             self.assertTrue(self._callFUT(CONNECTION) is CONNECTION)
 
 
-class Test_get_function(unittest2.TestCase):
+class Test_get_multi_function(unittest2.TestCase):
 
     def setUp(self):
         from gcloud.datastore._testing import _setup_defaults
@@ -170,25 +186,9 @@ def tearDown(self):
 
     def _callFUT(self, keys, missing=None, deferred=None,
                  connection=None, dataset_id=None):
-        from gcloud.datastore.api import get
-        return get(keys, missing=missing, deferred=deferred,
-                   connection=connection, dataset_id=dataset_id)
-
-    def _make_entity_pb(self, dataset_id, kind, integer_id,
-                        name=None, str_val=None):
-        from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
-
-        entity_pb = datastore_pb.Entity()
-        entity_pb.key.partition_id.dataset_id = dataset_id
-        path_element = entity_pb.key.path_element.add()
-        path_element.kind = kind
-        path_element.id = integer_id
-        if name is not None and str_val is not None:
-            prop = entity_pb.property.add()
-            prop.name = name
-            prop.value.string_value = str_val
-
-        return entity_pb
+        from gcloud.datastore.api import get_multi
+        return get_multi(keys, missing=missing, deferred=deferred,
+                         connection=connection, dataset_id=dataset_id)
 
     def test_wo_connection(self):
         from gcloud.datastore.key import Key
@@ -398,8 +398,7 @@ def test_hit(self):
         PATH = [{'kind': KIND, 'id': ID}]
 
         # Make a found entity pb to be returned from mock backend.
-        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
-                                         'foo', 'Foo')
+        entity_pb = _make_entity_pb(DATASET_ID, KIND, ID, 'foo', 'Foo')
 
         # Make a connection to return the entity pb.
         connection = _Connection(entity_pb)
@@ -426,8 +425,8 @@ def test_hit_multiple_keys_same_dataset(self):
         ID2 = 2345
 
         # Make a found entity pb to be returned from mock backend.
-        entity_pb1 = self._make_entity_pb(DATASET_ID, KIND, ID1)
-        entity_pb2 = self._make_entity_pb(DATASET_ID, KIND, ID2)
+        entity_pb1 = _make_entity_pb(DATASET_ID, KIND, ID1)
+        entity_pb2 = _make_entity_pb(DATASET_ID, KIND, ID2)
 
         # Make a connection to return the entity pbs.
         connection = _Connection(entity_pb1, entity_pb2)
@@ -469,8 +468,7 @@ def test_implicit_wo_transaction(self):
         PATH = [{'kind': KIND, 'id': ID}]
 
         # Make a found entity pb to be returned from mock backend.
-        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
-                                         'foo', 'Foo')
+        entity_pb = _make_entity_pb(DATASET_ID, KIND, ID, 'foo', 'Foo')
 
         # Make a connection to return the entity pb.
         CUSTOM_CONNECTION = _Connection(entity_pb)
@@ -507,8 +505,7 @@ def test_w_transaction(self):
         TRANSACTION = 'TRANSACTION'
 
         # Make a found entity pb to be returned from mock backend.
-        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
-                                         'foo', 'Foo')
+        entity_pb = _make_entity_pb(DATASET_ID, KIND, ID, 'foo', 'Foo')
 
         # Make a connection to return the entity pb.
         CUSTOM_CONNECTION = _Connection(entity_pb)
@@ -545,8 +542,7 @@ def test_max_loops(self):
         ID = 1234
 
         # Make a found entity pb to be returned from mock backend.
-        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
-                                         'foo', 'Foo')
+        entity_pb = _make_entity_pb(DATASET_ID, KIND, ID, 'foo', 'Foo')
 
         # Make a connection to return the entity pb.
         connection = _Connection(entity_pb)
@@ -566,6 +562,61 @@ def test_max_loops(self):
         self.assertEqual(deferred, [])
 
 
+class Test_get_function(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, key, missing=None, deferred=None,
+                 connection=None, dataset_id=None):
+        from gcloud.datastore.api import get
+        return get(key, missing=missing, deferred=deferred,
+                   connection=connection, dataset_id=dataset_id)
+
+    def test_hit(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+
+        # Make a found entity pb to be returned from mock backend.
+        entity_pb = _make_entity_pb(DATASET_ID, KIND, ID, 'foo', 'Foo')
+
+        # Make a connection to return the entity pb.
+        connection = _Connection(entity_pb)
+
+        key = Key(KIND, ID, dataset_id=DATASET_ID)
+        result = self._callFUT(key, connection=connection,
+                               dataset_id=DATASET_ID)
+        new_key = result.key
+
+        # Check the returned value is as expected.
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.dataset_id, DATASET_ID)
+        self.assertEqual(new_key.path, PATH)
+        self.assertEqual(list(result), ['foo'])
+        self.assertEqual(result['foo'], 'Foo')
+
+    def test_miss(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        connection = _Connection()
+        key = Key('Kind', 1234, dataset_id=DATASET_ID)
+        result = self._callFUT(key, connection=connection,
+                               dataset_id=DATASET_ID)
+        self.assertTrue(result is None)
+
+
 class Test_put_function(unittest2.TestCase):
 
     def setUp(self):

gcloud/datastore/test_dataset.py

@@ -52,9 +52,9 @@ def _get(*args, **kw):
         key = object()
 
         with _Monkey(MUT, get=_get):
-            dataset.get([key])
+            dataset.get(key)
 
-        self.assertEqual(_called_with[0][0], ([key],))
+        self.assertEqual(_called_with[0][0], (key,))
         self.assertTrue(_called_with[0][1]['missing'] is None)
         self.assertTrue(_called_with[0][1]['deferred'] is None)
         self.assertTrue(_called_with[0][1]['connection'] is None)
@@ -74,7 +74,50 @@ def _get(*args, **kw):
         key, missing, deferred = object(), [], []
 
         with _Monkey(MUT, get=_get):
-            dataset.get([key], missing, deferred)
+            dataset.get(key, missing, deferred)
+
+        self.assertEqual(_called_with[0][0], (key,))
+        self.assertTrue(_called_with[0][1]['missing'] is missing)
+        self.assertTrue(_called_with[0][1]['deferred'] is deferred)
+        self.assertTrue(_called_with[0][1]['connection'] is conn)
+        self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
+
+    def test_get_multi_defaults(self):
+        from gcloud.datastore import dataset as MUT
+        from gcloud._testing import _Monkey
+
+        _called_with = []
+
+        def _get_multi(*args, **kw):
+            _called_with.append((args, kw))
+
+        dataset = self._makeOne()
+        key = object()
+
+        with _Monkey(MUT, get_multi=_get_multi):
+            dataset.get_multi([key])
+
+        self.assertEqual(_called_with[0][0], ([key],))
+        self.assertTrue(_called_with[0][1]['missing'] is None)
+        self.assertTrue(_called_with[0][1]['deferred'] is None)
+        self.assertTrue(_called_with[0][1]['connection'] is None)
+        self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
+
+    def test_get_multi_explicit(self):
+        from gcloud.datastore import dataset as MUT
+        from gcloud._testing import _Monkey
+
+        _called_with = []
+
+        def _get_multi(*args, **kw):
+            _called_with.append((args, kw))
+
+        conn = object()
+        dataset = self._makeOne(connection=conn)
+        key, missing, deferred = object(), [], []
+
+        with _Monkey(MUT, get_multi=_get_multi):
+            dataset.get_multi([key], missing, deferred)
 
         self.assertEqual(_called_with[0][0], ([key],))
         self.assertTrue(_called_with[0][1]['missing'] is missing)