graphql-java · bbakerman · Apr 29, 2022 · Mar 16, 2022 · Mar 26, 2022 · Mar 26, 2022
diff --git a/src/main/java/graphql/analysis/MaxQueryComplexityInstrumentation.java b/src/main/java/graphql/analysis/MaxQueryComplexityInstrumentation.java
@@ -1,9 +1,14 @@
 package graphql.analysis;
 
+import graphql.ExecutionResult;
 import graphql.PublicApi;
 import graphql.execution.AbortExecutionException;
+import graphql.execution.ExecutionContext;
 import graphql.execution.instrumentation.InstrumentationContext;
+import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.SimpleInstrumentation;
+import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
+import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
 import graphql.validation.ValidationError;
 import org.slf4j.Logger;
@@ -12,10 +17,11 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
 import static graphql.Assert.assertNotNull;
-import static graphql.execution.instrumentation.SimpleInstrumentationContext.whenCompleted;
+import static graphql.execution.instrumentation.SimpleInstrumentationContext.noOp;
 import static java.util.Optional.ofNullable;
 
 /**
@@ -76,41 +82,52 @@ public MaxQueryComplexityInstrumentation(int maxComplexity, FieldComplexityCalcu
         this.maxQueryComplexityExceededFunction = maxQueryComplexityExceededFunction;
     }
 
+    @Override
+    public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
+        return new State();
+    }
+
     @Override
     public InstrumentationContext<List<ValidationError>> beginValidation(InstrumentationValidationParameters parameters) {
-        return whenCompleted((errors, throwable) -> {
-            if ((errors != null && errors.size() > 0) || throwable != null) {
-                return;
-            }
-            QueryTraverser queryTraverser = newQueryTraverser(parameters);
-
-            Map<QueryVisitorFieldEnvironment, Integer> valuesByParent = new LinkedHashMap<>();
-            queryTraverser.visitPostOrder(new QueryVisitorStub() {
-                @Override
-                public void visitField(QueryVisitorFieldEnvironment env) {
-                    int childsComplexity = valuesByParent.getOrDefault(env, 0);
-                    int value = calculateComplexity(env, childsComplexity);
-
-                    valuesByParent.compute(env.getParentEnvironment(), (key, oldValue) ->
-                            ofNullable(oldValue).orElse(0) + value
-                    );
-                }
-            });
-            int totalComplexity = valuesByParent.getOrDefault(null, 0);
-            if (log.isDebugEnabled()) {
-                log.debug("Query complexity: {}", totalComplexity);
-            }
-            if (totalComplexity > maxComplexity) {
-                QueryComplexityInfo queryComplexityInfo = QueryComplexityInfo.newQueryComplexityInfo()
-                        .complexity(totalComplexity)
-                        .instrumentationValidationParameters(parameters)
-                        .build();
-                boolean throwAbortException = maxQueryComplexityExceededFunction.apply(queryComplexityInfo);
-                if (throwAbortException) {
-                    throw mkAbortException(totalComplexity, maxComplexity);
-                }
+        State state = parameters.getInstrumentationState();
+        // for API backwards compatibility reasons we capture the validation parameters, so we can put them into QueryComplexityInfo
+        state.instrumentationValidationParameters.set(parameters);
+        return noOp();
+    }
+
+    @Override
+    public InstrumentationContext<ExecutionResult> beginExecuteOperation(InstrumentationExecuteOperationParameters instrumentationExecuteOperationParameters) {
+        State state = instrumentationExecuteOperationParameters.getInstrumentationState();
+        QueryTraverser queryTraverser = newQueryTraverser(instrumentationExecuteOperationParameters.getExecutionContext());
+
+        Map<QueryVisitorFieldEnvironment, Integer> valuesByParent = new LinkedHashMap<>();
+        queryTraverser.visitPostOrder(new QueryVisitorStub() {
+            @Override
+            public void visitField(QueryVisitorFieldEnvironment env) {
+                int childsComplexity = valuesByParent.getOrDefault(env, 0);
+                int value = calculateComplexity(env, childsComplexity);
+
+                valuesByParent.compute(env.getParentEnvironment(), (key, oldValue) ->
+                        ofNullable(oldValue).orElse(0) + value
+                );
             }
         });
+        int totalComplexity = valuesByParent.getOrDefault(null, 0);
+        if (log.isDebugEnabled()) {
+            log.debug("Query complexity: {}", totalComplexity);
+        }
+        if (totalComplexity > maxComplexity) {
+            QueryComplexityInfo queryComplexityInfo = QueryComplexityInfo.newQueryComplexityInfo()
+                    .complexity(totalComplexity)
+                    .instrumentationValidationParameters(state.instrumentationValidationParameters.get())
+                    .instrumentationExecuteOperationParameters(instrumentationExecuteOperationParameters)
+                    .build();
+            boolean throwAbortException = maxQueryComplexityExceededFunction.apply(queryComplexityInfo);
+            if (throwAbortException) {
+                throw mkAbortException(totalComplexity, maxComplexity);
+            }
+        }
+        return noOp();
     }
 
     /**
@@ -125,12 +142,12 @@ protected AbortExecutionException mkAbortException(int totalComplexity, int maxC
         return new AbortExecutionException("maximum query complexity exceeded " + totalComplexity + " > " + maxComplexity);
     }
 
-    QueryTraverser newQueryTraverser(InstrumentationValidationParameters parameters) {
+    QueryTraverser newQueryTraverser(ExecutionContext executionContext) {
         return QueryTraverser.newQueryTraverser()
-                .schema(parameters.getSchema())
-                .document(parameters.getDocument())
-                .operationName(parameters.getOperation())
-                .variables(parameters.getVariables())
+                .schema(executionContext.getGraphQLSchema())
+                .document(executionContext.getDocument())
+                .operationName(executionContext.getExecutionInput().getOperationName())
+                .variables(executionContext.getVariables())
                 .build();
     }
 
@@ -156,5 +173,8 @@ private FieldComplexityEnvironment convertEnv(QueryVisitorFieldEnvironment query
         );
     }
 
+    private static class State implements InstrumentationState {
+        AtomicReference<InstrumentationValidationParameters> instrumentationValidationParameters = new AtomicReference<>();
+    }
 
 }
diff --git a/src/main/java/graphql/analysis/MaxQueryDepthInstrumentation.java b/src/main/java/graphql/analysis/MaxQueryDepthInstrumentation.java
@@ -1,18 +1,18 @@
 package graphql.analysis;
 
+import graphql.ExecutionResult;
 import graphql.PublicApi;
 import graphql.execution.AbortExecutionException;
+import graphql.execution.ExecutionContext;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.SimpleInstrumentation;
-import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
-import graphql.validation.ValidationError;
+import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.List;
 import java.util.function.Function;
 
-import static graphql.execution.instrumentation.SimpleInstrumentationContext.whenCompleted;
+import static graphql.execution.instrumentation.SimpleInstrumentationContext.noOp;
 
 /**
  * Prevents execution if the query depth is greater than the specified maxDepth.
@@ -49,26 +49,22 @@ public MaxQueryDepthInstrumentation(int maxDepth, Function<QueryDepthInfo, Boole
     }
 
     @Override
-    public InstrumentationContext<List<ValidationError>> beginValidation(InstrumentationValidationParameters parameters) {
-        return whenCompleted((errors, throwable) -> {
-            if ((errors != null && errors.size() > 0) || throwable != null) {
-                return;
-            }
-            QueryTraverser queryTraverser = newQueryTraverser(parameters);
-            int depth = queryTraverser.reducePreOrder((env, acc) -> Math.max(getPathLength(env.getParentEnvironment()), acc), 0);
-            if (log.isDebugEnabled()) {
-                log.debug("Query depth info: {}", depth);
-            }
-            if (depth > maxDepth) {
-                QueryDepthInfo queryDepthInfo = QueryDepthInfo.newQueryDepthInfo()
-                        .depth(depth)
-                        .build();
-                boolean throwAbortException = maxQueryDepthExceededFunction.apply(queryDepthInfo);
-                if (throwAbortException) {
-                    throw mkAbortException(depth, maxDepth);
-                }
+    public InstrumentationContext<ExecutionResult> beginExecuteOperation(InstrumentationExecuteOperationParameters parameters) {
+        QueryTraverser queryTraverser = newQueryTraverser(parameters.getExecutionContext());
+        int depth = queryTraverser.reducePreOrder((env, acc) -> Math.max(getPathLength(env.getParentEnvironment()), acc), 0);
+        if (log.isDebugEnabled()) {
+            log.debug("Query depth info: {}", depth);
+        }
+        if (depth > maxDepth) {
+            QueryDepthInfo queryDepthInfo = QueryDepthInfo.newQueryDepthInfo()
+                    .depth(depth)
+                    .build();
+            boolean throwAbortException = maxQueryDepthExceededFunction.apply(queryDepthInfo);
+            if (throwAbortException) {
+                throw mkAbortException(depth, maxDepth);
             }
-        });
+        }
+        return noOp();
     }
 
     /**
@@ -83,12 +79,12 @@ protected AbortExecutionException mkAbortException(int depth, int maxDepth) {
         return new AbortExecutionException("maximum query depth exceeded " + depth + " > " + maxDepth);
     }
 
-    QueryTraverser newQueryTraverser(InstrumentationValidationParameters parameters) {
+    QueryTraverser newQueryTraverser(ExecutionContext executionContext) {
         return QueryTraverser.newQueryTraverser()
-                .schema(parameters.getSchema())
-                .document(parameters.getDocument())
-                .operationName(parameters.getOperation())
-                .variables(parameters.getVariables())
+                .schema(executionContext.getGraphQLSchema())
+                .document(executionContext.getDocument())
+                .operationName(executionContext.getExecutionInput().getOperationName())
+                .variables(executionContext.getVariables())
                 .build();
     }
 

diff --git a/src/main/java/graphql/analysis/QueryComplexityInfo.java b/src/main/java/graphql/analysis/QueryComplexityInfo.java
@@ -1,6 +1,7 @@
 package graphql.analysis;
 
 import graphql.PublicApi;
+import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
 
 /**
@@ -10,11 +11,13 @@
 public class QueryComplexityInfo {
 
     private final int complexity;
-    private InstrumentationValidationParameters instrumentationValidationParameters;
+    private final InstrumentationValidationParameters instrumentationValidationParameters;
+    private final InstrumentationExecuteOperationParameters instrumentationExecuteOperationParameters;
 
-    private QueryComplexityInfo(int complexity, InstrumentationValidationParameters parameters) {
-        this.complexity = complexity;
-        this.instrumentationValidationParameters = parameters;
+    private QueryComplexityInfo(Builder builder) {
+        this.complexity = builder.complexity;
+        this.instrumentationValidationParameters = builder.instrumentationValidationParameters;
+        this.instrumentationExecuteOperationParameters = builder.instrumentationExecuteOperationParameters;
     }
 
     /**
@@ -35,6 +38,15 @@ public InstrumentationValidationParameters getInstrumentationValidationParameter
         return instrumentationValidationParameters;
     }
 
+    /**
+     * This returns the instrumentation execute operation parameters.
+     *
+     * @return the instrumentation execute operation parameters.
+     */
+    public InstrumentationExecuteOperationParameters getInstrumentationExecuteOperationParameters() {
+        return instrumentationExecuteOperationParameters;
+    }
+
     @Override
     public String toString() {
         return "QueryComplexityInfo{" +
@@ -54,6 +66,7 @@ public static class Builder {
 
         private int complexity;
         private InstrumentationValidationParameters instrumentationValidationParameters;
+        private InstrumentationExecuteOperationParameters instrumentationExecuteOperationParameters;
 
         private Builder() {
         }
@@ -62,6 +75,7 @@ private Builder() {
          * The query complexity.
          *
          * @param complexity the query complexity
+         *
          * @return this builder
          */
         public Builder complexity(int complexity) {
@@ -73,18 +87,24 @@ public Builder complexity(int complexity) {
          * The instrumentation validation parameters.
          *
          * @param parameters the instrumentation validation parameters.
+         *
          * @return this builder
          */
         public Builder instrumentationValidationParameters(InstrumentationValidationParameters parameters) {
             this.instrumentationValidationParameters = parameters;
             return this;
         }
 
+        public Builder instrumentationExecuteOperationParameters(InstrumentationExecuteOperationParameters instrumentationExecuteOperationParameters) {
+            this.instrumentationExecuteOperationParameters = instrumentationExecuteOperationParameters;
+            return this;
+        }
+
         /**
          * @return a built {@link QueryComplexityInfo} object
          */
         public QueryComplexityInfo build() {
-            return new QueryComplexityInfo(complexity, instrumentationValidationParameters);
+            return new QueryComplexityInfo(this);
         }
     }
 }