def _datastore_run_query(query): """Run a query in Datastore. Args: query (query.QueryOptions): The query spec. Returns: tasklets.Future: """ query_pb = _query_to_protobuf(query) partition_id = entity_pb2.PartitionId(project_id=query.project, namespace_id=query.namespace) read_options = _datastore_api.get_read_options( query, default_read_consistency=_datastore_api.EVENTUAL) request = datastore_pb2.RunQueryRequest( project_id=query.project, partition_id=partition_id, query=query_pb, read_options=read_options, ) response = yield _datastore_api.make_call("RunQuery", request, timeout=query.timeout) log.debug(response) return response
def test_other_error(stub): api = stub.return_value future = tasklets.Future() api.foo.future.return_value = future class DummyException(Exception): pass try: raise DummyException("Have to raise in order to get traceback") except Exception as error: future.set_exception(error) request = object() with pytest.raises(DummyException): _api.make_call("foo", request, retries=0).result()
def test_no_retries(stub, _retry): api = stub.return_value future = tasklets.Future() api.foo.future.return_value = future _retry.retry_async.return_value = mock.Mock(return_value=future) future.set_result("bar") request = object() assert _api.make_call("foo", request, retries=0).result() == "bar" _retry.retry_async.assert_not_called()
def test_explicit_timeout(stub, _retry): api = stub.return_value future = tasklets.Future() api.foo.future.return_value = future _retry.retry_async.return_value = mock.Mock(return_value=future) future.set_result("bar") request = object() call = _api.make_call("foo", request, retries=0, timeout=20) assert call.result() == "bar" api.foo.future.assert_called_once_with(request, timeout=20)
def test_grpc_error(stub): api = stub.return_value future = tasklets.Future() api.foo.future.return_value = future class DummyError(grpc.Call, Exception): def code(self): return grpc.StatusCode.UNAVAILABLE def details(self): return "Where is the devil in?" try: raise DummyError("Have to raise in order to get traceback") except Exception as error: future.set_exception(error) request = object() with pytest.raises(core_exceptions.ServiceUnavailable): _api.make_call("foo", request, retries=0).result()
def test_explicit_retries(stub, _retry): api = stub.return_value future = tasklets.Future() api.foo.future.return_value = future _retry.retry_async.return_value = mock.Mock(return_value=future) future.set_result("bar") request = object() assert _api.make_call("foo", request, retries=4).result() == "bar" _retry.retry_async.assert_called_once() tasklet = _retry.retry_async.call_args[0][0] assert tasklet().result() == "bar" retries = _retry.retry_async.call_args[1]["retries"] assert retries == 4
def _run_query(project_id, namespace, query_pb): """Run a query in Datastore. Will potentially repeat the query to get all results. Args: project_id (str): The project/app id of the Datastore instance. namespace (str): The namespace to which to restrict results. query_pb (query_pb2.Query): The query protocol buffer representation. Returns: tasklets.Future: List[Tuple[query_pb2.EntityResult.ResultType, query_pb2.EntityResult]]: The raw query results. """ results = [] partition_id = entity_pb2.PartitionId(project_id=project_id, namespace_id=namespace) while True: # See what results we get from the backend request = datastore_pb2.RunQueryRequest(project_id=project_id, partition_id=partition_id, query=query_pb) response = yield _datastore_api.make_call("RunQuery", request) log.debug(response) batch = response.batch results.extend(((batch.entity_result_type, result) for result in batch.entity_results)) # Did we get all of them? if batch.more_results != MORE_RESULTS_TYPE_NOT_FINISHED: break # Still some results left to fetch. Update cursors and try again. query_pb.start_cursor = batch.end_cursor return results
def _datastore_run_query(query): """Run a query in Datastore. Args: query (query.QueryOptions): The query spec. Returns: tasklets.Future: """ query_pb = _query_to_protobuf(query) partition_id = entity_pb2.PartitionId(project_id=query.project, namespace_id=query.namespace) read_options = _datastore_api.get_read_options(query) request = datastore_pb2.RunQueryRequest( project_id=query.project, partition_id=partition_id, query=query_pb, read_options=read_options, ) response = yield _datastore_api.make_call("RunQuery", request, timeout=query.timeout) utils.logging_debug(log, response) raise tasklets.Return(response)