Add federation support

See gh-864

platform/build.gradle

@@ -32,6 +32,8 @@ dependencies {
 		api("jakarta.validation:jakarta.validation-api:3.0.2")
 		api("jakarta.persistence:jakarta.persistence-api:3.1.0")
 
+		api("com.apollographql.federation:federation-graphql-java-support:4.3.0")
+
 		api("com.google.code.findbugs:jsr305:3.0.2")
 
 		api("org.assertj:assertj-core:3.24.2")

spring-graphql/build.gradle

@@ -32,6 +32,8 @@ dependencies {
 
 	compileOnly 'com.fasterxml.jackson.core:jackson-databind'
 
+	compileOnly 'com.apollographql.federation:federation-graphql-java-support'
+
 	testImplementation 'org.junit.jupiter:junit-jupiter'
 	testImplementation 'org.assertj:assertj-core'
 	testImplementation 'org.mockito:mockito-core'
@@ -69,6 +71,7 @@ dependencies {
 	testImplementation 'com.fasterxml.jackson.core:jackson-databind'
 	testImplementation("com.fasterxml.jackson.datatype:jackson-datatype-jsr310")
 	testImplementation 'org.apache.tomcat.embed:tomcat-embed-el:10.0.21'
+	testImplementation 'com.apollographql.federation:federation-graphql-java-support'
 
 	testRuntimeOnly 'org.apache.logging.log4j:log4j-core'
 	testRuntimeOnly 'org.apache.logging.log4j:log4j-slf4j-impl'

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntitiesDataFetcher.java

@@ -0,0 +1,158 @@
+/*
+ * Copyright 2002-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.data.federation;
+
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletionException;
+
+import com.apollographql.federation.graphqljava._Entity;
+import graphql.GraphQLError;
+import graphql.GraphqlErrorBuilder;
+import graphql.execution.DataFetcherResult;
+import graphql.execution.ExecutionStepInfo;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DelegatingDataFetchingEnvironment;
+import reactor.core.publisher.Mono;
+
+import org.springframework.graphql.data.method.annotation.support.HandlerDataFetcherExceptionResolver;
+import org.springframework.graphql.execution.ErrorType;
+import org.springframework.lang.Nullable;
+
+/**
+ * DataFetcher that handles the "_entities" query by invoking
+ * {@link EntityHandlerMethod}s.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.3
+ * @see com.apollographql.federation.graphqljava.SchemaTransformer#fetchEntities(DataFetcher)
+ */
+final class EntitiesDataFetcher implements DataFetcher<Mono<DataFetcherResult<List<Object>>>> {
+
+	private final Map<String, EntityHandlerMethod> handlerMethods;
+
+	private final HandlerDataFetcherExceptionResolver exceptionResolver;
+
+
+	public EntitiesDataFetcher(
+			Map<String, EntityHandlerMethod> handlerMethods, HandlerDataFetcherExceptionResolver resolver) {
+
+		this.handlerMethods = new LinkedHashMap<>(handlerMethods);
+		this.exceptionResolver = resolver;
+	}
+
+
+	@Override
+	public Mono<DataFetcherResult<List<Object>>> get(DataFetchingEnvironment environment) {
+		List<Map<String, Object>> representations = environment.getArgument(_Entity.argumentName);
+
+		List<Mono<Object>> monoList = new ArrayList<>();
+		for (int index = 0; index < representations.size(); index++) {
+			Map<String, Object> map = representations.get(index);
+			if (!(map.get("__typename") instanceof String typename)) {
+				Exception ex = new RepresentationException(map, "Missing \"__typename\" argument");
+				monoList.add(resolveException(ex, environment, null, index));
+				continue;
+			}
+			EntityHandlerMethod handlerMethod = this.handlerMethods.get(typename);
+			if (handlerMethod == null) {
+				Exception ex = new RepresentationException(map, "No entity fetcher");
+				monoList.add(resolveException(ex, environment, null, index));
+				continue;
+			}
+			monoList.add(invokeResolver(environment, handlerMethod, map, index));
+		}
+		return Mono.zip(monoList, Arrays::asList).map(EntitiesDataFetcher::toDataFetcherResult);
+	}
+
+	private Mono<Object> invokeResolver(
+			DataFetchingEnvironment env, EntityHandlerMethod handlerMethod, Map<String, Object> map, int index) {
+
+		return handlerMethod.getEntity(env, map, index)
+				.switchIfEmpty(Mono.error(new RepresentationNotResolvedException(map, handlerMethod)))
+				.onErrorResume(ex -> resolveException(ex, env, handlerMethod, index));
+	}
+
+	private Mono<Object> resolveException(
+			Throwable ex, DataFetchingEnvironment env, @Nullable EntityHandlerMethod handlerMethod, int index) {
+
+		Throwable theEx = (ex instanceof CompletionException ? ex.getCause() : ex);
+		DataFetchingEnvironment theEnv = new EntityDataFetchingEnvironment(env, index);
+		Object handler = (handlerMethod != null ? handlerMethod.getBean() : null);
+
+		return this.exceptionResolver.resolveException(theEx, theEnv, handler)
+				.map(ErrorContainer::new)
+				.switchIfEmpty(Mono.fromCallable(() -> createDefaultError(theEx, theEnv)))
+				.cast(Object.class);
+	}
+
+	private ErrorContainer createDefaultError(Throwable ex, DataFetchingEnvironment env) {
+
+		ErrorType errorType = (ex instanceof RepresentationException representationEx ?
+				representationEx.getErrorType() : ErrorType.INTERNAL_ERROR);
+
+		return new ErrorContainer(GraphqlErrorBuilder.newError(env)
+				.errorType(errorType)
+				.message(ex.getMessage())
+				.build());
+	}
+
+	private static DataFetcherResult<List<Object>> toDataFetcherResult(List<Object> entities) {
+		List<GraphQLError> errors = new ArrayList<>();
+		for (int i = 0; i < entities.size(); i++) {
+			Object entity = entities.get(i);
+			if (entity instanceof ErrorContainer errorContainer) {
+				errors.addAll(errorContainer.errors());
+				entities.set(i, null);
+			}
+		}
+		return DataFetcherResult.<List<Object>>newResult().data(entities).errors(errors).build();
+	}
+
+
+	private static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
+
+		private final ExecutionStepInfo executionStepInfo;
+
+		public EntityDataFetchingEnvironment(DataFetchingEnvironment env, int index) {
+			super(env);
+			this.executionStepInfo = ExecutionStepInfo.newExecutionStepInfo(env.getExecutionStepInfo())
+					.path(env.getExecutionStepInfo().getPath().segment(index))
+					.build();
+		}
+
+		@Override
+		public ExecutionStepInfo getExecutionStepInfo() {
+			return this.executionStepInfo;
+		}
+	}
+
+
+	private record ErrorContainer(List<GraphQLError> errors) {
+
+		ErrorContainer(GraphQLError error) {
+			this(Collections.singletonList(error));
+		}
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2002-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.data.federation;
+
+import java.util.Map;
+
+import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DelegatingDataFetchingEnvironment;
+
+import org.springframework.core.ResolvableType;
+import org.springframework.graphql.data.GraphQlArgumentBinder;
+import org.springframework.graphql.data.method.annotation.Argument;
+import org.springframework.graphql.data.method.annotation.support.ArgumentMethodArgumentResolver;
+import org.springframework.validation.BindException;
+
+/**
+ * Resolver for a method parameter annotated with {@link Argument @Argument}.
+ * On {@code @EntityMapping} methods, the raw argument value is obtained from
+ * the "representation" input map for the entity with entries that identify
+ * the entity uniquely.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.3
+ */
+final class EntityArgumentMethodArgumentResolver extends ArgumentMethodArgumentResolver {
+
+
+	EntityArgumentMethodArgumentResolver(GraphQlArgumentBinder argumentBinder) {
+		super(argumentBinder);
+	}
+
+
+	@Override
+	protected Object doBind(
+			DataFetchingEnvironment environment, String name, ResolvableType targetType) throws BindException {
+
+		if (environment instanceof EntityDataFetchingEnvironment entityEnv) {
+			Map<String, Object> entityMap = entityEnv.getRepresentation();
+			Object rawValue = entityMap.get(name);
+			boolean isOmitted = !entityMap.containsKey(name);
+			return getArgumentBinder().bind(name, rawValue, isOmitted, targetType);
+		}
+
+		throw new IllegalStateException("Expected decorated DataFetchingEnvironment");
+	}
+
+	/**
+	 * Wrap the environment in order to also expose the entity representation map.
+	 */
+	public static DataFetchingEnvironment wrap(DataFetchingEnvironment env, Map<String, Object> representation) {
+		return new EntityDataFetchingEnvironment(env, representation);
+	}
+
+
+	private static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
+
+		private final Map<String, Object> representation;
+
+		EntityDataFetchingEnvironment(DataFetchingEnvironment env, Map<String, Object> representation) {
+			super(env);
+			this.representation = representation;
+		}
+
+		public Map<String, Object> getRepresentation() {
+			return this.representation;
+		}
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2002-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.data.federation;
+
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+
+import graphql.schema.DataFetchingEnvironment;
+import reactor.core.publisher.Mono;
+
+import org.springframework.graphql.data.method.HandlerMethod;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
+import org.springframework.graphql.data.method.annotation.support.DataFetcherHandlerMethodSupport;
+import org.springframework.lang.Nullable;
+
+/**
+ * Invokable controller method to fetch a federated entity.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.3
+ */
+final class EntityHandlerMethod extends DataFetcherHandlerMethodSupport {
+
+	public EntityHandlerMethod(
+			HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers,
+			@Nullable Executor executor) {
+
+		super(handlerMethod, resolvers, executor);
+	}
+
+
+	public Mono<Object> getEntity(
+			DataFetchingEnvironment environment, Map<String, Object> representation, int index) {
+
+		Object[] args;
+		try {
+			environment = EntityArgumentMethodArgumentResolver.wrap(environment, representation);
+			args = getMethodArgumentValues(environment, representation);
+		}
+		catch (Throwable ex) {
+			return Mono.error(ex);
+		}
+
+		Object result = doInvoke(environment.getGraphQlContext(), args);
+
+		if (result instanceof Mono<?> mono) {
+			return mono.cast(Object.class);
+		}
+		else if (result instanceof CompletableFuture<?> future) {
+			return Mono.fromFuture(future);
+		}
+		else {
+			return Mono.justOrEmpty(result);
+		}
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityMapping.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2002-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.data.federation;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import org.springframework.core.annotation.AliasFor;
+
+/**
+ * Annotation for mapping a handler method to a federated schema type.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.3
+ */
+@Target({ElementType.TYPE, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface EntityMapping {
+
+	/**
+	 * Customize the name of the entity to map to.
+	 * <p>By default, if not specified, this is initialized from the method name,
+	 * with the first letter changed to upper case via {@link Character#toUpperCase}.
+	 */
+	@AliasFor("value")
+	String name() default "";
+
+	/**
+	 * Effectively an alias for {@link #name()}.
+	 */
+	@AliasFor("name")
+	String value() default "";
+
+}