Provide consistent support for dynamically registered DataFetchers

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/DgsCodeRegistryBuilder.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2024 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.netflix.graphql.dgs
+
+import com.netflix.graphql.dgs.internal.DataFetcherResultProcessor
+import graphql.schema.*
+
+/**
+ * Utility wrapper for [GraphQLCodeRegistry.Builder] which provides
+ * a consistent registration mechanism of DataFetchers similar to the annotation-based approach.
+ * Can be used as a first parameter of a [DgsCodeRegistry] annotated method.
+ */
+class DgsCodeRegistryBuilder(
+    private val dataFetcherResultProcessors: List<DataFetcherResultProcessor>,
+    private val graphQLCodeRegistry: GraphQLCodeRegistry.Builder
+) {
+
+    fun dataFetcher(coordinates: FieldCoordinates, dataFetcher: DataFetcher<*>): DgsCodeRegistryBuilder {
+        val wrapped = DataFetcherFactories.wrapDataFetcher(dataFetcher) { dfe, result ->
+            result?.let {
+                val env = DgsDataFetchingEnvironment(dfe)
+                dataFetcherResultProcessors.find { it.supportsType(result) }?.process(result, env) ?: result
+            }
+        }
+
+        graphQLCodeRegistry.dataFetcher(coordinates, wrapped)
+        return this
+    }
+
+    fun hasDataFetcher(coordinates: FieldCoordinates): Boolean {
+        return graphQLCodeRegistry.hasDataFetcher(coordinates)
+    }
+
+    fun getDataFetcher(coordinates: FieldCoordinates, fieldDefinition: GraphQLFieldDefinition): DataFetcher<*> {
+        return graphQLCodeRegistry.getDataFetcher(coordinates, fieldDefinition)
+    }
+}

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsSchemaProvider.kt

@@ -41,8 +41,6 @@ import graphql.language.TypeName
 import graphql.language.UnionTypeDefinition
 import graphql.parser.MultiSourceReader
 import graphql.schema.Coercing
-import graphql.schema.DataFetcher
-import graphql.schema.DataFetcherFactories
 import graphql.schema.DataFetcherFactory
 import graphql.schema.FieldCoordinates
 import graphql.schema.GraphQLCodeRegistry
@@ -187,12 +185,14 @@ class DgsSchemaProvider(
         val runtimeWiringBuilder =
             RuntimeWiring.newRuntimeWiring().codeRegistry(codeRegistryBuilder).fieldVisibility(fieldVisibility)
 
+        val dgsCodeRegistryBuilder = DgsCodeRegistryBuilder(dataFetcherResultProcessors, codeRegistryBuilder)
+
         dgsComponents.asSequence()
             .mapNotNull { dgsComponent -> invokeDgsTypeDefinitionRegistry(dgsComponent, mergedRegistry) }
             .fold(mergedRegistry) { a, b -> a.merge(b) }
         findScalars(applicationContext, runtimeWiringBuilder)
         findDirectives(applicationContext, runtimeWiringBuilder)
-        findDataFetchers(dgsComponents, codeRegistryBuilder, mergedRegistry)
+        findDataFetchers(dgsComponents, dgsCodeRegistryBuilder, mergedRegistry)
         findTypeResolvers(dgsComponents, runtimeWiringBuilder, mergedRegistry)
 
         dgsComponents.forEach { dgsComponent ->
@@ -259,17 +259,26 @@ class DgsSchemaProvider(
         codeRegistryBuilder: GraphQLCodeRegistry.Builder,
         registry: TypeDefinitionRegistry
     ) {
+        val dgsCodeRegistryBuilder = DgsCodeRegistryBuilder(dataFetcherResultProcessors, codeRegistryBuilder)
+
         dgsComponent.annotatedMethods<DgsCodeRegistry>()
             .forEach { method ->
-                if (method.returnType != GraphQLCodeRegistry.Builder::class.java) {
-                    throw InvalidDgsConfigurationException("Method annotated with @DgsCodeRegistry must have return type GraphQLCodeRegistry.Builder")
+                if (method.returnType != GraphQLCodeRegistry.Builder::class.java && method.returnType != DgsCodeRegistryBuilder::class.java) {
+                    throw InvalidDgsConfigurationException("Method annotated with @DgsCodeRegistry must have return type GraphQLCodeRegistry.Builder or DgsCodeRegistryBuilder")
                 }
 
-                if (method.parameterCount != 2 || method.parameterTypes[0] != GraphQLCodeRegistry.Builder::class.java || method.parameterTypes[1] != TypeDefinitionRegistry::class.java) {
-                    throw InvalidDgsConfigurationException("Method annotated with @DgsCodeRegistry must accept the following arguments: GraphQLCodeRegistry.Builder, TypeDefinitionRegistry. ${dgsComponent.javaClass.name}.${method.name} has the following arguments: ${method.parameterTypes.joinToString()}")
+                if (method.parameterCount != 2 ||
+                    method.parameterTypes[0] != method.returnType || // Check that the first argument is of type DgsCodeRegistryBuilder or GraphQLCodeRegistry.Builder and the return type is the same
+                    method.parameterTypes[1] != TypeDefinitionRegistry::class.java
+                ) {
+                    throw InvalidDgsConfigurationException("Method annotated with @DgsCodeRegistry must accept the following arguments: GraphQLCodeRegistry.Builder or DgsCodeRegistryBuilder, and TypeDefinitionRegistry. ${dgsComponent.javaClass.name}.${method.name} has the following arguments: ${method.parameterTypes.joinToString()}")
                 }
 
-                ReflectionUtils.invokeMethod(method, dgsComponent.instance, codeRegistryBuilder, registry)
+                if (method.returnType == DgsCodeRegistryBuilder::class.java) {
+                    ReflectionUtils.invokeMethod(method, dgsComponent.instance, dgsCodeRegistryBuilder, registry)
+                } else if (method.returnType == GraphQLCodeRegistry.Builder::class.java) {
+                    ReflectionUtils.invokeMethod(method, dgsComponent.instance, codeRegistryBuilder, registry)
+                }
             }
     }
 
@@ -290,7 +299,7 @@ class DgsSchemaProvider(
 
     private fun findDataFetchers(
         dgsComponents: Collection<DgsBean>,
-        codeRegistryBuilder: GraphQLCodeRegistry.Builder,
+        codeRegistryBuilder: DgsCodeRegistryBuilder,
         typeDefinitionRegistry: TypeDefinitionRegistry
     ) {
         dgsComponents.forEach { dgsComponent ->
@@ -323,7 +332,7 @@ class DgsSchemaProvider(
 
     private fun registerDataFetcher(
         typeDefinitionRegistry: TypeDefinitionRegistry,
-        codeRegistryBuilder: GraphQLCodeRegistry.Builder,
+        codeRegistryBuilder: DgsCodeRegistryBuilder,
         dgsComponent: DgsBean,
         method: Method,
         dgsDataAnnotation: MergedAnnotation<DgsData>,
@@ -376,12 +385,9 @@ class DgsSchemaProvider(
                         // if we have a datafetcher explicitly defined for a parentType/field, use that and do not
                         // register the base implementation for interfaces
                         if (!codeRegistryBuilder.hasDataFetcher(FieldCoordinates.coordinates(implType.name, field))) {
-                            val dataFetcher =
-                                createBasicDataFetcher(method, dgsComponent.instance)
-                            codeRegistryBuilder.dataFetcher(
-                                FieldCoordinates.coordinates(implType.name, field),
-                                dataFetcher
-                            )
+                            val dataFetcher = methodDataFetcherFactory.createDataFetcher(dgsComponent.instance, method)
+                            codeRegistryBuilder.dataFetcher(FieldCoordinates.coordinates(implType.name, field), dataFetcher)
+
                             dataFetcherTracingInstrumentationEnabled["${implType.name}.$field"] =
                                 enableTracingInstrumentation
                             dataFetcherMetricsInstrumentationEnabled["${implType.name}.$field"] =
@@ -391,12 +397,9 @@ class DgsSchemaProvider(
                 }
                 is UnionTypeDefinition -> {
                     type.memberTypes.asSequence().filterIsInstance<TypeName>().forEach { memberType ->
-                        val dataFetcher =
-                            createBasicDataFetcher(method, dgsComponent.instance)
-                        codeRegistryBuilder.dataFetcher(
-                            FieldCoordinates.coordinates(memberType.name, field),
-                            dataFetcher
-                        )
+                        val dataFetcher = methodDataFetcherFactory.createDataFetcher(dgsComponent.instance, method)
+                        codeRegistryBuilder.dataFetcher(FieldCoordinates.coordinates(memberType.name, field), dataFetcher)
+
                         dataFetcherTracingInstrumentationEnabled["${memberType.name}.$field"] = enableTracingInstrumentation
                         dataFetcherMetricsInstrumentationEnabled["${memberType.name}.$field"] = enableMetricsInstrumentation
                     }
@@ -410,11 +413,9 @@ class DgsSchemaProvider(
                             matchingField.inputValueDefinitions.asSequence().map { it.name }.toSet()
                         )
                     }
-                    val dataFetcher = createBasicDataFetcher(method, dgsComponent.instance)
-                    codeRegistryBuilder.dataFetcher(
-                        FieldCoordinates.coordinates(parentType, field),
-                        dataFetcher
-                    )
+
+                    val dataFetcher = methodDataFetcherFactory.createDataFetcher(dgsComponent.instance, method)
+                    codeRegistryBuilder.dataFetcher(FieldCoordinates.coordinates(parentType, field), dataFetcher)
                 }
                 else -> {
                     throw InvalidDgsConfigurationException(
@@ -569,17 +570,6 @@ class DgsSchemaProvider(
         return null
     }
 
-    private fun createBasicDataFetcher(method: Method, dgsComponent: Any): DataFetcher<Any?> {
-        val dataFetcher = methodDataFetcherFactory.createDataFetcher(dgsComponent, method)
-
-        return DataFetcherFactories.wrapDataFetcher(dataFetcher) { dfe, result ->
-            result?.let {
-                val env = DgsDataFetchingEnvironment(dfe)
-                dataFetcherResultProcessors.find { it.supportsType(result) }?.process(result, env) ?: result
-            }
-        }
-    }
-
     private fun findTypeResolvers(
         dgsComponents: Collection<DgsBean>,
         runtimeWiringBuilder: RuntimeWiring.Builder,

graphql-dgs/src/test/kotlin/com/netflix/graphql/dgs/DgsSchemaProviderTest.kt

@@ -20,6 +20,7 @@ import com.netflix.graphql.dgs.exceptions.DataFetcherSchemaMismatchException
 import com.netflix.graphql.dgs.exceptions.InvalidDgsConfigurationException
 import com.netflix.graphql.dgs.exceptions.InvalidTypeResolverException
 import com.netflix.graphql.dgs.exceptions.NoSchemaFoundException
+import com.netflix.graphql.dgs.internal.DataFetcherResultProcessor
 import com.netflix.graphql.dgs.internal.DefaultInputObjectMapper
 import com.netflix.graphql.dgs.internal.DgsSchemaProvider
 import com.netflix.graphql.dgs.internal.kotlin.test.Show
@@ -73,7 +74,8 @@ internal class DgsSchemaProviderTest {
         typeDefinitionRegistry: TypeDefinitionRegistry? = null,
         schemaLocations: List<String> = listOf(DgsSchemaProvider.DEFAULT_SCHEMA_LOCATION),
         componentFilter: ((Any) -> Boolean)? = null,
-        schemaWiringValidationEnabled: Boolean = true
+        schemaWiringValidationEnabled: Boolean = true,
+        dataFetcherResultProcessors: List<DataFetcherResultProcessor> = emptyList()
     ): DgsSchemaProvider {
         return DgsSchemaProvider(
             applicationContext = applicationContext,
@@ -87,7 +89,8 @@ internal class DgsSchemaProviderTest {
                 )
             ),
             componentFilter = componentFilter,
-            schemaWiringValidationEnabled = schemaWiringValidationEnabled
+            schemaWiringValidationEnabled = schemaWiringValidationEnabled,
+            dataFetcherResultProcessors = dataFetcherResultProcessors
         )
     }
 
@@ -468,6 +471,7 @@ internal class DgsSchemaProviderTest {
     fun allowMergingStaticAndDynamicSchema() {
         @DgsComponent
         class CodeRegistryComponent {
+            // Result should not be processed by DataFetcherResultProcessors
             @DgsCodeRegistry
             fun registry(
                 codeRegistryBuilder: GraphQLCodeRegistry.Builder,
@@ -477,29 +481,66 @@ internal class DgsSchemaProviderTest {
                 val coordinates = FieldCoordinates.coordinates("Query", "myField")
                 return codeRegistryBuilder.dataFetcher(coordinates, df)
             }
+
+            // Result should be processed by DataFetcherResultProcessors
+            @DgsCodeRegistry
+            fun dgsProcessedRegistry(
+                codeRegistryBuilder: DgsCodeRegistryBuilder,
+                @Suppress("unused_parameter") registry: TypeDefinitionRegistry?
+            ): DgsCodeRegistryBuilder {
+                val df = DataFetcher { "Runtime added field" }
+                val coordinates = FieldCoordinates.coordinates("Query", "myProcessedField")
+                return codeRegistryBuilder.dataFetcher(coordinates, df)
+            }
         }
 
         contextRunner.withBeans(HelloFetcher::class, CodeRegistryComponent::class).run { context ->
             val typeDefinitionRegistry = TypeDefinitionRegistry()
             val objectTypeExtensionDefinition = ObjectTypeExtensionDefinition.newObjectTypeExtensionDefinition()
                 .name("Query")
-                .fieldDefinition(
-                    FieldDefinition.newFieldDefinition()
-                        .name("myField")
-                        .type(TypeName("String")).build()
-                )
-                .build()
+                .fieldDefinitions(
+                    listOf(
+                        FieldDefinition.newFieldDefinition()
+                            .name("myField")
+                            .type(TypeName("String")).build(),
+                        FieldDefinition.newFieldDefinition()
+                            .name("myProcessedField")
+                            .type(TypeName("String")).build()
+                    )
+                ).build()
+
+            val processor = object : DataFetcherResultProcessor {
+                override fun supportsType(originalResult: Any): Boolean {
+                    return true
+                }
+
+                override fun process(originalResult: Any, dfe: DgsDataFetchingEnvironment): Any {
+                    // Avoid processing other results apart from the one for this test
+                    if (originalResult != "Runtime added field") {
+                        return originalResult
+                    }
+                    return originalResult as String + " [suffixFromProcessor]"
+                }
+            }
 
             typeDefinitionRegistry.add(objectTypeExtensionDefinition)
-            val schema = schemaProvider(applicationContext = context, typeDefinitionRegistry = typeDefinitionRegistry).schema()
+            val schema = schemaProvider(
+                applicationContext = context,
+                typeDefinitionRegistry = typeDefinitionRegistry,
+                dataFetcherResultProcessors = listOf(processor)
+            ).schema()
             val build = GraphQL.newGraphQL(schema).build()
             assertHello(build)
 
-            val executionResult2 = build.execute("{myField}")
-            assertTrue(executionResult2.isDataPresent)
-
-            val data = executionResult2.getData<Map<String, *>>()
+            val executionResult = build.execute("{myField}")
+            assertTrue(executionResult.isDataPresent)
+            val data = executionResult.getData<Map<String, *>>()
             assertEquals("Runtime added field", data["myField"])
+
+            val processedExecutionResult = build.execute("{myProcessedField}")
+            assertTrue(processedExecutionResult.isDataPresent)
+            val processedData = processedExecutionResult.getData<Map<String, *>>()
+            assertEquals("Runtime added field [suffixFromProcessor]", processedData["myProcessedField"])
         }
 
         contextRunner.withBeans(HelloFetcher::class, CodeRegistryComponent::class).run { context ->