Skip to content

Commit

Permalink
✨ Allow to limit the GraphQL queries execution by setting a time limit
Browse files Browse the repository at this point in the history
  • Loading branch information
ujibang committed Nov 15, 2023
1 parent d82dcb4 commit 984ca9f
Show file tree
Hide file tree
Showing 5 changed files with 114 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -276,9 +276,9 @@ mongo:
# TTL in milliseconds; specify a value < 0 to never expire cached entries
schema-cache-ttl: 60000

# Time limit in milliseconds for processing queries on the server (without network latency). 0 means no time limit
# The time limit in milliseconds for processing queries. Set to 0 for no time limit.
query-time-limit: 0
# Time limit in milliseconds for processing aggregations on the server (without network latency). 0 means no time limit
# The time limit in milliseconds for processing aggregations. Set to 0 for no time limit.
aggregation-time-limit: 0

changeStreamActivator:
Expand Down Expand Up @@ -306,6 +306,8 @@ graphql:
default-limit: 100
# max-limit is the maximum value for a Query limit
max-limit: 1000
# The time limit in milliseconds for processing queries. Set to 0 for no time limit.
query-time-limit: 0
verbose: false

cacheInvalidator:
Expand Down
6 changes: 4 additions & 2 deletions core/src/main/resources/restheart-default-config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -275,9 +275,9 @@ mongo:
# TTL in milliseconds; specify a value < 0 to never expire cached entries
schema-cache-ttl: 60000

# Time limit in milliseconds for processing queries on the server (without network latency). 0 means no time limit
# The time limit in milliseconds for processing queries. Set to 0 for no time limit.
query-time-limit: 0
# Time limit in milliseconds for processing aggregations on the server (without network latency). 0 means no time limit
# The time limit in milliseconds for processing aggregations. Set to 0 for no time limit.
aggregation-time-limit: 0

# Deprecated: it will be removed in RH v8.0
Expand All @@ -298,6 +298,8 @@ graphql:
default-limit: 100
# max-limit is the maximum value for a Query limit
max-limit: 1000
# The time limit in milliseconds for processing queries. Set to 0 for no time limit.
query-time-limit: 0
verbose: false

# Proxied resources - expose exrernal API with RESTHeart acting as a reverese proxy
Expand Down
21 changes: 17 additions & 4 deletions graphql/src/main/java/org/restheart/graphql/GraphQLService.java
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@
import org.restheart.graphql.datafetchers.GraphQLDataFetcher;
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;
Expand Down Expand Up @@ -90,6 +92,7 @@ public class GraphQLService implements Service<GraphQLRequest, GraphQLResponse>
public static final Boolean DEFAULT_VERBOSE = false;
public static final int DEFAULT_DEFAULT_LIMIT = 100;
public static final int DEFAULT_MAX_LIMIT = 1_000;
public static final int DEFAULT_QUERY_TIME_LIMIT = 0; // disabled

private static final Logger LOGGER = LoggerFactory.getLogger(GraphQLService.class);

Expand All @@ -99,6 +102,7 @@ public class GraphQLService implements Service<GraphQLRequest, GraphQLResponse>
private Boolean verbose = DEFAULT_VERBOSE;
private int defaultLimit = DEFAULT_DEFAULT_LIMIT;
private int maxLimit = DEFAULT_MAX_LIMIT;
private int queryTimeLimit = DEFAULT_QUERY_TIME_LIMIT;

@Inject("mclient")
private MongoClient mclient;
Expand All @@ -115,8 +119,9 @@ public void init()throws ConfigurationException, NoSuchFieldException, IllegalAc
this.verbose = argOrDefault(config, "verbose", DEFAULT_VERBOSE);
this.verbose = argOrDefault(config, "verbose", DEFAULT_VERBOSE);

this.defaultLimit = argOrDefault(config, "default-limit", 100);
this.maxLimit = argOrDefault(config, "max-limit", 1000);
this.defaultLimit = argOrDefault(config, "default-limit", DEFAULT_DEFAULT_LIMIT);
this.maxLimit = argOrDefault(config, "max-limit", DEFAULT_MAX_LIMIT);
this.queryTimeLimit = argOrDefault(config, "query-time-limit", DEFAULT_QUERY_TIME_LIMIT);

AppDefinitionLoadingCache.setTTL(argOrDefault(config, "app-def-cache-ttl", 1_000));

Expand Down Expand Up @@ -190,8 +195,12 @@ public void handle(GraphQLRequest req, GraphQLResponse res) throws Exception {
}

var dispatcherInstrumentation = new DataLoaderDispatcherInstrumentation(dispatcherInstrumentationOptions);
var maxQueryTimeInstrumentation = new MaxQueryTimeInstrumentation(this.queryTimeLimit);

this.gql = GraphQL.newGraphQL(graphQLApp.getExecutableSchema()).instrumentation(dispatcherInstrumentation).build();
this.gql = GraphQL.newGraphQL(graphQLApp.getExecutableSchema())
.instrumentation(dispatcherInstrumentation)
.instrumentation(maxQueryTimeInstrumentation)
.build();

try {
var result = this.gql.execute(inputBuilder.build());
Expand All @@ -204,7 +213,11 @@ 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) {
res.setStatusCode(HttpStatus.SC_BAD_REQUEST);
if (result.getErrors().stream().anyMatch(e -> e instanceof QueryTimeoutException)) {
res.setStatusCode(HttpStatus.SC_REQUEST_TIMEOUT);
} else {
res.setStatusCode(HttpStatus.SC_BAD_REQUEST);
}
}

if (this.verbose) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
package org.restheart.graphql.instrumentation;


import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import graphql.ExecutionResult;
import graphql.execution.AbortExecutionException;
import graphql.execution.instrumentation.InstrumentationContext;
import graphql.execution.instrumentation.InstrumentationState;
import static graphql.execution.instrumentation.SimpleInstrumentationContext.noOp;
import graphql.execution.instrumentation.SimplePerformantInstrumentation;
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
import graphql.execution.instrumentation.parameters.InstrumentationFieldParameters;

/**
* Aborts the execution if the query time is greater than the specified queryTimeLimit.
*/
public class MaxQueryTimeInstrumentation extends SimplePerformantInstrumentation {

private static final Logger LOGGER = LoggerFactory.getLogger(MaxQueryTimeInstrumentation.class);

private final int queryTimeLimit; // in ms

/**
* Creates a new instrumentation that aborts the execution if the query time is greater than the specified queryTimeLimit
*
* @param queryTimeLimit max allowed query time, otherwise execution will be aborted. Set queryTimeLimit <= 0 to disable it.
*/
public MaxQueryTimeInstrumentation(int queryTimeLimit) {
this.queryTimeLimit = queryTimeLimit;
}

@Override
public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
return new State(System.currentTimeMillis());
}

@Override
public InstrumentationContext<ExecutionResult> beginField(InstrumentationFieldParameters parameters, InstrumentationState rawState) {
if (this.queryTimeLimit <= 0) {
return noOp();
}

State state = InstrumentationState.ofState(rawState);
var elapsed = System.currentTimeMillis() - state.startTime;


if (elapsed > this.queryTimeLimit) {
LOGGER.debug("Exceeded query time limit of {} while attempting to fetch field '{}'. Current query elapsed time: {}.", this.queryTimeLimit, parameters.getExecutionStepInfo().getPath(), elapsed);
throw mkAbortException(this.queryTimeLimit);
} else {
LOGGER.trace("Fetching data for field '{}' initiated. Current elapsed query time: {}.", parameters.getExecutionStepInfo().getPath(), elapsed);
return noOp();
}
}

/**
* Generate the exception with error message
*
* @param maxTime the maximum time allowed
*
* @return an instance of AbortExecutionException
*/
protected AbortExecutionException mkAbortException(long maxTime) {
return new QueryTimeoutException("Maximum query time limit of " + maxTime + "ms exceeded");
}

private static record State(long startTime) implements InstrumentationState {
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
/*
* Click nbfs://nbhost/SystemFileSystem/Templates/Licenses/license-default.txt to change this license
* Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template
*/

package org.restheart.graphql.instrumentation;

import graphql.execution.AbortExecutionException;

/**
*
* @author uji
*/
public class QueryTimeoutException extends AbortExecutionException {
public QueryTimeoutException(String message) {
super(message);
}
}

0 comments on commit 984ca9f

Please sign in to comment.