def test_allocate_ids_empty(self):
from google.cloud.datastore_v1.proto import datastore_pb2
project = "PROJECT"
rsp_pb = datastore_pb2.AllocateIdsResponse()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb.SerializeToString())]
)
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
# Make request.
ds_api = self._make_one(client)
response = ds_api.allocate_ids(project, [])
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
self.assertEqual(list(response.keys), [])
uri = _build_expected_url(client._base_url, project, "allocateIds")
request = _verify_protobuf_call(http, uri, datastore_pb2.AllocateIdsRequest())
self.assertEqual(list(request.keys), [])
def allocate_ids(self,
project_id,
keys,
retry=google.api_core.gapic_v1.method.DEFAULT,
timeout=google.api_core.gapic_v1.method.DEFAULT,
metadata=None):
"""
Allocates IDs for the given keys, which is useful for referencing an entity
before it is inserted.
Example:
>>> from google.cloud import datastore_v1
>>>
>>> client = datastore_v1.DatastoreClient()
>>>
>>> project_id = ''
>>> keys = []
>>>
>>> response = client.allocate_ids(project_id, keys)
Args:
project_id (str): The ID of the project against which to make the request.
keys (list[Union[dict, ~google.cloud.datastore_v1.types.Key]]): A list of keys with incomplete key paths for which to allocate IDs.
No key may be reserved/read-only.
If a dict is provided, it must be of the same form as the protobuf
message :class:`~google.cloud.datastore_v1.types.Key`
retry (Optional[google.api_core.retry.Retry]): A retry object used
to retry requests. If ``None`` is specified, requests will not
be retried.
timeout (Optional[float]): The amount of time, in seconds, to wait
for the request to complete. Note that if ``retry`` is
specified, the timeout applies to each individual attempt.
metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata
that is provided to the method.
Returns:
A :class:`~google.cloud.datastore_v1.types.AllocateIdsResponse` instance.
Raises:
google.api_core.exceptions.GoogleAPICallError: If the request
failed for any reason.
google.api_core.exceptions.RetryError: If the request failed due
to a retryable error and retry attempts failed.
ValueError: If the parameters are invalid.
"""
if metadata is None:
metadata = []
metadata = list(metadata)
request = datastore_pb2.AllocateIdsRequest(
project_id=project_id,
keys=keys,
)
return self._allocate_ids(
request, retry=retry, timeout=timeout, metadata=metadata)
def allocate_ids(self, project, key_pbs):
"""Perform an ``allocateIds`` request.
:type project: str
:param project: The project to connect to. This is
usually your project name in the cloud console.
:type key_pbs: list of :class:`.entity_pb2.Key`
:param key_pbs: The keys for which the backend should allocate IDs.
:rtype: :class:`.datastore_pb2.AllocateIdsResponse`
:returns: The returned protobuf response object.
"""
request_pb = _datastore_pb2.AllocateIdsRequest(keys=key_pbs)
return _rpc(self.client._http, project, 'allocateIds',
self.client._base_url, request_pb,
_datastore_pb2.AllocateIdsResponse)
def _datastore_allocate_ids(keys, retries=None):
"""Calls ``AllocateIds`` on Datastore.
Args:
keys (List[google.cloud.datastore_v1.entity_pb2.Key]): List of
incomplete keys to allocate.
retries (int): Number of times to potentially retry the call. If
:data:`None` is passed, will use :data:`_retry._DEFAULT_RETRIES`.
If :data:`0` is passed, the call is attempted only once.
Returns:
tasklets.Tasklet: A future for
:class:`google.cloud.datastore_v1.datastore_pb2.AllocateIdsResponse`
"""
client = context_module.get_context().client
request = datastore_pb2.AllocateIdsRequest(project_id=client.project,
keys=keys)
return make_call("AllocateIds", request, retries=retries)
def test_allocate_ids(self):
# Setup Expected Response
expected_response = {}
expected_response = datastore_pb2.AllocateIdsResponse(
**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
client = datastore_v1.DatastoreClient(channel=channel)
# Setup Request
project_id = 'projectId-1969970175'
keys = []
response = client.allocate_ids(project_id, keys)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = datastore_pb2.AllocateIdsRequest(
project_id=project_id, keys=keys)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_allocate_ids_non_empty(self):
from google.cloud.datastore_v1.proto import datastore_pb2
project = "PROJECT"
before_key_pbs = [
_make_key_pb(project, id_=None),
_make_key_pb(project, id_=None),
]
after_key_pbs = [_make_key_pb(project), _make_key_pb(project, id_=2345)]
rsp_pb = datastore_pb2.AllocateIdsResponse()
rsp_pb.keys.add().CopyFrom(after_key_pbs[0])
rsp_pb.keys.add().CopyFrom(after_key_pbs[1])
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb.SerializeToString())]
)
client_info = _make_client_info()
client = mock.Mock(
_http=http,
_base_url="test.invalid",
_client_info=client_info,
spec=["_http", "_base_url", "_client_info"],
)
# Make request.
ds_api = self._make_one(client)
response = ds_api.allocate_ids(project, before_key_pbs)
# Check the result and verify the callers.
self.assertEqual(list(response.keys), after_key_pbs)
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(client._base_url, project, "allocateIds")
request = _verify_protobuf_call(http, uri, datastore_pb2.AllocateIdsRequest())
self.assertEqual(len(request.keys), len(before_key_pbs))
for key_before, key_after in zip(before_key_pbs, request.keys):
self.assertEqual(key_before, key_after)
def allocate_ids(
self,
project_id,
keys,
retry=google.api_core.gapic_v1.method.DEFAULT,
timeout=google.api_core.gapic_v1.method.DEFAULT,
metadata=None,
):
"""
Allocates IDs for the given keys, which is useful for referencing an entity
before it is inserted.
Example:
>>> from google.cloud import datastore_v1
>>>
>>> client = datastore_v1.DatastoreClient()
>>>
>>> # TODO: Initialize `project_id`:
>>> project_id = ''
>>>
>>> # TODO: Initialize `keys`:
>>> keys = []
>>>
>>> response = client.allocate_ids(project_id, keys)
Args:
project_id (str): The ID of the project against which to make the request.
keys (list[Union[dict, ~google.cloud.datastore_v1.types.Key]]): A list of keys with incomplete key paths for which to allocate IDs.
No key may be reserved/read-only.
If a dict is provided, it must be of the same form as the protobuf
message :class:`~google.cloud.datastore_v1.types.Key`
retry (Optional[google.api_core.retry.Retry]): A retry object used
to retry requests. If ``None`` is specified, requests will not
be retried.
timeout (Optional[float]): The amount of time, in seconds, to wait
for the request to complete. Note that if ``retry`` is
specified, the timeout applies to each individual attempt.
metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata
that is provided to the method.
Returns:
A :class:`~google.cloud.datastore_v1.types.AllocateIdsResponse` instance.
Raises:
google.api_core.exceptions.GoogleAPICallError: If the request
failed for any reason.
google.api_core.exceptions.RetryError: If the request failed due
to a retryable error and retry attempts failed.
ValueError: If the parameters are invalid.
"""
# Wrap the transport method to add retry and timeout logic.
if "allocate_ids" not in self._inner_api_calls:
self._inner_api_calls[
"allocate_ids"
] = google.api_core.gapic_v1.method.wrap_method(
self.transport.allocate_ids,
default_retry=self._method_configs["AllocateIds"].retry,
default_timeout=self._method_configs["AllocateIds"].timeout,
client_info=self._client_info,
)
request = datastore_pb2.AllocateIdsRequest(project_id=project_id, keys=keys)
if metadata is None:
metadata = []
metadata = list(metadata)
try:
routing_header = [("project_id", project_id)]
except AttributeError:
pass
else:
routing_metadata = google.api_core.gapic_v1.routing_header.to_grpc_metadata(
routing_header
)
metadata.append(routing_metadata)
return self._inner_api_calls["allocate_ids"](
request, retry=retry, timeout=timeout, metadata=metadata
)
