feature: batch support for resolving federation requests (ExpediaGroup#306)

* feature: batch support for resolving federation requests

NOTE: This is a non-backwards compatible change as FederatedTypeResolver resolve signature is changed to process list of representations at once. This allows resolver to either instantiate entities one by one or use some batch logic.

Update _entities query resolver logic to resolve federation requests in batch mode. Specified federated object representations are grouped by the underlying typename and asynchronously resolved in batches. Batch results are then merged into a single list that preserves the original order of entities. If entity cannot be resolved, NULL will be returned instead. Any exceptions thrown while attempting to resolve the federated types will be returned together with partially resolved data.

* additional test for invalid result size

Author: dariuszkuc

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/FederatedSchemaGeneratorHooks.kt

@@ -3,6 +3,8 @@ package com.expedia.graphql.federation
 import com.expedia.graphql.annotations.GraphQLName
 import com.expedia.graphql.extensions.print
 import com.expedia.graphql.federation.directives.FieldSet
+import com.expedia.graphql.federation.execution.EntityResolver
+import com.expedia.graphql.federation.execution.FederatedTypeRegistry
 import com.expedia.graphql.federation.types.ANY_SCALAR_TYPE
 import com.expedia.graphql.federation.types.FIELD_SET_SCALAR_TYPE
 import com.expedia.graphql.federation.types.SERVICE_FIELD_DEFINITION
@@ -71,14 +73,7 @@ open class FederatedSchemaGeneratorHooks(private val federatedTypeRegistry: Fede
                 .replace("type ${originalQuery.name}", "type ${originalQuery.name} @extends")
                 .trim()
             federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, SERVICE_FIELD_DEFINITION.name), DataFetcher { _Service(sdl) })
-            federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, entityField.name), DataFetcher {
-                val representations: List<Map<String, Any>> = it.getArgument("representations")
-                representations.map { representation ->
-                    val type = representation["__typename"]?.toString() ?: throw FederationException("invalid _entity query - missing __typename in the representation, representation=$representation")
-                    val resolver = federatedTypeRegistry.getFederatedResolver(type) ?: throw FederationException("Federation exception - cannot find resolver for $type")
-                    resolver.resolve(representation)
-                }.toList()
-            })
+            federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, entityField.name), EntityResolver(federatedTypeRegistry))
             federatedCodeRegistry.typeResolver("_Entity") { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObjectName()) }
             federatedSchema.additionalType(ANY_SCALAR_TYPE)
         }

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/FederatedSchemaValidator.kt

@@ -1,6 +1,7 @@
 package com.expedia.graphql.federation
 
 import com.expedia.graphql.federation.directives.FieldSet
+import com.expedia.graphql.federation.exception.InvalidFederatedSchema
 import graphql.schema.GraphQLDirective
 import graphql.schema.GraphQLDirectiveContainer
 import graphql.schema.GraphQLFieldDefinition

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/FederatedTypeResolver.kt

@@ -0,0 +1,24 @@
+package com.expedia.graphql.federation.exception
+
+import graphql.ErrorClassification
+import graphql.ErrorType
+import graphql.GraphQLError
+import graphql.language.SourceLocation
+
+/**
+ * GraphQLError returned whenever exception occurs while trying to resolve federated entity representation.
+ */
+class FederatedRequestFailure(private val errorMessage: String, private val error: Exception? = null) : GraphQLError {
+    override fun getMessage(): String = errorMessage
+
+    override fun getErrorType(): ErrorClassification = ErrorType.DataFetchingException
+
+    override fun getLocations(): List<SourceLocation> = listOf(SourceLocation(-1, -1))
+
+    override fun getExtensions(): Map<String, Any>? =
+        if (error != null) {
+            mapOf("error" to error)
+        } else {
+            null
+        }
+}

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/FederationException.kt

@@ -0,0 +1,17 @@
+package com.expedia.graphql.federation.exception
+
+import graphql.ErrorClassification
+import graphql.ErrorType
+import graphql.GraphQLError
+import graphql.language.SourceLocation
+
+/**
+ * GraphQLError returned when federated representation cannot be resolved by any type resolvers.
+ */
+class InvalidFederatedRequest(private val errorMessage: String) : GraphQLError {
+    override fun getMessage(): String = errorMessage
+
+    override fun getErrorType(): ErrorClassification = ErrorType.ValidationError
+
+    override fun getLocations(): List<SourceLocation> = listOf(SourceLocation(-1, -1))
+}

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/exception/FederatedRequestFailure.kt

@@ -1,4 +1,4 @@
-package com.expedia.graphql.federation
+package com.expedia.graphql.federation.exception
 
 import com.expedia.graphql.exceptions.GraphQLKotlinException
 

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/exception/InvalidFederatedRequest.kt

@@ -0,0 +1,90 @@
+package com.expedia.graphql.federation.execution
+
+import com.expedia.graphql.federation.exception.FederatedRequestFailure
+import com.expedia.graphql.federation.exception.InvalidFederatedRequest
+import graphql.GraphQLError
+import graphql.execution.DataFetcherResult
+import graphql.schema.DataFetcher
+import graphql.schema.DataFetchingEnvironment
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.future.asCompletableFuture
+import java.util.concurrent.CompletableFuture
+
+/**
+ * Federated _entities query resolver.
+ */
+open class EntityResolver(private val federatedTypeRegistry: FederatedTypeRegistry) : DataFetcher<CompletableFuture<DataFetcherResult<List<Any?>>>> {
+
+    /**
+     * 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 indexedBatchRequestsByType = representations.withIndex().groupBy { it.value["__typename"].toString() }
+        return GlobalScope.async {
+            val data = mutableListOf<Any?>()
+            val errors = mutableListOf<GraphQLError>()
+            indexedBatchRequestsByType.map { (typeName, indexedRequests) ->
+                async {
+                    resolveType(typeName, indexedRequests)
+                }
+            }.awaitAll()
+                .flatten()
+                .sortedBy { it.first }
+                .forEach {
+                    val result = it.second
+                    if (result is GraphQLError) {
+                        data.add(null)
+                        errors.add(result)
+                    } else {
+                        data.add(result)
+                    }
+                }
+            DataFetcherResult.newResult<List<Any?>>()
+                .data(data)
+                .errors(errors)
+                .build()
+        }.asCompletableFuture()
+    }
+
+    private suspend fun resolveType(typeName: String, indexedRequests: List<IndexedValue<Map<String, Any>>>): List<Pair<Int, Any?>> {
+        val indices = indexedRequests.map { it.index }
+        val batch = indexedRequests.map { it.value }
+        val results = resolveBatch(typeName, batch)
+        return if (results.size != indices.size) {
+            indices.map {
+                it to FederatedRequestFailure("Federation batch request for $typeName generated different number of results than requested, representations=${indices.size}, results=${results.size}")
+            }
+        } else {
+            indices.zip(results)
+        }
+    }
+
+    @Suppress("TooGenericExceptionCaught")
+    private suspend fun resolveBatch(typeName: String, batch: List<Map<String, Any>>): List<Any?> {
+        val resolver = federatedTypeRegistry.getFederatedResolver(typeName)
+        return if (resolver != null) {
+            try {
+                resolver.resolve(batch)
+            } catch (e: Exception) {
+                batch.map {
+                    FederatedRequestFailure("Exception was thrown while trying to resolve federated type, representation=$it", e)
+                }
+            }
+        } else {
+            batch.map {
+                InvalidFederatedRequest("Unable to resolve federated type, representation=$it")
+            }
+        }
+    }
+}

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/InvalidFederatedSchema.kt renamed to graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/exception/InvalidFederatedSchema.kt

@@ -1,4 +1,4 @@
-package com.expedia.graphql.federation
+package com.expedia.graphql.federation.execution
 
 /**
  * Simple registry that holds mapping of all registered federated GraphQL types and their corresponding resolvers.

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/execution/EntityResolver.kt

@@ -0,0 +1,17 @@
+package com.expedia.graphql.federation.execution
+
+/**
+ * Resolver used to retrieve target federated types.
+ */
+interface FederatedTypeResolver<out T> {
+
+    /**
+     * Resolves underlying federated types based on the passed in _entities query representations. Entities
+     * need to be resolved in the same order they were specified by the list of representations. Each passed
+     * in representation should either be resolved to a target entity OR NULL if entity cannot be resolved.
+     *
+     * @param representations _entity query representations that are required to instantiate the target type
+     * @return list of the target federated type instances
+     */
+    suspend fun resolve(representations: List<Map<String, Any>>): List<T?>
+}

graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/FederatedTypeRegistry.kt renamed to graphql-kotlin-federation/src/main/kotlin/com/expedia/graphql/federation/execution/FederatedTypeRegistry.kt

@@ -2,11 +2,10 @@ package com.expedia.graphql.federation
 
 import com.expedia.graphql.TopLevelObject
 import com.expedia.graphql.extensions.print
+import com.expedia.graphql.federation.execution.FederatedTypeRegistry
 import graphql.schema.GraphQLUnionType
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
-import test.data.queries.federated.Book
-import test.data.queries.federated.User
 import test.data.queries.simple.SimpleQuery
 import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
@@ -84,26 +83,9 @@ class FederatedSchemaGeneratorTest {
 
     @Test
     fun `verify can generate federated schema`() {
-        val bookResolver = object : FederatedTypeResolver<Book> {
-            override fun resolve(keys: Map<String, Any>): Book {
-                val book = Book(keys["id"].toString())
-                keys["weight"]?.toString()?.toDoubleOrNull()?.let {
-                    book.weight = it
-                }
-                return book
-            }
-        }
-        val userResolver = object : FederatedTypeResolver<User> {
-            override fun resolve(keys: Map<String, Any>): User {
-                val id = keys["userId"].toString().toInt()
-                val name = keys["name"].toString()
-                return User(id, name)
-            }
-        }
-
         val config = FederatedSchemaGeneratorConfig(
             supportedPackages = listOf("test.data.queries.federated"),
-            hooks = FederatedSchemaGeneratorHooks(FederatedTypeRegistry(mapOf("Book" to bookResolver, "User" to userResolver)))
+            hooks = FederatedSchemaGeneratorHooks(FederatedTypeRegistry(emptyMap()))
         )
 
         val schema = toFederatedSchema(config)