Coverage for src/lexigram/graphql/security/depth.py: 97%
69 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 04:37 +0800
1"""Query depth limiting.
3This module provides depth limiting for GraphQL queries
4to prevent deeply nested queries that could impact performance.
5"""
7from __future__ import annotations
9from typing import TYPE_CHECKING
11from strawberry.extensions import SchemaExtension
13from lexigram.graphql.exceptions import QueryTooDeepError
14from lexigram.logging import get_logger
16if TYPE_CHECKING:
17 from collections.abc import Iterator
19 from graphql.language.ast import (
20 DocumentNode,
21 FragmentDefinitionNode,
22 SelectionSetNode,
23 )
26logger = get_logger(__name__)
29class DepthLimitValidator:
30 """Validate query depth against a limit.
32 Analyzes GraphQL queries to determine their depth and
33 validates against a configured maximum depth.
35 Example:
36 ```python
37 validator = DepthLimitValidator(max_depth=10)
39 # Validate a parsed document
40 depth = validator.get_depth(document)
41 validator.validate(document) # Raises if too deep
42 ```
43 """
45 def __init__(
46 self,
47 max_depth: int = 10,
48 ignore_introspection: bool = True,
49 ) -> None:
50 """Initialize the validator.
52 Args:
53 max_depth: Maximum allowed query depth.
54 ignore_introspection: Skip depth check for introspection queries.
55 """
56 self._max_depth = max_depth
57 self._ignore_introspection = ignore_introspection
59 @property
60 def max_depth(self) -> int:
61 """Get maximum depth limit."""
62 return self._max_depth
64 def get_depth(self, document: DocumentNode) -> int:
65 """Calculate the maximum depth of a query.
67 Args:
68 document: Parsed GraphQL document.
70 Returns:
71 Maximum depth of the query.
72 """
73 max_depth = 0
75 for definition in document.definitions:
76 if hasattr(definition, "selection_set"):
77 depth = self._calculate_depth(definition.selection_set, 0, document)
78 max_depth = max(max_depth, depth)
80 return max_depth
82 def _calculate_depth(
83 self,
84 selection_set: SelectionSetNode | None,
85 current_depth: int,
86 document: DocumentNode | None = None,
87 *,
88 _visited_fragments: frozenset[str] | None = None,
89 ) -> int:
90 """Recursively calculate depth, traversing fragment spreads.
92 Handles three AST node types:
93 - ``Field``: increments depth and recurses into its selection set.
94 - ``InlineFragment``: transparent wrapper; recurses at the same depth.
95 - ``FragmentSpread``: looks up the named fragment definition from
96 *document* and recurses at the same depth. Already-visited
97 fragment names are tracked via *_visited_fragments* to prevent
98 infinite recursion from cyclically-defined fragments.
100 Args:
101 selection_set: Selection set to analyze.
102 current_depth: Current depth level.
103 document: Full parsed document, required for ``FragmentSpread``
104 resolution. Pass ``None`` only in contexts where the document
105 is unavailable — fragment spreads will then be skipped.
106 _visited_fragments: Names of fragments already being traversed in
107 the current call chain (cycle guard).
109 Returns:
110 Maximum depth from this point.
111 """
112 if selection_set is None:
113 return current_depth
115 from graphql.language.ast import (
116 FragmentDefinitionNode,
117 FragmentSpreadNode,
118 InlineFragmentNode,
119 )
121 visited = _visited_fragments or frozenset()
122 max_depth = current_depth
124 for selection in selection_set.selections:
125 # Skip introspection fields early
126 if self._ignore_introspection and hasattr(selection, "name"):
127 name = (
128 selection.name.value
129 if hasattr(selection.name, "value")
130 else str(selection.name)
131 )
132 if name.startswith("__"):
133 continue
135 if isinstance(selection, FragmentSpreadNode):
136 # FragmentSpread: locate the named fragment and recurse.
137 # The spread itself does not add a depth level — only the
138 # fields inside the fragment do.
139 if document is None:
140 continue
141 frag_name = selection.name.value
142 if frag_name in visited:
143 continue # cycle guard
144 frag_def: FragmentDefinitionNode | None = next(
145 (
146 d
147 for d in document.definitions
148 if isinstance(d, FragmentDefinitionNode)
149 and d.name.value == frag_name
150 ),
151 None,
152 )
153 if frag_def is not None:
154 depth = self._calculate_depth(
155 frag_def.selection_set,
156 current_depth,
157 document,
158 _visited_fragments=visited | {frag_name},
159 )
160 max_depth = max(max_depth, depth)
162 elif isinstance(selection, InlineFragmentNode):
163 # InlineFragment: transparent type guard; recurse at same depth.
164 if selection.selection_set:
165 depth = self._calculate_depth(
166 selection.selection_set,
167 current_depth,
168 document,
169 _visited_fragments=visited,
170 )
171 max_depth = max(max_depth, depth)
173 elif hasattr(selection, "selection_set") and selection.selection_set:
174 # Field: increment depth for each nested selection set.
175 depth = self._calculate_depth(
176 selection.selection_set,
177 current_depth + 1,
178 document,
179 _visited_fragments=visited,
180 )
181 max_depth = max(max_depth, depth)
183 return max_depth
185 def validate(self, document: DocumentNode) -> None:
186 """Validate query depth.
188 Args:
189 document: Parsed GraphQL document.
191 Raises:
192 DepthLimitError: If query exceeds depth limit.
193 """
194 depth = self.get_depth(document)
196 if depth > self._max_depth:
197 raise QueryTooDeepError(
198 f"Query depth {depth} exceeds maximum allowed depth of {self._max_depth}",
199 )
201 logger.debug("Query depth %d is within limit %d", depth, self._max_depth)
204class DepthLimitExtension(SchemaExtension):
205 """Strawberry extension for query depth limiting.
207 Add this extension to your schema to automatically
208 validate query depth before execution.
210 Example:
211 ```python
212 from lexigram.graphql.security import DepthLimitExtension
214 schema = strawberry.Schema(
215 query=Query,
216 extensions=[DepthLimitExtension(max_depth=10)],
217 )
218 ```
219 """
221 def __init__(
222 self,
223 max_depth: int = 10,
224 ignore_introspection: bool = True,
225 ) -> None:
226 """Initialize the extension.
228 Args:
229 max_depth: Maximum allowed query depth.
230 ignore_introspection: Skip introspection queries.
231 """
232 self._validator = DepthLimitValidator(
233 max_depth=max_depth,
234 ignore_introspection=ignore_introspection,
235 )
237 def on_validate(self) -> Iterator[None]:
238 """Hook called during operation validation."""
239 execution_context = self.execution_context
241 # Validate depth
242 if execution_context.graphql_document:
243 try:
244 self._validator.validate(execution_context.graphql_document)
245 except QueryTooDeepError as e:
246 logger.warning("Query rejected: %s", e)
247 raise
249 yield
252def create_depth_limit(
253 max_depth: int = 10,
254 ignore_introspection: bool = True,
255) -> DepthLimitExtension:
256 """Create a depth limit extension (convenience function).
258 Args:
259 max_depth: Maximum allowed depth.
260 ignore_introspection: Skip introspection queries.
262 Returns:
263 Configured extension.
265 Example:
266 ```python
267 schema = strawberry.Schema(
268 query=Query,
269 extensions=[create_depth_limit(5)],
270 )
271 ```
272 """
273 return DepthLimitExtension(
274 max_depth=max_depth,
275 ignore_introspection=ignore_introspection,
276 )
279__all__ = [
280 "DepthLimitExtension",
281 "DepthLimitValidator",
282 "create_depth_limit",
283]