def _next_page_helper(self, txn_id=None): from google.api_core import page_iterator from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore_v1.proto import entity_pb2 from google.cloud.datastore_v1.proto 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_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( 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 test_get_multi_miss_w_missing(self): from google.cloud.datastore_v1.proto import entity_pb2 from google.cloud.datastore_v1.proto 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, [key_pb], read_options=read_options)
def test_lookup_multiple_keys_empty_response(self): from google.cloud.datastore_v1.proto 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 = _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.lookup(project, [key_pb1, key_pb2], read_options=read_options) # 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) 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_eventually_consistent(): options = _api.get_read_options( _options.ReadOptions(read_consistency=_api.EVENTUAL) ) assert options == datastore_pb2.ReadOptions( read_consistency=datastore_pb2.ReadOptions.EVENTUAL )
def _get_read_options(options): """Get the read options for a request. Args: options (Dict[str, Any]): The options for the request. For example, ``{"read_consistency": EVENTUAL}``. May contain options unrelated to creating a :class:`datastore_pb2.ReadOptions` instance, which will be ignored. Returns: datastore_pb2.ReadOptions: The options instance for passing to the Datastore gRPC API. Raises: ValueError: When ``read_consistency`` is set to ``EVENTUAL`` and there is a transaction. """ state = _runstate.current() transaction = options.get("transaction", state.transaction) read_consistency = options.get("read_consistency") if read_consistency is None: read_consistency = options.get("read_policy") # Legacy NDB if transaction is not None and read_consistency is EVENTUAL: raise ValueError( "read_consistency must be EVENTUAL when in transaction") return datastore_pb2.ReadOptions(read_consistency=read_consistency, transaction=transaction)
def test_lookup_multiple_keys_w_deferred(self): from google.cloud.datastore_v1.proto 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, [key_pb1, key_pb2], read_options=read_options) # 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(self): from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore.key import Key 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) result, = client.get_multi([key]) 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() ds_api.lookup.assert_called_once_with( self.PROJECT, [key.to_protobuf()], read_options=read_options, )
def test_idle_callback(_datastore_lookup, entity_pb2, context): class MockKey: def __init__(self, key=None): self.key = key def ParseFromString(self, key): self.key = key rpc = tasklets.Future("_datastore_lookup") _datastore_lookup.return_value = rpc entity_pb2.Key = MockKey eventloop = mock.Mock(spec=("queue_rpc", "run")) with context.new(eventloop=eventloop).use() as context: batch = _api._LookupBatch(_options.ReadOptions()) batch.lookup_callback = mock.Mock() batch.todo.update({"foo": ["one", "two"], "bar": ["three"]}) batch.idle_callback() called_with = _datastore_lookup.call_args[0] called_with_keys = set( (mock_key.key for mock_key in called_with[0]) ) assert called_with_keys == set(["foo", "bar"]) called_with_options = called_with[1] assert called_with_options == datastore_pb2.ReadOptions() rpc.set_result(None) batch.lookup_callback.assert_called_once_with(rpc)
def test_eventually_consistent_legacy(runstate): options = _api._get_read_options( {"read_policy": _api.EVENTUAL_CONSISTENCY} ) assert options == datastore_pb2.ReadOptions( read_consistency=datastore_pb2.ReadOptions.EVENTUAL )
def test_get_multi_hit_multiple_keys_same_project(self): from google.cloud.datastore_v1.proto 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, [key1.to_protobuf(), key2.to_protobuf()], read_options=read_options, )
def test_get_multi_miss_w_deferred(self): from google.cloud.datastore_v1.proto 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, [key_pb], read_options=read_options, )
def get_read_options(options, default_read_consistency=None): """Get the read options for a request. Args: options (_options.ReadOptions): The options for the request. May contain options unrelated to creating a :class:`datastore_pb2.ReadOptions` instance, which will be ignored. default_read_consistency: Use this value for ``read_consistency`` if neither ``transaction`` nor ``read_consistency`` are otherwise specified. Returns: datastore_pb2.ReadOptions: The options instance for passing to the Datastore gRPC API. Raises: ValueError: When ``read_consistency`` is set to ``EVENTUAL`` and there is a transaction. """ transaction = _get_transaction(options) read_consistency = options.read_consistency if transaction is None: if read_consistency is None: read_consistency = default_read_consistency elif read_consistency is EVENTUAL: raise ValueError( "read_consistency must not be EVENTUAL when in transaction") return datastore_pb2.ReadOptions(read_consistency=read_consistency, transaction=transaction)
def test_default_w_transaction(self): from google.cloud.datastore_v1.proto 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_eventual_wo_transaction(self): from google.cloud.datastore_v1.proto 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_get_multi_w_deferred_from_backend_but_not_passed(self): from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore_v1.proto 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, [key2_pb], read_options=read_options, ) ds_api.lookup.assert_any_call( self.PROJECT, [key1_pb, key2_pb], read_options=read_options, )
def test_get_multi_miss(self): from google.cloud.datastore_v1.proto 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, [key.to_protobuf()], read_options=read_options)
def test_run_query_w_eventual_no_transaction(self): from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore_v1.proto import entity_pb2 from google.cloud.datastore_v1.proto import query_pb2 project = "PROJECT" kind = "Nonesuch" cursor = b"\x00" query_pb = self._make_query_pb(kind) partition_id = entity_pb2.PartitionId(project_id=project) read_options = datastore_pb2.ReadOptions( read_consistency=datastore_pb2.ReadOptions.EVENTUAL ) 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 = _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.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) self.assertEqual(request.read_options, read_options)
def test_it(_datastore_api): query = query_module.QueryOptions(project="testing", namespace="") query_pb = _datastore_query._query_to_protobuf(query) _datastore_api.make_call.return_value = utils.future_result("foo") read_options = datastore_pb2.ReadOptions() request = datastore_pb2.RunQueryRequest( project_id="testing", partition_id=entity_pb2.PartitionId(project_id="testing", namespace_id=""), query=query_pb, read_options=read_options, ) _datastore_api.get_read_options.return_value = read_options assert _datastore_query._datastore_run_query(query).result() == "foo" _datastore_api.make_call.assert_called_once_with("RunQuery", request, timeout=None) _datastore_api.get_read_options.assert_called_once_with( query, default_read_consistency=_datastore_api.EVENTUAL)
def test_run_query_w_namespace_nonempty_result(self): from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore_v1.proto import entity_pb2 from google.cloud.datastore_v1.proto 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_lookup_single_key_nonempty_response(self): from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore_v1.proto 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_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) 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_run_query_wo_eventual_w_transaction(self): from google.cloud.datastore_v1.proto import datastore_pb2 from google.cloud.datastore_v1.proto import entity_pb2 from google.cloud.datastore_v1.proto 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 = _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) self.assertEqual(request.read_options, read_options)
def test_idle_callback(_datastore_lookup, entity_pb2, runstate): class MockKey: def __init__(self, key=None): self.key = key def ParseFromString(self, key): self.key = key entity_pb2.Key = MockKey runstate.eventloop = mock.Mock(spec=("queue_rpc", "run")) batch = _api._LookupBatch({}) batch.todo.update({"foo": ["one", "two"], "bar": ["three"]}) batch.idle_callback() called_with = _datastore_lookup.call_args[0] called_with_keys = set((mock_key.key for mock_key in called_with[0])) assert called_with_keys == set(["foo", "bar"]) called_with_options = called_with[1] assert called_with_options == datastore_pb2.ReadOptions() rpc = _datastore_lookup.return_value runstate.eventloop.queue_rpc.assert_called_once_with( rpc, batch.lookup_callback )
def test_get_multi_hit_w_transaction(self): from google.cloud.datastore_v1.proto 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, [key.to_protobuf()], read_options=read_options)
def test_eventually_consistent(runstate): options = _api._get_read_options({"read_consistency": _api.EVENTUAL}) assert options == datastore_pb2.ReadOptions( read_consistency=datastore_pb2.ReadOptions.EVENTUAL )
def test_args_override_transaction(runstate): runstate.transaction = b"txfoo" options = _api._get_read_options({"transaction": b"txbar"}) assert options == datastore_pb2.ReadOptions(transaction=b"txbar")
def test_no_args_transaction(runstate): runstate.transaction = b"txfoo" options = _api._get_read_options({}) assert options == datastore_pb2.ReadOptions(transaction=b"txfoo")
def test_no_args_no_transaction(runstate): assert _api._get_read_options({}) == datastore_pb2.ReadOptions()
def test_args_override_transaction(context): with context.new(transaction=b"txfoo").use(): options = _api.get_read_options( _options.ReadOptions(transaction=b"txbar") ) assert options == datastore_pb2.ReadOptions(transaction=b"txbar")
def test_no_args_no_transaction(): assert ( _api.get_read_options(_options.ReadOptions()) == datastore_pb2.ReadOptions() )