Support directives on the schema

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

@@ -1,7 +1,9 @@
 package graphql.kickstart.tools
 
 import graphql.language.Description
+import graphql.language.Directive
 import graphql.language.SchemaDefinition
+import graphql.language.SchemaExtensionDefinition
 import graphql.language.TypeName
 
 /**
@@ -11,27 +13,38 @@ internal class RootTypeInfo private constructor(
     private val queryType: TypeName?,
     private val mutationType: TypeName?,
     private val subscriptionType: TypeName?,
-    private val description: Description?
+    private val description: Description?,
+    private val directives: List<Directive>
 ) {
     companion object {
         const val DEFAULT_QUERY_NAME = "Query"
         const val DEFAULT_MUTATION_NAME = "Mutation"
         const val DEFAULT_SUBSCRIPTION_NAME = "Subscription"
 
         fun fromSchemaDefinitions(definitions: List<SchemaDefinition>): RootTypeInfo {
-            val queryType = definitions.lastOrNull()?.operationTypeDefinitions?.find { it.name == "query" }?.typeName
-            val mutationType = definitions.lastOrNull()?.operationTypeDefinitions?.find { it.name == "mutation" }?.typeName
-            val subscriptionType = definitions.lastOrNull()?.operationTypeDefinitions?.find { it.name == "subscription" }?.typeName
-            val description = definitions.lastOrNull()?.description
+            // SchemaExtensionDefinition is a subclass of SchemaDefinition, so `definitions` contains the extensions too
+            val schemaDefinition = definitions.lastOrNull { it !is SchemaExtensionDefinition }
+            val extensionDefinitions = definitions.filterIsInstance<SchemaExtensionDefinition>()
+            // the schema definition comes first, then its extensions
+            val allDefinitions = listOfNotNull(schemaDefinition) + extensionDefinitions
 
-            return RootTypeInfo(queryType, mutationType, subscriptionType, description)
+            val operationTypes = allDefinitions.flatMap { it.operationTypeDefinitions }.associate { it.name to it.typeName }
+
+            return RootTypeInfo(
+                operationTypes["query"],
+                operationTypes["mutation"],
+                operationTypes["subscription"],
+                schemaDefinition?.description,
+                allDefinitions.flatMap { it.directives }
+            )
         }
     }
 
     fun getQueryName() = queryType?.name ?: DEFAULT_QUERY_NAME
     fun getMutationName() = mutationType?.name ?: DEFAULT_MUTATION_NAME
     fun getSubscriptionName() = subscriptionType?.name ?: DEFAULT_SUBSCRIPTION_NAME
     fun getDescription() = description?.content
+    fun getDirectives() = directives
 
     fun isMutationRequired() = mutationType != null
     fun isSubscriptionRequired() = subscriptionType != null

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

@@ -5,15 +5,38 @@ import graphql.schema.*
 /**
  * @author Andrew Potter
  */
-data class SchemaObjects(
+data class SchemaObjects @JvmOverloads constructor(
     val query: GraphQLObjectType,
     val mutation: GraphQLObjectType?,
     val subscription: GraphQLObjectType?,
     val dictionary: Set<GraphQLType>,
     val directives: Set<GraphQLDirective>,
     val codeRegistryBuilder: GraphQLCodeRegistry.Builder,
-    val description: String?
+    val description: String?,
+    // directives applied to the schema itself, e.g. `extend schema @link(...)`
+    val schemaAppliedDirectives: List<GraphQLAppliedDirective> = emptyList()
 ) {
+    // keeps the copy() signature from before schemaAppliedDirectives was added, for binary compatibility
+    @Deprecated("Kept for binary compatibility", level = DeprecationLevel.HIDDEN)
+    fun copy(
+        query: GraphQLObjectType = this.query,
+        mutation: GraphQLObjectType? = this.mutation,
+        subscription: GraphQLObjectType? = this.subscription,
+        dictionary: Set<GraphQLType> = this.dictionary,
+        directives: Set<GraphQLDirective> = this.directives,
+        codeRegistryBuilder: GraphQLCodeRegistry.Builder = this.codeRegistryBuilder,
+        description: String? = this.description
+    ) = copy(
+        query = query,
+        mutation = mutation,
+        subscription = subscription,
+        dictionary = dictionary,
+        directives = directives,
+        codeRegistryBuilder = codeRegistryBuilder,
+        description = description,
+        schemaAppliedDirectives = schemaAppliedDirectives
+    )
+
     // TODO change dictionary to Set<GraphQLNamedType> in the next major version and remove this cast
     @Suppress("UNCHECKED_CAST")
     private fun namedDictionary(): Set<GraphQLNamedType> = dictionary as Set<GraphQLNamedType>
@@ -29,6 +52,7 @@ data class SchemaObjects(
             .subscription(subscription)
             .additionalTypes(namedDictionary())
             .additionalDirectives(directives)
+            .withSchemaAppliedDirectives(schemaAppliedDirectives)
             .codeRegistry(codeRegistryBuilder.build())
             .build()
     }

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

@@ -1,6 +1,7 @@
 package graphql.kickstart.tools
 
 import graphql.Directives
+import graphql.Scalars
 import graphql.introspection.Introspection
 import graphql.introspection.Introspection.DirectiveLocation.INPUT_FIELD_DEFINITION
 import graphql.kickstart.tools.directive.DirectiveWiringHelper
@@ -104,7 +105,10 @@ class SchemaParser internal constructor(
         val additionalObjects = objects.filter { o -> o != query && o != subscription && o != mutation }
 
         val types = (additionalObjects.toSet() as Set<GraphQLType>) + inputObjects + enums + interfaces + unions
-        return SchemaObjects(query, mutation, subscription, types, schemaDirectives, codeRegistryBuilder, rootInfo.getDescription())
+        // 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()).toList()
+        return SchemaObjects(query, mutation, subscription, types, schemaDirectives, codeRegistryBuilder, rootInfo.getDescription(), schemaAppliedDirectives)
     }
 
     /**
@@ -352,7 +356,8 @@ class SchemaParser internal constructor(
         return directives.map { directive ->
             val graphQLDirective = schemaDirectives.find { d -> d.name == directive.name }
                 ?: BUILT_IN_DIRECTIVES[directive.name]
-                ?: throw SchemaError("Found applied directive ${directive.name} without corresponding directive definition.")
+                ?: if (options.allowUndeclaredDirectives) return@map buildUndeclaredAppliedDirective(directive)
+                else throw SchemaError("Found applied directive ${directive.name} without corresponding directive definition.")
             val graphQLArguments = graphQLDirective.arguments.associateBy { it.name }
 
             GraphQLAppliedDirective.newDirective()
@@ -394,6 +399,36 @@ class SchemaParser internal constructor(
         }.toTypedArray()
     }
 
+    private fun buildUndeclaredAppliedDirective(directive: Directive): GraphQLAppliedDirective {
+        return GraphQLAppliedDirective.newDirective()
+            .name(directive.name)
+            .definition(directive)
+            .comparatorRegistry(runtimeWiring.comparatorRegistry)
+            .apply {
+                directive.arguments.forEach { arg ->
+                    argument(GraphQLAppliedDirectiveArgument.newArgument()
+                        .name(arg.name)
+                        .type(guessDirectiveArgumentType(directive, arg.name, arg.value))
+                        .valueLiteral(arg.value)
+                        .build()
+                    )
+                }
+            }
+            .build()
+    }
+
+    // there's no directive definition to look the type up in, so it's guessed from the value
+    private fun guessDirectiveArgumentType(directive: Directive, argumentName: String, value: Value<*>): GraphQLInputType {
+        return when (value) {
+            is StringValue, is NullValue -> Scalars.GraphQLString
+            is IntValue -> Scalars.GraphQLInt
+            is FloatValue -> Scalars.GraphQLFloat
+            is BooleanValue -> Scalars.GraphQLBoolean
+            is ArrayValue -> GraphQLList(value.values.firstOrNull()?.let { guessDirectiveArgumentType(directive, argumentName, it) } ?: Scalars.GraphQLString)
+            else -> throw SchemaError("Can't guess the type of argument ${directive.name}#$argumentName of undeclared directive ${directive.name}, please declare the directive.")
+        }
+    }
+
     // TODO remove this once directives are fully replaced with applied directives
     private fun buildDirectives(
         directives: List<Directive>,
@@ -408,7 +443,9 @@ class SchemaParser internal constructor(
                 names.add(directive.name)
                 val graphQLDirective = this.schemaDirectives.find { d -> d.name == directive.name }
                     ?: BUILT_IN_DIRECTIVES[directive.name]
-                    ?: throw SchemaError("Found applied directive ${directive.name} without corresponding directive definition.")
+                    // graphql-java rejects legacy directives without a definition, so undeclared ones are only kept as applied directives
+                    ?: if (options.allowUndeclaredDirectives) continue
+                    else throw SchemaError("Found applied directive ${directive.name} without corresponding directive definition.")
                 val graphQLArguments = graphQLDirective.arguments.associateBy { it.name }
                 output.add(
                     GraphQLDirective.newDirective()

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

@@ -36,7 +36,8 @@ data class SchemaParserOptions internal constructor(
     val typeDefinitionFactories: List<TypeDefinitionFactory>,
     val fieldVisibility: GraphqlFieldVisibility?,
     val includeUnusedTypes: Boolean,
-    val useCommentsForDescriptions: Boolean
+    val useCommentsForDescriptions: Boolean,
+    val allowUndeclaredDirectives: Boolean
 ) {
     companion object {
         @JvmStatic
@@ -67,6 +68,7 @@ data class SchemaParserOptions internal constructor(
         private var fieldVisibility: GraphqlFieldVisibility? = null
         private var includeUnusedTypes = false
         private var useCommentsForDescriptions = true
+        private var allowUndeclaredDirectives = false
 
         fun contextClass(contextClass: Class<*>) = this.apply {
             this.contextClass = contextClass
@@ -153,6 +155,14 @@ data class SchemaParserOptions internal constructor(
             this.useCommentsForDescriptions = useCommentsForDescriptions
         }
 
+        /**
+         * Allows applying directives that have no definition in the schema, e.g. the Apollo Federation directives imported through `@link`.
+         * They're only added as applied directives, with argument types guessed from their values.
+         */
+        fun allowUndeclaredDirectives(allowUndeclaredDirectives: Boolean) = this.apply {
+            this.allowUndeclaredDirectives = allowUndeclaredDirectives
+        }
+
         @ExperimentalCoroutinesApi
         fun build(): SchemaParserOptions {
             val coroutineContextProvider = coroutineContextProvider
@@ -194,7 +204,8 @@ data class SchemaParserOptions internal constructor(
                 typeDefinitionFactories,
                 fieldVisibility,
                 includeUnusedTypes,
-                useCommentsForDescriptions
+                useCommentsForDescriptions,
+                allowUndeclaredDirectives
             )
         }
     }

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

@@ -7,6 +7,8 @@ import graphql.relay.SimpleListConnection
 import graphql.schema.*
 import graphql.schema.idl.SchemaDirectiveWiring
 import graphql.schema.idl.SchemaDirectiveWiringEnvironment
+import graphql.schema.idl.SchemaPrinter
+import org.junit.Assert.assertThrows
 import org.junit.Test
 
 class DirectiveTest {
@@ -304,6 +306,149 @@ class DirectiveTest {
         assert((schema.getType("Book") as GraphQLObjectType).getField("name").isDeprecated)
     }
 
+    @Test
+    fun `should apply directives on the schema and its extensions`() {
+        val schema = SchemaParser.newParser()
+            .schemaString(
+                """
+                directive @link(url: String!, for: LinkPurpose, import: [String]) repeatable on SCHEMA
+                directive @contact(name: String!) on SCHEMA
+
+                enum LinkPurpose {
+                    SECURITY
+                    EXECUTION
+                }
+
+                extend schema @link(url: "https://specs.apollo.dev/federation/v2.3", import: ["@key", "@shareable"])
+
+                schema @contact(name: "books-team") {
+                    query: Query
+                }
+
+                extend schema @link(url: "https://specs.apollo.dev/link/v1.0", for: SECURITY)
+
+                type Query {
+                    books: [Book!]
+                }
+
+                type Book {
+                    id: Int!
+                    name: String!
+                }
+                """)
+            .resolvers(QueryResolver())
+            .dictionary(LinkPurpose::class)
+            .build()
+            .makeExecutableSchema()
+
+        assertEquals(schema.schemaAppliedDirectives.map { it.name }, listOf("contact", "link", "link"))
+        assertEquals(
+            schema.getSchemaAppliedDirectives("link").map { it.getArgument("url")?.getValue<String>() },
+            listOf("https://specs.apollo.dev/federation/v2.3", "https://specs.apollo.dev/link/v1.0")
+        )
+        assertEquals(
+            schema.getSchemaAppliedDirectives("link").first().getArgument("import")?.getValue<List<String>>(),
+            listOf("@key", "@shareable")
+        )
+        assertEquals(schema.getSchemaAppliedDirectives("link").last().getArgument("for")?.getValue<LinkPurpose>(), LinkPurpose.SECURITY)
+
+        val printed = SchemaPrinter(SchemaPrinter.Options.defaultOptions().includeSchemaDefinition(true)).print(schema)
+        assert(printed.contains("""schema @contact(name : "books-team") @link(import : ["@key", "@shareable"], url : "https://specs.apollo.dev/federation/v2.3") @link(for : SECURITY, url : "https://specs.apollo.dev/link/v1.0"){""")) {
+            printed
+        }
+    }
+
+    @Test
+    fun `should fail on undeclared schema directive`() {
+        val error = assertThrows(SchemaError::class.java) {
+            SchemaParser.newParser()
+                .schemaString(
+                    """
+                    extend schema @link(url: "https://specs.apollo.dev/federation/v2.3")
+
+                    type Query {
+                        books: [Book!]
+                    }
+
+                    type Book {
+                        id: Int!
+                        name: String!
+                    }
+                    """)
+                .resolvers(QueryResolver())
+                .build()
+                .makeExecutableSchema()
+        }
+
+        assertEquals(error.message, "Found applied directive link without corresponding directive definition.")
+    }
+
+    @Test
+    fun `should allow undeclared directives when enabled`() {
+        val schema = SchemaParser.newParser()
+            .schemaString(
+                """
+                extend schema @link(url: "https://specs.apollo.dev/federation/v2.3", import: ["@key", "@shareable"])
+
+                type Query {
+                    books: [Book!] @shareable
+                }
+
+                type Book @key(fields: "id", resolvable: true) @custom(weight: 1.5, limit: 3, tag: null) {
+                    id: Int!
+                    name: String!
+                }
+                """)
+            .resolvers(QueryResolver())
+            .options(SchemaParserOptions.newOptions().allowUndeclaredDirectives(true).build())
+            .build()
+            .makeExecutableSchema()
+
+        val link = schema.getSchemaAppliedDirective("link")
+        assertEquals(link.getArgument("import")?.getValue<List<String>>(), listOf("@key", "@shareable"))
+        assertNotNull(schema.queryType.getField("books").getAppliedDirective("shareable"))
+
+        val book = schema.getObjectType("Book")!!
+        val key = book.getAppliedDirective("key")
+        assertEquals(key.getArgument("fields")?.getValue<String>(), "id")
+        assertEquals(key.getArgument("resolvable")?.getValue<Boolean>(), true)
+        val custom = book.getAppliedDirective("custom")
+        assertEquals(custom.getArgument("weight")?.getValue<Double>(), 1.5)
+        assertEquals(custom.getArgument("limit")?.getValue<Int>(), 3)
+        assertNull(custom.getArgument("tag")?.getValue<String>())
+        // graphql-java rejects legacy directives without a definition
+        assert(book.directives.isEmpty())
+
+        val printed = SchemaPrinter(SchemaPrinter.Options.defaultOptions().includeSchemaDefinition(true)).print(schema)
+        assert(printed.contains("""type Book @custom(limit : 3, tag : null, weight : 1.5) @key(fields : "id", resolvable : true) {""")) {
+            printed
+        }
+    }
+
+    @Test
+    fun `should fail on undeclared directive with an argument whose type can't be guessed`() {
+        val error = assertThrows(SchemaError::class.java) {
+            SchemaParser.newParser()
+                .schemaString(
+                    """
+                    type Query {
+                        books: [Book!] @policy(purpose: SECURITY)
+                    }
+
+                    type Book {
+                        id: Int!
+                        name: String!
+                    }
+                    """)
+                .resolvers(QueryResolver())
+                .options(SchemaParserOptions.newOptions().allowUndeclaredDirectives(true).build())
+                .build()
+                .makeExecutableSchema()
+        }
+
+        assertEquals(error.message, "Can't guess the type of argument policy#purpose of undeclared directive policy, please declare the directive.")
+    }
+
     private class BookQueryResolver : GraphQLQueryResolver {
         fun book(filter: BookFilter): Book? = null
     }
@@ -324,6 +469,11 @@ class DirectiveTest {
         val name: String
     )
 
+    private enum class LinkPurpose {
+        SECURITY,
+        EXECUTION
+    }
+
     private enum class AllowedState {
         ALLOWED,
         DISALLOWED