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)