Allow resolver methods to take both the context and the environment

Author: oryan-block

src/main/kotlin/graphql/kickstart/tools/resolver/FieldResolverScanner.kt

@@ -25,7 +25,7 @@ internal class FieldResolverScanner(val options: SchemaParserOptions) {
 
     private val log = LoggerFactory.getLogger(javaClass)
 
-    private val allowedLastArgumentTypes = listOfNotNull(DataFetchingEnvironment::class.java, GraphQLContext::class.java, options.contextClass)
+    private val allowedTrailingParameterTypes = listOfNotNull(DataFetchingEnvironment::class.java, GraphQLContext::class.java, options.contextClass)
 
     private val methodsByNameCache = mutableMapOf<Pair<Class<out Any>, Boolean>, Map<String, List<Method>>>()
 
@@ -171,10 +171,9 @@ internal class FieldResolverScanner(val options: SchemaParserOptions) {
         }
 
         val methodParameterCount = method.parameterCountWithoutContinuation()
-        val methodLastParameter = method.parameterTypes.getOrNull(methodParameterCount - 1)
 
-        val correctParameterCount = methodParameterCount == requiredCount ||
-            (methodParameterCount == (requiredCount + 1) && allowedLastArgumentTypes.contains(methodLastParameter))
+        val correctParameterCount = methodParameterCount >= requiredCount &&
+            method.parameterTypes.slice(requiredCount until methodParameterCount).all { it in allowedTrailingParameterTypes }
         return correctParameterCount && appropriateFirstParameter
     }
 
@@ -203,7 +202,7 @@ internal class FieldResolverScanner(val options: SchemaParserOptions) {
         val sourceLocation = field.sourceLocation?.let { "$sourceName:${it.line}" } ?: "<unknown>"
 
         return "No method${if (scannedProperties) " or field" else ""} found as defined in schema $sourceLocation with any of the following signatures " +
-            "(with or without one of $allowedLastArgumentTypes as the last argument), in priority order:\n${signatures.joinToString("\n  ")}" +
+            "(optionally followed by any of $allowedTrailingParameterTypes), in priority order:\n${signatures.joinToString("\n  ")}" +
             if (isSubscription) "\n\nNote that a Subscription data fetcher must return a Publisher of events" else ""
     }
 

src/main/kotlin/graphql/kickstart/tools/resolver/MethodFieldResolver.kt

@@ -47,7 +47,6 @@ internal class MethodFieldResolver(
 
     private val isSuspendFunction = method.isSuspendFunction()
     private val numberOfParameters = method.parameterCountWithoutContinuation()
-    private val hasAdditionalParameter = numberOfParameters == (field.inputValueDefinitions.size + getIndexOffset() + 1)
 
     internal val objectMapper: ObjectMapper by lazy { options.objectMapperProvider.provide(field) }
 
@@ -103,11 +102,10 @@ internal class MethodFieldResolver(
             }
         }
 
-        // Add DataFetchingEnvironment/Context argument
-        if (this.hasAdditionalParameter) {
-            // suspend functions have a trailing Continuation parameter
-            when (this.method.parameterTypes[numberOfParameters - 1]) {
-                null -> throw ResolverError("Expected at least one argument but got none, this is most likely a bug with graphql-java-tools")
+        // Add DataFetchingEnvironment/Context arguments, stopping before the Continuation parameter of a suspend function
+        val firstTrailingParameter = field.inputValueDefinitions.size + getIndexOffset()
+        method.parameterTypes.slice(firstTrailingParameter until numberOfParameters).forEach { parameterType ->
+            when (parameterType) {
                 options.contextClass -> args.add { environment ->
                     val context: Any? = environment.graphQlContext[options.contextClass]
                     if (context != null) {

src/test/kotlin/graphql/kickstart/tools/MethodFieldResolverDataFetcherTest.kt

@@ -278,6 +278,20 @@ class MethodFieldResolverDataFetcherTest {
         assertEquals(resolver.get(createEnvironment(DataClass(), context = context)), true)
     }
 
+    @Test
+    fun `data fetcher passes custom context and environment if method has both as extra arguments`() {
+        val customContext = ContextClass()
+        val context = GraphQLContext.of(mapOf(ContextClass::class.java to customContext))
+        val options = SchemaParserOptions.newOptions().contextClass(ContextClass::class).build()
+        val resolver = createFetcher("active", options = options, resolver = object : GraphQLResolver<DataClass> {
+            fun isActive(dataClass: DataClass, ctx: ContextClass, env: DataFetchingEnvironment): Boolean {
+                return ctx == customContext && env.graphQlContext == context
+            }
+        })
+
+        assertEquals(resolver.get(createEnvironment(DataClass(), context = context)), true)
+    }
+
     @Test
     fun `data fetcher passes context to suspend function if method has extra argument and context is specified`() {
         val context = GraphQLContext.newContext().build()

src/test/kotlin/graphql/kickstart/tools/SchemaParserTest.kt

@@ -1103,7 +1103,7 @@ class SchemaParserTest {
         }
 
         val expected = """
-            No method or field found as defined in schema <unknown>:3 with any of the following signatures (with or without one of [interface graphql.schema.DataFetchingEnvironment, class graphql.GraphQLContext] as the last argument), in priority order:
+            No method or field found as defined in schema <unknown>:3 with any of the following signatures (optionally followed by any of [interface graphql.schema.DataFetchingEnvironment, class graphql.GraphQLContext]), in priority order:
 
               graphql.kickstart.tools.SchemaParserTest${"$"}parser should verify subscription resolver return type${"$"}Subscription.onItemCreated()
               graphql.kickstart.tools.SchemaParserTest${"$"}parser should verify subscription resolver return type${"$"}Subscription.getOnItemCreated()
@@ -1143,7 +1143,7 @@ class SchemaParserTest {
         }
 
         val expected = """
-            No method or field found as defined in schema <unknown>:3 with any of the following signatures (with or without one of [interface graphql.schema.DataFetchingEnvironment, class graphql.GraphQLContext] as the last argument), in priority order:
+            No method or field found as defined in schema <unknown>:3 with any of the following signatures (optionally followed by any of [interface graphql.schema.DataFetchingEnvironment, class graphql.GraphQLContext]), in priority order:
 
               graphql.kickstart.tools.SchemaParserTest${"$"}parser should verify subscription resolver generic future return type${"$"}Subscription.onItemCreated()
               graphql.kickstart.tools.SchemaParserTest${"$"}parser should verify subscription resolver generic future return type${"$"}Subscription.getOnItemCreated()