Datastore: expose 'reserveIds' API.

datastore/google/cloud/datastore/client.py

@@ -648,3 +648,30 @@ def do_something(entity):
         if "namespace" not in kwargs:
             kwargs["namespace"] = self.namespace
         return Query(self, **kwargs)
+
+    def reserve_ids(self, complete_key, num_ids):
+        """Reserve a list of IDs from a complete key.
+
+        :type complete_key: :class:`google.cloud.datastore.key.Key`
+        :param complete_key: Partial key to use as base for reserved IDs.
+
+        :type num_ids: int
+        :param num_ids: The number of IDs to reserve.
+
+        :rtype: class:`NoneType`
+        :returns: None
+        :raises: :class:`ValueError` if `complete_key`` is not a
+                 Complete key.
+        """
+        if complete_key.is_partial:
+            raise ValueError(("Key is not Complete.", complete_key))
+
+        if not isinstance(num_ids, int):
+            raise ValueError(("num_ids is not a valid integer.", num_ids))
+
+        complete_key_pb = complete_key.to_protobuf()
+        complete_key_pbs = [complete_key_pb] * num_ids
+
+        self._datastore_api.reserve_ids(complete_key.project, complete_key_pbs)
+
+        return None

datastore/tests/unit/test_client.py

@@ -860,12 +860,42 @@ def test_allocate_ids_w_partial_key(self):
         # Check the IDs returned.
         self.assertEqual([key._id for key in result], list(range(num_ids)))
 
-    def test_allocate_ids_with_completed_key(self):
+    def test_allocate_ids_w_completed_key(self):
         creds = _make_credentials()
         client = self._make_one(credentials=creds)
 
-        COMPLETE_KEY = _Key(self.PROJECT)
-        self.assertRaises(ValueError, client.allocate_ids, COMPLETE_KEY, 2)
+        complete_key = _Key(self.PROJECT)
+        self.assertRaises(ValueError, client.allocate_ids, complete_key, 2)
+
+    def test_reserve_ids_w_completed_key(self):
+        num_ids = 2
+        creds = _make_credentials()
+        client = self._make_one(credentials=creds, _use_grpc=False)
+        complete_key = _Key(self.PROJECT)
+        reserve_ids = mock.Mock()
+        ds_api = mock.Mock(reserve_ids=reserve_ids, spec=["reserve_ids"])
+        client._datastore_api_internal = ds_api
+        self.assertTrue(not complete_key.is_partial)
+        client.reserve_ids(complete_key, num_ids)
+        expected_keys = [complete_key.to_protobuf()] * num_ids
+        reserve_ids.assert_called_once_with(self.PROJECT, expected_keys)
+
+    def test_reserve_ids_w_partial_key(self):
+        num_ids = 2
+        incomplete_key = _Key(self.PROJECT)
+        incomplete_key._id = None
+        creds = _make_credentials()
+        client = self._make_one(credentials=creds)
+        with self.assertRaises(ValueError):
+            client.reserve_ids(incomplete_key, num_ids)
+
+    def test_reserve_ids_w_wrong_num_ids(self):
+        num_ids = "2"
+        complete_key = _Key(self.PROJECT)
+        creds = _make_credentials()
+        client = self._make_one(credentials=creds)
+        with self.assertRaises(ValueError):
+            client.reserve_ids(complete_key, num_ids)
 
     def test_key_w_project(self):
         KIND = "KIND"