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

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

@@ -41,7 +41,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.extensions.addDirectivesIfNotPresent
 import com.expediagroup.graphql.generator.federation.types.ANY_SCALAR_TYPE
@@ -68,7 +68,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 scalarDefinitionRegex = "(^\".+\"$[\\r\\n])?^scalar (_FieldSet|_Any)$[\\r\\n]*".toRegex(setOf(RegexOption.MULTILINE, RegexOption.IGNORE_CASE))
     private val emptyQueryRegex = "^type Query @extends \\s*\\{\\s*}\\s*".toRegex(setOf(RegexOption.MULTILINE, RegexOption.IGNORE_CASE))
     private val serviceFieldRegex = "\\s*_service: _Service!".toRegex(setOf(RegexOption.MULTILINE, RegexOption.IGNORE_CASE))
@@ -166,7 +169,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()) }
             federatedSchemaBuilder.additionalType(ANY_SCALAR_TYPE)
         }

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

@@ -0,0 +1,107 @@
+/*
+ * 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.FederatedTypeResolverExecutor
+import com.expediagroup.graphql.generator.federation.execution.resolverexecutor.ResolvableEntity
+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<TypeResolver>
+) : DataFetcher<CompletableFuture<DataFetcherResult<List<Any?>>>> {
+
+    constructor(vararg resolvers: TypeResolver) : this(resolvers.toList())
+    /**
+     * Pre-compute resolvers by typename so, we don't have to search on every request
+     */
+    private val resolversByType: Map<String, TypeResolver> = resolvers.associateBy(TypeResolver::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<FederatedTypeResolver<*>>>()
+
+        representations.withIndex()
+            .groupBy { (_, representation) -> representation[TYPENAME_FIELD].toString() }
+            .forEach { (typeName, indexedRequests) ->
+                when (val resolver = resolversByType[typeName]) {
+                    is FederatedTypePromiseResolver<*> -> {
+                        entitiesWithPromiseResolver += ResolvableEntity(typeName, indexedRequests, resolver)
+                    }
+                    is FederatedTypeResolver<*> -> {
+                        entitiesWithSuspendResolver += ResolvableEntity(typeName, indexedRequests, resolver)
+                    }
+                    null -> {
+                        representationsWithoutResolver += indexedRequests
+                    }
+                }
+            }
+
+        val noResolverErrors: CompletableFuture<List<Map<Int, Any?>>> = CompletableFuture.completedFuture(
+            listOf(
+                representationsWithoutResolver.map(IndexedValue<Map<String, Any>>::index).zip(
+                    representationsWithoutResolver.map(IndexedValue<Map<String, Any>>::value).map { representation ->
+                        InvalidFederatedRequest("Unable to resolve federated type, representation=$representation")
+                    }
+                ).toMap()
+            )
+        )
+
+        val allPromises: List<CompletableFuture<List<Map<Int, Any?>>>> = listOf(
+            FederatedTypePromiseResolverExecutor.execute(entitiesWithPromiseResolver, env),
+            FederatedTypeResolverExecutor.execute(entitiesWithSuspendResolver, env),
+            noResolverErrors
+        )
+
+        return CompletableFuture.allOf(
+            *allPromises.toTypedArray()
+        ).thenApply {
+            allPromises.asSequence()
+                .map(CompletableFuture<List<Map<Int, Any?>>>::join)
+                .flatten()
+                .map(Map<Int, Any?>::toList)
+                .flatten()
+                .sortedBy(Pair<Int, Any?>::first)
+                .map(Pair<Int, Any?>::second)
+                .toDataFetcherResult()
+        }
+    }
+}

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

@@ -0,0 +1,31 @@
+/*
+ * 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 graphql.schema.DataFetchingEnvironment
+import java.util.concurrent.CompletableFuture
+
+interface FederatedTypePromiseResolver<T> : TypeResolver {
+    /**
+     * Resolves underlying federated types by returning a CompletableFuture
+     *
+     * @param environment DataFetchingEnvironment for executing this query
+     * @param representations _entity query representations that are required to instantiate the target type
+     * @return promise of list of the target federated type instances
+     */
+    fun resolve(environment: DataFetchingEnvironment, representations: List<Map<String, Any>>): CompletableFuture<List<T?>>
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 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.
@@ -18,21 +18,12 @@ package com.expediagroup.graphql.generator.federation.execution
 
 import graphql.schema.DataFetchingEnvironment
 
-/**
- * Resolver used to retrieve target federated types.
- */
-interface FederatedTypeResolver<out T> {
+interface FederatedTypeResolver<out T> : TypeResolver {
 
-    /**
-     * This is the GraphQL name of the type [T]. It is used when running the resolvers and inspecting the
-     * GraphQL "__typename" property during the entities requests
-     */
-    val typeName: String
+    override val typeName: String
 
     /**
-     * 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.
+     * Resolves underlying federated types by using suspending functions
      *
      * @param environment DataFetchingEnvironment for executing this query
      * @param representations _entity query representations that are required to instantiate the target type

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

@@ -0,0 +1,31 @@
+/*
+ * 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
+
+/**
+ * Abstraction that provides a convenient way to resolve 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.
+ */
+sealed interface TypeResolver {
+    /**
+     * This is the GraphQL name of the type. It is used when running the resolvers and inspecting the
+     * GraphQL "__typename" property during the entities requests
+     */
+    val typeName: String
+}

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

@@ -0,0 +1,59 @@
+/*
+ * 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.resolverexecutor
+
+import com.expediagroup.graphql.generator.federation.exception.FederatedRequestFailure
+import com.expediagroup.graphql.generator.federation.execution.FederatedTypePromiseResolver
+import graphql.schema.DataFetchingEnvironment
+import java.util.concurrent.CompletableFuture
+
+object FederatedTypePromiseResolverExecutor : TypeResolverExecutor<FederatedTypePromiseResolver<*>> {
+    override fun execute(
+        resolvableEntities: List<ResolvableEntity<FederatedTypePromiseResolver<*>>>,
+        environment: DataFetchingEnvironment
+    ): CompletableFuture<List<Map<Int, Any?>>> {
+        val futures: List<CompletableFuture<Map<Int, Any?>>> = resolvableEntities.map { resolvableEntity ->
+            resolveEntity(resolvableEntity, environment)
+        }
+        return CompletableFuture.allOf(
+            *futures.toTypedArray()
+        ).thenApply {
+            futures.map(CompletableFuture<Map<Int, Any?>>::join)
+        }
+    }
+
+    @Suppress("TooGenericExceptionCaught")
+    private fun resolveEntity(
+        resolvableEntity: ResolvableEntity<FederatedTypePromiseResolver<*>>,
+        environment: DataFetchingEnvironment,
+    ): CompletableFuture<Map<Int, Any?>> {
+        val indexes = resolvableEntity.indexedRepresentations.map(IndexedValue<Map<String, Any>>::index)
+        val representations = resolvableEntity.indexedRepresentations.map(IndexedValue<Map<String, Any>>::value)
+        val resultsPromise = try {
+            resolvableEntity.resolver.resolve(environment, representations)
+        } catch (e: Exception) {
+            CompletableFuture.completedFuture(
+                representations.map {
+                    FederatedRequestFailure("Exception was thrown while trying to resolve federated type, representation=$it", e)
+                }
+            )
+        }
+        return resultsPromise.thenApply { results ->
+            handleResults(resolvableEntity.typeName, indexes, results)
+        }
+    }
+}