Merge pull request #895 from graphql-java-kickstart/bugfix/875

Run directive wirings on the assembled schema

Author: oryan-block

src/main/kotlin/graphql/kickstart/tools/SchemaParser.kt

@@ -71,7 +71,8 @@ class SchemaParser internal constructor(
     private val codeRegistryBuilder = GraphQLCodeRegistry.newCodeRegistry()
     private val directiveWiringHelper = DirectiveWiringHelper(runtimeWiring, codeRegistryBuilder)
 
-    private lateinit var schemaDirectives: SchemaDirectives
+    // by name, in the order they're declared
+    private lateinit var schemaDirectives: MutableMap<String, GraphQLDirective>
 
     /**
      * Parses the given schema with respect to the given dictionary and returns GraphQL objects.
@@ -88,21 +89,14 @@ class SchemaParser internal constructor(
 
         // Create GraphQL objects
         val inputObjects: MutableList<GraphQLInputObjectType> = mutableListOf()
-        createDirectives(inputObjects)
+        createDirectives()
         inputObjectDefinitions.forEach {
             if (inputObjects.none { io -> io.name == it.name }) {
                 inputObjects.add(createInputObject(it, inputObjects, mutableSetOf()))
             }
         }
         val interfaces = interfaceDefinitions.map { createInterfaceObject(it, inputObjects) }
         val objects = objectDefinitions.map { createObject(it, interfaces, inputObjects) }
-        if (options.unmappedInputFields != UnmappedInputFields.IGNORE) {
-            val checker = UnmappedInputFieldChecker(inputObjects.associateBy { it.name }, options.unmappedInputFields)
-            objectDefinitions.forEach { type ->
-                fieldResolversByType[type].orEmpty().values.filterIsInstance<MethodFieldResolver>().forEach { checker.check(type.name, it) }
-            }
-            checker.report()
-        }
         val unions = unionDefinitions.map { createUnionObject(it, objects) }
         val enums = enumDefinitions.map { createEnumObject(it) }
 
@@ -128,7 +122,33 @@ class SchemaParser internal constructor(
         // Directives applied to the schema itself, e.g. `extend schema @link(...)`. Only GraphQLAppliedDirectives are built here, not
         // the deprecated GraphQLDirective ones: graphql-java doesn't resolve type references in those, so enum arguments would fail.
         val schemaAppliedDirectives = buildAppliedDirectives(rootInfo.getDirectives(), SCHEMA).toList()
-        return SchemaObjects(query, mutation, subscription, types, schemaDirectives.all, codeRegistryBuilder, rootInfo.getDescription(), schemaAppliedDirectives)
+        val schemaObjects = SchemaObjects(query, mutation, subscription, types, schemaDirectives.values.toSet(), codeRegistryBuilder, rootInfo.getDescription(), schemaAppliedDirectives)
+        // in the order the types were created
+        val wiredObjects = if (directiveWiringHelper.isWiringRequired) wire(schemaObjects, inputObjects + interfaces + objects + unions + enums) else schemaObjects
+        if (options.unmappedInputFields != UnmappedInputFields.IGNORE) {
+            val wiredInputObjects = wiredObjects.dictionary.filterIsInstance<GraphQLInputObjectType>().associateBy { it.name }
+            val checker = UnmappedInputFieldChecker(wiredInputObjects, options.unmappedInputFields)
+            objectDefinitions.forEach { type ->
+                fieldResolversByType[type].orEmpty().values.filterIsInstance<MethodFieldResolver>().forEach { checker.check(type.name, it) }
+            }
+            checker.report()
+        }
+        return wiredObjects
+    }
+
+    /**
+     * Runs the directive wirings on the types once the schema is assembled, so they see resolved types, and swaps in the types they change.
+     */
+    private fun wire(schemaObjects: SchemaObjects, types: List<GraphQLNamedType>): SchemaObjects {
+        // building the schema resolves the type references in these types
+        val wired = directiveWiringHelper.wireTypes(schemaObjects.toSchema(), types)
+        if (wired.isEmpty()) return schemaObjects
+        return schemaObjects.copy(
+            query = wired[schemaObjects.query.name] as GraphQLObjectType? ?: schemaObjects.query,
+            mutation = schemaObjects.mutation?.let { wired[it.name] as GraphQLObjectType? ?: it },
+            subscription = schemaObjects.subscription?.let { wired[it.name] as GraphQLObjectType? ?: it },
+            dictionary = schemaObjects.dictionary.mapTo(linkedSetOf()) { wired[(it as GraphQLNamedType).name] ?: it }
+        )
     }
 
     /**
@@ -178,7 +198,7 @@ class SchemaParser internal constructor(
                 }
             }
 
-        return directiveWiringHelper.wireObject(builder.build())
+        return builder.build()
     }
 
     private fun createInputObject(definition: InputObjectTypeDefinition, inputObjects: List<GraphQLInputObjectType>,
@@ -214,7 +234,7 @@ class SchemaParser internal constructor(
                 }
             }
 
-        return directiveWiringHelper.wireInputObject(builder.build())
+        return builder.build()
     }
 
     private fun createEnumObject(definition: EnumTypeDefinition): GraphQLEnumType {
@@ -259,7 +279,7 @@ class SchemaParser internal constructor(
                 }
             }
 
-        return directiveWiringHelper.wireEnum(builder.build())
+        return builder.build()
     }
 
     private fun createInterfaceObject(interfaceDefinition: InterfaceTypeDefinition, inputObjects: List<GraphQLInputObjectType>): GraphQLInterfaceType {
@@ -286,7 +306,7 @@ class SchemaParser internal constructor(
                 }
             }
 
-        return directiveWiringHelper.wireInterFace(builder.build())
+        return builder.build()
     }
 
     private fun createUnionObject(definition: UnionTypeDefinition, types: List<GraphQLObjectType>): GraphQLUnionType {
@@ -304,7 +324,7 @@ class SchemaParser internal constructor(
                 getLeafUnionObjects(definition, types).forEach { possibleType(it) }
             }
 
-        return directiveWiringHelper.wireUnion(builder.build())
+        return builder.build()
     }
 
     private fun getLeafUnionObjects(definition: UnionTypeDefinition, types: List<GraphQLObjectType>): List<GraphQLObjectType> {
@@ -353,7 +373,7 @@ class SchemaParser internal constructor(
             .description(getDocumentation(fieldDefinition, options))
             .definition(fieldDefinition)
             .apply { getDeprecated(fieldDefinition.directives)?.let { deprecate(it) } }
-            .type(determineOutputType(fieldDefinition.type, inputObjects))
+            .type(determineOutputType(fieldDefinition.type))
             .withAppliedDirectives(*buildAppliedDirectives(fieldDefinition.directives, FIELD_DEFINITION))
             .withDirectives(*buildDirectives(fieldDefinition.directives, FIELD_DEFINITION))
             .apply {
@@ -376,11 +396,10 @@ class SchemaParser internal constructor(
             .build()
     }
 
-    private fun createDirectives(inputObjects: MutableList<GraphQLInputObjectType>) {
-        schemaDirectives = SchemaDirectives(inputObjects)
-        directiveDefinitions.forEach { schemaDirectives.find(it.name) }
+    private fun createDirectives() {
+        schemaDirectives = directiveDefinitions.associateTo(linkedMapOf()) { it.name to createDirective(it) }
         // because the arguments can have directives too, we attach them only after the directives themselves are created
-        schemaDirectives.update { d ->
+        schemaDirectives.replaceAll { _, d ->
             val arguments = d.arguments.map { a -> a.transform {
                 it.withAppliedDirectives(*buildAppliedDirectives(a.definition!!.directives, ARGUMENT_DEFINITION))
                     .withDirectives(*buildDirectives(a.definition!!.directives, ARGUMENT_DEFINITION))
@@ -389,33 +408,7 @@ class SchemaParser internal constructor(
         }
     }
 
-    /**
-     * The directive definitions. The input objects used by their arguments need the directives applied to them, so each directive
-     * is first built with type references to input objects, then rebuilt with the input objects the first time it's looked up.
-     * Only a directive applied within its own argument types gets its first build there.
-     */
-    private inner class SchemaDirectives(private val inputObjects: MutableList<GraphQLInputObjectType>) {
-        // by name, in the order they're declared
-        private val byName = directiveDefinitions.associateTo(linkedMapOf()) { definition ->
-            definition.name to createDirective(definition) { determineType(GraphQLInputType::class, it, permittedTypesForInputObject, inputObjects) as GraphQLInputType }
-        }
-        private val unfinished = directiveDefinitions.associateByTo(mutableMapOf()) { it.name }
-        // shared, so that an input object that's still being built isn't built again for another directive
-        private val referencingInputObjects = mutableSetOf<String>()
-
-        val all: Set<GraphQLDirective> get() = byName.values.toSet()
-
-        fun find(name: String): GraphQLDirective? {
-            unfinished.remove(name)?.let { definition ->
-                byName[name] = createDirective(definition) { determineInputType(it, inputObjects, referencingInputObjects) }
-            }
-            return byName[name]
-        }
-
-        fun update(transform: (GraphQLDirective) -> GraphQLDirective) = byName.replaceAll { _, directive -> transform(directive) }
-    }
-
-    private fun createDirective(definition: DirectiveDefinition, determineArgumentType: (Type<*>) -> GraphQLInputType): GraphQLDirective {
+    private fun createDirective(definition: DirectiveDefinition): GraphQLDirective {
         val locations = definition.directiveLocations.map { Introspection.DirectiveLocation.valueOf(it.name) }.toTypedArray()
 
         return GraphQLDirective.newDirective()
@@ -427,18 +420,19 @@ class SchemaParser internal constructor(
             .repeatable(definition.isRepeatable)
             .apply {
                 definition.inputValueDefinitions.forEach { argumentDefinition ->
-                    argument(createDirectiveArgument(argumentDefinition, determineArgumentType))
+                    argument(createDirectiveArgument(argumentDefinition))
                 }
             }
             .build()
     }
 
-    private fun createDirectiveArgument(definition: InputValueDefinition, determineArgumentType: (Type<*>) -> GraphQLInputType): GraphQLArgument {
+    private fun createDirectiveArgument(definition: InputValueDefinition): GraphQLArgument {
         return GraphQLArgument.newArgument()
             .name(definition.name)
             .definition(definition)
             .description(getDocumentation(definition, options))
-            .type(determineArgumentType(definition.type))
+            // type references, so that an input object used by a directive applied to it doesn't contain itself
+            .type(determineType(GraphQLInputType::class, definition.type, permittedTypesForInputObject) as GraphQLInputType)
             .apply { getDeprecated(definition.directives)?.let { deprecate(it) } }
             .apply { definition.defaultValue?.let { defaultValueLiteral(it) } }
             .build()
@@ -508,7 +502,7 @@ class SchemaParser internal constructor(
             }
         }
 
-    private fun findDirective(name: String): GraphQLDirective? = schemaDirectives.find(name) ?: BUILT_IN_DIRECTIVES[name]
+    private fun findDirective(name: String): GraphQLDirective? = schemaDirectives[name] ?: BUILT_IN_DIRECTIVES[name]
 
     private fun buildUndeclaredAppliedDirective(directive: Directive): GraphQLAppliedDirective {
         return GraphQLAppliedDirective.newDirective()
@@ -597,22 +591,18 @@ class SchemaParser internal constructor(
     private fun missingArgumentsWithDefault(directive: Directive, graphQLDirective: GraphQLDirective): List<GraphQLArgument> =
         graphQLDirective.arguments.filter { it.hasSetDefaultValue() && directive.getArgument(it.name) == null }
 
-    private fun determineOutputType(typeDefinition: Type<*>, inputObjects: List<GraphQLInputObjectType>) =
-        determineType(GraphQLOutputType::class, typeDefinition, permittedTypesForObject, inputObjects) as GraphQLOutputType
+    private fun determineOutputType(typeDefinition: Type<*>) =
+        determineType(GraphQLOutputType::class, typeDefinition, permittedTypesForObject) as GraphQLOutputType
 
+    // named types other than scalars are type references
     private fun <T : Any> determineType(
         expectedType: KClass<T>,
         typeDefinition: Type<*>,
-        allowedTypeReferences: Set<String>,
-        inputObjects: List<GraphQLInputObjectType>
+        allowedTypeReferences: Set<String>
     ): GraphQLType =
         when (typeDefinition) {
-            is ListType -> GraphQLList(determineType(expectedType, typeDefinition.type, allowedTypeReferences, inputObjects))
-            is NonNullType -> GraphQLNonNull(determineType(expectedType, typeDefinition.type, allowedTypeReferences, inputObjects))
-            is InputObjectTypeDefinition -> {
-                log.info("Create input object")
-                createInputObject(typeDefinition, inputObjects, mutableSetOf())
-            }
+            is ListType -> GraphQLList(determineType(expectedType, typeDefinition.type, allowedTypeReferences))
+            is NonNullType -> GraphQLNonNull(determineType(expectedType, typeDefinition.type, allowedTypeReferences))
             is TypeName -> {
                 val scalarType = customScalars[typeDefinition.name]
                     ?: GRAPHQL_SCALARS[typeDefinition.name]
@@ -622,7 +612,7 @@ class SchemaParser internal constructor(
                     if (!allowedTypeReferences.contains(typeDefinition.name)) {
                         throw unexpectedTypeError(typeDefinition, expectedType)
                     }
-                    inputObjects.find { it.name == typeDefinition.name } ?: GraphQLTypeReference(typeDefinition.name)
+                    GraphQLTypeReference(typeDefinition.name)
                 }
             }
             else -> throw SchemaError("Unknown type: $typeDefinition")

src/main/kotlin/graphql/kickstart/tools/directive/DirectiveWiringHelper.kt

@@ -16,31 +16,91 @@ internal class DirectiveWiringHelper(
 ) {
     private val schemaDirectiveParameters = Parameters(runtimeWiring, codeRegistryBuilder)
 
-    fun wireObject(objectType: GraphQLObjectType): GraphQLObjectType {
+    val isWiringRequired = runtimeWiring.registeredDirectiveWiring.isNotEmpty() || runtimeWiring.directiveWiring.isNotEmpty()
+
+    /**
+     * Runs the wirings on the given types of the assembled schema. If they change any, returns every type by name with type references to
+     * the others, so that the next build resolves them to the wired types.
+     */
+    fun wireTypes(schema: GraphQLSchema, types: List<GraphQLNamedType>): Map<String, GraphQLNamedType> {
+        val changed = linkedMapOf<String, GraphQLNamedType>()
+        types.forEach { type ->
+            val wired = when (type) {
+                is GraphQLObjectType -> wireObject(type)
+                is GraphQLInterfaceType -> wireInterFace(type)
+                is GraphQLInputObjectType -> wireInputObject(type)
+                is GraphQLEnumType -> wireEnum(type)
+                is GraphQLUnionType -> wireUnion(type)
+                else -> type
+            }
+            if (wired !== type) changed[type.name] = wired
+        }
+        if (changed.isEmpty()) return changed
+        val references = TypeReferences(schema)
+        return types.associate { it.name to references.inType(changed[it.name] ?: it) }
+    }
+
+    /**
+     * Swaps the schema's own types that a type refers to for type references, like they were before the schema was built.
+     */
+    private class TypeReferences(private val schema: GraphQLSchema) {
+        // the types wirings build in code, rebuilt once each
+        private val built = IdentityHashMap<GraphQLType, GraphQLType>()
+
+        @Suppress("UNCHECKED_CAST")
+        fun <T : GraphQLNamedType> inType(type: T): T = when (type) {
+            is GraphQLObjectType -> type.transform { it.clearFields().fields(type.fields.map(::inField)).replaceInterfaces(of(type.interfaces)) }
+            is GraphQLInterfaceType -> type.transform { it.clearFields().fields(type.fields.map(::inField)).replaceInterfacesOrReferences(of(type.interfaces)) }
+            is GraphQLUnionType -> type.transform { it.replacePossibleTypes(of(type.types)) }
+            is GraphQLInputObjectType -> type.transform { it.clearFields().fields(type.fields.map(::inInputField)) }
+            else -> type
+        } as T
+
+        private fun inField(field: GraphQLFieldDefinition) = field.transform {
+            it.type(of(field.type) as GraphQLOutputType)
+                .clearArguments()
+                .arguments(field.arguments.map { argument -> argument.transform { a -> a.type(of(argument.type) as GraphQLInputType) } })
+        }
+
+        private fun inInputField(field: GraphQLInputObjectField) = field.transform { it.type(of(field.type) as GraphQLInputType) }
+
+        private fun of(type: GraphQLType): GraphQLType = when (type) {
+            is GraphQLNonNull -> GraphQLNonNull.nonNull(of(type.wrappedType))
+            is GraphQLList -> GraphQLList.list(of(type.wrappedType))
+            // scalars aren't wired, and a reference only finds a scalar another field uses
+            is GraphQLScalarType -> type
+            is GraphQLNamedType -> if (schema.typeMap[type.name] === type) GraphQLTypeReference.typeRef(type.name) else built.getOrPut(type) { inType(type) }
+            else -> type
+        }
+
+        private fun of(types: List<GraphQLNamedOutputType>) = types.map { of(it) as GraphQLNamedOutputType }
+    }
+
+    private fun wireObject(objectType: GraphQLObjectType): GraphQLObjectType {
         return wireFields(objectType)
             .let { fields -> if (objectType.fields != fields) objectType.transform { it.clearFields().fields(fields) } else objectType }
             .let { wireDirectives(WiringWrapper(it, OBJECT, SchemaDirectiveWiring::onObject)) }
     }
 
-    fun wireInterFace(interfaceType: GraphQLInterfaceType): GraphQLInterfaceType {
+    private fun wireInterFace(interfaceType: GraphQLInterfaceType): GraphQLInterfaceType {
         return wireFields(interfaceType)
             .let { fields -> if (interfaceType.fields != fields) interfaceType.transform { it.clearFields().fields(fields) } else interfaceType }
             .let { wireDirectives(WiringWrapper(it, INTERFACE, SchemaDirectiveWiring::onInterface)) }
     }
 
-    fun wireInputObject(inputObjectType: GraphQLInputObjectType): GraphQLInputObjectType {
+    private fun wireInputObject(inputObjectType: GraphQLInputObjectType): GraphQLInputObjectType {
         return wireInputFields(inputObjectType)
             .let { fields -> if (inputObjectType.fields != fields) inputObjectType.transform { it.clearFields().fields(fields) } else inputObjectType }
             .let { wireDirectives(WiringWrapper(it, INPUT_OBJECT, SchemaDirectiveWiring::onInputObjectType)) }
     }
 
-    fun wireEnum(enumType: GraphQLEnumType): GraphQLEnumType {
+    private fun wireEnum(enumType: GraphQLEnumType): GraphQLEnumType {
         return wireEnumValues(enumType)
             .let { values -> if (enumType.values != values) enumType.transform { it.clearValues().values(values) } else enumType }
             .let { wireDirectives(WiringWrapper(it, ENUM, SchemaDirectiveWiring::onEnum)) }
     }
 
-    fun wireUnion(unionType: GraphQLUnionType): GraphQLUnionType {
+    private fun wireUnion(unionType: GraphQLUnionType): GraphQLUnionType {
         return wireDirectives(WiringWrapper(unionType, UNION, SchemaDirectiveWiring::onUnion))
     }