66import graphql .execution .instrumentation .SimpleInstrumentation ;
77import graphql .execution .instrumentation .parameters .InstrumentationValidationParameters ;
88import graphql .validation .ValidationError ;
9+ import org .slf4j .Logger ;
10+ import org .slf4j .LoggerFactory ;
911
1012import java .util .LinkedHashMap ;
1113import java .util .List ;
1214import java .util .Map ;
15+ import java .util .function .Function ;
1316
1417import static graphql .Assert .assertNotNull ;
1518import static graphql .execution .instrumentation .SimpleInstrumentationContext .whenCompleted ;
1619import 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
2228public 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+ }
0 commit comments