Support DataLoader argument in `@EntityMapping` methods

[rstoyanchev]

@EntityMapping methods support a method signature with a list of argument id's as input and returning a list of entities as output. This allows loading all entities at once much like @QueryMapping methods.

Users migrating from GraphQL Java Kickstart previously used a DataLoader in entity resolvers, and would like to be able to do the same. We should explore the option to support injecting a DataLoader argument. This could be as simple as adding DataLoaderMethodArgumentResolver to the list of resolvers, but need to experiment.

This is based on discussion below #1023 (comment).

[xcq1]

Hey, really cool that you follow up on this. Let me answer your open questions from the other issue:

Aside from re-using your existing code, I'm curious, in relation to data loader caching, do you actually have cases where a federated type appears again in nested parts of the graph, and is that even supported in federation

From what I've gathered about federation it is a perfect fit for the data loader interface, since the federation gateway will keep a list of ids of all types while traversing the graph. Since I already have that data loader for my sub-graph, it only makes sense to me to directly pipe the federation request to the data loader as there should be no fundamental difference whether the query tries to fetch information available exclusively from my subgraph or from a list of federated references.

Off the top of my head, I couldn't find a query that really requests both parts at the same time, but it would certainly be conceivable. Consider for example a graph of users and business objects when information about the currently logged-in user is requested and some business objects which can reference any user as creator.

Let me show you a shortened version of my the implementation I ended up with:

@Controller
class UserBatchQueryAndFederatedController(
    registry: BatchLoaderRegistry,
    private val userRepository: UserRepository,
    // ...
) {
    init {
        registry.forTypePair(UserId::class.java, UserGraphQL::class.java).registerMappedBatchLoader { keys, _ ->
            // ...
            val users = userRepository.findAllByIds(keys)
            Mono.just(
                users.associate { it.userId to it.toGraphQL() }.toMap()
            )
        }
    }

    @QueryMapping
    fun user(@Argument id: String, dataLoader: DataLoader<UserId, UserGraphQL>): CompletableFuture<UserGraphQL?> =
        dataLoader.load(UserId(id))

    @EntityMapping
    fun user(
        @Argument idList: List<String>,
        env: DataFetchingEnvironment
    ): CompletableFuture<List<UserGraphQL>> {
        val dataLoader = env.getDataLoader<UserId, UserGraphQL>(UserGraphQL::class.java.name) ?: error("No dataloader for User")
        return dataLoader.loadMany(idList.map { UserId(it) })
    }
}

As you can see QueryMapping and EntityMapping seem very similar, but for EntityMapping I have to manually unravel the automatically created data loader since injecting it would fail with a runtime exception.

It occurs to me that support for DataLoader could be as simple as adding the DataLoaderMethodArgumentResolver to the list of argument resolvers. At the moment it's not very convenient to do that, but it is possible by overriding FederationSchemaFactory#initArgumentResolvers. Let us know, if you do try it.

That sounds great, but I couldn't get it to work quickly. FederationSchemaFactory is final, so it's not straightforward to override the method in application code.