Moves Dataset.allocate_ids to connection module.

Addresses fifth part of googleapis#477.

gcloud/datastore/__init__.py

@@ -48,7 +48,7 @@
 
 from gcloud import credentials
 from gcloud.datastore import _implicit_environ
-from gcloud.datastore.connection import Connection
+from gcloud.datastore import connection as connection_module
 
 
 SCOPE = ('https://www.googleapis.com/auth/datastore ',
@@ -104,7 +104,7 @@ def get_connection():
     """
     implicit_credentials = credentials.get_credentials()
     scoped_credentials = implicit_credentials.create_scoped(SCOPE)
-    return Connection(credentials=scoped_credentials)
+    return connection_module.Connection(credentials=scoped_credentials)
 
 
 def get_dataset(dataset_id):
@@ -156,18 +156,3 @@ def get_entities(keys):
     :returns: The requested entities.
     """
     return _require_dataset().get_entities(keys)
-
-
-def allocate_ids(incomplete_key, num_ids):
-    """Allocates a list of IDs from a partial key.
-
-    :type incomplete_key: A :class:`gcloud.datastore.key.Key`
-    :param incomplete_key: The partial key to use as base for allocated IDs.
-
-    :type num_ids: A :class:`int`.
-    :param num_ids: The number of IDs to allocate.
-
-    :rtype: list of :class:`gcloud.datastore.key.Key`
-    :returns: The (complete) keys allocated with `incomplete_key` as root.
-    """
-    return _require_dataset().allocate_ids(incomplete_key, num_ids)

gcloud/datastore/connection.py

@@ -14,13 +14,13 @@
 
 """Connections to gcloud datastore API servers."""
 
-from gcloud import connection
+from gcloud import connection as base_connection
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore import helpers
 from gcloud.datastore.dataset import Dataset
 
 
-class Connection(connection.Connection):
+class Connection(base_connection.Connection):
     """A connection to the Google Cloud Datastore via the Protobuf API.
 
     This class should understand only the basic types (and protobufs)
@@ -125,7 +125,7 @@ def build_api_url(cls, dataset_id, method, base_url=None,
             api_version=(api_version or cls.API_VERSION),
             dataset_id=dataset_id, method=method)
 
-    def transaction(self, transaction=connection.Connection._EMPTY):
+    def transaction(self, transaction=base_connection.Connection._EMPTY):
         """Getter/setter for the connection's transaction object.
 
         :type transaction: :class:`gcloud.datastore.transaction.Transaction`,
@@ -575,3 +575,35 @@ def _copy_deferred_keys(lookup_request, lookup_response):
         lookup_request.key.remove(old_key)
     for def_key in lookup_response.deferred:
         lookup_request.key.add().CopyFrom(def_key)
+
+
+def allocate_ids(incomplete_key, num_ids, connection, dataset_id):
+    """Allocates a list of IDs from a partial key.
+
+    :type incomplete_key: A :class:`gcloud.datastore.key.Key`
+    :param incomplete_key: Partial key to use as base for allocated IDs.
+
+    :type num_ids: A :class:`int`.
+    :param num_ids: The number of IDs to allocate.
+
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: The connection used to allocate IDs.
+
+    :type dataset_id: :class:`str`.
+    :param dataset_id: The ID of the dataset used to allocate.
+
+    :rtype: list of :class:`gcloud.datastore.key.Key`
+    :returns: The (complete) keys allocated with `incomplete_key` as root.
+    :raises: `ValueError` if `incomplete_key` is not a partial key.
+    """
+    if not incomplete_key.is_partial:
+        raise ValueError(('Key is not partial.', incomplete_key))
+
+    incomplete_key_pb = incomplete_key.to_protobuf()
+    incomplete_key_pbs = [incomplete_key_pb] * num_ids
+
+    allocated_key_pbs = connection.allocate_ids(dataset_id, incomplete_key_pbs)
+    allocated_ids = [allocated_key_pb.path_element[-1].id
+                     for allocated_key_pb in allocated_key_pbs]
+    return [incomplete_key.completed_key(allocated_id)
+            for allocated_id in allocated_ids]

gcloud/datastore/dataset.py

@@ -112,29 +112,3 @@ def get_entities(self, keys, missing=None, deferred=None):
             entities.append(helpers.entity_from_protobuf(
                 entity_pb, dataset=self))
         return entities
-
-    def allocate_ids(self, incomplete_key, num_ids):
-        """Allocates a list of IDs from a partial key.
-
-        :type incomplete_key: A :class:`gcloud.datastore.key.Key`
-        :param incomplete_key: Partial key to use as base for allocated IDs.
-
-        :type num_ids: A :class:`int`.
-        :param num_ids: The number of IDs to allocate.
-
-        :rtype: list of :class:`gcloud.datastore.key.Key`
-        :returns: The (complete) keys allocated with `incomplete_key` as root.
-        :raises: `ValueError` if `incomplete_key` is not a partial key.
-        """
-        if not incomplete_key.is_partial:
-            raise ValueError(('Key is not partial.', incomplete_key))
-
-        incomplete_key_pb = incomplete_key.to_protobuf()
-        incomplete_key_pbs = [incomplete_key_pb] * num_ids
-
-        allocated_key_pbs = self.connection().allocate_ids(
-            self.id(), incomplete_key_pbs)
-        allocated_ids = [allocated_key_pb.path_element[-1].id
-                         for allocated_key_pb in allocated_key_pbs]
-        return [incomplete_key.completed_key(allocated_id)
-                for allocated_id in allocated_ids]

gcloud/datastore/test___init__.py

@@ -181,19 +181,3 @@ def test_get_entities(self):
         with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
             result = gcloud.datastore.get_entities(DUMMY_KEYS)
         self.assertTrue(result == DUMMY_VALS)
-
-    def test_allocate_ids(self):
-        import gcloud.datastore
-        from gcloud.datastore import _implicit_environ
-        from gcloud.datastore.key import Key
-        from gcloud.datastore.test_entity import _Dataset
-        from gcloud._testing import _Monkey
-
-        CUSTOM_DATASET = _Dataset()
-        NUM_IDS = 2
-        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
-            INCOMPLETE_KEY = Key('KIND')
-            result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
-
-        # Check the IDs returned.
-        self.assertEqual([key.id for key in result], range(1, NUM_IDS + 1))

gcloud/datastore/test_connection.py

@@ -1139,6 +1139,38 @@ def mutation(self):
         self.assertEqual(len(mutation.delete), 1)
 
 
+class Test_allocate_ids_function(unittest2.TestCase):
+
+    def _callFUT(self, incomplete_key, num_ids, connection, dataset_id):
+        from gcloud.datastore.connection import allocate_ids
+        return allocate_ids(incomplete_key, num_ids, connection, dataset_id)
+
+    def test_allocate_ids(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_dataset import _Connection
+
+        DATASET_ID = 'DATASET'
+        INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
+        CONNECTION = _Connection()
+        NUM_IDS = 2
+        result = self._callFUT(INCOMPLETE_KEY, NUM_IDS,
+                               CONNECTION, DATASET_ID)
+
+        # Check the IDs returned match.
+        self.assertEqual([key.id for key in result], range(NUM_IDS))
+
+        # Check connection is called correctly.
+        self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
+        self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
+
+    def test_allocate_ids_with_complete(self):
+        from gcloud.datastore.test_entity import _Key
+
+        COMPLETE_KEY = _Key()
+        self.assertRaises(ValueError, self._callFUT,
+                          COMPLETE_KEY, 2, None, None)
+
+
 class Http(object):
 
     _called_with = None

gcloud/datastore/test_dataset.py

@@ -108,31 +108,6 @@ def test_get_entities_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
-    def test_allocate_ids(self):
-        from gcloud.datastore.key import Key
-
-        DATASET_ID = 'DATASET'
-        INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
-        CONNECTION = _Connection()
-        NUM_IDS = 2
-        DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
-        result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
-
-        # Check the IDs returned match.
-        self.assertEqual([key.id for key in result], range(NUM_IDS))
-
-        # Check connection is called correctly.
-        self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
-        self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
-
-    def test_allocate_ids_with_complete(self):
-        from gcloud.datastore.test_entity import _Key
-
-        COMPLETE_KEY = _Key()
-        DATASET = self._makeOne(None)
-        self.assertRaises(ValueError, DATASET.allocate_ids,
-                          COMPLETE_KEY, 2)
-
 
 class _Connection(object):
     _called_with = None

gcloud/datastore/test_entity.py

@@ -274,9 +274,6 @@ def connection(self):
     def get_entities(self, keys):
         return [self.get(key) for key in keys]
 
-    def allocate_ids(self, incomplete_key, num_ids):
-        return [incomplete_key.completed_key(i + 1) for i in range(num_ids)]
-
 
 class _Connection(object):
     _transaction = _saved = _deleted = None

regression/datastore.py

@@ -18,6 +18,7 @@
 import unittest2
 
 from gcloud import datastore
+from gcloud.datastore.connection import allocate_ids
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 from gcloud.datastore.query import Query
@@ -48,7 +49,9 @@ class TestDatastoreAllocateIDs(TestDatastore):
     def test_allocate_ids(self):
         incomplete_key = Key('Kind')
         num_ids = 10
-        allocated_keys = datastore.allocate_ids(incomplete_key, num_ids)
+        connection = datastore.get_connection()
+        allocated_keys = allocate_ids(incomplete_key, num_ids,
+                                      connection, DATASET_ID)
         self.assertEqual(len(allocated_keys), num_ids)
 
         unique_ids = set()