Apply directives declared on the schema and its extensions

Directives on `schema @foo { ... }` and `extend schema @foo` were never
read, so e.g. Apollo Federation v2's `extend schema @link(...)` was
missing from the GraphQLSchema and from federation-jvm's _service.sdl.
They're now added as schema applied directives, the schema definition's
first, then its extensions'.

RootTypeInfo also took the last SchemaDefinition, which can be an
extension, so `schema { query: MyQuery }` followed by `extend schema @foo`
lost the custom root type. Operation types are now merged across the
definition and its extensions.

Only applied directives are set on the schema: graphql-java doesn't
resolve type references in the deprecated legacy schema directives, so
an enum argument would fail validation.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: oryan-block

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

@@ -104,7 +104,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 applied directives are built here:
+        // graphql-java doesn't resolve type references in the deprecated legacy schema directives, so enum arguments would fail.
+        val schemaAppliedDirectives = buildAppliedDirectives(rootInfo.getDirectives()).toList()
+        return SchemaObjects(query, mutation, subscription, types, schemaDirectives, codeRegistryBuilder, rootInfo.getDescription(), schemaAppliedDirectives)
     }
 
     /**

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,83 @@ 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.")
+    }
+
     private class BookQueryResolver : GraphQLQueryResolver {
         fun book(filter: BookFilter): Book? = null
     }
@@ -324,6 +403,11 @@ class DirectiveTest {
         val name: String
     )
 
+    private enum class LinkPurpose {
+        SECURITY,
+        EXECUTION
+    }
+
     private enum class AllowedState {
         ALLOWED,
         DISALLOWED

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

@@ -436,7 +436,7 @@ class SchemaClassScannerTest {
                 directive @key(fields: FieldSet!, resolvable: Boolean = true) repeatable on OBJECT | INTERFACE
                 directive @extends on OBJECT | INTERFACE
                 directive @external on FIELD_DEFINITION | OBJECT
-                directive @link(url: String!, as: String, for: link__Purpose) repeatable on SCHEMA
+                directive @link(url: String!, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
 
                 extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@shareable"])
 
@@ -465,7 +465,7 @@ class SchemaClassScannerTest {
             .options(SchemaParserOptions.newOptions().includeUnusedTypes(true).build())
             .dictionary(User::class)
             .dictionary("link__Purpose", LinkPurpose::class)
-            .scalars(fieldSetScalar)
+            .scalars(fieldSetScalar, linkImportScalar)
             .build()
             .makeExecutableSchema()
 
@@ -488,6 +488,16 @@ class SchemaClassScannerTest {
         })
         .build()
 
+    private val linkImportScalar: GraphQLScalarType = GraphQLScalarType.newScalar()
+        .name("link__Import")
+        .coercing(object : Coercing<String, String> {
+            override fun serialize(input: Any, context: GraphQLContext, locale: Locale) = input.toString()
+            override fun parseValue(input: Any, context: GraphQLContext, locale: Locale) = input.toString()
+            override fun parseLiteral(input: Value<*>, variables: CoercedVariables, context: GraphQLContext, locale: Locale) =
+                input.toString()
+        })
+        .build()
+
     class Product {
         var name: String? = null
     }

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

@@ -628,6 +628,47 @@ class SchemaParserTest {
         assertNull(schema.description)
     }
 
+    @Test
+    fun `parser should use root types declared in schema extensions`() {
+        val schema = SchemaParser.newParser()
+            .schemaString(
+                """
+                "This is a schema level description"
+                schema {
+                    query: SubstituteQuery
+                }
+
+                extend schema {
+                    mutation: SubstituteMutation
+                }
+
+                directive @contact(name: String!) on SCHEMA
+
+                extend schema @contact(name: "books-team")
+
+                type SubstituteQuery {
+                    query: String
+                }
+
+                type SubstituteMutation {
+                    mutation: String
+                }
+                """)
+            .resolvers(
+                object : GraphQLQueryResolver {
+                    fun query(): String? = null
+                },
+                object : GraphQLMutationResolver {
+                    fun mutation(): String? = null
+                })
+            .build()
+            .makeExecutableSchema()
+
+        assertEquals(schema.queryType.name, "SubstituteQuery")
+        assertEquals(schema.mutationType?.name, "SubstituteMutation")
+        assertEquals(schema.description, "This is a schema level description")
+    }
+
     enum class EnumType {
         TEST
     }