graphql-java · bbakerman · Nov 28, 2022 · Nov 28, 2022 · Nov 30, 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,7 +1,6 @@
 package graphql.analysis.values;
 
 import com.google.common.collect.ImmutableList;
-import graphql.Assert;
 import graphql.PublicApi;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.DataFetchingEnvironmentImpl;
@@ -23,6 +22,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;
 
 /**
@@ -125,20 +126,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 +161,16 @@ 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;
     }
 
     private static Object visitPreOrderImpl(Object coercedValue, GraphQLInputType startingInputType, InputElements containingElements, ValueVisitor visitor) {
@@ -166,13 +189,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 +240,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,7 @@
 package graphql.analysis.values;
 
 import graphql.PublicSpi;
+import graphql.schema.GraphQLArgument;
 import graphql.schema.GraphQLEnumType;
 import graphql.schema.GraphQLInputObjectField;
 import graphql.schema.GraphQLInputObjectType;
@@ -122,4 +123,19 @@ 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;
+    }
+
 }
diff --git a/src/test/groovy/graphql/analysis/values/ValueTraverserTest.groovy b/src/test/groovy/graphql/analysis/values/ValueTraverserTest.groovy
@@ -1,11 +1,13 @@
 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.GraphQLArgument
 import graphql.schema.GraphQLEnumType
 import graphql.schema.GraphQLFieldDefinition
 import graphql.schema.GraphQLFieldsContainer
@@ -17,6 +19,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 +296,78 @@ 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 a null changes"() {
         def sdl = """
             type Query {