Refactor federation code for readability (#887)

smyrick · Shane Myrick · web-flow · commit a41fc35b6047 · 2020-09-21T12:38:40.000-05:00
Co-authored-by: Shane Myrick &lt;accounts@shanemyrick.com&gt;
diff --git a/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/federation/FederatedSchemaGeneratorHooks.kt b/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/federation/FederatedSchemaGeneratorHooks.kt
@@ -30,6 +30,7 @@ import com.expediagroup.graphql.federation.execution.EntityResolver
 import com.expediagroup.graphql.federation.execution.FederatedTypeResolver
 import com.expediagroup.graphql.federation.extensions.addDirectivesIfNotPresent
 import com.expediagroup.graphql.federation.types.ANY_SCALAR_TYPE
+import com.expediagroup.graphql.federation.types.ENTITY_UNION_NAME
 import com.expediagroup.graphql.federation.types.FIELD_SET_SCALAR_TYPE
 import com.expediagroup.graphql.federation.types.SERVICE_FIELD_DEFINITION
 import com.expediagroup.graphql.federation.types._Service
@@ -83,52 +84,71 @@ open class FederatedSchemaGeneratorHooks(private val resolvers: List<FederatedTy
             .field(SERVICE_FIELD_DEFINITION)
             .withDirective(EXTENDS_DIRECTIVE_TYPE)
 
-        /**
-         * Register the data fetcher for the SDL returned by _service field.
-         *
-         * It should NOT contain:
-         *   - default schema definition
-         *   - empty Query type
-         *   - any directive definitions
-         *   - any custom directives
-         *   - new federated scalars
-         */
-        val sdl = originalSchema.print(
-            includeDefaultSchemaDefinition = false,
-            includeDirectiveDefinitions = false,
-            includeDirectivesFilter = customDirectivePredicate
-        ).replace(scalarDefinitionRegex, "")
-            .replace(emptyQueryRegex, "")
-            .trim()
+        // Register the data fetcher for the _service query
+        val sdl = getFederatedServiceSdl(originalSchema)
         federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, SERVICE_FIELD_DEFINITION.name), DataFetcher { _Service(sdl) })
 
-        val entityTypeNames = originalSchema.allTypesAsList
-            .asSequence()
-            .filterIsInstance<GraphQLObjectType>()
-            .filter { type -> type.getDirective(KEY_DIRECTIVE_NAME) != null }
-            .map { it.name }
-            .toSet()
-
-        // Add the _entities field to the query
+        // Add the _entities field to the query and register all the _Entity union types
+        val entityTypeNames = getFederatedEntities(originalSchema)
         if (entityTypeNames.isNotEmpty()) {
             val entityField = generateEntityFieldDefinition(entityTypeNames)
             federatedQuery.field(entityField)
 
             federatedCodeRegistry.dataFetcher(FieldCoordinates.coordinates(originalQuery.name, entityField.name), EntityResolver(resolvers))
-            federatedCodeRegistry.typeResolver("_Entity") { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObjectName()) }
+            federatedCodeRegistry.typeResolver(ENTITY_UNION_NAME) { env: TypeResolutionEnvironment -> env.schema.getObjectType(env.getObjectName()) }
             federatedSchemaBuilder.additionalType(ANY_SCALAR_TYPE)
         }
 
         return federatedSchemaBuilder.query(federatedQuery.build())
             .codeRegistry(federatedCodeRegistry.build())
     }
 
-    // skip validation for empty query type - federation will add _service query
+    /**
+     * Skip validation for empty query type - federation will add _service query
+     */
     override fun didGenerateQueryObject(type: GraphQLObjectType): GraphQLObjectType = type
-}
 
-private fun TypeResolutionEnvironment.getObjectName(): String? {
-    val kClass = this.getObject<Any>().javaClass.kotlin
-    return kClass.findAnnotation<GraphQLName>()?.value
-        ?: kClass.simpleName
+    /**
+     * Get the modified SDL returned by _service field
+     *
+     * It should NOT contain:
+     *   - default schema definition
+     *   - empty Query type
+     *   - any directive definitions
+     *   - any custom directives
+     *   - new federated scalars
+     *
+     * See the federation spec for more details:
+     * https://www.apollographql.com/docs/apollo-server/federation/federation-spec/#query_service
+     */
+    private fun getFederatedServiceSdl(schema: GraphQLSchema): String {
+        return schema.print(
+            includeDefaultSchemaDefinition = false,
+            includeDirectiveDefinitions = false,
+            includeDirectivesFilter = customDirectivePredicate
+        ).replace(scalarDefinitionRegex, "")
+            .replace(emptyQueryRegex, "")
+            .trim()
+    }
+
+    /**
+     * Get all the federation entities in the _Entity union, aka all the types with the @key directive.
+     *
+     * See the federation spec:
+     * https://www.apollographql.com/docs/apollo-server/federation/federation-spec/#union-_entity
+     */
+    private fun getFederatedEntities(originalSchema: GraphQLSchema): Set<String> {
+        return originalSchema.allTypesAsList
+            .asSequence()
+            .filterIsInstance<GraphQLObjectType>()
+            .filter { type -> type.getDirective(KEY_DIRECTIVE_NAME) != null }
+            .map { it.name }
+            .toSet()
+    }
+
+    private fun TypeResolutionEnvironment.getObjectName(): String? {
+        val kClass = this.getObject<Any>().javaClass.kotlin
+        return kClass.findAnnotation<GraphQLName>()?.value
+            ?: kClass.simpleName
+    }
 }
diff --git a/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/federation/execution/EntityResolver.kt b/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/federation/execution/EntityResolver.kt
@@ -20,11 +20,12 @@ import graphql.GraphQLError
 import graphql.execution.DataFetcherResult
 import graphql.schema.DataFetcher
 import graphql.schema.DataFetchingEnvironment
+import java.util.concurrent.CompletableFuture
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.future.future
-import java.util.concurrent.CompletableFuture
 
 private const val TYPENAME_FIELD = "__typename"
 private const val REPRESENTATIONS = "representations"
@@ -51,16 +52,13 @@ open class EntityResolver(resolvers: List<FederatedTypeResolver<*>>) : DataFetch
      */
     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_FIELD].toString() }
+
         return GlobalScope.future {
             val data = mutableListOf<Any?>()
             val errors = mutableListOf<GraphQLError>()
-            indexedBatchRequestsByType.map { (typeName, indexedRequests) ->
-                async {
-                    resolveType(env, typeName, indexedRequests, resolverMap)
-                }
-            }.awaitAll()
+
+            resolveRequests(indexedBatchRequestsByType, env)
                 .flatten()
                 .sortedBy { it.first }
                 .forEach {
@@ -72,10 +70,21 @@ open class EntityResolver(resolvers: List<FederatedTypeResolver<*>>) : DataFetch
                         data.add(result)
                     }
                 }
+
             DataFetcherResult.newResult<List<Any?>>()
                 .data(data)
                 .errors(errors)
                 .build()
         }
     }
+
+    private suspend fun resolveRequests(indexedBatchRequestsByType: Map<String, List<IndexedValue<Map<String, Any>>>>, env: DataFetchingEnvironment): List<List<Pair<Int, Any?>>> {
+        return coroutineScope {
+            indexedBatchRequestsByType.map { (typeName, indexedRequests) ->
+                async {
+                    resolveType(env, typeName, indexedRequests, resolverMap)
+                }
+            }.awaitAll()
+        }
+    }
 }
diff --git a/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/federation/types/_Entity.kt b/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/federation/types/_Entity.kt
@@ -23,6 +23,9 @@ import graphql.schema.GraphQLNonNull
 import graphql.schema.GraphQLTypeReference
 import graphql.schema.GraphQLUnionType
 
+internal const val ENTITY_UNION_NAME = "_Entity"
+internal const val ENTITIES_FIELD_NAME = "_entities"
+
 /**
  * Generates union of all types that use the  @key directive, including both types native to the schema and extended types.
  */
@@ -47,14 +50,14 @@ internal fun generateEntityFieldDefinition(federatedTypes: Set<String>): GraphQL
     val graphQLType = GraphQLNonNull(
         GraphQLList(
             GraphQLUnionType.newUnionType()
-                .name("_Entity")
+                .name(ENTITY_UNION_NAME)
                 .possibleTypes(*possibleTypes)
                 .build()
         )
     )
 
     return GraphQLFieldDefinition.newFieldDefinition()
-        .name("_entities")
+        .name(ENTITIES_FIELD_NAME)
         .description("Union of all types that use the @key directive, including both types native to the schema and extended types")
         .argument(graphQLArgument)
         .type(graphQLType)
diff --git a/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/federation/FederatedSchemaGeneratorTest.kt b/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/federation/FederatedSchemaGeneratorTest.kt
@@ -21,11 +21,12 @@ import com.expediagroup.graphql.extensions.print
 import com.expediagroup.graphql.federation.data.queries.simple.NestedQuery
 import com.expediagroup.graphql.federation.data.queries.simple.SimpleQuery
 import com.expediagroup.graphql.federation.directives.KEY_DIRECTIVE_NAME
+import com.expediagroup.graphql.federation.types.ENTITY_UNION_NAME
 import graphql.schema.GraphQLUnionType
-import org.junit.jupiter.api.Assertions.assertEquals
-import org.junit.jupiter.api.Test
 import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
 
 private val FEDERATED_SDL =
     """
@@ -137,7 +138,7 @@ class FederatedSchemaGeneratorTest {
         assertNotNull(productType)
         assertNotNull(productType.getDirective(KEY_DIRECTIVE_NAME))
 
-        val entityUnion = schema.getType("_Entity") as? GraphQLUnionType
+        val entityUnion = schema.getType(ENTITY_UNION_NAME) as? GraphQLUnionType
         assertNotNull(entityUnion)
         assertTrue(entityUnion.types.contains(productType))
     }
diff --git a/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/federation/execution/FederatedQueryResolverTest.kt b/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/federation/execution/FederatedQueryResolverTest.kt
@@ -19,13 +19,14 @@ package com.expediagroup.graphql.federation.execution
 import com.expediagroup.graphql.federation.data.BookResolver
 import com.expediagroup.graphql.federation.data.UserResolver
 import com.expediagroup.graphql.federation.data.federatedTestSchema
+import com.expediagroup.graphql.federation.types.ENTITIES_FIELD_NAME
 import graphql.ExecutionInput
 import graphql.GraphQL
-import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
+import org.junit.jupiter.api.Test
 
 const val FEDERATED_QUERY =
 """
@@ -65,7 +66,7 @@ class FederatedQueryResolverTest {
         val result = graphQL.executeAsync(executionInput).get().toSpecification()
 
         assertNotNull(result["data"] as? Map<*, *>) { data ->
-            assertNotNull(data["_entities"] as? List<*>) { entities ->
+            assertNotNull(data[ENTITIES_FIELD_NAME] as? List<*>) { entities ->
                 assertFalse(entities.isEmpty())
                 assertEquals(representations.size, entities.size)
                 assertNotNull(entities[0] as? Map<*, *>) { user ->
diff --git a/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/federation/types/EntityTest.kt b/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/federation/types/EntityTest.kt
@@ -18,11 +18,11 @@ package com.expediagroup.graphql.federation.types
 
 import com.expediagroup.graphql.extensions.unwrapType
 import graphql.schema.GraphQLUnionType
-import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
+import org.junit.jupiter.api.Test
 
 class EntityTest {
 
@@ -37,14 +37,14 @@ class EntityTest {
     fun `generateEntityFieldDefinition should return a valid type on a single set`() {
         val result = generateEntityFieldDefinition(setOf("MyType"))
         assertNotNull(result)
-        assertEquals(expected = "_entities", actual = result.name)
+        assertEquals(expected = ENTITIES_FIELD_NAME, actual = result.name)
         assertFalse(result.description.isNullOrEmpty())
         assertEquals(expected = 1, actual = result.arguments.size)
 
         val graphQLUnionType = result.type.unwrapType() as? GraphQLUnionType
 
         assertNotNull(graphQLUnionType)
-        assertEquals(expected = "_Entity", actual = graphQLUnionType.name)
+        assertEquals(expected = ENTITY_UNION_NAME, actual = graphQLUnionType.name)
         assertEquals(expected = 1, actual = graphQLUnionType.types.size)
         assertEquals(expected = "MyType", actual = graphQLUnionType.types.first().name)
     }
