Coverage for src/lexigram/graphql/security/complexity.py: 81%
124 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 04:37 +0800
1"""Query complexity analysis for GraphQL.
3This module provides complexity scoring to prevent expensive queries
4from being executed. Fields can be annotated with costs, and lists
5can be multiplied based on pagination arguments.
6"""
8from __future__ import annotations
10from collections.abc import Iterator
11from dataclasses import dataclass, field
12from typing import TYPE_CHECKING, Any
14from strawberry.extensions import SchemaExtension
16from lexigram.graphql import constants as const
17from lexigram.logging import get_logger
19if TYPE_CHECKING:
20 from collections.abc import Callable
22 from graphql import GraphQLSchema
23 from graphql.language.ast import DocumentNode, SelectionSetNode
25logger = get_logger(__name__)
28# G-27 FIX:
29# Registry keyed by the *resolver function name* (which maps to the GraphQL field name
30# under Strawberry's default snake_to_camel convention) so that ComplexityAnalyzer can
31# look up costs purely from the query's field names.
32#
33# Keys are lowercase function names; callers must register consistently.
34_FIELD_COSTS_BY_NAME: dict[str, float] = {}
35_LIST_COST_MULTIPLIERS_BY_NAME: dict[str, Callable[..., float]] = {}
38def cost(value: float) -> Callable:
39 """Decorator to set the cost of a field resolver.
41 Sets ``func.__graphql_cost__`` on the resolver and registers the cost
42 in the module-level lookup table keyed by ``func.__name__`` so that
43 :class:`ComplexityAnalyzer` can resolve the cost from the GraphQL
44 field name during query analysis.
46 Args:
47 value: The cost value for this field.
49 Returns:
50 Decorator function.
52 Example:
53 @cost(10)
54 def resolve_expensive_field(self, info):
55 return expensive_operation()
56 """
58 def decorator(func: Callable) -> Callable:
59 func.__graphql_cost__ = value # type: ignore[attr-defined]
60 _FIELD_COSTS_BY_NAME[func.__name__] = value
61 return func
63 return decorator
66def list_cost(multiplier: Callable[..., float]) -> Callable:
67 """Decorator to set the cost multiplier for a list field.
69 The multiplier function receives the field's arguments and should
70 return a multiplier value based on pagination arguments.
72 Sets ``func.__graphql_list_cost_multiplier__`` on the resolver and
73 registers the multiplier in the module-level lookup table so that
74 :class:`ComplexityAnalyzer` can use it from the field name.
76 Args:
77 multiplier: Function that calculates the list multiplier.
79 Returns:
80 Decorator function.
82 Example:
83 @list_cost(lambda first=10, last=10, **_: (first or 10) * (last or 10))
84 def resolve_list(self, info, **kwargs):
85 return items
86 """
88 def decorator(func: Callable) -> Callable:
89 func.__graphql_list_cost_multiplier__ = multiplier # type: ignore[attr-defined]
90 _LIST_COST_MULTIPLIERS_BY_NAME[func.__name__] = multiplier
91 return func
93 return decorator
96@dataclass
97class ComplexityResult:
98 """Result of complexity analysis.
100 Attributes:
101 complexity: Total complexity score.
102 field_count: Number of fields in the query.
103 list_count: Number of list fields.
104 warnings: Any warnings generated during analysis.
105 """
107 complexity: float = 0
108 field_count: int = 0
109 list_count: int = 0
110 warnings: list[str] = field(default_factory=list)
113class ComplexityAnalyzer:
114 """Analyze and score query complexity.
116 Calculates complexity scores based on:
117 - Base cost per field (default: 1)
118 - List cost multiplier (default: 10)
119 - Field-level cost annotations (@cost decorator)
120 - List pagination arguments (first/last)
122 Example:
123 analyzer = ComplexityAnalyzer(max_complexity=1000)
124 result = analyzer.analyze(document, schema)
126 if result.complexity > 1000:
127 raise QueryTooComplexError(f"Query complexity {result.complexity} exceeds limit")
128 """
130 def __init__(
131 self,
132 max_complexity: int = 1000,
133 default_field_cost: float = 1.0,
134 default_list_cost: float = 10.0,
135 ) -> None:
136 """Initialize the analyzer.
138 Args:
139 max_complexity: Maximum allowed complexity score.
140 default_field_cost: Default cost for a scalar field.
141 default_list_cost: Default cost multiplier for list fields.
142 """
143 self._max_complexity = max_complexity
144 self._default_field_cost = default_field_cost
145 self._default_list_cost = default_list_cost
147 @property
148 def max_complexity(self) -> int:
149 """Get maximum complexity limit."""
150 return self._max_complexity
152 def analyze(
153 self,
154 document: DocumentNode,
155 schema: GraphQLSchema | None = None,
156 ) -> ComplexityResult:
157 """Analyze query complexity.
159 Args:
160 document: Parsed GraphQL document.
161 schema: Optional GraphQL schema for type-aware analysis.
163 Returns:
164 Complexity result with score and metadata.
165 """
166 result = ComplexityResult()
167 total_complexity = 0.0
169 for definition in document.definitions:
170 if hasattr(definition, "selection_set") and definition.selection_set:
171 complexity = self._analyze_selection_set(
172 definition.selection_set,
173 result,
174 schema,
175 depth=0,
176 )
177 total_complexity += complexity
179 result.complexity = total_complexity
180 return result
182 def _analyze_selection_set(
183 self,
184 selection_set: SelectionSetNode,
185 result: ComplexityResult,
186 schema: GraphQLSchema | None,
187 depth: int,
188 ) -> float:
189 """Recursively analyze a selection set.
191 Args:
192 selection_set: The selection set to analyze.
193 result: Result object to accumulate metrics.
194 schema: Optional GraphQL schema.
195 depth: Current depth in the query.
197 Returns:
198 Complexity score for this selection set.
199 """
200 total_cost = 0.0
202 for selection in selection_set.selections:
203 if not hasattr(selection, "name"):
204 continue
206 field_name = (
207 selection.name.value
208 if hasattr(selection.name, "value")
209 else str(selection.name)
210 )
212 # Skip introspection fields
213 if field_name.startswith("__"):
214 continue
216 # Check if this is a list field (has arguments like first, last)
217 is_list = self._is_list_field(selection, schema)
219 # Calculate base cost
220 field_cost = self._get_field_cost(field_name, is_list)
222 # Get list multiplier if applicable
223 multiplier = 1.0
224 if is_list:
225 multiplier = self._get_list_multiplier(selection)
226 result.list_count += 1
228 # Calculate child complexity if there's a nested selection
229 child_cost = 0.0
230 if hasattr(selection, "selection_set") and selection.selection_set:
231 child_cost = self._analyze_selection_set(
232 selection.selection_set,
233 result,
234 schema,
235 depth + 1,
236 )
238 # Combine costs
239 field_total = (field_cost + child_cost) * multiplier
240 total_cost += field_total
241 result.field_count += 1
243 return total_cost
245 def _is_list_field(
246 self,
247 field_node: Any,
248 schema: GraphQLSchema | None,
249 ) -> bool:
250 """Check if a field returns a list type.
252 Args:
253 field_node: The field node to check.
254 schema: Optional schema for type info.
256 Returns:
257 True if the field returns a list.
258 """
259 # Check for pagination arguments commonly used with lists
260 args: list[Any] = list(field_node.arguments or [])
261 arg_names = {arg.name.value for arg in args}
263 list_args = {"first", "last", "limit", "after", "before"}
264 if list_args & arg_names:
265 return True
267 # If we have schema, we could check the return type
268 # but for now, we'll rely on argument presence
269 return False
271 def _get_field_cost(self, field_name: str, is_list: bool = False) -> float:
272 """Get the cost for a field.
274 Checks the registry populated by the :func:`cost` decorator first,
275 using ``field_name`` as the key (matching ``func.__name__``). Falls
276 back to the configured default costs.
278 Args:
279 field_name: GraphQL field name.
280 is_list: Whether the field returns a list.
282 Returns:
283 Cost value for the field.
284 """
285 # G-27 FIX: look up by field name in the decorator-populated registry.
286 if field_name in _FIELD_COSTS_BY_NAME:
287 return _FIELD_COSTS_BY_NAME[field_name]
289 if is_list:
290 return self._default_list_cost
292 return self._default_field_cost
294 def _get_list_multiplier(self, field_node: Any) -> float:
295 """Get the multiplier for a list field based on its arguments.
297 Args:
298 field_node: The list field node.
300 Returns:
301 Multiplier value (defaults to 10 if no arguments found).
302 """
303 args: list[Any] = list(field_node.arguments or [])
304 args_dict = {}
305 for arg in args:
306 if hasattr(arg, "value"):
307 # Handle both name and value attributes
308 arg_name = (
309 arg.name.value if hasattr(arg.name, "value") else str(arg.name)
310 )
311 if hasattr(arg.value, "value"):
312 args_dict[arg_name] = arg.value.value
313 else:
314 args_dict[arg_name] = arg.value
316 # Check for first/last/limit
317 first = args_dict.get("first")
318 last = args_dict.get("last")
319 limit = args_dict.get("limit")
321 if first is not None:
322 return float(first)
323 if last is not None:
324 return float(last)
325 if limit is not None:
326 return float(limit)
328 # Default multiplier if no pagination args
329 return self._default_list_cost
331 def check(
332 self, document: DocumentNode, schema: GraphQLSchema | None = None
333 ) -> tuple[bool, ComplexityResult]:
334 """Check query complexity and return a detailed result.
336 Unlike :meth:`validate`, this method never raises — it returns a
337 ``(is_valid, result)`` tuple so callers can inspect the details.
339 Args:
340 document: Parsed GraphQL document.
341 schema: Optional GraphQL schema.
343 Returns:
344 Tuple of (is_valid, result).
345 """
346 result = self.analyze(document, schema)
347 is_valid = result.complexity <= self._max_complexity
349 if not is_valid:
350 result.warnings.append(
351 f"Query complexity {result.complexity} exceeds maximum of {self._max_complexity}",
352 )
354 return is_valid, result
356 def validate(
357 self, document: DocumentNode, schema: GraphQLSchema | None = None
358 ) -> None:
359 """Validate query complexity; raise if the limit is exceeded.
361 Conforms to :class:`~lexigram.contracts.graphql.ValidationRule` so
362 instances of :class:`ComplexityAnalyzer` are usable wherever a
363 validator is expected (alongside :class:`DepthLimitValidator` and
364 :class:`AliasLimitValidator`).
366 Args:
367 document: Parsed GraphQL document.
368 schema: Optional GraphQL schema.
370 Raises:
371 QueryTooComplexError: If the query exceeds :attr:`max_complexity`.
372 """
373 is_valid, result = self.check(document, schema)
374 if not is_valid:
375 from lexigram.graphql.exceptions import QueryTooComplexError
377 msg = (
378 result.warnings[0]
379 if result.warnings
380 else f"Query complexity {result.complexity} exceeds maximum of {self._max_complexity}"
381 )
382 raise QueryTooComplexError(msg)
385class ComplexityLimitExtension(SchemaExtension):
386 """Enforce query complexity on every operation.
388 Wraps :class:`ComplexityAnalyzer` at the schema-extension boundary so the
389 limit applies to every execution path through the schema.
390 """
392 def __init__(
393 self,
394 max_complexity: int = const.DEFAULT_MAX_COMPLEXITY,
395 default_field_cost: float = 1.0,
396 default_list_cost: float = 10.0,
397 ) -> None:
398 self._analyzer = ComplexityAnalyzer(
399 max_complexity=max_complexity,
400 default_field_cost=default_field_cost,
401 default_list_cost=default_list_cost,
402 )
404 def on_validate(self) -> Iterator[None]:
405 """Reject the operation when complexity exceeds the configured maximum."""
406 execution_context = self.execution_context
407 if execution_context.graphql_document:
408 self._analyzer.validate(execution_context.graphql_document)
409 yield
412# Convenience function for quick analysis
413def analyze_complexity(
414 query: str,
415 max_complexity: int = 1000,
416) -> ComplexityResult:
417 """Analyze complexity of a GraphQL query.
419 Args:
420 query: GraphQL query string.
421 max_complexity: Maximum allowed complexity.
423 Returns:
424 Complexity result.
425 """
426 from graphql import parse
428 analyzer = ComplexityAnalyzer(max_complexity=max_complexity)
429 document = parse(query)
430 return analyzer.analyze(document)
433__all__ = [
434 "ComplexityAnalyzer",
435 "ComplexityLimitExtension",
436 "ComplexityResult",
437 "analyze_complexity",
438 "cost",
439 "list_cost",
440]