graphql-java · bbakerman · Dec 6, 2022 · Nov 28, 2022 · Nov 28, 2022 · Nov 30, 2022
diff --git a/src/main/java/graphql/analysis/values/ValueTraverser.java b/src/main/java/graphql/analysis/values/ValueTraverser.java
@@ -1,10 +1,11 @@
 package graphql.analysis.values;
 
 import com.google.common.collect.ImmutableList;
-import graphql.Assert;
 import graphql.PublicApi;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.DataFetchingEnvironmentImpl;
+import graphql.schema.GraphQLAppliedDirective;
+import graphql.schema.GraphQLAppliedDirectiveArgument;
 import graphql.schema.GraphQLArgument;
 import graphql.schema.GraphQLEnumType;
 import graphql.schema.GraphQLFieldDefinition;
@@ -23,6 +24,8 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static graphql.Assert.assertShouldNeverHappen;
+import static graphql.Assert.assertTrue;
 import static graphql.analysis.values.ValueVisitor.ABSENCE_SENTINEL;
 
 /**
@@ -51,10 +54,10 @@ private static class InputElements implements ValueVisitor.InputElements {
         private final List<GraphQLInputSchemaElement> unwrappedInputElements;
         private final GraphQLInputValueDefinition lastElement;
 
-        private InputElements(GraphQLInputValueDefinition startElement) {
+        private InputElements(GraphQLInputSchemaElement startElement) {
             this.inputElements = ImmutableList.of(startElement);
             this.unwrappedInputElements = ImmutableList.of(startElement);
-            this.lastElement = startElement;
+            this.lastElement = startElement instanceof GraphQLInputValueDefinition ? (GraphQLInputValueDefinition) startElement : null;
         }
 
         private InputElements(ImmutableList<GraphQLInputSchemaElement> inputElements) {
@@ -66,7 +69,7 @@ private InputElements(ImmutableList<GraphQLInputSchemaElement> inputElements) {
             List<GraphQLInputValueDefinition> inputValDefs = unwrappedInputElements.stream()
                     .filter(it -> it instanceof GraphQLInputValueDefinition)
                     .map(GraphQLInputValueDefinition.class::cast).collect(Collectors.toList());
-            this.lastElement = inputValDefs.get(inputValDefs.size() - 1);
+            this.lastElement = inputValDefs.isEmpty() ? null : inputValDefs.get(inputValDefs.size() - 1);
         }
 
 
@@ -125,20 +128,33 @@ public static Map<String, Object> visitPreOrder(Map<String, Object> coercedArgum
             String key = fieldArgument.getName();
             Object argValue = coercedArgumentValues.get(key);
             InputElements inputElements = new InputElements(fieldArgument);
-            Object newValue = visitPreOrderImpl(argValue, fieldArgument.getType(), inputElements, visitor);
+            Object newValue = visitor.visitArgumentValue(argValue, fieldArgument, inputElements);
             if (hasChanged(newValue, argValue)) {
                 if (!copied) {
                     coercedArgumentValues = new LinkedHashMap<>(coercedArgumentValues);
                     copied = true;
                 }
                 setNewValue(coercedArgumentValues, key, newValue);
             }
+            if (newValue != ABSENCE_SENTINEL) {
+                newValue = visitPreOrderImpl(argValue, fieldArgument.getType(), inputElements, visitor);
+                if (hasChanged(newValue, argValue)) {
+                    if (!copied) {
+                        coercedArgumentValues = new LinkedHashMap<>(coercedArgumentValues);
+                        copied = true;
+                    }
+                    setNewValue(coercedArgumentValues, key, newValue);
+                }
+            }
         }
         return coercedArgumentValues;
     }
 
     /**
-     * This will visit a single argument of a {@link GraphQLArgument} and if the visitor changes the value, it will return a new argument
+     * This will visit a single argument of a {@link GraphQLArgument} and if the visitor changes the value, it will return a new argument value
+     * <p>
+     * Note you cannot return the ABSENCE_SENTINEL from this method as its makes no sense to be somehow make the argument disappear.  Use
+     * {@link #visitPreOrder(Map, GraphQLFieldDefinition, ValueVisitor)} say to remove arguments in the fields map of arguments.
      *
      * @param coercedArgumentValue the starting coerced argument value
      * @param argument             the argument definition
@@ -147,7 +163,40 @@ public static Map<String, Object> visitPreOrder(Map<String, Object> coercedArgum
      * @return the same value if nothing changes or a new value if the visitor changes anything
      */
     public static Object visitPreOrder(Object coercedArgumentValue, GraphQLArgument argument, ValueVisitor visitor) {
-        return visitPreOrderImpl(coercedArgumentValue, argument.getType(), new InputElements(argument), visitor);
+        InputElements inputElements = new InputElements(argument);
+        Object newValue = visitor.visitArgumentValue(coercedArgumentValue, argument, inputElements);
+        if (newValue == ABSENCE_SENTINEL) {
+            assertShouldNeverHappen("It makes no sense to return the ABSENCE_SENTINEL during the visitPreOrder GraphQLArgument method");
+        }
+        newValue = visitPreOrderImpl(newValue, argument.getType(), inputElements, visitor);
+        if (newValue == ABSENCE_SENTINEL) {
+            assertShouldNeverHappen("It makes no sense to return the ABSENCE_SENTINEL during the visitPreOrder GraphQLArgument method");
+        }
+        return newValue;
+    }
+
+    /**
+     * This will visit a single argument of a {@link GraphQLAppliedDirective} and if the visitor changes the value, it will return a new argument value
+     * <p>
+     * Note you cannot return the ABSENCE_SENTINEL from this method as its makes no sense to be somehow make the argument disappear.
+     *
+     * @param coercedArgumentValue the starting coerced argument value
+     * @param argument             the applied argument
+     * @param visitor              the visitor to use
+     *
+     * @return the same value if nothing changes or a new value if the visitor changes anything
+     */
+    public static Object visitPreOrder(Object coercedArgumentValue, GraphQLAppliedDirectiveArgument argument, ValueVisitor visitor) {
+        InputElements inputElements = new InputElements(argument);
+        Object newValue = visitor.visitAppliedDirectiveArgumentValue(coercedArgumentValue, argument, inputElements);
+        if (newValue == ABSENCE_SENTINEL) {
+            assertShouldNeverHappen("It makes no sense to return the ABSENCE_SENTINEL during the visitPreOrder GraphQLAppliedDirectiveArgument method");
+        }
+        newValue = visitPreOrderImpl(newValue, argument.getType(), inputElements, visitor);
+        if (newValue == ABSENCE_SENTINEL) {
+            assertShouldNeverHappen("It makes no sense to return the ABSENCE_SENTINEL during the visitPreOrder GraphQLAppliedDirectiveArgument method");
+        }
+        return newValue;
     }
 
     private static Object visitPreOrderImpl(Object coercedValue, GraphQLInputType startingInputType, InputElements containingElements, ValueVisitor visitor) {
@@ -166,13 +215,13 @@ private static Object visitPreOrderImpl(Object coercedValue, GraphQLInputType st
         } else if (inputType instanceof GraphQLEnumType) {
             return visitor.visitEnumValue(coercedValue, (GraphQLEnumType) inputType, containingElements);
         } else {
-            return Assert.assertShouldNeverHappen("ValueTraverser can only be called on full materialised schemas");
+            return assertShouldNeverHappen("ValueTraverser can only be called on full materialised schemas");
         }
     }
 
     private static Object visitObjectValue(Object coercedValue, GraphQLInputObjectType inputObjectType, InputElements containingElements, ValueVisitor visitor) {
         if (coercedValue != null) {
-            Assert.assertTrue(coercedValue instanceof Map, () -> "A input object type MUST have an Map<String,Object> value");
+            assertTrue(coercedValue instanceof Map, () -> "A input object type MUST have an Map<String,Object> value");
         }
         @SuppressWarnings("unchecked")
         Map<String, Object> map = (Map<String, Object>) coercedValue;
@@ -217,7 +266,7 @@ private static Object visitObjectValue(Object coercedValue, GraphQLInputObjectTy
 
     private static Object visitListValue(Object coercedValue, GraphQLList listInputType, InputElements containingElements, ValueVisitor visitor) {
         if (coercedValue != null) {
-            Assert.assertTrue(coercedValue instanceof List, () -> "A list type MUST have an List value");
+            assertTrue(coercedValue instanceof List, () -> "A list type MUST have an List value");
         }
         @SuppressWarnings("unchecked")
         List<Object> list = (List<Object>) coercedValue;

diff --git a/src/main/java/graphql/analysis/values/ValueVisitor.java b/src/main/java/graphql/analysis/values/ValueVisitor.java
@@ -1,6 +1,8 @@
 package graphql.analysis.values;
 
 import graphql.PublicSpi;
+import graphql.schema.GraphQLAppliedDirectiveArgument;
+import graphql.schema.GraphQLArgument;
 import graphql.schema.GraphQLEnumType;
 import graphql.schema.GraphQLInputObjectField;
 import graphql.schema.GraphQLInputObjectType;
@@ -122,4 +124,33 @@ interface InputElements {
     default @Nullable List<Object> visitListValue(@Nullable List<Object> coercedValue, GraphQLList listInputType, InputElements inputElements) {
         return coercedValue;
     }
+
+
+    /**
+     * This is called when a {@link GraphQLArgument} is encountered
+     *
+     * @param coercedValue    the value that is in coerced form
+     * @param graphQLArgument the {@link GraphQLArgument} in play
+     * @param inputElements   the elements that lead to this value and type
+     *
+     * @return the same value or a new value
+     */
+    default @Nullable Object visitArgumentValue(@Nullable Object coercedValue, GraphQLArgument graphQLArgument, InputElements inputElements) {
+        return coercedValue;
+    }
+
+
+    /**
+     * This is called when a {@link GraphQLAppliedDirectiveArgument} is encountered
+     *
+     * @param coercedValue                    the value that is in coerced form
+     * @param graphQLAppliedDirectiveArgument the {@link GraphQLAppliedDirectiveArgument} in play
+     * @param inputElements                   the elements that lead to this value and type
+     *
+     * @return the same value or a new value
+     */
+    default @Nullable Object visitAppliedDirectiveArgumentValue(@Nullable Object coercedValue, GraphQLAppliedDirectiveArgument graphQLAppliedDirectiveArgument, InputElements inputElements) {
+        return coercedValue;
+    }
+
 }
diff --git a/src/main/java/graphql/schema/GraphQLAppliedDirectiveArgument.java b/src/main/java/graphql/schema/GraphQLAppliedDirectiveArgument.java
@@ -26,7 +26,7 @@
  * You can think of them as 'instances' of {@link GraphQLArgument}, when applied to a directive on a schema element
  */
 @PublicApi
-public class GraphQLAppliedDirectiveArgument implements GraphQLNamedSchemaElement {
+public class GraphQLAppliedDirectiveArgument implements GraphQLNamedSchemaElement, GraphQLInputSchemaElement {
 
     private final String name;
     private final InputValueWithState value;

diff --git a/src/test/groovy/graphql/analysis/values/ValueTraverserTest.groovy b/src/test/groovy/graphql/analysis/values/ValueTraverserTest.groovy
@@ -1,11 +1,14 @@
 package graphql.analysis.values
 
+import graphql.AssertException
 import graphql.ExecutionInput
 import graphql.GraphQL
 import graphql.TestUtil
 import graphql.schema.DataFetcher
 import graphql.schema.DataFetchingEnvironment
 import graphql.schema.DataFetchingEnvironmentImpl
+import graphql.schema.GraphQLAppliedDirectiveArgument
+import graphql.schema.GraphQLArgument
 import graphql.schema.GraphQLEnumType
 import graphql.schema.GraphQLFieldDefinition
 import graphql.schema.GraphQLFieldsContainer
@@ -17,6 +20,7 @@ import graphql.schema.GraphQLNamedSchemaElement
 import graphql.schema.GraphQLScalarType
 import graphql.schema.idl.SchemaDirectiveWiring
 import graphql.schema.idl.SchemaDirectiveWiringEnvironment
+import org.jetbrains.annotations.Nullable
 import spock.lang.Specification
 
 import static graphql.schema.idl.RuntimeWiring.newRuntimeWiring
@@ -293,6 +297,154 @@ class ValueTraverserTest extends Specification {
         ]
     }
 
+    def "can visit arguments and change things"() {
+        def sdl = """
+            type Query {
+                field(arg1 : Input!, arg2 : Input, removeArg : Input) : String
+            }
+
+            input Input {
+                name : String
+                age : Int
+                input : Input
+            }
+        """
+        def schema = TestUtil.schema(sdl)
+
+        def fieldDef = schema.getObjectType("Query").getFieldDefinition("field")
+        def argValues = [
+                arg1     :
+                        [name: "Tess", age: 42],
+                arg2     :
+                        [name: "Tom", age: 24],
+                removeArg:
+                        [name: "Gone-ski", age: 99],
+        ]
+        def visitor = new ValueVisitor() {
+            @Override
+            Object visitArgumentValue(@Nullable Object coercedValue, GraphQLArgument graphQLArgument, ValueVisitor.InputElements inputElements) {
+                if (graphQLArgument.name == "arg2") {
+                    return [name: "Harry Potter", age: 54]
+                }
+                if (graphQLArgument.name == "removeArg") {
+                    return ABSENCE_SENTINEL
+                }
+                return coercedValue
+            }
+        }
+        when:
+        def actual = ValueTraverser.visitPreOrder(argValues, fieldDef, visitor)
+
+        def expected = [
+                arg1:
+                        [name: "Tess", age: 42],
+                arg2:
+                        [name: "Harry Potter", age: 54]
+        ]
+        then:
+        actual == expected
+
+
+        // can change a DFE arguments
+        when:
+        def startingDFE = DataFetchingEnvironmentImpl.newDataFetchingEnvironment().fieldDefinition(fieldDef).arguments(argValues).build()
+        def newDFE = ValueTraverser.visitPreOrder(startingDFE, visitor)
+
+        then:
+        newDFE.getArguments() == expected
+        newDFE.getFieldDefinition() == fieldDef
+
+        // can change a single arguments
+        when:
+        def newValues = ValueTraverser.visitPreOrder(argValues['arg2'], fieldDef.getArgument("arg2"), visitor)
+
+        then:
+        newValues == [name: "Harry Potter", age: 54]
+
+        // catches non sense states
+        when:
+        ValueTraverser.visitPreOrder([:], fieldDef.getArgument("removeArg"), visitor)
+
+        then:
+        thrown(AssertException.class)
+    }
+
+    def "can handle applied directive arguments"() {
+        def sdl = """
+            directive @d(
+                arg1 :  Input
+                arg2 :  Input
+                removeArg :  Input
+            ) on FIELD_DEFINITION
+
+            type Query {
+                field : String @d(
+                arg1:
+                        {name: "Tom Riddle", age: 42}
+                arg2:
+                        {name: "Ron Weasley", age: 42}
+                removeArg:
+                        {name: "Ron Weasley", age: 42}
+                )
+            }
+
+            input Input {
+                name : String
+                age : Int
+                input : Input
+            }
+        """
+        def schema = TestUtil.schema(sdl)
+
+        def fieldDef = schema.getObjectType("Query").getFieldDefinition("field")
+        def appliedDirective = fieldDef.getAppliedDirective("d")
+        def visitor = new ValueVisitor() {
+
+            @Override
+            Object visitScalarValue(@Nullable Object coercedValue, GraphQLScalarType inputType, ValueVisitor.InputElements inputElements) {
+                if (coercedValue == "Tom Riddle") {
+                    return "Happy Potter"
+                }
+                return coercedValue
+            }
+
+            @Override
+            Object visitAppliedDirectiveArgumentValue(@Nullable Object coercedValue, GraphQLAppliedDirectiveArgument graphQLAppliedDirectiveArgument, ValueVisitor.InputElements inputElements) {
+                if (graphQLAppliedDirectiveArgument.name == "arg2") {
+                    return [name: "Harry Potter", age: 54]
+                }
+                if (graphQLAppliedDirectiveArgument.name == "removeArg") {
+                    return ABSENCE_SENTINEL
+                }
+                return coercedValue
+            }
+        }
+
+
+        def appliedDirectiveArgument = appliedDirective.getArgument("arg1")
+        when:
+        def actual = ValueTraverser.visitPreOrder(appliedDirectiveArgument.getValue(), appliedDirectiveArgument, visitor)
+
+        then:
+        actual == [name: "Happy Potter", age: 42]
+
+        when:
+        appliedDirectiveArgument = appliedDirective.getArgument("arg2")
+        actual = ValueTraverser.visitPreOrder(appliedDirectiveArgument.getValue(), appliedDirectiveArgument, visitor)
+
+        then:
+        actual == [name: "Harry Potter", age: 54]
+
+
+        // catches non sense states
+        when:
+        appliedDirectiveArgument = appliedDirective.getArgument("removeArg")
+        ValueTraverser.visitPreOrder(appliedDirectiveArgument.getValue(), appliedDirectiveArgument, visitor)
+
+        then:
+        thrown(AssertException.class)
+    }
+
     def "can handle a null changes"() {
         def sdl = """
             type Query {