dataloader_with_filtering.groovy

package graphql.execution.instrumentation.dataloader

import graphql.ExecutionInput
import graphql.TestUtil
import graphql.schema.DataFetcher
import org.dataloader.BatchLoader
import org.dataloader.DataLoader
import org.dataloader.DataLoaderOptions
import org.dataloader.DataLoaderRegistry
import spock.lang.Specification

import java.util.concurrent.CompletableFuture
import java.util.concurrent.CompletionStage

import static graphql.schema.idl.RuntimeWiring.newRuntimeWiring
import static graphql.schema.idl.TypeRuntimeWiring.newTypeWiring

class DataLoaderFilteringTest extends Specification {

    class Blog {
        String id
        String name
        String purpose

        @Override
        String toString() {
            return id + name;
        }
    }

    class Article {
        String id
        String blogId
        String title

        @Override
        String toString() {
            return id + ":" + blogId + ":" + title;
        }
    }

    class ArticleSearchKey {
        String blogId
        String searchText
    }


    def blogData = [
            new Blog(id: 1, name: "Java", purpose: "Java stuff"),
            new Blog(id: 2, name: "Reactive", purpose: "Rx stuff")
    ]

    def articleData = [
            new Article(id: 101, blogId: 1, title: "Changes in the LTR policy of the JDK"),
            new Article(id: 102, blogId: 1, title: "CPU impact of the JVM Memory Model"),

            new Article(id: 200, blogId: 2, title: "So you want to use less CPU"),
            new Article(id: 201, blogId: 2, title: "Why reactive systems use less CPU")
    ]

    CompletionStage<List<List<Article>>> getArticles(List<ArticleSearchKey> searchKeys) {
        //
        // this could be turned into SQL where its something like
        //
        // select from article where blog in (:theSourceBlogs) and title like 'foo'
        //
        // by de-duping the searchText into a set say
        //
        def blogArticles = []
        for (ArticleSearchKey searchKey : searchKeys) {
            def list = []
            for (Article article : articleData) {
                if (searchKey.blogId == article.blogId) {
                    if (article.title.toLowerCase().contains(searchKey.searchText.toLowerCase())) {
                        list.add(article)
                    }
                }
            }
            blogArticles.add(list)
        }
        return CompletableFuture.completedFuture(blogArticles)
    }


    def batchLoaderWithContext = new BatchLoader<ArticleSearchKey, List<Article>>() {

        @Override
        CompletionStage<List<List<Article>>> load(List<ArticleSearchKey> keys) {
            return getArticles(keys)
        }
    }


    DataFetcher articleFetcher = { env ->
        Blog blog = env.getSource()
        def whereTitleContains = env.getArgument("whereTitleContains")
        def searchKey = new ArticleSearchKey(blogId: blog.id, searchText: whereTitleContains)
        return env.getDataLoader("articleLoader").load(searchKey)
    }

    DataFetcher blogFetcher = { env ->
        return blogData
    }

    def spec = '''
        type Query {
            blogs : [Blog]
        }
        
        type Blog {
            id : ID
            name : String
            purpose : String
            articles(whereTitleContains : String) : [Article]
        }
        
        type Article {
            id : ID
            blogId : ID
            title :String
        }
    '''

    def runtimeWiring = newRuntimeWiring()
            .type(newTypeWiring("Blog").dataFetcher("articles", articleFetcher))
            .type(newTypeWiring("Query").dataFetcher("blogs", blogFetcher))
            .build()

    def graphQL = TestUtil.graphQL(spec, runtimeWiring).build()


    def "dataloader with filtering in the second level"() {

        // disable caching to show it always gets the full results
        DataLoaderOptions options = DataLoaderOptions.newOptions().setCachingEnabled(false)

        DataLoader<ArticleSearchKey, List<Article>> articleLoader = DataLoader.newDataLoader(batchLoaderWithContext, options)
        DataLoaderRegistry registry = new DataLoaderRegistry()
        registry.register("articleLoader", articleLoader)


        when:
        def executionResult = graphQL.execute(ExecutionInput.newExecutionInput()
                .dataLoaderRegistry(registry)
                .query('''
        query {
          blogs {
            id
            name
            articlesWithCPUInTitle: articles(whereTitleContains: "CPU") {
                blogId
                title
            }
            articlesWithReactiveInTitle: articles(whereTitleContains: "reactive") {
                blogId
                title
            }
          }
        }
            '''))
        then:
        executionResult.errors.isEmpty()
        println ""
        println executionResult.data
        println ""
    }

}