From 9fc2d3b9d713e9e6fa7811164b43df750b7e1d25 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Mon, 18 Sep 2023 17:42:09 +0100
Subject: [PATCH] ConnectionTypeDefinitionConfigurer supports type extensions

Closes gh-816
---
 .../ConnectionTypeDefinitionConfigurer.java   |  6 +++-
 ...nnectionTypeDefinitionConfigurerTests.java | 30 +++++++++++++++----
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurer.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurer.java
index 27625a207..fc9ae6ae0 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurer.java
+++ b/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;
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurerTests.java
index 5287e2786..0473ba87d 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurerTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/ConnectionTypeDefinitionConfigurerTests.java
@@ -45,7 +45,29 @@ 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 ->
@@ -53,7 +75,8 @@ void connectionTypeGeneration() {
 
 		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) {