Use the type name for the connection as defined in the schema

Author: oliemansm

src/main/kotlin/com/coxautodev/graphql/tools/relay/RelayConnectionFactory.kt

@@ -1,12 +1,15 @@
 package com.coxautodev.graphql.tools.relay
 
 import com.coxautodev.graphql.tools.TypeDefinitionFactory
+import graphql.language.Argument
+import graphql.language.Comment
 import graphql.language.Definition
 import graphql.language.Directive
 import graphql.language.FieldDefinition
 import graphql.language.ListType
 import graphql.language.NonNullType
 import graphql.language.ObjectTypeDefinition
+import graphql.language.SourceLocation
 import graphql.language.StringValue
 import graphql.language.TypeDefinition
 import graphql.language.TypeName
@@ -35,32 +38,32 @@ class RelayConnectionFactory : TypeDefinitionFactory {
         return definitions
     }
 
-    private fun findConnectionDirectives(definitions: List<Definition<*>>): List<Directive> {
+    private fun findConnectionDirectives(definitions: List<Definition<*>>): List<DirectiveWithField> {
         return definitions.filterIsInstance<ObjectTypeDefinition>()
                 .flatMap { it.fieldDefinitions }
-                .flatMap { it.directives }
+                .flatMap { it.directivesWithField() }
                 .filter { it.name == "connection" }
     }
 
-    private fun createDefinitions(directive: Directive): List<ObjectTypeDefinition> {
+    private fun createDefinitions(directive: DirectiveWithField): List<ObjectTypeDefinition> {
         val definitions = mutableListOf<ObjectTypeDefinition>()
-        definitions.add(createEdgeDefinition(directive.forTypeName()))
-        definitions.add(createConnectionDefinition(directive.forTypeName()))
+        definitions.add(createConnectionDefinition(directive.getTypeName()))
+        definitions.add(createEdgeDefinition(directive.getTypeName(), directive.forTypeName()))
         return definitions.toList()
     }
 
     private fun createConnectionDefinition(type: String): ObjectTypeDefinition =
             ObjectTypeDefinition.newObjectTypeDefinition()
-                    .name(type + "Connection")
+                    .name(type)
                     .fieldDefinition(FieldDefinition("edges", ListType(TypeName(type + "Edge"))))
                     .fieldDefinition(FieldDefinition("pageInfo", TypeName("PageInfo")))
                     .build()
 
-    private fun createEdgeDefinition(type: String): ObjectTypeDefinition =
+    private fun createEdgeDefinition(connectionType: String, nodeType: String): ObjectTypeDefinition =
             ObjectTypeDefinition.newObjectTypeDefinition()
-                    .name(type + "Edge")
+                    .name(connectionType + "Edge")
                     .fieldDefinition(FieldDefinition("cursor", TypeName("String")))
-                    .fieldDefinition(FieldDefinition("node", TypeName(type)))
+                    .fieldDefinition(FieldDefinition("node", TypeName(nodeType)))
                     .build()
 
     private fun createPageInfo(): ObjectTypeDefinition =
@@ -74,4 +77,16 @@ class RelayConnectionFactory : TypeDefinitionFactory {
         return (this.getArgument("for").value as StringValue).value
     }
 
+    private fun Directive.withField(field: FieldDefinition): DirectiveWithField {
+        return DirectiveWithField(field, this.name, this.arguments, this.sourceLocation, this.comments)
+    }
+
+    private fun FieldDefinition.directivesWithField(): List<DirectiveWithField> {
+        return this.directives.map { it.withField(this) }
+    }
+
+    class DirectiveWithField(val field: FieldDefinition, name: String, arguments: List<Argument>, sourceLocation: SourceLocation, comments: List<Comment>) : Directive(name, arguments, sourceLocation, comments) {
+        fun getTypeName() = (field.type as TypeName).name
+    }
+
 }

src/test/resources/RelayConnection.graphqls

@@ -1,5 +1,5 @@
 type Query {
-    users(first: Int, after: String): UserConnection @connection(for: "User")
+    users(first: Int, after: String): UserRelayConnection @connection(for: "User")
 }
 
 type User {