feat: add 'retry'/'timeout' args to 'Client.allocate_ids'

Toward #3

google/cloud/datastore/client.py

@@ -644,7 +644,7 @@ def delete_multi(self, keys, retry=None, timeout=None):
         if not in_batch:
             current.commit(retry=retry, timeout=timeout)
 
-    def allocate_ids(self, incomplete_key, num_ids):
+    def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
         """Allocate a list of IDs from a partial key.
 
         :type incomplete_key: :class:`google.cloud.datastore.key.Key`
@@ -653,6 +653,17 @@ def allocate_ids(self, incomplete_key, num_ids):
         :type num_ids: int
         :param num_ids: The number of IDs to allocate.
 
+        :type retry: :class:`google.api_core.retry.Retry`
+        :param retry:
+            A retry object used to retry requests. If ``None`` is specified,
+            requests will be retried using a default configuration.
+
+        :type timeout: float
+        :param timeout:
+            Time, in seconds, to wait for the request to complete.
+            Note that if ``retry`` is specified, the timeout applies
+            to each individual attempt.
+
         :rtype: list of :class:`google.cloud.datastore.key.Key`
         :returns: The (complete) keys allocated with ``incomplete_key`` as
                   root.
@@ -665,8 +676,16 @@ def allocate_ids(self, incomplete_key, num_ids):
         incomplete_key_pb = incomplete_key.to_protobuf()
         incomplete_key_pbs = [incomplete_key_pb] * num_ids
 
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         response_pb = self._datastore_api.allocate_ids(
-            incomplete_key.project, incomplete_key_pbs
+            incomplete_key.project, incomplete_key_pbs, **kwargs
         )
         allocated_ids = [
             allocated_key_pb.path[-1].id for allocated_key_pb in response_pb.keys

tests/unit/test_client.py

@@ -920,6 +920,36 @@ def test_allocate_ids_w_partial_key(self):
         # Check the IDs returned.
         self.assertEqual([key._id for key in result], list(range(num_ids)))
 
+        expected_keys = [incomplete_key.to_protobuf()] * num_ids
+        alloc_ids.assert_called_once_with(self.PROJECT, expected_keys)
+
+    def test_allocate_ids_w_partial_key_w_retry_w_timeout(self):
+        num_ids = 2
+
+        incomplete_key = _Key(self.PROJECT)
+        incomplete_key._id = None
+        retry = mock.Mock()
+        timeout = 100000
+
+        creds = _make_credentials()
+        client = self._make_one(credentials=creds, _use_grpc=False)
+        allocated = mock.Mock(keys=[_KeyPB(i) for i in range(num_ids)], spec=["keys"])
+        alloc_ids = mock.Mock(return_value=allocated, spec=[])
+        ds_api = mock.Mock(allocate_ids=alloc_ids, spec=["allocate_ids"])
+        client._datastore_api_internal = ds_api
+
+        result = client.allocate_ids(
+            incomplete_key, num_ids, retry=retry, timeout=timeout
+        )
+
+        # Check the IDs returned.
+        self.assertEqual([key._id for key in result], list(range(num_ids)))
+
+        expected_keys = [incomplete_key.to_protobuf()] * num_ids
+        alloc_ids.assert_called_once_with(
+            self.PROJECT, expected_keys, retry=retry, timeout=timeout
+        )
+
     def test_allocate_ids_w_completed_key(self):
         creds = _make_credentials()
         client = self._make_one(credentials=creds)