def _next_page_helper(self, txn_id=None):
from google.cloud.iterator import Page
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import query_pb2
from google.cloud.datastore.query import Query
more_enum = query_pb2.QueryResultBatch.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)
iterator = self._make_one(query, client)
page = iterator._next_page()
self.assertIsInstance(page, 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(
project, partition_id, read_options, query=empty_query)
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.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(self):
"""Get the next page in the iterator.
:rtype: :class:`~google.cloud.iterator.Page`
:returns: The next page in the iterator (or :data:`None` if
there are no pages left).
"""
if not self._more_results:
return None
query_pb = self._build_protobuf()
transaction = self.client.current_transaction
if transaction is None:
read_options = _datastore_pb2.ReadOptions()
else:
read_options = _datastore_pb2.ReadOptions(
transaction=transaction.id)
partition_id = _entity_pb2.PartitionId(
project_id=self._query.project, namespace_id=self._query.namespace)
response_pb = self.client._datastore_api.run_query(
self._query.project,
partition_id,
read_options,
query=query_pb,
)
entity_pbs = self._process_query_results(response_pb)
return page_iterator.Page(self, entity_pbs, self._item_to_value)
def test_eventual_wo_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
read_options = self._call_fut(True, None)
expected = datastore_pb2.ReadOptions(
read_consistency=datastore_pb2.ReadOptions.EVENTUAL)
self.assertEqual(read_options, expected)
def test_default_w_transaction(self):
from google.cloud.proto.datastore.v1 import 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_multiple_keys_empty_response(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
project = 'PROJECT'
key_pb1 = _make_key_pb(project)
key_pb2 = _make_key_pb(project, id_=2345)
rsp_pb = datastore_pb2.LookupResponse()
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
ds_api = self._make_one(client)
response = ds_api.lookup(project, read_options, [key_pb1, key_pb2])
# Check the result and verify the callers.
self.assertEqual(response, rsp_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)
cw = http._called_with
_verify_protobuf_call(self, cw, uri)
request = datastore_pb2.LookupRequest()
request.ParseFromString(cw['body'])
self.assertEqual(list(request.keys), [key_pb1, key_pb2])
self.assertEqual(request.read_options, read_options)
def test_get_multi_hit_w_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.datastore.key import Key
txn_id = b'123'
kind = 'Kind'
id_ = 1234
path = [{'kind': kind, 'id': id_}]
# Make a found entity pb to be returned from mock backend.
entity_pb = _make_entity_pb(self.PROJECT, kind, id_, 'foo', 'Foo')
# Make a connection to return the entity pb.
creds = _make_credentials()
client = self._make_one(credentials=creds)
lookup_response = _make_lookup_response(results=[entity_pb])
ds_api = _make_datastore_api(lookup_response=lookup_response)
client._datastore_api_internal = ds_api
key = Key(kind, id_, project=self.PROJECT)
txn = client.transaction()
txn._id = txn_id
result, = client.get_multi([key], transaction=txn)
new_key = result.key
# Check the returned value is as expected.
self.assertIsNot(new_key, key)
self.assertEqual(new_key.project, self.PROJECT)
self.assertEqual(new_key.path, path)
self.assertEqual(list(result), ['foo'])
self.assertEqual(result['foo'], 'Foo')
read_options = datastore_pb2.ReadOptions(transaction=txn_id)
ds_api.lookup.assert_called_once_with(self.PROJECT, read_options,
[key.to_protobuf()])
def test_lookup_multiple_keys_w_deferred(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
project = 'PROJECT'
key_pb1 = _make_key_pb(project)
key_pb2 = _make_key_pb(project, id_=2345)
rsp_pb = datastore_pb2.LookupResponse()
rsp_pb.deferred.add().CopyFrom(key_pb1)
rsp_pb.deferred.add().CopyFrom(key_pb2)
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb.SerializeToString())])
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
ds_api = self._make_one(client)
response = ds_api.lookup(project, read_options, [key_pb1, key_pb2])
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(client._base_url, project, 'lookup')
self.assertEqual(len(response.found), 0)
self.assertEqual(len(response.missing), 0)
self.assertEqual(list(response.deferred), [key_pb1, key_pb2])
request = _verify_protobuf_call(http, uri,
datastore_pb2.LookupRequest())
self.assertEqual(list(request.keys), [key_pb1, key_pb2])
self.assertEqual(request.read_options, read_options)
def test_get_multi_hit_multiple_keys_same_project(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.datastore.key import Key
kind = 'Kind'
id1 = 1234
id2 = 2345
# Make a found entity pb to be returned from mock backend.
entity_pb1 = _make_entity_pb(self.PROJECT, kind, id1)
entity_pb2 = _make_entity_pb(self.PROJECT, kind, id2)
# Make a connection to return the entity pbs.
creds = _make_credentials()
client = self._make_one(credentials=creds)
lookup_response = _make_lookup_response(
results=[entity_pb1, entity_pb2])
ds_api = _make_datastore_api(lookup_response=lookup_response)
client._datastore_api_internal = ds_api
key1 = Key(kind, id1, project=self.PROJECT)
key2 = Key(kind, id2, project=self.PROJECT)
retrieved1, retrieved2 = client.get_multi([key1, key2])
# Check values match.
self.assertEqual(retrieved1.key.path, key1.path)
self.assertEqual(dict(retrieved1), {})
self.assertEqual(retrieved2.key.path, key2.path)
self.assertEqual(dict(retrieved2), {})
read_options = datastore_pb2.ReadOptions()
ds_api.lookup.assert_called_once_with(
self.PROJECT, read_options,
[key1.to_protobuf(), key2.to_protobuf()])
def test_get_multi_miss_w_missing(self):
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.datastore.key import Key
KIND = 'Kind'
ID = 1234
# Make a missing entity pb to be returned from mock backend.
missed = entity_pb2.Entity()
missed.key.partition_id.project_id = self.PROJECT
path_element = missed.key.path.add()
path_element.kind = KIND
path_element.id = ID
creds = _make_credentials()
client = self._make_one(credentials=creds)
# Set missing entity on mock connection.
lookup_response = _make_lookup_response(missing=[missed])
ds_api = _make_datastore_api(lookup_response=lookup_response)
client._datastore_api_internal = ds_api
key = Key(KIND, ID, project=self.PROJECT)
missing = []
entities = client.get_multi([key], missing=missing)
self.assertEqual(entities, [])
key_pb = key.to_protobuf()
self.assertEqual([missed.key.to_protobuf() for missed in missing],
[key_pb])
read_options = datastore_pb2.ReadOptions()
ds_api.lookup.assert_called_once_with(self.PROJECT, read_options,
[key_pb])
def test_get_multi_miss(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.datastore.key import Key
creds = _make_credentials()
client = self._make_one(credentials=creds)
ds_api = _make_datastore_api()
client._datastore_api_internal = ds_api
key = Key('Kind', 1234, project=self.PROJECT)
results = client.get_multi([key])
self.assertEqual(results, [])
read_options = datastore_pb2.ReadOptions()
ds_api.lookup.assert_called_once_with(self.PROJECT, read_options,
[key.to_protobuf()])
def test_get_multi_w_deferred_from_backend_but_not_passed(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.datastore.entity import Entity
from google.cloud.datastore.key import Key
key1 = Key('Kind', project=self.PROJECT)
key1_pb = key1.to_protobuf()
key2 = Key('Kind', 2345, project=self.PROJECT)
key2_pb = key2.to_protobuf()
entity1_pb = entity_pb2.Entity()
entity1_pb.key.CopyFrom(key1_pb)
entity2_pb = entity_pb2.Entity()
entity2_pb.key.CopyFrom(key2_pb)
creds = _make_credentials()
client = self._make_one(credentials=creds)
# Mock up two separate requests. Using an iterable as side_effect
# allows multiple return values.
lookup_response1 = _make_lookup_response(results=[entity1_pb],
deferred=[key2_pb])
lookup_response2 = _make_lookup_response(results=[entity2_pb])
ds_api = _make_datastore_api()
ds_api.lookup = mock.Mock(
side_effect=[lookup_response1, lookup_response2], spec=[])
client._datastore_api_internal = ds_api
missing = []
found = client.get_multi([key1, key2], missing=missing)
self.assertEqual(len(found), 2)
self.assertEqual(len(missing), 0)
# Check the actual contents on the response.
self.assertIsInstance(found[0], Entity)
self.assertEqual(found[0].key.path, key1.path)
self.assertEqual(found[0].key.project, key1.project)
self.assertIsInstance(found[1], Entity)
self.assertEqual(found[1].key.path, key2.path)
self.assertEqual(found[1].key.project, key2.project)
self.assertEqual(ds_api.lookup.call_count, 2)
read_options = datastore_pb2.ReadOptions()
ds_api.lookup.assert_any_call(self.PROJECT, read_options, [key2_pb])
ds_api.lookup.assert_any_call(self.PROJECT, read_options,
[key1_pb, key2_pb])
def test_lookup_multiple_keys_w_deferred(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud import _http as connection_module
from google.cloud.datastore._http import _CLIENT_INFO
project = 'PROJECT'
key_pb1 = _make_key_pb(project)
key_pb2 = _make_key_pb(project, id_=2345)
rsp_pb = datastore_pb2.LookupResponse()
rsp_pb.deferred.add().CopyFrom(key_pb1)
rsp_pb.deferred.add().CopyFrom(key_pb2)
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
ds_api = self._make_one(client)
response = ds_api.lookup(project, read_options, [key_pb1, key_pb2])
# Check the result and verify the callers.
self.assertEqual(response, rsp_pb)
uri = _build_expected_url(client._base_url, project, 'lookup')
self.assertEqual(len(response.found), 0)
self.assertEqual(len(response.missing), 0)
self.assertEqual(list(response.deferred), [key_pb1, key_pb2])
cw = http._called_with
_verify_protobuf_call(self, cw, uri)
self.assertEqual(cw['uri'], uri)
self.assertEqual(cw['method'], 'POST')
expected_headers = {
'Content-Type': 'application/x-protobuf',
'User-Agent': connection_module.DEFAULT_USER_AGENT,
'Content-Length': str(len(cw['body'])),
connection_module.CLIENT_INFO_HEADER: _CLIENT_INFO,
}
self.assertEqual(cw['headers'], expected_headers)
request = datastore_pb2.LookupRequest()
request.ParseFromString(cw['body'])
self.assertEqual(list(request.keys), [key_pb1, key_pb2])
self.assertEqual(request.read_options, read_options)
def test_run_query_w_namespace_nonempty_result(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import 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.FULL,
entity_results=[
query_pb2.EntityResult(entity=entity_pb2.Entity()),
],
more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb.SerializeToString())])
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# 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)
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)
self.assertEqual(request.query, query_pb)
def test_run_query_wo_eventual_w_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import query_pb2
project = 'PROJECT'
kind = 'Nonesuch'
cursor = b'\x00'
transaction = b'TRANSACTION'
query_pb = self._make_query_pb(kind)
partition_id = entity_pb2.PartitionId(project_id=project)
read_options = datastore_pb2.ReadOptions(transaction=transaction)
rsp_pb = datastore_pb2.RunQueryResponse(
batch=query_pb2.QueryResultBatch(
entity_result_type=query_pb2.EntityResult.FULL,
end_cursor=cursor,
more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS,
))
# Create mock HTTP and client with response.
http = Http({'status': '200'}, rsp_pb.SerializeToString())
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# 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)
uri = _build_expected_url(client._base_url, project, 'runQuery')
cw = http._called_with
_verify_protobuf_call(self, cw, uri)
request = datastore_pb2.RunQueryRequest()
request.ParseFromString(cw['body'])
self.assertEqual(request.partition_id, partition_id)
self.assertEqual(request.query, query_pb)
self.assertEqual(request.read_options, read_options)
def test_get_multi_miss_w_deferred(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.datastore.key import Key
key = Key('Kind', 1234, project=self.PROJECT)
key_pb = key.to_protobuf()
# Set deferred entity on mock connection.
creds = _make_credentials()
client = self._make_one(credentials=creds)
lookup_response = _make_lookup_response(deferred=[key_pb])
ds_api = _make_datastore_api(lookup_response=lookup_response)
client._datastore_api_internal = ds_api
deferred = []
entities = client.get_multi([key], deferred=deferred)
self.assertEqual(entities, [])
self.assertEqual([def_key.to_protobuf() for def_key in deferred],
[key_pb])
read_options = datastore_pb2.ReadOptions()
ds_api.lookup.assert_called_once_with(self.PROJECT, read_options,
[key_pb])
def test_lookup_single_key_nonempty_response(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
from google.cloud.proto.datastore.v1 import entity_pb2
project = 'PROJECT'
key_pb = _make_key_pb(project)
rsp_pb = datastore_pb2.LookupResponse()
entity = entity_pb2.Entity()
entity.key.CopyFrom(key_pb)
rsp_pb.found.add(entity=entity)
read_options = datastore_pb2.ReadOptions()
# Create mock HTTP and client with response.
http = _make_requests_session(
[_make_response(content=rsp_pb.SerializeToString())])
client = mock.Mock(_http=http,
_base_url='test.invalid',
spec=['_http', '_base_url'])
# Make request.
ds_api = self._make_one(client)
response = ds_api.lookup(project, read_options, [key_pb])
# Check the result and verify the callers.
self.assertEqual(response, rsp_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])
self.assertEqual(request.read_options, read_options)
def test_default_wo_transaction(self):
from google.cloud.proto.datastore.v1 import datastore_pb2
read_options = self._call_fut(False, None)
expected = datastore_pb2.ReadOptions()
self.assertEqual(read_options, expected)
