feat: CompletableFuture in FederatedTypeResolver, standardize entities resolver with other subgraph implementations

examples/client/server/src/main/kotlin/com/expediagroup/graphql/examples/client/server/CustomFederatedHooks.kt

@@ -11,7 +11,7 @@ import java.util.UUID
 import kotlin.reflect.KType
 
 @Component
-class CustomFederatedHooks(resolvers: List<FederatedTypeResolver<*>>) : FederatedSchemaGeneratorHooks(resolvers, true) {
+class CustomFederatedHooks(resolvers: List<FederatedTypeResolver>) : FederatedSchemaGeneratorHooks(resolvers, true) {
     override fun willGenerateGraphQLType(type: KType): GraphQLType? = when (type.classifier) {
         UUID::class -> graphqlUUIDType
         ULocale::class -> graphqlULocaleType

examples/client/server/src/main/kotlin/com/expediagroup/graphql/examples/client/server/FederatedQuery.kt

@@ -10,7 +10,7 @@ import com.expediagroup.graphql.generator.federation.directives.OverrideDirectiv
 import com.expediagroup.graphql.generator.federation.directives.ProvidesDirective
 import com.expediagroup.graphql.generator.federation.directives.ShareableDirective
 import com.expediagroup.graphql.generator.federation.directives.TagDirective
-import com.expediagroup.graphql.generator.federation.execution.FederatedTypeResolver
+import com.expediagroup.graphql.generator.federation.execution.FederatedTypeSuspendResolver
 import com.expediagroup.graphql.generator.scalars.ID
 import com.expediagroup.graphql.server.operations.Query
 import graphql.schema.DataFetchingEnvironment
@@ -105,28 +105,24 @@ data class User(
 )
 
 @Component
-class ProductsResolver : FederatedTypeResolver<Product> {
+class ProductsResolver : FederatedTypeSuspendResolver<Product> {
     override val typeName: String = "Product"
 
     override suspend fun resolve(
         environment: DataFetchingEnvironment,
-        representations: List<Map<String, Any>>
-    ): List<Product?> = representations.map {
-        Product.byReference(it)
-    }
+        representation: Map<String, Any>
+    ): Product? = Product.byReference(representation)
 }
 
 @Component
-class UserResolver : FederatedTypeResolver<User> {
+class UserResolver : FederatedTypeSuspendResolver<User> {
     override val typeName: String = "User"
 
     override suspend fun resolve(
         environment: DataFetchingEnvironment,
-        representations: List<Map<String, Any>>
-    ): List<User?> {
-        return representations.map {
-            val email = it["email"]?.toString() ?: throw RuntimeException("invalid entity reference")
-            User(email = email, name = "default", totalProductsCreated = 1337)
-        }
+        representation: Map<String, Any>
+    ): User? {
+        val email = representation["email"]?.toString() ?: throw RuntimeException("invalid entity reference")
+        return User(email = email, name = "default", totalProductsCreated = 1337)
     }
 }

examples/federation/base-app/src/main/kotlin/com/expediagroup/graphql/examples/federation/base/hooks/CustomFederationSchemaGeneratorHooks.kt

@@ -30,7 +30,7 @@ import kotlin.reflect.KType
 /**
  * Schema generator hook that adds additional scalar types.
  */
-class CustomFederationSchemaGeneratorHooks(resolvers: List<FederatedTypeResolver<*>>) : FederatedSchemaGeneratorHooks(resolvers) {
+class CustomFederationSchemaGeneratorHooks(resolvers: List<FederatedTypeResolver>) : FederatedSchemaGeneratorHooks(resolvers) {
 
     /**
      * Register additional GraphQL scalar types.

examples/federation/extend-app/src/main/kotlin/com/expediagroup/graphql/examples/federation/extend/schema/WidgetResolver.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 Expedia, Inc
+ * Copyright 2022 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,24 +17,27 @@
 package com.expediagroup.graphql.examples.federation.extend.schema
 
 import com.expediagroup.graphql.examples.federation.extend.service.RandomNumberService
-import com.expediagroup.graphql.generator.federation.execution.FederatedTypeResolver
+import com.expediagroup.graphql.generator.federation.execution.FederatedTypeSuspendResolver
 import graphql.schema.DataFetchingEnvironment
 import org.springframework.stereotype.Component
 
 @Component
-class WidgetResolver(private val randomNumberService: RandomNumberService) : FederatedTypeResolver<Widget> {
+class WidgetResolver(private val randomNumberService: RandomNumberService) : FederatedTypeSuspendResolver<Widget> {
     override val typeName: String = "Widget"
 
     @Suppress("UNCHECKED_CAST")
-    override suspend fun resolve(environment: DataFetchingEnvironment, representations: List<Map<String, Any>>): List<Widget?> = representations.map {
-        // Extract the 'id' from the other service
-        val id = it["id"]?.toString()?.toIntOrNull() ?: throw InvalidWidgetIdException()
-        val listOfValues = it["listOfValues"] as? List<Int>
+    override suspend fun resolve(
+        environment: DataFetchingEnvironment,
+        representation: Map<String, Any>
+    ): Widget? {
+        // Extract the 'id' from the representation map provided by the other service
+        val id = representation["id"]?.toString()?.toIntOrNull() ?: throw InvalidWidgetIdException()
+        val listOfValues = representation["listOfValues"] as? List<Int>
 
         // If we needed to construct a Widget which has data from other APIs,
         // this is the place where we could call them with the widget id
         val valueFromExtend = randomNumberService.getInt()
-        Widget(id, listOfValues, valueFromExtend)
+        return Widget(id, listOfValues, valueFromExtend)
     }
 
     class InvalidWidgetIdException : RuntimeException()

generator/graphql-kotlin-federation/README.md

@@ -121,13 +121,16 @@ type _Service {
 
 ### Extended Schema (Reviews Subgraph)
 
-Extended federated GraphQL schemas provide additional functionality to the types already exposed by other GraphQL services. In the example below, `Product` type is extended to add new `reviews` field to it. Primary key needed to instantiate the `Product` type (i.e. `id`) has to match the `@key` definition on the base type. Since primary keys are defined on the base type and are only referenced from the extended type, all of the fields that are part of the field set specified in `@key` directive have to be marked as `@external`.
+Extended federated GraphQL schemas provide additional functionality to the types already exposed by other GraphQL services.
+In the example below, `Product` type is extended to add new `reviews` field to it. Primary key needed to instantiate
+the `Product` type (i.e. `id`) has to match the `@key` definition on the base type.
+Since primary keys are defined on the base type and are only referenced from the extended type,
+all the fields that are part of the field set specified in `@key` directive have to be marked as `@external`.
 
 ```kotlin
 @KeyDirective(fields = FieldSet("id"))
 @ExtendsDirective
 data class Product(@ExternalDirective val id: Int) {
-
     fun reviews(): List<Review> {
         // returns list of product reviews
     }
@@ -136,9 +139,12 @@ data class Product(@ExternalDirective val id: Int) {
 data class Review(val reviewId: String, val text: String)
 
 // Generate the schema
-val productResolver = object: FederatedTypeResolver<Product> {
-    override fun resolve(keys: Map<String, Any>): Product {
-        val id = keys["id"]?.toString()?.toIntOrNull()
+val productResolver = object: FederatedTypeSuspendResolver<Product> {
+    override fun resolve(
+        environment: DataFetchingEnvironment,
+        representation: Map<String, Any>
+    ): Product {
+        val id = representation["id"]?.toString()?.toIntOrNull()
         // instantiate product using id
     }
 }

generator/graphql-kotlin-federation/build.gradle.kts

@@ -2,10 +2,14 @@ description = "Federated GraphQL schema generator"
 
 val junitVersion: String by project
 val federationGraphQLVersion: String by project
+val reactorVersion: String by project
+val reactorExtensionsVersion: String by project
 
 dependencies {
     api(project(path = ":graphql-kotlin-schema-generator"))
     api("com.apollographql.federation:federation-graphql-java-support:$federationGraphQLVersion")
+    testImplementation("io.projectreactor.kotlin:reactor-kotlin-extensions:$reactorExtensionsVersion")
+    testImplementation("io.projectreactor:reactor-core:$reactorVersion")
     testImplementation("org.junit.jupiter:junit-jupiter-params:$junitVersion")
 }
 

generator/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/generator/federation/FederatedSchemaGeneratorHooks.kt

@@ -40,7 +40,7 @@ import com.expediagroup.graphql.generator.federation.directives.SHAREABLE_DIRECT
 import com.expediagroup.graphql.generator.federation.directives.TAG_DIRECTIVE_TYPE
 import com.expediagroup.graphql.generator.federation.directives.appliedLinkDirective
 import com.expediagroup.graphql.generator.federation.exception.IncorrectFederatedDirectiveUsage
-import com.expediagroup.graphql.generator.federation.execution.EntityResolver
+import com.expediagroup.graphql.generator.federation.execution.EntitiesDataFetcher
 import com.expediagroup.graphql.generator.federation.execution.FederatedTypeResolver
 import com.expediagroup.graphql.generator.federation.types.ANY_SCALAR_TYPE
 import com.expediagroup.graphql.generator.federation.types.ENTITY_UNION_NAME
@@ -67,7 +67,10 @@ import kotlin.reflect.full.findAnnotation
 /**
  * Hooks for generating federated GraphQL schema.
  */
-open class FederatedSchemaGeneratorHooks(private val resolvers: List<FederatedTypeResolver<*>>, private val optInFederationV2: Boolean = false) : SchemaGeneratorHooks {
+open class FederatedSchemaGeneratorHooks(
+    private val resolvers: List<FederatedTypeResolver>,
+    private val optInFederationV2: Boolean = false
+) : SchemaGeneratorHooks {
     private val validator = FederatedSchemaValidator()
 
     private val federationV2OnlyDirectiveNames: Set<String> = setOf(
@@ -158,7 +161,7 @@ open class FederatedSchemaGeneratorHooks(private val resolvers: List<FederatedTy
             val entityField = generateEntityFieldDefinition(entityTypeNames)
             federatedQuery.field(entityField)
 
-            federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, entityField.name), EntityResolver(resolvers))
+            federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, entityField.name), EntitiesDataFetcher(resolvers))
             federatedCodeRegistry.typeResolver(ENTITY_UNION_NAME) { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObjectName()) }
 
             builder.query(federatedQuery)

generator/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcher.kt

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2022 Expedia, Inc
+ *
+ * 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 com.expediagroup.graphql.generator.federation.execution
+
+import com.expediagroup.graphql.generator.federation.exception.InvalidFederatedRequest
+import com.expediagroup.graphql.generator.federation.execution.resolverexecutor.FederatedTypePromiseResolverExecutor
+import com.expediagroup.graphql.generator.federation.execution.resolverexecutor.FederatedTypeSuspendResolverExecutor
+import com.expediagroup.graphql.generator.federation.execution.resolverexecutor.ResolvableEntity
+import com.expediagroup.graphql.generator.federation.extensions.collectAll
+import com.expediagroup.graphql.generator.federation.extensions.toDataFetcherResult
+import graphql.execution.DataFetcherResult
+import graphql.schema.DataFetcher
+import graphql.schema.DataFetchingEnvironment
+import java.util.concurrent.CompletableFuture
+
+private const val TYPENAME_FIELD = "__typename"
+private const val REPRESENTATIONS = "representations"
+
+/**
+ * Federated _entities field data fetcher.
+ */
+open class EntitiesDataFetcher(
+    resolvers: List<FederatedTypeResolver>
+) : DataFetcher<CompletableFuture<DataFetcherResult<List<Any?>>>> {
+
+    constructor(vararg resolvers: FederatedTypeResolver) : this(resolvers.toList())
+    /**
+     * Pre-compute resolvers by typename so, we don't have to search on every request
+     */
+    private val resolversByType: Map<String, FederatedTypeResolver> = resolvers.associateBy(FederatedTypeResolver::typeName)
+
+    /**
+     * Resolves entities based on the passed in representations argument. Entities are resolved in the same order
+     * they are specified in the list of representations. If target representation cannot be resolved, NULL will
+     * be returned instead.
+     *
+     * Representations are grouped by the underlying typename and each batch is resolved asynchronously before merging
+     * the results back into a single list that preserves the original order.
+     *
+     * @return list of resolved nullable entities
+     */
+    override fun get(env: DataFetchingEnvironment): CompletableFuture<DataFetcherResult<List<Any?>>> {
+        val representations: List<Map<String, Any>> = env.getArgument(REPRESENTATIONS)
+
+        val representationsWithoutResolver = mutableListOf<IndexedValue<Map<String, Any>>>()
+        val entitiesWithPromiseResolver = mutableListOf<ResolvableEntity<FederatedTypePromiseResolver<*>>>()
+        val entitiesWithSuspendResolver = mutableListOf<ResolvableEntity<FederatedTypeSuspendResolver<*>>>()
+
+        representations.withIndex()
+            .groupBy { (_, representation) -> representation[TYPENAME_FIELD].toString() }
+            .forEach { (typeName, indexedRequests) ->
+                when (val resolver = resolversByType[typeName]) {
+                    is FederatedTypePromiseResolver<*> -> {
+                        entitiesWithPromiseResolver += ResolvableEntity(typeName, indexedRequests, resolver)
+                    }
+                    is FederatedTypeSuspendResolver<*> -> {
+                        entitiesWithSuspendResolver += ResolvableEntity(typeName, indexedRequests, resolver)
+                    }
+                    null -> {
+                        representationsWithoutResolver += indexedRequests
+                    }
+                }
+            }
+
+        val noResolverErrors: CompletableFuture<List<Map<Int, Any?>>> = CompletableFuture.completedFuture(
+            listOf(
+                representationsWithoutResolver.associateBy(IndexedValue<Map<String, Any>>::index) { (_, representation) ->
+                    InvalidFederatedRequest("Unable to resolve federated type, representation=$representation")
+                }
+            )
+        )
+
+        val promises: List<CompletableFuture<List<Map<Int, Any?>>>> = listOf(
+            FederatedTypePromiseResolverExecutor.execute(entitiesWithPromiseResolver, env),
+            FederatedTypeSuspendResolverExecutor.execute(entitiesWithSuspendResolver, env),
+            noResolverErrors
+        )
+
+        return promises
+            .collectAll()
+            .thenApply { results ->
+                results.asSequence()
+                    .flatten()
+                    .map(Map<Int, Any?>::toList)
+                    .flatten()
+                    .sortedBy(Pair<Int, Any?>::first)
+                    .map(Pair<Int, Any?>::second)
+                    .toDataFetcherResult()
+            }
+    }
+}