✨ Refactor graphql query time limit code

#487
SoftInstigate · Nov 15, 2023 · 7f95635 · 7f95635
1 parent 984ca9f
commit 7f95635
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/...nstrumentation/QueryTimeoutException.java → ...graphql/GraphQLQueryTimeoutException.java b/...nstrumentation/QueryTimeoutException.java → ...graphql/GraphQLQueryTimeoutException.java
@@ -3,16 +3,16 @@
  * Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template
  */
 
-package org.restheart.graphql.instrumentation;
+package org.restheart.graphql;
 
 import graphql.execution.AbortExecutionException;
 
 /**
  *
  * @author uji
  */
-public class QueryTimeoutException extends AbortExecutionException {
-    public QueryTimeoutException(String message) {
+public class GraphQLQueryTimeoutException extends AbortExecutionException {
+    public GraphQLQueryTimeoutException(String message) {
         super(message);
     }
 }
diff --git a/graphql/src/main/java/org/restheart/graphql/GraphQLService.java b/graphql/src/main/java/org/restheart/graphql/GraphQLService.java
@@ -45,7 +45,6 @@
 import org.restheart.graphql.dataloaders.AggregationBatchLoader;
 import org.restheart.graphql.dataloaders.QueryBatchLoader;
 import org.restheart.graphql.instrumentation.MaxQueryTimeInstrumentation;
-import org.restheart.graphql.instrumentation.QueryTimeoutException;
 import org.restheart.graphql.models.AggregationMapping;
 import org.restheart.graphql.models.GraphQLApp;
 import org.restheart.graphql.models.QueryMapping;
@@ -74,6 +73,8 @@
 import graphql.ExecutionResultImpl;
 import graphql.GraphQL;
 import graphql.GraphQLError;
+import graphql.execution.instrumentation.ChainedInstrumentation;
+import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentationOptions;
 import graphql.language.Document;
@@ -194,12 +195,12 @@ public void handle(GraphQLRequest req, GraphQLResponse res) throws Exception {
             dispatcherInstrumentationOptions = dispatcherInstrumentationOptions.includeStatistics(true);
         }
 
-        var dispatcherInstrumentation = new DataLoaderDispatcherInstrumentation(dispatcherInstrumentationOptions);
-        var maxQueryTimeInstrumentation = new MaxQueryTimeInstrumentation(this.queryTimeLimit);
+        var chainedInstrumentations = new ArrayList<Instrumentation>();
+        chainedInstrumentations.add(new MaxQueryTimeInstrumentation(this.queryTimeLimit));
+        chainedInstrumentations.add(new DataLoaderDispatcherInstrumentation(dispatcherInstrumentationOptions));
 
         this.gql = GraphQL.newGraphQL(graphQLApp.getExecutableSchema())
-            .instrumentation(dispatcherInstrumentation)
-            .instrumentation(maxQueryTimeInstrumentation)
+            .instrumentation(new ChainedInstrumentation(chainedInstrumentations))
             .build();
 
         try {
@@ -213,7 +214,7 @@ public void handle(GraphQLRequest req, GraphQLResponse res) throws Exception {
             //  The graphql specification specifies:
             //  If an error was encountered during the execution that prevented a valid response, the data entry in the response should be null."
             if (result.getErrors() != null && !result.getErrors().isEmpty() && result.getData() == null) {
-                if (result.getErrors().stream().anyMatch(e -> e instanceof QueryTimeoutException)) {
+                if (result.getErrors().stream().anyMatch(e -> e instanceof GraphQLQueryTimeoutException)) {
                     res.setStatusCode(HttpStatus.SC_REQUEST_TIMEOUT);
                 } else {
                     res.setStatusCode(HttpStatus.SC_BAD_REQUEST);

diff --git a/graphql/src/main/java/org/restheart/graphql/instrumentation/MaxQueryTimeInstrumentation.java b/graphql/src/main/java/org/restheart/graphql/instrumentation/MaxQueryTimeInstrumentation.java
@@ -1,6 +1,7 @@
 package org.restheart.graphql.instrumentation;
 
 
+import org.restheart.graphql.GraphQLQueryTimeoutException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -63,7 +64,7 @@ public InstrumentationContext<ExecutionResult> beginField(InstrumentationFieldPa
      * @return an instance of AbortExecutionException
      */
     protected AbortExecutionException mkAbortException(long maxTime) {
-        return new QueryTimeoutException("Maximum query time limit of " + maxTime + "ms exceeded");
+        return new GraphQLQueryTimeoutException("Maximum query time limit of " + maxTime + "ms exceeded");
     }
 
     private static record State(long startTime) implements InstrumentationState {