<commit_before>import logging
from graphql.language.ast import FragmentSpread


class QueryCostException(Exception):
    pass


log = logging.getLogger('graphapi')


def _get_counts(info, fragments):
    multiplier = 1
    inner_multiplier = 0

    if isinstance(info, FragmentSpread):
        for selection in fragments[info.name.value].selection_set.selections:
            inner_multiplier += _get_counts(selection, fragments)
    else:
        # the multiplier is either 1 or the number of elements returned
        for argument in info.arguments:
            if argument.name.value in ('first', 'last'):
                multiplier = int(argument.value.value)

        # count up how many multi-nodes inside
        if info.selection_set:
            for selection in info.selection_set.selections:
                inner_multiplier += _get_counts(selection, fragments)

    # if this wasn't a multi-node, this counts as one node
    if inner_multiplier == 0:
        inner_multiplier = 1

    return multiplier * inner_multiplier


class QueryProtectionMiddleware(object):
    def __init__(self, max_cost=5000):
        self.max_cost = max_cost

    def resolve(self, next, root, info, **args):
        if root is None:
            count = _get_counts(info.field_asts[0], info.fragments)
            log.debug(f'graphql query name={info.field_name} asts={info.field_asts} cost={count}')
            if count > self.max_cost:
                raise QueryCostException(
                    f'Query Cost is too high ({count}), limit is {self.max_cost}'
                )<commit_msg>Issue 173 - allow "first" and "last" args to be variables.<commit_after>import logging
from graphql.language.ast import FragmentSpread


class QueryCostException(Exception):
    pass


log = logging.getLogger('graphapi')


def _get_counts(info, fragments, args={}):
    multiplier = 1
    inner_multiplier = 0

    if isinstance(info, FragmentSpread):
        for selection in fragments[info.name.value].selection_set.selections:
            inner_multiplier += _get_counts(selection, fragments)
    else:
        # the multiplier is either 1 or the number of elements returned
        multiplier = args.get('first', args.get('last', 1))

        # count up how many multi-nodes inside
        if info.selection_set:
            for selection in info.selection_set.selections:
                inner_multiplier += _get_counts(selection, fragments)

    # if this wasn't a multi-node, this counts as one node
    if inner_multiplier == 0:
        inner_multiplier = 1

    return multiplier * inner_multiplier


class QueryProtectionMiddleware(object):
    def __init__(self, max_cost=5000):
        self.max_cost = max_cost

    def resolve(self, next, root, info, **args):
        if root is None:
            count = _get_counts(info.field_asts[0], info.fragments, args)
            log.debug(f'graphql query name={info.field_name} asts={info.field_asts} cost={count}')
            if count > self.max_cost:
                raise QueryCostException(
                    f'Query Cost is too high ({count}), limit is {self.max_cost}'
                )
{'fraction_non_alphanumeric': 0.0884418901660281, 'fraction_numerical': 0.007662835249042145, 'mean_word_length': 2.690223792697291, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6306712', 'n_tokens_mistral': 915, 'n_tokens_neox': 873, 'n_words': 250}