Skip to content

Commit 25378bd

Browse files
Marcel Overdijkbbakerman
authored andcommitted
Added query complexity and depth logging + added 'simulation' option (graphql-java#1282)
* Added query complexity and depth logging + added 'simulation' option * Added query complexity|depth info objects * Added support for custom actions when max query complexity or depth is exceeded * Updated javadoc
1 parent 5e2c94f commit 25378bd

6 files changed

Lines changed: 279 additions & 9 deletions

File tree

‎src/main/java/graphql/analysis/MaxQueryComplexityInstrumentation.java‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,32 +6,50 @@
66
import graphql.execution.instrumentation.SimpleInstrumentation;
77
import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
88
import graphql.validation.ValidationError;
9+
import org.slf4j.Logger;
10+
import org.slf4j.LoggerFactory;
911

1012
import java.util.LinkedHashMap;
1113
import java.util.List;
1214
import java.util.Map;
15+
import java.util.function.Function;
1316

1417
import static graphql.Assert.assertNotNull;
1518
import static graphql.execution.instrumentation.SimpleInstrumentationContext.whenCompleted;
1619
import static java.util.Optional.ofNullable;
1720

1821
/**
19-
* Prevents execution if the query complexity is greater than the specified maxComplexity
22+
* Prevents execution if the query complexity is greater than the specified maxComplexity.
23+
*
24+
* Use the {@link Function<QueryComplexityInfo, Boolean>} parameter to supply a function to perform a custom action when the max complexity
25+
* is exceeded. If the function returns {@code true} a {@link AbortExecutionException} is thrown.
2026
*/
2127
@PublicApi
2228
public class MaxQueryComplexityInstrumentation extends SimpleInstrumentation {
2329

30+
private static final Logger log = LoggerFactory.getLogger(MaxQueryComplexityInstrumentation.class);
2431

2532
private final int maxComplexity;
2633
private final FieldComplexityCalculator fieldComplexityCalculator;
34+
private final Function<QueryComplexityInfo, Boolean> maxQueryComplexityExceededFunction;
2735

2836
/**
2937
* new Instrumentation with default complexity calculator which is `1 + childComplexity`
3038
*
3139
* @param maxComplexity max allowed complexity, otherwise execution will be aborted
3240
*/
3341
public MaxQueryComplexityInstrumentation(int maxComplexity) {
34-
this(maxComplexity, (env, childComplexity) -> 1 + childComplexity);
42+
this(maxComplexity, (queryComplexityInfo) -> true);
43+
}
44+
45+
/**
46+
* new Instrumentation with default complexity calculator which is `1 + childComplexity`
47+
*
48+
* @param maxComplexity max allowed complexity, otherwise execution will be aborted
49+
* @param maxQueryComplexityExceededFunction the function to perform when the max complexity is exceeded
50+
*/
51+
public MaxQueryComplexityInstrumentation(int maxComplexity, Function<QueryComplexityInfo, Boolean> maxQueryComplexityExceededFunction) {
52+
this(maxComplexity, (env, childComplexity) -> 1 + childComplexity, maxQueryComplexityExceededFunction);
3553
}
3654

3755
/**
@@ -41,11 +59,23 @@ public MaxQueryComplexityInstrumentation(int maxComplexity) {
4159
* @param fieldComplexityCalculator custom complexity calculator
4260
*/
4361
public MaxQueryComplexityInstrumentation(int maxComplexity, FieldComplexityCalculator fieldComplexityCalculator) {
62+
this(maxComplexity, fieldComplexityCalculator, (queryComplexityInfo) -> true);
63+
}
64+
65+
/**
66+
* new Instrumentation with custom complexity calculator
67+
*
68+
* @param maxComplexity max allowed complexity, otherwise execution will be aborted
69+
* @param fieldComplexityCalculator custom complexity calculator
70+
* @param maxQueryComplexityExceededFunction the function to perform when the max complexity is exceeded
71+
*/
72+
public MaxQueryComplexityInstrumentation(int maxComplexity, FieldComplexityCalculator fieldComplexityCalculator,
73+
Function<QueryComplexityInfo, Boolean> maxQueryComplexityExceededFunction) {
4474
this.maxComplexity = maxComplexity;
4575
this.fieldComplexityCalculator = assertNotNull(fieldComplexityCalculator, "calculator can't be null");
76+
this.maxQueryComplexityExceededFunction = maxQueryComplexityExceededFunction;
4677
}
4778

48-
4979
@Override
5080
public InstrumentationContext<List<ValidationError>> beginValidation(InstrumentationValidationParameters parameters) {
5181
return whenCompleted((errors, throwable) -> {
@@ -66,11 +96,16 @@ public void visitField(QueryVisitorFieldEnvironment env) {
6696
);
6797
}
6898
});
69-
7099
int totalComplexity = valuesByParent.getOrDefault(null, 0);
71-
100+
log.debug("Query complexity: {}", totalComplexity);
72101
if (totalComplexity > maxComplexity) {
73-
throw mkAbortException(totalComplexity, maxComplexity);
102+
QueryComplexityInfo queryComplexityInfo = QueryComplexityInfo.newQueryComplexityInfo()
103+
.complexity(totalComplexity)
104+
.build();
105+
boolean throwAbortException = maxQueryComplexityExceededFunction.apply(queryComplexityInfo);
106+
if (throwAbortException) {
107+
throw mkAbortException(totalComplexity, maxComplexity);
108+
}
74109
}
75110
});
76111
}
@@ -119,4 +154,4 @@ private FieldComplexityEnvironment convertEnv(QueryVisitorFieldEnvironment query
119154
}
120155

121156

122-
}
157+
}

‎src/main/java/graphql/analysis/MaxQueryDepthInstrumentation.java‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,46 @@
66
import graphql.execution.instrumentation.SimpleInstrumentation;
77
import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
88
import graphql.validation.ValidationError;
9+
import org.slf4j.Logger;
10+
import org.slf4j.LoggerFactory;
911

1012
import java.util.List;
13+
import java.util.function.Function;
1114

1215
import static graphql.execution.instrumentation.SimpleInstrumentationContext.whenCompleted;
1316

1417
/**
15-
* Prevents execution if the query depth is greater than the specified maxDepth
18+
* Prevents execution if the query depth is greater than the specified maxDepth.
19+
*
20+
* Use the {@link Function<QueryDepthInfo, Boolean>} parameter to supply a function to perform a custom action when the max depth is
21+
* exceeded. If the function returns {@code true} a {@link AbortExecutionException} is thrown.
1622
*/
1723
@PublicApi
1824
public class MaxQueryDepthInstrumentation extends SimpleInstrumentation {
1925

26+
private static final Logger log = LoggerFactory.getLogger(MaxQueryDepthInstrumentation.class);
27+
2028
private final int maxDepth;
29+
private final Function<QueryDepthInfo, Boolean> maxQueryDepthExceededFunction;
2130

31+
/**
32+
* Creates a new instrumentation that tracks the query depth.
33+
*
34+
* @param maxDepth max allowed depth, otherwise execution will be aborted
35+
*/
2236
public MaxQueryDepthInstrumentation(int maxDepth) {
37+
this(maxDepth, (queryDepthInfo) -> true);
38+
}
39+
40+
/**
41+
* Creates a new instrumentation that tracks the query depth.
42+
*
43+
* @param maxDepth max allowed depth, otherwise execution will be aborted
44+
* @param maxQueryDepthExceededFunction the function to perform when the max depth is exceeded
45+
*/
46+
public MaxQueryDepthInstrumentation(int maxDepth, Function<QueryDepthInfo, Boolean> maxQueryDepthExceededFunction) {
2347
this.maxDepth = maxDepth;
48+
this.maxQueryDepthExceededFunction = maxQueryDepthExceededFunction;
2449
}
2550

2651
@Override
@@ -31,8 +56,15 @@ public InstrumentationContext<List<ValidationError>> beginValidation(Instrumenta
3156
}
3257
QueryTraversal queryTraversal = newQueryTraversal(parameters);
3358
int depth = queryTraversal.reducePreOrder((env, acc) -> Math.max(getPathLength(env.getParentEnvironment()), acc), 0);
59+
log.debug("Query depth info: {}", depth);
3460
if (depth > maxDepth) {
35-
throw mkAbortException(depth, maxDepth);
61+
QueryDepthInfo queryDepthInfo = QueryDepthInfo.newQueryDepthInfo()
62+
.depth(depth)
63+
.build();
64+
boolean throwAbortException = maxQueryDepthExceededFunction.apply(queryDepthInfo);
65+
if (throwAbortException) {
66+
throw mkAbortException(depth, maxDepth);
67+
}
3668
}
3769
});
3870
}
Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package graphql.analysis;
2+
3+
import graphql.PublicApi;
4+
5+
/**
6+
* The query complexity info.
7+
*/
8+
@PublicApi
9+
public class QueryComplexityInfo {
10+
11+
private final int complexity;
12+
13+
private QueryComplexityInfo(int complexity) {
14+
this.complexity = complexity;
15+
}
16+
17+
/**
18+
* This returns the query complexity.
19+
*
20+
* @return the query complexity
21+
*/
22+
public int getComplexity() {
23+
return complexity;
24+
}
25+
26+
@Override
27+
public String toString() {
28+
return "QueryComplexityInfo{" +
29+
"complexity=" + complexity +
30+
'}';
31+
}
32+
33+
/**
34+
* @return a new {@link QueryComplexityInfo} builder
35+
*/
36+
public static Builder newQueryComplexityInfo() {
37+
return new Builder();
38+
}
39+
40+
@PublicApi
41+
public static class Builder {
42+
43+
private int complexity;
44+
45+
private Builder() {
46+
}
47+
48+
/**
49+
* The query complexity.
50+
*
51+
* @param complexity the query complexity
52+
* @return this builder
53+
*/
54+
public Builder complexity(int complexity) {
55+
this.complexity = complexity;
56+
return this;
57+
}
58+
59+
/**
60+
* @return a built {@link QueryComplexityInfo} object
61+
*/
62+
public QueryComplexityInfo build() {
63+
return new QueryComplexityInfo(complexity);
64+
}
65+
}
66+
}
Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package graphql.analysis;
2+
3+
import graphql.PublicApi;
4+
5+
/**
6+
* The query depth info.
7+
*/
8+
@PublicApi
9+
public class QueryDepthInfo {
10+
11+
private final int depth;
12+
13+
private QueryDepthInfo(int depth) {
14+
this.depth = depth;
15+
}
16+
17+
/**
18+
* This returns the query depth.
19+
*
20+
* @return the query depth
21+
*/
22+
public int getDepth() {
23+
return depth;
24+
}
25+
26+
@Override
27+
public String toString() {
28+
return "QueryDepthInfo{" +
29+
"depth=" + depth +
30+
'}';
31+
}
32+
33+
/**
34+
* @return a new {@link QueryDepthInfo} builder
35+
*/
36+
public static Builder newQueryDepthInfo() {
37+
return new Builder();
38+
}
39+
40+
@PublicApi
41+
public static class Builder {
42+
43+
private int depth;
44+
45+
private Builder() {
46+
}
47+
48+
/**
49+
* The query depth.
50+
*
51+
* @param depth the depth complexity
52+
* @return this builder
53+
*/
54+
public Builder depth(int depth) {
55+
this.depth = depth;
56+
return this;
57+
}
58+
59+
/**
60+
* @return a built {@link QueryDepthInfo} object
61+
*/
62+
public QueryDepthInfo build() {
63+
return new QueryDepthInfo(depth);
64+
}
65+
}
66+
}

‎src/test/groovy/graphql/analysis/MaxQueryComplexityInstrumentationTest.groovy‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import graphql.validation.ValidationError
1111
import graphql.validation.ValidationErrorType
1212
import spock.lang.Specification
1313

14+
import java.util.function.Function
15+
1416
class MaxQueryComplexityInstrumentationTest extends Specification {
1517

1618
Document createQuery(String query) {
@@ -154,6 +156,39 @@ class MaxQueryComplexityInstrumentationTest extends Specification {
154156

155157
}
156158

159+
def "custom max query complexity exceeded function"() {
160+
given:
161+
def schema = TestUtil.schema("""
162+
type Query{
163+
foo: Foo
164+
bar: String
165+
}
166+
type Foo {
167+
scalar: String
168+
foo: Foo
169+
}
170+
""")
171+
def query = createQuery("""
172+
{f2: foo {scalar foo{scalar}} f1: foo { foo {foo {foo {foo{foo{scalar}}}}}} }
173+
""")
174+
Boolean test = false
175+
Function<QueryComplexityInfo, Boolean> maxQueryComplexityExceededFunction = new Function<QueryComplexityInfo, Boolean>() {
176+
@Override
177+
Boolean apply(final QueryComplexityInfo queryComplexityInfo) {
178+
test = true
179+
return false
180+
}
181+
}
182+
MaxQueryComplexityInstrumentation queryComplexityInstrumentation = new MaxQueryComplexityInstrumentation(10, maxQueryComplexityExceededFunction)
183+
ExecutionInput executionInput = Mock(ExecutionInput)
184+
InstrumentationValidationParameters validationParameters = new InstrumentationValidationParameters(executionInput, query, schema, null)
185+
InstrumentationContext instrumentationContext = queryComplexityInstrumentation.beginValidation(validationParameters)
186+
when:
187+
instrumentationContext.onCompleted(null, null)
188+
then:
189+
test == true
190+
notThrown(Exception)
191+
}
157192
}
158193

159194

0 commit comments

Comments
 (0)