diff --git a/aiida/backends/tests/restapi.py b/aiida/backends/tests/restapi.py index 53a432ae95..3bcda52b6b 100644 --- a/aiida/backends/tests/restapi.py +++ b/aiida/backends/tests/restapi.py @@ -727,6 +727,22 @@ def test_calculation_input_filters(self): result_node_type="data", result_name="inputs") + def test_calculation_iotree(self): + """ + Get filtered inputs list for given calculations + """ + node_uuid = self.get_dummy_data()["calculations"][1]["uuid"] + url = self.get_url_prefix() + '/calculations/' + str( + node_uuid) + '/io/tree?in_limit=1&out_limit=1' + with self.app.test_client() as client: + rv = client.get(url) + response = json.loads(rv.data) + self.assertEqual(len(response["data"]["nodes"]), 3) + self.assertEqual(len(response["data"]["edges"]),2) + RESTApiTestCase.compare_extra_response_data(self, "calculations", + url, + response, uuid=node_uuid) + ############### calculation attributes ############# def test_calculation_attributes(self): """ diff --git a/aiida/restapi/common/utils.py b/aiida/restapi/common/utils.py index fc6169c340..a479a869ab 100644 --- a/aiida/restapi/common/utils.py +++ b/aiida/restapi/common/utils.py @@ -509,6 +509,10 @@ def build_translator_parameters(self, field_list): filename = None rtype = None + # io tree limit parameters + tree_in_limit = None + tree_out_limit = None + ## Count how many time a key has been used for the filters and check if # reserved keyword # have been used twice, @@ -575,6 +579,12 @@ def build_translator_parameters(self, field_list): raise RestInputValidationError( "You cannot specify rtype more than " "once") + if 'in_limit' in field_counts.keys() and field_counts['in_limit'] > 1: + raise RestInputValidationError( + "You cannot specify in_limit more than once") + if 'out_limit' in field_counts.keys() and field_counts['out_limit'] > 1: + raise RestInputValidationError( + "You cannot specify out_limit more than once") ## Extract results for field in field_list: @@ -675,6 +685,20 @@ def build_translator_parameters(self, field_list): "only assignment operator '=' " "is permitted after 'rtype'") + elif field[0] == 'in_limit': + if field[1] == '=': + tree_in_limit = field[2] + else: + raise RestInputValidationError( + "only assignment operator '=' is permitted after 'in_limit'") + + elif field[0] == 'out_limit': + if field[1] == '=': + tree_out_limit = field[2] + else: + raise RestInputValidationError( + "only assignment operator '=' is permitted after 'out_limit'") + else: ## Construct the filter entry. @@ -707,7 +731,7 @@ def build_translator_parameters(self, field_list): # limit = self.LIMIT_DEFAULT return (limit, offset, perpage, orderby, filters, alist, nalist, elist, - nelist, downloadformat, visformat, filename, rtype) + nelist, downloadformat, visformat, filename, rtype, tree_in_limit, tree_out_limit) def parse_query_string(self, query_string): """ diff --git a/aiida/restapi/resources.py b/aiida/restapi/resources.py index 6dbf20cef0..1f72cf9fe7 100644 --- a/aiida/restapi/resources.py +++ b/aiida/restapi/resources.py @@ -122,9 +122,11 @@ def get(self, id=None, page=None): ## Parse request (resource_type, page, id, query_type) = self.utils.parse_path( path, parse_pk_uuid=self.parse_pk_uuid) + + # pylint: disable=unused-variable (limit, offset, perpage, orderby, filters, _alist, _nalist, _elist, - _nelist, _downloadformat, _visformat, _filename, - _rtype) = self.utils.parse_query_string(query_string) + _nelist, _downloadformat, _visformat, _filename, _rtype, tree_in_limit, + tree_out_limit) = self.utils.parse_query_string(query_string) ## Validate request self.utils.validate_request( @@ -225,8 +227,8 @@ def get(self, id=None, page=None): path, parse_pk_uuid=self.parse_pk_uuid) (limit, offset, perpage, orderby, filters, alist, nalist, elist, nelist, - downloadformat, visformat, filename, - rtype) = self.utils.parse_query_string(query_string) + downloadformat, visformat, filename, rtype, tree_in_limit, + tree_out_limit) = self.utils.parse_query_string(query_string) ## Validate request self.utils.validate_request( @@ -249,8 +251,8 @@ def get(self, id=None, page=None): ## Treat the statistics elif query_type == "statistics": (limit, offset, perpage, orderby, filters, alist, nalist, elist, - nelist, downloadformat, visformat, filename, - rtype) = self.utils.parse_query_string(query_string) + nelist, downloadformat, visformat, filename, rtype, tree_in_limit, + tree_out_limit) = self.utils.parse_query_string(query_string) headers = self.utils.build_headers(url=request.url, total_count=0) if filters: usr = filters["user"]["=="] @@ -258,10 +260,11 @@ def get(self, id=None, page=None): usr = None results = self.trans.get_statistics(usr) - # TODO Might need to be improved + # TODO improve the performance of tree endpoint by getting the data from database faster + # TODO add pagination for this endpoint (add default max limit) elif query_type == "tree": headers = self.utils.build_headers(url=request.url, total_count=0) - results = self.trans.get_io_tree(id) + results = self.trans.get_io_tree(id, tree_in_limit, tree_out_limit) else: ## Initialize the translator self.trans.set_query( diff --git a/aiida/restapi/translator/node.py b/aiida/restapi/translator/node.py index 91caa97f30..ac4d48376e 100644 --- a/aiida/restapi/translator/node.py +++ b/aiida/restapi/translator/node.py @@ -519,7 +519,7 @@ def get_statistics(self, user_pk=None): return qmanager.get_creation_statistics(user_pk=user_pk) - def get_io_tree(self, uuid_pattern): + def get_io_tree(self, uuid_pattern, tree_in_limit, tree_out_limit): from aiida.orm.querybuilder import QueryBuilder from aiida.orm.node import Node @@ -569,12 +569,15 @@ def get_node_shape(ntype): }) nodeCount += 1 + # get all inputs qb = QueryBuilder() qb.append(Node, tag="main", project=['*'], filters=self._id_filter) qb.append(Node, tag="in", project=['*'], edge_project=['label', 'type'], input_of='main') + if tree_in_limit is not None: + qb.limit(tree_in_limit) input_node_pks = {} @@ -625,6 +628,8 @@ def get_node_shape(ntype): filters=self._id_filter) qb.append(Node, tag="out", project=['*'], edge_project=['label', 'type'], output_of='main') + if tree_out_limit is not None: + qb.limit(tree_out_limit) output_node_pks = {} @@ -668,5 +673,16 @@ def get_node_shape(ntype): "label": linktype }) + # count total no of nodes + qb = QueryBuilder() + qb.append(Node, tag="main", project=['id'], filters=self._id_filter) + qb.append(Node, tag="in", project=['id'], input_of='main') + no_of_inputs = qb.count() + + qb = QueryBuilder() + qb.append(Node, tag="main", project=['id'], filters=self._id_filter) + qb.append(Node, tag="out", project=['id'], output_of='main') + no_of_outputs = qb.count() + - return {"nodes": nodes, "edges": edges} + return {"nodes": nodes, "edges": edges, "no_of_incomings": no_of_inputs, "no_of_outgoings": no_of_outputs}