fix: allow @link on schema type (#246)

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

@@ -1,5 +1,6 @@
 package com.apollographql.federation.graphqljava;
 
+import com.apollographql.federation.graphqljava.exceptions.UnsupportedLinkImportException;
 import graphql.language.Argument;
 import graphql.language.ArrayValue;
 import graphql.language.AstTransformer;
@@ -10,6 +11,7 @@
 import graphql.language.ObjectValue;
 import graphql.language.SDLNamedDefinition;
 import graphql.language.ScalarTypeDefinition;
+import graphql.language.SchemaDefinition;
 import graphql.language.StringValue;
 import graphql.language.TypeDefinition;
 import graphql.language.TypeName;
@@ -22,12 +24,12 @@
 import graphql.schema.idl.TypeDefinitionRegistry;
 import java.io.File;
 import java.io.Reader;
-import java.util.AbstractMap;
-import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -220,115 +222,95 @@ private static RuntimeWiring addScalarsToRuntimeWiring(
   }
 
   /**
-   * Looks for an @link extension and gets the import from it
+   * Looks for a @link directive and parses imports information from it.
    *
    * @return - null if this typeDefinitionRegistry is not Fed2 - the list of imports and potential
    *     renames else
    */
   private static @Nullable Map<String, String> fed2DirectiveImports(
       TypeDefinitionRegistry typeDefinitionRegistry) {
-    List<Directive> linkDirectives =
-        typeDefinitionRegistry.getSchemaExtensionDefinitions().stream()
-            .flatMap(
-                schemaExtensionDefinition ->
-                    schemaExtensionDefinition.getDirectives().stream()
-                        .filter(directive -> directive.getName().equals("link")))
-            .filter(
-                directive -> {
-                  Optional<Argument> arg =
-                      directive.getArguments().stream()
-                          .filter(argument -> argument.getName().equals("url"))
-                          .findFirst();
-
-                  if (!arg.isPresent()) {
-                    return false;
-                  }
-
-                  Value value = arg.get().getValue();
-                  if (!(value instanceof StringValue)) {
-                    return false;
-                  }
-
-                  StringValue stringValue = (StringValue) value;
-                  return stringValue.getValue().equals("https://specs.apollo.dev/federation/v2.0");
-                })
-            .collect(Collectors.toList());
-
-    if (linkDirectives.isEmpty()) {
+    List<Directive> federationLinkDirectives =
+        typeDefinitionRegistry
+            .schemaDefinition()
+            .map(Federation::getFederationLinkDirective)
+            .map(Collections::singletonList)
+            .orElseGet(
+                () ->
+                    typeDefinitionRegistry.getSchemaExtensionDefinitions().stream()
+                        .map(Federation::getFederationLinkDirective)
+                        .filter(Objects::nonNull)
+                        .collect(Collectors.toList()));
+
+    if (federationLinkDirectives.isEmpty()) {
       return null;
+    } else {
+      Map<String, String> imports = new HashMap<>();
+      federationLinkDirectives.forEach(directive -> imports.putAll(parseLinkImports(directive)));
+
+      imports.put("@link", "@link");
+      return imports;
     }
+  }
 
-    Map<String, String> imports =
-        linkDirectives.stream()
-            .flatMap(
-                directive -> {
-                  Optional<Argument> arg =
-                      directive.getArguments().stream()
-                          .filter(argument -> argument.getName().equals("import"))
-                          .findFirst();
-
-                  if (!arg.isPresent()) {
-                    return Stream.empty();
-                  }
-
-                  Value value = arg.get().getValue();
-                  if (!(value instanceof ArrayValue)) {
-                    return Stream.empty();
-                  }
-
-                  ArrayValue arrayValue = (ArrayValue) value;
-
-                  List<Map.Entry<String, String>> entries = new ArrayList<>();
-
-                  for (Value imp : arrayValue.getValues()) {
-                    if (imp instanceof StringValue) {
-                      String name = ((StringValue) imp).getValue();
-                      entries.add(new AbstractMap.SimpleEntry(name, name));
-                    } else if (imp instanceof ObjectValue) {
-                      ObjectValue objectValue = (ObjectValue) imp;
-                      Optional<ObjectField> nameField =
-                          objectValue.getObjectFields().stream()
-                              .filter(field -> field.getName().equals("name"))
-                              .findFirst();
-                      Optional<ObjectField> asField =
-                          objectValue.getObjectFields().stream()
-                              .filter(field -> field.getName().equals("as"))
-                              .findFirst();
-
-                      if (!nameField.isPresent()) {
-                        throw new RuntimeException("Unsupported import: " + imp);
-                      }
-
-                      Value nameValue = nameField.get().getValue();
-                      if (!(nameValue instanceof StringValue)) {
-                        throw new RuntimeException("Unsupported import: " + imp);
-                      }
-
-                      String as;
-                      if (!asField.isPresent()) {
-                        as = null;
-                      } else {
-                        Value asValue = asField.get().getValue();
-                        if (!(asValue instanceof StringValue)) {
-                          throw new RuntimeException("Unsupported import: " + imp);
-                        }
-                        as = ((StringValue) asValue).getValue();
-                      }
-
-                      entries.add(
-                          new AbstractMap.SimpleEntry(((StringValue) nameValue).getValue(), as));
-                    } else {
-                      throw new RuntimeException("Unsupported import: " + imp.toString());
-                    }
-                  }
-
-                  return entries.stream();
-                })
-            .collect(
-                Collectors.toMap(
-                    Map.Entry::getKey, Map.Entry::getValue, (value1, value2) -> value2));
-
-    imports.put("@link", "@link");
+  private static @Nullable Directive getFederationLinkDirective(SchemaDefinition schemaDefinition) {
+    return schemaDefinition.getDirectives("link").stream()
+        .filter(
+            directive -> {
+              Argument urlArgument = directive.getArgument("url");
+              if (urlArgument != null && urlArgument.getValue() instanceof StringValue) {
+                StringValue value = (StringValue) urlArgument.getValue();
+                return "https://specs.apollo.dev/federation/v2.0".equals(value.getValue());
+              } else {
+                return false;
+              }
+            })
+        .findAny()
+        .orElse(null);
+  }
+
+  private static Map<String, String> parseLinkImports(Directive linkDirective) {
+    final Map<String, String> imports = new HashMap<>();
+
+    final Argument importArgument = linkDirective.getArgument("import");
+    if (importArgument != null && importArgument.getValue() instanceof ArrayValue) {
+      final ArrayValue linkImports = (ArrayValue) importArgument.getValue();
+      for (Value importedDefinition : linkImports.getValues()) {
+        if (importedDefinition instanceof StringValue) {
+          // no rename
+          final String name = ((StringValue) importedDefinition).getValue();
+          imports.put(name, name);
+        } else if (importedDefinition instanceof ObjectValue) {
+          // renamed import
+          final ObjectValue importedObjectValue = (ObjectValue) importedDefinition;
+
+          final Optional<ObjectField> nameField =
+              importedObjectValue.getObjectFields().stream()
+                  .filter(field -> field.getName().equals("name"))
+                  .findFirst();
+          final Optional<ObjectField> renameAsField =
+              importedObjectValue.getObjectFields().stream()
+                  .filter(field -> field.getName().equals("as"))
+                  .findFirst();
+
+          if (!nameField.isPresent() || !(nameField.get().getValue() instanceof StringValue)) {
+            throw new UnsupportedLinkImportException(importedObjectValue);
+          }
+          final String name = ((StringValue) nameField.get().getValue()).getValue();
+
+          if (!renameAsField.isPresent()) {
+            imports.put(name, name);
+          } else {
+            final Value renamedAsValue = renameAsField.get().getValue();
+            if (!(renamedAsValue instanceof StringValue)) {
+              throw new UnsupportedLinkImportException(importedObjectValue);
+            }
+            imports.put(name, ((StringValue) renamedAsValue).getValue());
+          }
+        } else {
+          throw new UnsupportedLinkImportException(importedDefinition);
+        }
+      }
+    }
     return imports;
   }
 }

graphql-java-support/src/main/java/com/apollographql/federation/graphqljava/exceptions/UnsupportedLinkImportException.java

@@ -0,0 +1,17 @@
+package com.apollographql.federation.graphqljava.exceptions;
+
+import graphql.language.Value;
+
+/**
+ * Exception thrown when processing invalid `@link` import definitions.
+ *
+ * <p>Unsupported imports: - specifying object import without specifying String name - specifying
+ * object rename that is not a String - attempting to import definition that is not a String nor an
+ * object definition
+ */
+public class UnsupportedLinkImportException extends RuntimeException {
+
+  public UnsupportedLinkImportException(Value importedDefinition) {
+    super("Unsupported import: " + importedDefinition);
+  }
+}

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

@@ -190,6 +190,16 @@ public void verifyWeCannotRenameInaccessibleDirective() {
         () -> verifyFederationTransformation("schemas/renamedInaccessibleImport.graphql", true));
   }
 
+  @Test
+  public void verifyFederationV2Transformation_renames() {
+    verifyFederationTransformation("schemas/renamedImports.graphql", true);
+  }
+
+  @Test
+  public void verifyFederationV2Transformation_linkOnSchema() {
+    verifyFederationTransformation("schemas/schemaImport.graphql", true);
+  }
+
   private GraphQLSchema verifyFederationTransformation(
       String schemaFileName, boolean isFederationV2) {
     final RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring().build();

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

@@ -0,0 +1,33 @@
+extend schema
+@link(url: "https://specs.apollo.dev/federation/v2.0",
+    import: [{ name: "@key", as: "@myKey" }, { name: "@shareable" }, "@provides", "@external", "@tag", "@extends", "@override", "@inaccessible"])
+
+type Product @myKey(fields: "id") @myKey(fields: "sku package") @myKey(fields: "sku variation { id }") {
+    id: ID!
+    sku: String
+    package: String
+    variation: ProductVariation
+    dimensions: ProductDimension
+    createdBy: User @provides(fields: "totalProductsCreated")
+    notes: String @tag(name: "internal")
+}
+
+type ProductVariation {
+    id: ID!
+}
+
+type ProductDimension @shareable {
+    size: String
+    weight: Float
+    unit: String @inaccessible
+}
+
+type Query {
+    product(id: ID!): Product
+}
+
+type User @myKey(fields: "email") {
+    email: ID!
+    name: String @shareable @override(from: "users")
+    totalProductsCreated: Int @external
+}

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

@@ -0,0 +1,62 @@
+schema @link(import : [{name : "@key", as : "@myKey"}, {name : "@shareable"}, "@provides", "@external", "@tag", "@extends", "@override", "@inaccessible"], url : "https://specs.apollo.dev/federation/v2.0"){
+  query: Query
+}
+
+directive @extends on OBJECT | INTERFACE
+
+directive @external on OBJECT | FIELD_DEFINITION
+
+directive @federation__requires(fields: federation__FieldSet!) on FIELD_DEFINITION
+
+directive @inaccessible on SCALAR | OBJECT | FIELD_DEFINITION | ARGUMENT_DEFINITION | INTERFACE | UNION | ENUM | ENUM_VALUE | INPUT_OBJECT | INPUT_FIELD_DEFINITION
+
+directive @link(as: String, import: [link__Import], url: String!) repeatable on SCHEMA
+
+directive @myKey(fields: federation__FieldSet!, resolvable: Boolean = true) repeatable on OBJECT | INTERFACE
+
+directive @override(from: String!) on FIELD_DEFINITION
+
+directive @provides(fields: federation__FieldSet!) on FIELD_DEFINITION
+
+directive @shareable on OBJECT | FIELD_DEFINITION
+
+directive @tag(name: String!) repeatable on SCALAR | OBJECT | FIELD_DEFINITION | ARGUMENT_DEFINITION | INTERFACE | UNION | ENUM | ENUM_VALUE | INPUT_OBJECT | INPUT_FIELD_DEFINITION
+
+type Product @myKey(fields : "id", resolvable : true) @myKey(fields : "sku package", resolvable : true) @myKey(fields : "sku variation { id }", resolvable : true) {
+  createdBy: User @provides(fields : "totalProductsCreated")
+  dimensions: ProductDimension
+  id: ID!
+  notes: String @tag(name : "internal")
+  package: String
+  sku: String
+  variation: ProductVariation
+}
+
+type ProductDimension @shareable {
+  size: String
+  unit: String @inaccessible
+  weight: Float
+}
+
+type ProductVariation {
+  id: ID!
+}
+
+type Query {
+  _service: _Service!
+  product(id: ID!): Product
+}
+
+type User @myKey(fields : "email", resolvable : true) {
+  email: ID!
+  name: String @override(from : "users") @shareable
+  totalProductsCreated: Int @external
+}
+
+type _Service {
+  sdl: String!
+}
+
+scalar federation__FieldSet
+
+scalar link__Import

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

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