def get_read_options(eventual, transaction_id):
"""Validate rules for read options, and assign to the request.
Helper method for ``lookup()`` and ``run_query``.
:type eventual: bool
:param eventual: Flag indicating if ``EVENTUAL`` or ``STRONG``
consistency should be used.
:type transaction_id: bytes
:param transaction_id: A transaction identifier (may be null).
:rtype: :class:`.datastore_pb2.ReadOptions`
:returns: The read options corresponding to the inputs.
:raises: :class:`ValueError` if ``eventual`` is ``True`` and the
``transaction_id`` is not ``None``.
"""
if transaction_id is None:
if eventual:
return datastore_pb2.ReadOptions(
read_consistency=datastore_pb2.ReadOptions.ReadConsistency.EVENTUAL
)
else:
return datastore_pb2.ReadOptions()
else:
if eventual:
raise ValueError("eventual must be False when in a transaction")
else:
return datastore_pb2.ReadOptions(transaction=transaction_id)
def _next_page_helper(self, txn_id=None, retry=None, timeout=None):
from google.api_core import page_iterator
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
from google.cloud.datastore.query import Query
more_enum = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED
result = _make_query_response([], b"", more_enum, 0)
project = "prujekt"
ds_api = _make_datastore_api(result)
if txn_id is None:
client = _Client(project, datastore_api=ds_api)
else:
transaction = mock.Mock(id=txn_id, spec=["id"])
client = _Client(project,
datastore_api=ds_api,
transaction=transaction)
query = Query(client)
kwargs = {}
if retry is not None:
kwargs["retry"] = retry
if timeout is not None:
kwargs["timeout"] = timeout
iterator = self._make_one(query, client, **kwargs)
page = iterator._next_page()
self.assertIsInstance(page, page_iterator.Page)
self.assertIs(page._parent, iterator)
partition_id = entity_pb2.PartitionId(project_id=project)
if txn_id is None:
read_options = datastore_pb2.ReadOptions()
else:
read_options = datastore_pb2.ReadOptions(transaction=txn_id)
empty_query = query_pb2.Query()
ds_api.run_query.assert_called_once_with(
request={
"project_id": project,
"partition_id": partition_id,
"read_options": read_options,
"query": empty_query,
},
**kwargs,
)
def test_default_w_transaction(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
txn_id = b"123abc-easy-as"
read_options = self._call_fut(False, txn_id)
expected = datastore_pb2.ReadOptions(transaction=txn_id)
self.assertEqual(read_options, expected)
def test_lookup_single_key_empty_response_w_transaction(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
project = "PROJECT"
transaction = b"TRANSACTION"
key_pb = _make_key_pb(project)
rsp_pb = datastore_pb2.LookupResponse()
read_options = datastore_pb2.ReadOptions(transaction=transaction)
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._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.lookup(project, [key_pb], read_options=read_options)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "lookup")
self.assertEqual(len(response.found), 0)
self.assertEqual(len(response.missing), 0)
self.assertEqual(len(response.deferred), 0)
request = _verify_protobuf_call(http, uri,
datastore_pb2.LookupRequest())
self.assertEqual(list(request.keys), [key_pb._pb])
self.assertEqual(request.read_options, read_options._pb)
def test_eventual_wo_transaction(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
read_options = self._call_fut(True, None)
expected = datastore_pb2.ReadOptions(
read_consistency=datastore_pb2.ReadOptions.ReadConsistency.EVENTUAL
)
self.assertEqual(read_options, expected)
def test__next_page_w_skipped_lt_offset(self):
from google.api_core import page_iterator
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
from google.cloud.datastore.query import Query
project = "prujekt"
skipped_1 = 100
skipped_cursor_1 = b"DEADBEEF"
skipped_2 = 50
skipped_cursor_2 = b"FACEDACE"
more_enum = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED
result_1 = _make_query_response([], b"", more_enum, skipped_1)
result_1.batch.skipped_cursor = skipped_cursor_1
result_2 = _make_query_response([], b"", more_enum, skipped_2)
result_2.batch.skipped_cursor = skipped_cursor_2
ds_api = _make_datastore_api(result_1, result_2)
client = _Client(project, datastore_api=ds_api)
query = Query(client)
offset = 150
iterator = self._make_one(query, client, offset=offset)
page = iterator._next_page()
self.assertIsInstance(page, page_iterator.Page)
self.assertIs(page._parent, iterator)
partition_id = entity_pb2.PartitionId(project_id=project)
read_options = datastore_pb2.ReadOptions()
query_1 = query_pb2.Query(offset=offset)
query_2 = query_pb2.Query(start_cursor=skipped_cursor_1,
offset=(offset - skipped_1))
expected_calls = [
mock.call(
request={
"project_id": project,
"partition_id": partition_id,
"read_options": read_options,
"query": query,
}) for query in [query_1, query_2]
]
self.assertEqual(ds_api.run_query.call_args_list, expected_calls)
def test_run_query_w_namespace_nonempty_result(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
project = "PROJECT"
kind = "Kind"
namespace = "NS"
query_pb = self._make_query_pb(kind)
partition_id = entity_pb2.PartitionId(project_id=project,
namespace_id=namespace)
read_options = datastore_pb2.ReadOptions()
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(
entity_result_type=query_pb2.EntityResult.ResultType.FULL,
entity_results=[
query_pb2.EntityResult(entity=entity_pb2.Entity())
],
more_results=query_pb2.QueryResultBatch.MoreResultsType.
NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._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.run_query(project,
partition_id,
read_options,
query=query_pb)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "runQuery")
request = _verify_protobuf_call(http, uri,
datastore_pb2.RunQueryRequest())
self.assertEqual(request.partition_id, partition_id._pb)
self.assertEqual(request.query, query_pb._pb)
def test_lookup_multiple_keys_w_missing(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
project = "PROJECT"
key_pb1 = _make_key_pb(project)
key_pb2 = _make_key_pb(project, id_=2345)
rsp_pb = datastore_pb2.LookupResponse()
er_1 = rsp_pb._pb.missing.add()
er_1.entity.key.CopyFrom(key_pb1._pb)
er_2 = rsp_pb._pb.missing.add()
er_2.entity.key.CopyFrom(key_pb2._pb)
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._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.lookup(project, [key_pb1, key_pb2],
read_options=read_options)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "lookup")
self.assertEqual(len(response.found), 0)
self.assertEqual(len(response.deferred), 0)
missing_keys = [result.entity.key for result in response.missing]
self.assertEqual(missing_keys, [key_pb1._pb, key_pb2._pb])
request = _verify_protobuf_call(http, uri,
datastore_pb2.LookupRequest())
self.assertEqual(list(request.keys), [key_pb1._pb, key_pb2._pb])
self.assertEqual(request.read_options, read_options._pb)
def test_lookup_single_key_nonempty_response(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
project = "PROJECT"
key_pb = _make_key_pb(project)
rsp_pb = datastore_pb2.LookupResponse()
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key_pb._pb)
rsp_pb._pb.found.add(entity=entity._pb)
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb._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.lookup(project, [key_pb], read_options=read_options)
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "lookup")
self.assertEqual(len(response.found), 1)
self.assertEqual(len(response.missing), 0)
self.assertEqual(len(response.deferred), 0)
found = response.found[0].entity
self.assertEqual(found.key.path[0].kind, "Kind")
self.assertEqual(found.key.path[0].id, 1234)
request = _verify_protobuf_call(http, uri,
datastore_pb2.LookupRequest())
self.assertEqual(list(request.keys), [key_pb._pb])
self.assertEqual(request.read_options, read_options._pb)
def test_default_wo_transaction(self):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
read_options = self._call_fut(False, None)
expected = datastore_pb2.ReadOptions()
self.assertEqual(read_options, expected)
def _run_query_helper(
self,
read_consistency=None,
transaction=None,
namespace=None,
found=0,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore_v1.types import query as query_pb2
project = "PROJECT"
kind = "Nonesuch"
query_pb = self._make_query_pb(kind)
partition_kw = {"project_id": project}
if namespace is not None:
partition_kw["namespace_id"] = namespace
partition_id = entity_pb2.PartitionId(**partition_kw)
options_kw = {}
if read_consistency is not None:
options_kw["read_consistency"] = read_consistency
if transaction is not None:
options_kw["transaction"] = transaction
read_options = datastore_pb2.ReadOptions(**options_kw)
cursor = b"\x00"
batch_kw = {
"entity_result_type":
query_pb2.EntityResult.ResultType.FULL,
"end_cursor":
cursor,
"more_results":
query_pb2.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS,
}
if found:
batch_kw["entity_results"] = [
query_pb2.EntityResult(entity=entity_pb2.Entity())
] * found
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(**batch_kw))
http = _make_requests_session(
[_make_response(content=rsp_pb._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"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"partition_id": partition_id,
"read_options": read_options,
"query": query_pb,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.run_query(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
uri = _build_expected_url(client._base_url, project, "runQuery")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.RunQueryRequest(),
retry=retry,
timeout=timeout,
)
self.assertEqual(request.partition_id, partition_id._pb)
self.assertEqual(request.query, query_pb._pb)
self.assertEqual(request.read_options, read_options._pb)
def _lookup_multiple_helper(
self,
found=0,
missing=0,
deferred=0,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
project = "PROJECT"
key_pb1 = _make_key_pb(project)
key_pb2 = _make_key_pb(project, id_=2345)
keys = [key_pb1, key_pb2]
read_options = datastore_pb2.ReadOptions()
rsp_pb = datastore_pb2.LookupResponse()
found_keys = []
for i_found in range(found):
key = keys[i_found]
found_keys.append(key._pb)
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key._pb)
rsp_pb._pb.found.add(entity=entity._pb)
missing_keys = []
for i_missing in range(missing):
key = keys[i_missing]
missing_keys.append(key._pb)
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key._pb)
rsp_pb._pb.missing.add(entity=entity._pb)
deferred_keys = []
for i_deferred in range(deferred):
key = keys[i_deferred]
deferred_keys.append(key._pb)
rsp_pb._pb.deferred.append(key._pb)
http = _make_requests_session(
[_make_response(content=rsp_pb._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"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"keys": keys,
"read_options": read_options,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.lookup(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
self.assertEqual([found.entity.key for found in response.found],
found_keys)
self.assertEqual([missing.entity.key for missing in response.missing],
missing_keys)
self.assertEqual(list(response.deferred), deferred_keys)
uri = _build_expected_url(client._base_url, project, "lookup")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.LookupRequest(),
retry=retry,
timeout=timeout,
)
self.assertEqual(list(request.keys), [key_pb1._pb, key_pb2._pb])
self.assertEqual(request.read_options, read_options._pb)
def _lookup_single_helper(
self,
read_consistency=None,
transaction=None,
empty=True,
retry=None,
timeout=None,
):
from google.cloud.datastore_v1.types import datastore as datastore_pb2
from google.cloud.datastore_v1.types import entity as entity_pb2
project = "PROJECT"
key_pb = _make_key_pb(project)
options_kw = {}
if read_consistency is not None:
options_kw["read_consistency"] = read_consistency
if transaction is not None:
options_kw["transaction"] = transaction
read_options = datastore_pb2.ReadOptions(**options_kw)
rsp_pb = datastore_pb2.LookupResponse()
if not empty:
entity = entity_pb2.Entity()
entity.key._pb.CopyFrom(key_pb._pb)
rsp_pb._pb.found.add(entity=entity._pb)
http = _make_requests_session(
[_make_response(content=rsp_pb._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"],
)
ds_api = self._make_one(client)
request = {
"project_id": project,
"keys": [key_pb],
"read_options": read_options,
}
kwargs = _make_retry_timeout_kwargs(retry, timeout, http)
response = ds_api.lookup(request=request, **kwargs)
self.assertEqual(response, rsp_pb._pb)
if empty:
self.assertEqual(len(response.found), 0)
else:
self.assertEqual(len(response.found), 1)
self.assertEqual(len(response.missing), 0)
self.assertEqual(len(response.deferred), 0)
uri = _build_expected_url(client._base_url, project, "lookup")
request = _verify_protobuf_call(
http,
uri,
datastore_pb2.LookupRequest(),
retry=retry,
timeout=timeout,
)
if retry is not None:
retry.assert_called_once_with(http.request)
self.assertEqual(list(request.keys), [key_pb._pb])
self.assertEqual(request.read_options, read_options._pb)
