diff --git a/config/dev.yml b/config/dev.yml --- a/config/dev.yml +++ b/config/dev.yml @@ -9,3 +9,5 @@ debug: yes server-type: asgi + +max_query_cost: 1000 diff --git a/config/staging.yml b/config/staging.yml --- a/config/staging.yml +++ b/config/staging.yml @@ -9,3 +9,5 @@ debug: yes server-type: wsgi + +max_query_cost: 1000 diff --git a/swh/graphql/errors/__init__.py b/swh/graphql/errors/__init__.py --- a/swh/graphql/errors/__init__.py +++ b/swh/graphql/errors/__init__.py @@ -3,7 +3,7 @@ # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information -from .errors import ObjectNotFoundError, PaginationError +from .errors import ObjectNotFoundError, PaginationError, QueryCostError from .handlers import format_error -__all__ = ["ObjectNotFoundError", "PaginationError", "format_error"] +__all__ = ["ObjectNotFoundError", "PaginationError", "QueryCostError", "format_error"] diff --git a/swh/graphql/errors/errors.py b/swh/graphql/errors/errors.py --- a/swh/graphql/errors/errors.py +++ b/swh/graphql/errors/errors.py @@ -16,3 +16,7 @@ def __init__(self, message, errors=None): # FIXME, log this error super().__init__(f"{self.msg}: {message}") + + +class QueryCostError(Exception): + """ """ diff --git a/swh/graphql/resolvers/base_node.py b/swh/graphql/resolvers/base_node.py --- a/swh/graphql/resolvers/base_node.py +++ b/swh/graphql/resolvers/base_node.py @@ -7,6 +7,7 @@ from collections import namedtuple from swh.graphql.errors import ObjectNotFoundError +from swh.graphql.utils import query_cost class BaseNode(ABC): @@ -21,6 +22,14 @@ self._node = self._get_node(node_data) # handle the errors, if any, after _node is set self._handle_node_errors() + self.cost = self._get_node_cost() + # validate and register the cost + query_cost.query_cost.add_cost(self.cost) + + def _get_node_cost(self): + # get the cost for this node + # Override for specific behaviour + return query_cost.get_node_cost(self) def _get_node(self, node_data): """ diff --git a/swh/graphql/utils/__init__.py b/swh/graphql/utils/__init__.py new file mode 100644 diff --git a/swh/graphql/utils/query_cost.py b/swh/graphql/utils/query_cost.py new file mode 100644 --- /dev/null +++ b/swh/graphql/utils/query_cost.py @@ -0,0 +1,36 @@ +# Copyright (C) 2022 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + +from swh.graphql.errors import QueryCostError + + +class QueryCost: + total = 0 + + def get_max_query_cost(self): + # Add logic to have different limits on user type + return 1000 + + def total_query_cost(self): + return self.total + + def add_cost(self, cost): + self.total += cost + max_query_cost = self.get_max_query_cost() + if self.total > max_query_cost: + # FIXME, return the already calculated data + raise QueryCostError(max_query_cost, self.total) + + def register_usage(self): + # FIXME, register the usage against the user + pass + + +query_cost = QueryCost() + + +def get_node_cost(node): + # FIXME, logic to calculate node cost + return 1