fix: fixed imports when both schema and extension are present (#258)

While introducing [Fed 2.1 support](#249) I added a validation to ensure only single federation `@link` import is present. That logic had a bug that if `schema` was present it never hit the else condition as `Optional` had a value in it (empty stream). Updated logic to correctly process both `schema` and its extensions.

Resolves: #257

Author: dariuszkuc

graphql-java-support/src/main/java/com/apollographql/federation/graphqljava/directives/LinkDirectiveProcessor.java

@@ -38,15 +38,19 @@ private LinkDirectiveProcessor() {}
    */
   public static @Nullable Stream<SDLNamedDefinition> loadFederationImportedDefinitions(
       TypeDefinitionRegistry typeDefinitionRegistry) {
-    List<Directive> federationLinkDirectives =
+
+    Stream<Directive> schemaLinkDirectives =
         typeDefinitionRegistry
             .schemaDefinition()
             .map(LinkDirectiveProcessor::getFederationLinkDirectives)
-            .orElseGet(
-                () ->
-                    typeDefinitionRegistry.getSchemaExtensionDefinitions().stream()
-                        .flatMap(LinkDirectiveProcessor::getFederationLinkDirectives))
-            .collect(Collectors.toList());
+            .orElse(Stream.empty());
+
+    Stream<Directive> extensionLinkDirectives =
+        typeDefinitionRegistry.getSchemaExtensionDefinitions().stream()
+            .flatMap(LinkDirectiveProcessor::getFederationLinkDirectives);
+
+    List<Directive> federationLinkDirectives =
+        Stream.concat(schemaLinkDirectives, extensionLinkDirectives).collect(Collectors.toList());
 
     if (federationLinkDirectives.isEmpty()) {
       return null;

graphql-java-support/src/test/java/com/apollographql/federation/graphqljava/FederationTest.java

@@ -235,6 +235,20 @@ public void verifyFederationV2Transformation_multipleFedLinks_throwsException()
         () -> Federation.transform(schemaSDL).fetchEntities(env -> null).build());
   }
 
+  @Test
+  public void verifySchemaCanBeExtended() {
+    verifyFederationTransformation("schemas/extendSchema.graphql", true);
+  }
+
+  @Test
+  public void
+      verifyFederationV2Transformation_multipleFedLinksSchemaAndExtension_throwsException() {
+    final String schemaSDL = FileUtils.readResource("schemas/multipleSchemaLinks.graphql");
+    assertThrows(
+        MultipleFederationLinksException.class,
+        () -> Federation.transform(schemaSDL).fetchEntities(env -> null).build());
+  }
+
   private GraphQLSchema verifyFederationTransformation(
       String schemaFileName, boolean isFederationV2) {
     final RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring().build();

graphql-java-support/src/test/resources/schemas/extendSchema.graphql

@@ -0,0 +1,14 @@
+extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key"])
+
+schema {
+  query: Query
+}
+
+type Product @key(fields: "id") {
+    id: ID!
+}
+
+type Query {
+    product(id: ID!): Product
+}
+

graphql-java-support/src/test/resources/schemas/extendSchema_federated.graphql

@@ -0,0 +1,45 @@
+schema @link(import : ["@key"], url : "https://specs.apollo.dev/federation/v2.0"){
+  query: Query
+}
+
+directive @federation__extends on OBJECT | INTERFACE
+
+directive @federation__external on OBJECT | FIELD_DEFINITION
+
+directive @federation__override(from: String!) on FIELD_DEFINITION
+
+directive @federation__provides(fields: federation__FieldSet!) on FIELD_DEFINITION
+
+directive @federation__requires(fields: federation__FieldSet!) on FIELD_DEFINITION
+
+directive @federation__shareable on OBJECT | FIELD_DEFINITION
+
+directive @inaccessible on SCALAR | OBJECT | FIELD_DEFINITION | ARGUMENT_DEFINITION | INTERFACE | UNION | ENUM | ENUM_VALUE | INPUT_OBJECT | INPUT_FIELD_DEFINITION
+
+directive @key(fields: federation__FieldSet!, resolvable: Boolean = true) repeatable on OBJECT | INTERFACE
+
+directive @link(as: String, import: [link__Import], url: String!) repeatable on SCHEMA
+
+directive @tag(name: String!) repeatable on SCALAR | OBJECT | FIELD_DEFINITION | ARGUMENT_DEFINITION | INTERFACE | UNION | ENUM | ENUM_VALUE | INPUT_OBJECT | INPUT_FIELD_DEFINITION
+
+union _Entity = Product
+
+type Product @key(fields : "id", resolvable : true) {
+  id: ID!
+}
+
+type Query {
+  _entities(representations: [_Any!]!): [_Entity]!
+  _service: _Service!
+  product(id: ID!): Product
+}
+
+type _Service {
+  sdl: String!
+}
+
+scalar _Any
+
+scalar federation__FieldSet
+
+scalar link__Import

graphql-java-support/src/test/resources/schemas/multipleSchemaLinks.graphql

@@ -0,0 +1,14 @@
+extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key"])
+
+schema @link(url: "https://specs.apollo.dev/federation/v2.1", import: ["@tag"]) {
+  query: Query
+}
+
+type Product @key(fields: "id") {
+    id: ID!
+    name: String! @tag(name: "example")
+}
+
+type Query {
+    product(id: ID!): Product
+}