ConnectionTypeDefinitionConfigurer supports type extensions

Closes gh-816

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurer.java

@@ -15,9 +15,11 @@
  */
 package org.springframework.graphql.execution;
 
+import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import graphql.language.FieldDefinition;
 import graphql.language.ImplementingTypeDefinition;
@@ -85,7 +87,9 @@ public void configure(TypeDefinitionRegistry registry) {
 	}
 
 	private static Set<String> findConnectionTypeNames(TypeDefinitionRegistry registry) {
-		return registry.types().values().stream()
+		return Stream.concat(
+						registry.types().values().stream(),
+						registry.objectTypeExtensions().values().stream().flatMap(Collection::stream))
 				.filter(definition -> definition instanceof ImplementingTypeDefinition)
 				.flatMap(definition -> {
 					ImplementingTypeDefinition<?> typeDefinition = (ImplementingTypeDefinition<?>) definition;

spring-graphql/src/test/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurerTests.java

@@ -45,15 +45,38 @@ public class ConnectionTypeDefinitionConfigurerTests {
 
 	@Test
 	void connectionTypeGeneration() {
+		GraphQlSetup graphQlSetup = GraphQlSetup.schemaResource(BookSource.paginationSchema);
+		testConnectionTypeGeneration(graphQlSetup);
+	}
+
+	@Test
+	void connectionTypeGenerationWithObjectExtension() {
+		String schema = """
+				type Query {
+					bookById(id:ID): Book
+				}
+				type Book {
+					id: ID
+					name: String
+				}
+				extend type Query {
+					books(first:Int, after:String): BookConnection
+				}
+				""";
+		GraphQlSetup graphQlSetup = GraphQlSetup.schemaContent(schema);
+		testConnectionTypeGeneration(graphQlSetup);
+	}
 
+	private static void testConnectionTypeGeneration(GraphQlSetup graphQlSetup) {
 		List<Book> books = BookSource.books();
 
 		DataFetcher<?> dataFetcher = environment ->
 				createConnection(books, book -> new DefaultConnectionCursor("book:" + book.getId()));
 
 		String document = BookSource.booksConnectionQuery("");
 
-		Mono<ExecutionGraphQlResponse> response = initGraphQlSetup()
+		Mono<ExecutionGraphQlResponse> response = graphQlSetup
+				.typeDefinitionConfigurer(new ConnectionTypeDefinitionConfigurer())
 				.dataFetcher("Query", "books", dataFetcher)
 				.toGraphQlService()
 				.execute(document);
@@ -78,11 +101,6 @@ void connectionTypeGeneration() {
 		);
 	}
 
-	private GraphQlSetup initGraphQlSetup() {
-		return GraphQlSetup.schemaResource(BookSource.paginationSchema)
-				.typeDefinitionConfigurer(new ConnectionTypeDefinitionConfigurer());
-	}
-
 	private static <N> Connection<N> createConnection(
 			List<N> nodes, Function<N, ConnectionCursor> cursorFunction) {