Filter objectTypeNames for non-batched hydration

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

@@ -4,6 +4,7 @@ import graphql.nadel.hints.AllDocumentVariablesHint
 import graphql.nadel.hints.LegacyOperationNamesHint
 import graphql.nadel.hints.NadelDeferSupportHint
 import graphql.nadel.hints.NadelExecuteOnEngineSchemaHint
+import graphql.nadel.hints.NadelHydrationFilterObjectTypesHint
 import graphql.nadel.hints.NadelSharedTypeRenamesHint
 import graphql.nadel.hints.NadelShortCircuitEmptyQueryHint
 import graphql.nadel.hints.NadelVirtualTypeSupportHint
@@ -18,6 +19,7 @@ data class NadelExecutionHints(
     val shortCircuitEmptyQuery: NadelShortCircuitEmptyQueryHint,
     val virtualTypeSupport: NadelVirtualTypeSupportHint,
     val executeOnEngineSchema: NadelExecuteOnEngineSchemaHint,
+    val hydrationFilterObjectTypes: NadelHydrationFilterObjectTypesHint,
 ) {
     /**
      * Returns a builder with the same field values as this object.
@@ -38,6 +40,7 @@ data class NadelExecutionHints(
         private var sharedTypeRenames = NadelSharedTypeRenamesHint { false }
         private var virtualTypeSupport = NadelVirtualTypeSupportHint { false }
         private var executeOnEngineSchema = NadelExecuteOnEngineSchemaHint { false }
+        private var hydrationFilterObjectTypes = NadelHydrationFilterObjectTypesHint { false }
 
         constructor()
 
@@ -47,6 +50,7 @@ data class NadelExecutionHints(
             newResultMergerAndNamespacedTypename = nadelExecutionHints.newResultMergerAndNamespacedTypename
             shortCircuitEmptyQuery = nadelExecutionHints.shortCircuitEmptyQuery
             executeOnEngineSchema = nadelExecutionHints.executeOnEngineSchema
+            hydrationFilterObjectTypes = nadelExecutionHints.hydrationFilterObjectTypes
         }
 
         fun legacyOperationNames(flag: LegacyOperationNamesHint): Builder {
@@ -89,6 +93,11 @@ data class NadelExecutionHints(
             return this
         }
 
+        fun hydrationFilterObjectTypes(flag: NadelHydrationFilterObjectTypesHint): Builder {
+            hydrationFilterObjectTypes = flag
+            return this
+        }
+
         fun build(): NadelExecutionHints {
             return NadelExecutionHints(
                 legacyOperationNames,
@@ -98,7 +107,8 @@ data class NadelExecutionHints(
                 sharedTypeRenames,
                 shortCircuitEmptyQuery,
                 virtualTypeSupport,
-                executeOnEngineSchema
+                executeOnEngineSchema,
+                hydrationFilterObjectTypes,
             )
         }
     }

lib/src/main/java/graphql/nadel/engine/NadelExecutionContext.kt

@@ -10,8 +10,6 @@ import graphql.nadel.hooks.NadelExecutionHooks
 import graphql.nadel.result.NadelResultTracker
 import graphql.normalized.ExecutableNormalizedOperation
 import kotlinx.coroutines.CoroutineScope
-import java.util.concurrent.CompletableFuture
-import java.util.concurrent.ConcurrentHashMap
 
 data class NadelExecutionContext internal constructor(
     val executionInput: ExecutionInput,
@@ -26,8 +24,6 @@ data class NadelExecutionContext internal constructor(
     internal val isPartitionedCall: Boolean = false,
     internal val executionCoroutine: CoroutineScope,
 ) {
-    private val serviceContexts = ConcurrentHashMap<String, CompletableFuture<Any?>>()
-
     val userContext: Any?
         get() {
             return executionInput.context

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

@@ -79,6 +79,15 @@ interface NadelGenericHydrationInstruction {
      */
     val backingFieldDef: GraphQLFieldDefinition
 
+    /**
+     * The names of the object types the backing field returns.
+     *
+     * Note: when virtual types are involved these are the virtual types.
+     *
+     * Could be empty if there are no object types involved e.g. returns scalar, enum etc.
+     */
+    val backingFieldReturnsObjectTypeNames: Set<String>
+
     /**
      * The container of the backing field in the overall schema referenced by [queryPathToBackingField].
      */
@@ -101,6 +110,7 @@ data class NadelHydrationFieldInstruction(
     override val sourceFields: List<NadelQueryPath>,
     override val backingFieldDef: GraphQLFieldDefinition,
     override val backingFieldContainer: GraphQLFieldsContainer,
+    override val backingFieldReturnsObjectTypeNames: Set<String>,
     override val condition: NadelHydrationCondition?,
     /**
      * Hydration can bring about virtual types.
@@ -124,6 +134,7 @@ data class NadelBatchHydrationFieldInstruction(
     override val sourceFields: List<NadelQueryPath>,
     override val backingFieldDef: GraphQLFieldDefinition,
     override val backingFieldContainer: GraphQLFieldsContainer,
+    override val backingFieldReturnsObjectTypeNames: Set<String>,
     override val condition: NadelHydrationCondition?,
     val batchSize: Int,
     val batchHydrationMatchStrategy: NadelBatchHydrationMatchStrategy,

lib/src/main/java/graphql/nadel/engine/transform/hydration/NadelHydrationFieldsBuilder.kt

@@ -1,5 +1,6 @@
 package graphql.nadel.engine.transform.hydration
 
+import graphql.nadel.NadelExecutionHints
 import graphql.nadel.Service
 import graphql.nadel.engine.NadelExecutionContext
 import graphql.nadel.engine.blueprint.NadelBatchHydrationFieldInstruction
@@ -40,7 +41,13 @@ internal object NadelHydrationFieldsBuilder {
                 makeBackingQueries(
                     instruction = instruction,
                     fieldArguments = args,
-                    fieldChildren = deepClone(virtualField.children),
+                    fieldChildren = deepClone(
+                        if (executionContext.hints.hydrationFilterObjectTypes()) {
+                            filterChildren(instruction, virtualField.children)
+                        } else {
+                            virtualField.children
+                        }
+                    ),
                     executionBlueprint = executionBlueprint,
                 )
             }
@@ -56,30 +63,37 @@ internal object NadelHydrationFieldsBuilder {
     }
 
     fun makeBatchBackingQueries(
+        executionHints: NadelExecutionHints,
         executionBlueprint: NadelOverallExecutionBlueprint,
         instruction: NadelBatchHydrationFieldInstruction,
         aliasHelper: NadelAliasHelper,
         virtualField: ExecutableNormalizedField,
         argBatches: List<Map<NadelHydrationArgument, NormalizedInputValue>>,
     ): List<ExecutableNormalizedField> {
-        val backingFieldOverallObjectTypeNames = getBackingFieldOverallObjectTypenames(instruction, executionBlueprint)
-        val fieldChildren = deepClone(fields = virtualField.children)
-            .mapNotNull { childField ->
-                val objectTypesAreNotReturnedByBackingField =
-                    backingFieldOverallObjectTypeNames.none { it in childField.objectTypeNames }
+        val fieldChildren = if (executionHints.hydrationFilterObjectTypes()) {
+            deepClone(fields = filterChildren(instruction, virtualField.children)) +
+                makeObjectIdFields(executionBlueprint, aliasHelper, instruction)
+        } else {
+            val backingFieldOverallObjectTypeNames =
+                getBackingFieldOverallObjectTypenames(instruction, executionBlueprint)
+            deepClone(fields = virtualField.children)
+                .mapNotNull { childField ->
+                    val objectTypesAreNotReturnedByBackingField =
+                        backingFieldOverallObjectTypeNames.none { it in childField.objectTypeNames }
 
-                if (objectTypesAreNotReturnedByBackingField) {
-                    null
-                } else {
-                    childField.toBuilder()
-                        .clearObjectTypesNames()
-                        .objectTypeNames(childField.objectTypeNames.filter { it in backingFieldOverallObjectTypeNames })
-                        .build()
+                    if (objectTypesAreNotReturnedByBackingField) {
+                        null
+                    } else {
+                        childField.toBuilder()
+                            .clearObjectTypesNames()
+                            .objectTypeNames(childField.objectTypeNames.filter { it in backingFieldOverallObjectTypeNames })
+                            .build()
+                    }
                 }
-            }
-            .let { children ->
-                children + makeObjectIdFields(executionBlueprint, aliasHelper, instruction)
-            }
+                .let { children ->
+                    children + makeObjectIdFields(executionBlueprint, aliasHelper, instruction)
+                }
+        }
 
         return argBatches.map { argBatch ->
             makeBackingQueries(
@@ -91,6 +105,29 @@ internal object NadelHydrationFieldsBuilder {
         }
     }
 
+    private fun filterChildren(
+        instruction: NadelGenericHydrationInstruction,
+        children: List<ExecutableNormalizedField>,
+    ): List<ExecutableNormalizedField> {
+        return children
+            .mapNotNull { childField ->
+                val legalObjectTypeNames =
+                    childField.objectTypeNames.intersect(instruction.backingFieldReturnsObjectTypeNames)
+
+                if (legalObjectTypeNames.isEmpty()) {
+                    null
+                } else if (legalObjectTypeNames == childField.objectTypeNames) {
+                    childField
+                } else {
+                    childField.toBuilder()
+                        .clearObjectTypesNames()
+                        .objectTypeNames(legalObjectTypeNames.toList())
+                        .build()
+                }
+            }
+    }
+
+    @Deprecated("Will be removed once NadelHydrationFilterObjectTypesHint is rolled out")
     private fun getBackingFieldOverallObjectTypenames(
         instruction: NadelBatchHydrationFieldInstruction,
         executionBlueprint: NadelOverallExecutionBlueprint,

lib/src/main/java/graphql/nadel/engine/transform/hydration/batch/NadelNewBatchHydrator.kt

@@ -16,7 +16,6 @@ import graphql.nadel.engine.transform.getInstructionsForNode
 import graphql.nadel.engine.transform.hydration.NadelHydrationFieldsBuilder
 import graphql.nadel.engine.transform.hydration.NadelHydrationUtil.getInstructionsToAddErrors
 import graphql.nadel.engine.transform.hydration.batch.NadelBatchHydrationTransform.State
-import graphql.nadel.engine.transform.hydration.batch.NadelNewBatchHydrator.SourceObjectMetadata
 import graphql.nadel.engine.transform.hydration.batch.indexing.NadelBatchHydrationIndexBasedIndexer
 import graphql.nadel.engine.transform.hydration.batch.indexing.NadelBatchHydrationIndexKey
 import graphql.nadel.engine.transform.hydration.batch.indexing.NadelBatchHydrationIndexer
@@ -444,6 +443,7 @@ internal class NadelNewBatchHydrator(
 
         val queries = NadelHydrationFieldsBuilder
             .makeBatchBackingQueries(
+                executionHints = executionContext.hints,
                 executionBlueprint = executionBlueprint,
                 instruction = instruction,
                 aliasHelper = aliasHelper,

lib/src/main/java/graphql/nadel/engine/util/GraphQLUtil.kt

@@ -60,13 +60,11 @@ import graphql.schema.GraphQLCodeRegistry
 import graphql.schema.GraphQLFieldDefinition
 import graphql.schema.GraphQLFieldsContainer
 import graphql.schema.GraphQLInputType
-import graphql.schema.GraphQLInterfaceType
 import graphql.schema.GraphQLNamedType
 import graphql.schema.GraphQLObjectType
 import graphql.schema.GraphQLSchema
 import graphql.schema.GraphQLType
 import graphql.schema.GraphQLTypeUtil
-import graphql.schema.GraphQLUnionType
 import graphql.schema.GraphQLUnmodifiedType
 import graphql.schema.idl.TypeUtil
 import kotlinx.coroutines.future.asDeferred
@@ -525,20 +523,31 @@ internal suspend fun NadelInstrumentation.beginExecute(
  * Turns GraphQL types to object types when possible e.g. finds concrete implementations
  * for interfaces, gets object types inside unions, and returns objects as is.
  */
-fun resolveObjectTypes(
+inline fun resolveObjectTypes(
     schema: GraphQLSchema,
     type: GraphQLType,
     onNotObjectType: (GraphQLType) -> Nothing,
 ): List<GraphQLObjectType> {
-    return when (val unwrappedType = type.unwrapAll()) {
-        is GraphQLObjectType -> listOf(unwrappedType)
-        is GraphQLUnionType -> unwrappedType.types.flatMap {
-            resolveObjectTypes(schema, type = it, onNotObjectType)
-        }
-
-        is GraphQLInterfaceType -> schema.getImplementations(unwrappedType)
-        else -> onNotObjectType(unwrappedType)
-    }
+    return type.unwrapAll()
+        .whenType(
+            enumType = onNotObjectType,
+            inputObjectType = onNotObjectType,
+            interfaceType = { interfaceType ->
+                schema.getImplementations(interfaceType)
+                    .map { objectImpl -> objectImpl }
+            },
+            objectType = {
+                listOf(it)
+            },
+            scalarType = onNotObjectType,
+            unionType = { unionType ->
+                unionType.types
+                    .map { memberType ->
+                        memberType as? GraphQLObjectType
+                            ?: throw IllegalArgumentException("Member type must be object but found ${memberType?.javaClass?.name}")
+                    }
+            },
+        )
 }
 
 /**

lib/src/main/java/graphql/nadel/hints/NadelHydrationFilterObjectTypesHint.kt

@@ -0,0 +1,5 @@
+package graphql.nadel.hints
+
+fun interface NadelHydrationFilterObjectTypesHint {
+    operator fun invoke(): Boolean
+}

lib/src/main/java/graphql/nadel/validation/hydration/NadelHydrationValidation.kt

@@ -9,6 +9,7 @@ import graphql.nadel.definition.hydration.NadelHydrationDefinition
 import graphql.nadel.definition.virtualType.hasVirtualTypeDefinition
 import graphql.nadel.engine.blueprint.NadelBatchHydrationFieldInstruction
 import graphql.nadel.engine.blueprint.NadelHydrationFieldInstruction
+import graphql.nadel.engine.blueprint.NadelVirtualTypeContext
 import graphql.nadel.engine.blueprint.hydration.NadelBatchHydrationMatchStrategy
 import graphql.nadel.engine.blueprint.hydration.NadelHydrationArgument
 import graphql.nadel.engine.blueprint.hydration.NadelHydrationCondition
@@ -25,6 +26,7 @@ import graphql.nadel.engine.util.isList
 import graphql.nadel.engine.util.isNonNull
 import graphql.nadel.engine.util.makeFieldCoordinates
 import graphql.nadel.engine.util.partitionCount
+import graphql.nadel.engine.util.resolveObjectTypes
 import graphql.nadel.engine.util.singleOfTypeOrNull
 import graphql.nadel.engine.util.startsWith
 import graphql.nadel.engine.util.unwrapAll
@@ -279,6 +281,7 @@ class NadelHydrationValidation internal constructor(
                 timeout = hydrationDefinition.timeout,
                 sourceFields = sourceFields,
                 backingFieldDef = backingField,
+                backingFieldReturnsObjectTypeNames = getReturnsObjectTypeNames(backingField, virtualTypeContext),
                 backingFieldContainer = backingFieldContainer,
                 condition = hydrationCondition,
                 virtualTypeContext = virtualTypeContext,
@@ -356,6 +359,7 @@ class NadelHydrationValidation internal constructor(
                 sourceFields = sourceFields,
                 backingFieldDef = backingField,
                 backingFieldContainer = backingFieldContainer,
+                backingFieldReturnsObjectTypeNames = getReturnsObjectTypeNames(backingField, null),
                 condition = hydrationCondition,
                 batchSize = hydrationDefinition.batchSize,
                 batchHydrationMatchStrategy = matchStrategy,
@@ -744,4 +748,21 @@ class NadelHydrationValidation internal constructor(
         }
         return false
     }
+
+    context(NadelValidationContext)
+    private fun getReturnsObjectTypeNames(
+        backingField: GraphQLFieldDefinition,
+        virtualTypeContext: NadelVirtualTypeContext?,
+    ): Set<String> {
+        val returnsObjectTypes = resolveObjectTypes(engineSchema, backingField.type) { return emptySet() }
+            .mapTo(mutableSetOf()) { it.name }
+
+        return if (virtualTypeContext == null) {
+            returnsObjectTypes
+        } else {
+            returnsObjectTypes.mapTo(mutableSetOf()) {
+                virtualTypeContext.backingTypeToVirtualType[it] ?: it
+            }
+        }
+    }
 }

test/src/test/kotlin/graphql/nadel/tests/legacy/new hydration/polymorphic hydrations/solitary polymorphic hydration snapshot.kt

@@ -64,8 +64,6 @@ public class `solitary polymorphic hydration snapshot` : TestSnapshot() {
                 | {
                 |   humanById(id: "HUMAN-0") {
                 |     __typename
-                |     __typename__type_filter__id: __typename
-                |     __typename__type_filter__breed: __typename
                 |     id
                 |     name
                 |   }
@@ -77,9 +75,7 @@ public class `solitary polymorphic hydration snapshot` : TestSnapshot() {
                 |   "data": {
                 |     "humanById": {
                 |       "__typename": "Human",
-                |       "__typename__type_filter__id": "Human",
                 |       "id": "HUMAN-0",
-                |       "__typename__type_filter__breed": "Human",
                 |       "name": "Fanny Longbottom"
                 |     }
                 |   }
@@ -94,8 +90,6 @@ public class `solitary polymorphic hydration snapshot` : TestSnapshot() {
                 | {
                 |   petById(id: "PET-0") {
                 |     __typename
-                |     __typename__type_filter__id: __typename
-                |     __typename__type_filter__name: __typename
                 |     breed
                 |     id
                 |   }
@@ -108,9 +102,7 @@ public class `solitary polymorphic hydration snapshot` : TestSnapshot() {
                 |     "petById": {
                 |       "__typename": "Pet",
                 |       "id": "PET-0",
-                |       "__typename__type_filter__id": "Pet",
-                |       "breed": "Akita",
-                |       "__typename__type_filter__name": "Pet"
+                |       "breed": "Akita"
                 |     }
                 |   }
                 | }

test/src/test/kotlin/graphql/nadel/tests/legacy/new hydration/polymorphic hydrations/solitary polymorphic hydration when hook returns null snapshot.kt

@@ -64,8 +64,6 @@ public class `solitary polymorphic hydration when hook returns null snapshot` :
                 | {
                 |   humanById(id: "HUMAN-0") {
                 |     __typename
-                |     __typename__type_filter__id: __typename
-                |     __typename__type_filter__breed: __typename
                 |     id
                 |     name
                 |   }
@@ -77,9 +75,7 @@ public class `solitary polymorphic hydration when hook returns null snapshot` :
                 |   "data": {
                 |     "humanById": {
                 |       "__typename": "Human",
-                |       "__typename__type_filter__id": "Human",
                 |       "id": "HUMAN-0",
-                |       "__typename__type_filter__breed": "Human",
                 |       "name": "Fanny Longbottom"
                 |     }
                 |   }