Fix top level hydration (#572)

* Fix top level hydration

* Add plain hydration

* Functional solution

lib/src/main/java/graphql/nadel/NextgenEngine.kt

@@ -29,7 +29,6 @@ import graphql.nadel.engine.transform.result.NadelResultTransformer
 import graphql.nadel.engine.util.MutableJsonMap
 import graphql.nadel.engine.util.beginExecute
 import graphql.nadel.engine.util.compileToDocument
-import graphql.nadel.engine.util.copy
 import graphql.nadel.engine.util.getOperationKind
 import graphql.nadel.engine.util.newExecutionResult
 import graphql.nadel.engine.util.newGraphQLError
@@ -47,9 +46,9 @@ import graphql.nadel.instrumentation.parameters.NadelInstrumentationOnErrorParam
 import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters.ChildStep.Companion.DocumentCompilation
 import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters.RootStep
 import graphql.nadel.instrumentation.parameters.child
-import graphql.nadel.schema.NadelDirectives.namespacedDirectiveDefinition
 import graphql.nadel.result.NadelResultMerger
 import graphql.nadel.result.NadelResultTracker
+import graphql.nadel.schema.NadelDirectives.namespacedDirectiveDefinition
 import graphql.nadel.util.OperationNameUtil
 import graphql.normalized.ExecutableNormalizedField
 import graphql.normalized.ExecutableNormalizedOperationFactory.createExecutableNormalizedOperationWithRawVariables
@@ -64,7 +63,6 @@ import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.future.asCompletableFuture
 import kotlinx.coroutines.future.asDeferred
@@ -285,11 +283,10 @@ internal class NextgenEngine(
                 field = topLevelField
             )
         }
-        val transformedQuery = queryTransform.result.single()
         val result: ServiceExecutionResult = timer.time(step = RootStep.ServiceExecution.child(service.name)) {
             executeService(
                 service = service,
-                transformedQuery = transformedQuery,
+                topLevelFields = queryTransform.result,
                 executionContext = executionContext,
                 serviceExecutionContext = serviceExecutionContext,
                 executionHydrationDetails = executionContext.hydrationDetails,
@@ -299,11 +296,11 @@ internal class NextgenEngine(
             executionContext.incrementalResultSupport.defer(
                 result.incrementalItemPublisher
                     .asFlow()
-                    .onEach {delayedIncrementalResult ->
+                    .onEach { delayedIncrementalResult ->
                         // Transform
                         delayedIncrementalResult.incremental
                             ?.filterIsInstance<DeferPayload>()
-                            ?.forEach {deferPayload ->
+                            ?.forEach { deferPayload ->
                                 resultTransformer
                                     .transform(
                                         executionContext = executionContext,
@@ -314,7 +311,8 @@ internal class NextgenEngine(
                                         service = service,
                                         result = result,
                                         deferPayload = deferPayload,
-                                    ) }
+                                    )
+                            }
                     }
             )
         }
@@ -338,7 +336,7 @@ internal class NextgenEngine(
 
     private suspend fun executeService(
         service: Service,
-        transformedQuery: ExecutableNormalizedField,
+        topLevelFields: List<ExecutableNormalizedField>,
         executionContext: NadelExecutionContext,
         serviceExecutionContext: NadelServiceExecutionContext,
         executionHydrationDetails: ServiceExecutionHydrationDetails? = null,
@@ -352,9 +350,9 @@ internal class NextgenEngine(
         val compileResult = timer.time(step = DocumentCompilation) {
             compileToDocument(
                 schema = service.underlyingSchema,
-                operationKind = transformedQuery.getOperationKind(engineSchema),
+                operationKind = topLevelFields.first().getOperationKind(engineSchema),
                 operationName = getOperationName(service, executionContext),
-                topLevelFields = listOf(transformedQuery),
+                topLevelFields = topLevelFields,
                 variablePredicate = jsonPredicate,
                 deferSupport = executionContext.hints.deferSupport(),
             )
@@ -370,9 +368,16 @@ internal class NextgenEngine(
             serviceContext = executionContext.getContextForService(service).await(),
             serviceExecutionContext = serviceExecutionContext,
             hydrationDetails = executionHydrationDetails,
-            executableNormalizedField = transformedQuery,
+            // Prefer non __typename field first, otherwise we just get first
+            executableNormalizedField = topLevelFields
+                .asSequence()
+                .filterNot {
+                    it.fieldName == TypeNameMetaFieldDef.name
+                }
+                .firstOrNull() ?: topLevelFields.first(),
         )
-        val serviceExecution = chooseServiceExecution(service, transformedQuery, executionContext.hints)
+
+        val serviceExecution = getServiceExecution(service, topLevelFields, executionContext.hints)
         val serviceExecResult = try {
             serviceExecution.execute(serviceExecParams)
                 .asDeferred()
@@ -408,39 +413,53 @@ internal class NextgenEngine(
             )
         }
 
-        val transformedData: MutableJsonMap = serviceExecResult.data.let { data ->
-            data.takeIf { transformedQuery.resultKey in data }
-                ?: mutableMapOf(transformedQuery.resultKey to null)
-        }
+        val transformedData: MutableJsonMap = serviceExecResult.data
+            .let { data ->
+                // Ensures data always has root fields as keys
+                topLevelFields
+                    .asSequence()
+                    .map {
+                        it.resultKey
+                    }
+                    .associateWithTo(mutableMapOf()) { resultKey ->
+                        data[resultKey]
+                    }
+            }
 
-        return when(serviceExecResult) {
+        return when (serviceExecResult) {
             is NadelServiceExecutionResultImpl -> serviceExecResult.copy(data = transformedData)
             is NadelIncrementalServiceExecutionResult -> serviceExecResult.copy(data = transformedData)
         }
     }
 
-    private fun chooseServiceExecution(
+    private fun getServiceExecution(
         service: Service,
-        transformedQuery: ExecutableNormalizedField,
+        topLevelFields: List<ExecutableNormalizedField>,
         hints: NadelExecutionHints,
     ): ServiceExecution {
-        return when {
-            hints.shortCircuitEmptyQuery(service) && onlyTopLevelTypenameField(transformedQuery) ->
-                engineSchemaIntrospectionService.serviceExecution
-            else -> service.serviceExecution
+        if (hints.shortCircuitEmptyQuery(service) && isOnlyTopLevelFieldTypename(topLevelFields)) {
+            return engineSchemaIntrospectionService.serviceExecution
         }
+
+        return service.serviceExecution
     }
 
-    private fun onlyTopLevelTypenameField(executableNormalizedField: ExecutableNormalizedField): Boolean {
-        if (executableNormalizedField.fieldName == TypeNameMetaFieldDef.name) {
+    private fun isOnlyTopLevelFieldTypename(topLevelFields: List<ExecutableNormalizedField>): Boolean {
+        val topLevelField = topLevelFields.singleOrNull() ?: return false
+
+        if (topLevelField.fieldName == TypeNameMetaFieldDef.name) {
             return true
         }
-        val operationType = engineSchema.getTypeAs<GraphQLObjectType>(executableNormalizedField.singleObjectTypeName)
-        val topLevelFieldDefinition = operationType.getField(executableNormalizedField.name)
+
+        val operationType = engineSchema.getTypeAs<GraphQLObjectType>(topLevelField.singleObjectTypeName)
+        val topLevelFieldDefinition = operationType.getField(topLevelField.name)
+
         return if (topLevelFieldDefinition.hasAppliedDirective(namespacedDirectiveDefinition.name)) {
-            executableNormalizedField.hasChildren()
-                && executableNormalizedField.children.all { it.name == TypeNameMetaFieldDef.name }
-        } else false
+            topLevelField.hasChildren()
+                && topLevelField.children.all { it.name == TypeNameMetaFieldDef.name }
+        } else {
+            false
+        }
     }
 
     private fun getDocumentVariablePredicate(hints: NadelExecutionHints, service: Service): VariablePredicate {

lib/src/main/java/graphql/nadel/engine/blueprint/NadelExecutionBlueprintFactory.kt

@@ -43,7 +43,6 @@ import graphql.schema.GraphQLFieldsContainer
 import graphql.schema.GraphQLObjectType
 import graphql.schema.GraphQLScalarType
 import graphql.schema.GraphQLSchema
-import graphql.schema.GraphQLType
 import java.math.BigInteger
 
 internal object NadelExecutionBlueprintFactory {
@@ -329,7 +328,7 @@ private class Factory(
                     return@mapNotNull null
                 }
 
-                val underlyingParentType = getUnderlyingType(hydratedFieldParentType)
+                val underlyingParentType = getUnderlyingType(hydratedFieldParentType, hydratedFieldDef)
                     ?: error("No underlying type for: ${hydratedFieldParentType.name}")
                 val fieldDefs = underlyingParentType.getFieldsAlong(inputValueDef.valueSource.queryPathToField.segments)
                 inputValueDef.takeIf {
@@ -410,7 +409,7 @@ private class Factory(
     private fun getBatchHydrationSourceFields(
         matchStrategy: NadelBatchHydrationMatchStrategy,
         hydrationArgs: List<NadelHydrationActorInputDef>,
-        condition: NadelHydrationCondition?
+        condition: NadelHydrationCondition?,
     ): List<NadelQueryPath> {
         val paths = (when (matchStrategy) {
             NadelBatchHydrationMatchStrategy.MatchIndex -> emptyList()
@@ -541,7 +540,7 @@ private class Factory(
                     val pathToField = argSourceType.pathToField
                     FieldResultValue(
                         queryPathToField = NadelQueryPath(pathToField),
-                        fieldDefinition = getUnderlyingType(hydratedFieldParentType)
+                        fieldDefinition = getUnderlyingType(hydratedFieldParentType, hydratedFieldDef)
                             ?.getFieldAt(pathToField)
                             ?: error("No field defined at: ${hydratedFieldParentType.name}.${pathToField.joinToString(".")}"),
                     )
@@ -561,11 +560,25 @@ private class Factory(
         }
     }
 
-    private fun <T : GraphQLType> getUnderlyingType(overallType: T): T? {
+    /**
+     * Gets the underlying type for an [GraphQLObjectType]
+     *
+     * The [childField] is there in case the [overallType] is an operation type.
+     * In that case we still need to know which service's operation type to return.
+     */
+    private fun getUnderlyingType(
+        overallType: GraphQLObjectType,
+        childField: GraphQLFieldDefinition,
+    ): GraphQLObjectType? {
         val renameInstruction = makeTypeRenameInstruction(overallType as? GraphQLDirectiveContainer ?: return null)
-        val service = definitionNamesToService[overallType.name]
-            ?: error("Unknown service for type: ${overallType.name}")
         val underlyingName = renameInstruction?.underlyingName ?: overallType.name
+
+        val fieldCoordinates = makeFieldCoordinates(overallType, childField)
+
+        val service = definitionNamesToService[overallType.name]
+            ?: coordinatesToService[fieldCoordinates]
+            ?: error("Unable to determine service for $fieldCoordinates")
+
         return service.underlyingSchema.getTypeAs(underlyingName)
     }
 

test/src/test/kotlin/graphql/nadel/tests/next/fixtures/batchHydration/BatchHydrationAtQueryTypeTest.kt

@@ -0,0 +1,62 @@
+package graphql.nadel.tests.next.fixtures.batchHydration
+
+import graphql.nadel.engine.util.strictAssociateBy
+import graphql.nadel.tests.next.NadelIntegrationTest
+
+class BatchHydrationAtQueryTypeTest : NadelIntegrationTest(
+    query = """
+        query {
+          myIssues {
+            title
+          }
+        }
+    """.trimIndent(),
+    services = listOf(
+        Service(
+            name = "issues",
+            overallSchema = """
+                type Query {
+                  issuesByIds(ids: [ID!]!): [Issue]
+                  myIssueKeys(limit: Int! = 25): [ID!] @hidden
+                  myIssues: [Issue]
+                    @hydrated(
+                      service: "issues"
+                      field: "issuesByIds"
+                      arguments: [{name: "ids", value: "$source.myIssueKeys"}]
+                      identifiedBy: "id"
+                    )
+                }
+                type Issue {
+                  id: ID!
+                  title: String
+                }
+            """.trimIndent(),
+            runtimeWiring = { wiring ->
+                data class Issue(
+                    val id: String,
+                    val title: String,
+                )
+
+                val issuesByIds = listOf(
+                    Issue(id = "hello", title = "Hello there"),
+                    Issue(id = "afternoon", title = "Good afternoon"),
+                    Issue(id = "bye", title = "Farewell"),
+                ).strictAssociateBy { it.id }
+
+                wiring
+                    .type("Query") { type ->
+                        type
+                            .dataFetcher("issuesByIds") { env ->
+                                env.getArgument<List<String>>("ids")
+                                    ?.map {
+                                        issuesByIds[it]
+                                    }
+                            }
+                            .dataFetcher("myIssueKeys") { env ->
+                                listOf("hello", "bye")
+                            }
+                    }
+            },
+        ),
+    ),
+)