Try to make federation work with reactive

smallrye · Sep 5, 2023 · 36883cd · 36883cd
1 parent 012a85c
commit 36883cd
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/server/implementation/src/main/java/io/smallrye/graphql/bootstrap/FederationDataFetcher.java b/server/implementation/src/main/java/io/smallrye/graphql/bootstrap/FederationDataFetcher.java
@@ -1,15 +1,17 @@
 package io.smallrye.graphql.bootstrap;
 
-import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
 
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
 
 import com.apollographql.federation.graphqljava._Entity;
 
+import graphql.execution.Async;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.DelegatingDataFetchingEnvironment;
@@ -20,7 +22,7 @@
 import graphql.schema.GraphQLObjectType;
 import graphql.schema.GraphQLOutputType;
 
-class FederationDataFetcher implements DataFetcher<List<Object>> {
+class FederationDataFetcher implements DataFetcher<CompletableFuture<List<Object>>> {
 
     private final GraphQLObjectType queryType;
     private final GraphQLCodeRegistry codeRegistry;
@@ -31,10 +33,11 @@ public FederationDataFetcher(GraphQLObjectType queryType, GraphQLCodeRegistry co
     }
 
     @Override
-    public List<Object> get(DataFetchingEnvironment environment) throws Exception {
-        return environment.<List<Map<String, Object>>> getArgument(_Entity.argumentName).stream()
-                .map(representations -> fetchEntities(environment, representations))
-                .collect(toList());
+    public CompletableFuture<List<Object>> get(DataFetchingEnvironment environment) throws Exception {
+        return sequence(environment.<List<Map<String, Object>>> getArgument(_Entity.argumentName).stream()
+                .map(representations -> fetchEntities(environment, representations)).map(Async::toCompletableFuture)
+                .collect(Collectors.toList()));
+
     }
 
     private Object fetchEntities(DataFetchingEnvironment env, Map<String, Object> representations) {
@@ -90,4 +93,11 @@ public <T> T getArgumentOrDefault(String name, T defaultValue) {
             throw new RuntimeException("can't fetch data from " + field, e);
         }
     }
+
+    static <T> CompletableFuture<List<T>> sequence(List<CompletableFuture<T>> com) {
+        return CompletableFuture.allOf(com.toArray(new CompletableFuture<?>[0]))
+                .thenApply(v -> com.stream()
+                        .map(CompletableFuture::join)
+                        .collect(Collectors.toList()));
+    }
 }