graphql-java · bbakerman · Dec 12, 2018 · Dec 11, 2018
diff --git a/src/main/java/graphql/GraphQL.java b/src/main/java/graphql/GraphQL.java
@@ -540,12 +540,14 @@ private PreparsedDocumentEntry parseAndValidate(ExecutionInput executionInput, G
     }
 
     private ParseResult parse(ExecutionInput executionInput, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState) {
-        InstrumentationContext<Document> parseInstrumentation = instrumentation.beginParse(new InstrumentationExecutionParameters(executionInput, graphQLSchema, instrumentationState));
+        InstrumentationExecutionParameters parameters = new InstrumentationExecutionParameters(executionInput, graphQLSchema, instrumentationState);
+        InstrumentationContext<Document> parseInstrumentation = instrumentation.beginParse(parameters);
 
         Parser parser = new Parser();
         Document document;
         try {
             document = parser.parseDocument(executionInput.getQuery());
+            document = instrumentation.instrumentDocument(document, parameters);
         } catch (ParseCancellationException e) {
             parseInstrumentation.onCompleted(null, e);
             return ParseResult.ofError(e);

diff --git a/src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java b/src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java
@@ -21,7 +21,6 @@
 import graphql.schema.GraphQLSchema;
 import graphql.validation.ValidationError;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -176,6 +175,15 @@ public ExecutionInput instrumentExecutionInput(ExecutionInput executionInput, In
         return executionInput;
     }
 
+    @Override
+    public Document instrumentDocument(Document document, InstrumentationExecutionParameters parameters) {
+        for (Instrumentation instrumentation : instrumentations) {
+            InstrumentationState state = getState(instrumentation, parameters.getInstrumentationState());
+            document = instrumentation.instrumentDocument(document, parameters.withNewState(state));
+        }
+        return document;
+    }
+
     @Override
     public GraphQLSchema instrumentSchema(GraphQLSchema schema, InstrumentationExecutionParameters parameters) {
         for (Instrumentation instrumentation : instrumentations) {

diff --git a/src/main/java/graphql/execution/instrumentation/Instrumentation.java b/src/main/java/graphql/execution/instrumentation/Instrumentation.java
@@ -166,6 +166,18 @@ default ExecutionInput instrumentExecutionInput(ExecutionInput executionInput, I
         return executionInput;
     }
 
+    /**
+     * This is called to instrument a {@link graphql.language.Document} before it is used allowing you to adjust the query AST if you so desire
+     *
+     * @param document   the document to be used
+     * @param parameters the parameters describing the execution
+     *
+     * @return a non null instrumented Document, the default is to return to the same object
+     */
+    default Document instrumentDocument(Document document, InstrumentationExecutionParameters parameters) {
+        return document;
+    }
+
     /**
      * This is called to instrument a {@link graphql.schema.GraphQLSchema} before it is used to parse, validate
      * and execute a query, allowing you to adjust what types are used.

diff --git a/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy b/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy
@@ -5,8 +5,12 @@ import graphql.GraphQL
 import graphql.StarWarsSchema
 import graphql.execution.AsyncExecutionStrategy
 import graphql.execution.batched.BatchedExecutionStrategy
+import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
 import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
+import graphql.language.AstPrinter
+import graphql.language.Document
+import graphql.parser.Parser
 import graphql.schema.DataFetcher
 import graphql.schema.DataFetchingEnvironment
 import graphql.schema.PropertyDataFetcher
@@ -279,4 +283,51 @@ class InstrumentationTest extends Specification {
 
         er.data == [artoo: [id: '2001'], r2d2: [name: 'R2-D2']]
     }
+
+    def "document can be intercepted by instrumentation and changed"() {
+
+        given:
+
+        def query = """
+        {
+            hero {
+                id
+            }
+        }
+        """
+        def newQuery = """
+        {
+            hero {
+                name
+            }
+        }
+        """
+
+        def instrumentation = new TestingInstrumentation() {
+
+            @Override
+            Document instrumentDocument(Document document, InstrumentationExecutionParameters parameters) {
+                this.capturedData["originalDoc"] = AstPrinter.printAst(document)
+                new Parser().parseDocument(newQuery);
+            }
+
+        }
+
+        def graphQL = GraphQL
+                .newGraphQL(StarWarsSchema.starWarsSchema)
+                .instrumentation(instrumentation)
+                .build()
+
+        when:
+        def er = graphQL.execute(query)
+
+        then:
+        er.data == [hero: [name: 'R2-D2']]
+        instrumentation.capturedData["originalDoc"] == '''query {
+  hero {
+    id
+  }
+}
+'''
+    }
 }
diff --git a/src/test/groovy/graphql/execution/instrumentation/TestingInstrumentation.groovy b/src/test/groovy/graphql/execution/instrumentation/TestingInstrumentation.groovy
@@ -26,6 +26,7 @@ class TestingInstrumentation implements Instrumentation {
     List<Throwable> throwableList = []
     List<DataFetchingEnvironment> dfInvocations = []
     List<Class> dfClasses = []
+    def capturedData = [:]
 
     @Override
     InstrumentationState createState() {