rest api: iotree limit (#2458)

* added tree_in_limit and tree_out_limit filters to tree endpoint of REST API
* tree endpoint now returns total no. of incoming and outgoing nodes

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):
         """

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):
         """

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,19 +251,20 @@ 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"]["=="]
             else:
                 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(

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}