package graphql.execution.instrumentation.dataloader

import com.github.javafaker.Faker

import graphql.ExecutionInput

import graphql.GraphQL

import graphql.TestUtil

import graphql.execution.instrumentation.dataloader.models.Company

import graphql.execution.instrumentation.dataloader.models.Person

import graphql.execution.instrumentation.dataloader.models.Product

import graphql.schema.DataFetcher

import graphql.schema.DataFetchingEnvironment

import graphql.schema.idl.RuntimeWiring

import org.dataloader.BatchLoader

import org.dataloader.DataLoader

import org.dataloader.DataLoaderFactory

import org.dataloader.DataLoaderRegistry

import spock.lang.Specification

import java.util.concurrent.CompletableFuture

import java.util.concurrent.CompletionStage

import java.util.stream.Collectors

import static graphql.schema.idl.RuntimeWiring.newRuntimeWiring

class PeopleCompaniesAndProductsDataLoaderTest extends Specification {

int personBatchLoadInvocationCount = 0

int personBatchLoadKeyCount = 0

int companyBatchLoadInvocationCount = 0

int companyBatchLoadKeyCount = 0

void setup() {

personBatchLoadInvocationCount = 0

personBatchLoadKeyCount = 0

companyBatchLoadInvocationCount = 0

companyBatchLoadKeyCount = 0

}

Faker faker = new Faker()

String schema = """

BatchLoader<Integer, Person> personBatchLoader = new BatchLoader<Integer, Person>() {

@Override

CompletionStage<List<Person>> load(List<Integer> keys) {

personBatchLoadInvocationCount += 1

personBatchLoadKeyCount += keys.size()

CompletableFuture.supplyAsync({

keys.stream()

.map({ id -> new Person(id, faker.name().fullName(), id + 200) })

.collect(Collectors.toList())

})

}

}

BatchLoader<Integer, Company> companyBatchLoader = new BatchLoader<Integer, Company>() {

@Override

CompletionStage<List<Company>> load(List<Integer> keys) {

companyBatchLoadInvocationCount += 1

companyBatchLoadKeyCount += keys.size()

CompletableFuture.supplyAsync({

return keys.stream()

.map({ id -> new Company(id, faker.company().name()) })

.collect(Collectors.toList())

})

}

}

DataFetcher productsDF = new DataFetcher() {

@Override

Object get(DataFetchingEnvironment environment) {

List<Product> products = new ArrayList<>();

for (int id = 0; id < 5; id++) {

List<Integer> madeBy = [id * 10001, id * 10002, id * 10003, id * 10004, id * 10005]

products.add(new Product(id.toString(), faker.commerce().productName(), id + 200, madeBy))

}

return products

}

}

DataFetcher suppliedByDF = new DataFetcher() {

@Override

Object get(DataFetchingEnvironment environment) {

Product source = environment.getSource()

DataLoaderRegistry dlRegistry = environment.getGraphQlContext().get("registry")

DataLoader<Integer, Person> personDL = dlRegistry.getDataLoader("person")

return personDL.load(source.getSuppliedById())

}

}

DataFetcher madeByDF = new DataFetcher() {

@Override

Object get(DataFetchingEnvironment environment) {

Product source = environment.getSource()

DataLoaderRegistry dlRegistry = environment.getGraphQlContext().get("registry")

DataLoader<Integer, Person> personDL = dlRegistry.getDataLoader("person")

return personDL.loadMany(source.getMadeByIds())

}

}

DataFetcher companyDF = new DataFetcher() {

@Override

Object get(DataFetchingEnvironment environment) {

Person source = environment.getSource()

DataLoaderRegistry dlRegistry = environment.getGraphQlContext().get("registry")

DataLoader<Integer, Company> companyDL = dlRegistry.getDataLoader("company")

return companyDL.load(source.getCompanyId())

}

}

RuntimeWiring runtimeWiring = newRuntimeWiring()

.type("QueryType", { builder -> builder.dataFetcher("products", productsDF) })

.type("Product", { builder -> builder.dataFetcher("suppliedBy", suppliedByDF) })

.type("Product", { builder -> builder.dataFetcher("madeBy", madeByDF) })

.type("Person", { builder -> builder.dataFetcher("company", companyDF) })

.build()

def graphQLSchema = TestUtil.schema(schema, runtimeWiring)

def query = """

private DataLoaderRegistry buildRegistry() {

DataLoader<Integer, Person> personDataLoader = DataLoaderFactory.newDataLoader(personBatchLoader)

DataLoader<Integer, Company> companyDataLoader = DataLoaderFactory.newDataLoader(companyBatchLoader)

DataLoaderRegistry registry = new DataLoaderRegistry()

registry.register("person", personDataLoader)

registry.register("company", companyDataLoader)

return registry

}

def "ensure performant loading with field tracking"() {

DataLoaderRegistry registry = buildRegistry()

GraphQL graphQL = GraphQL

.newGraphQL(graphQLSchema)

.build()

when:

ExecutionInput executionInput = ExecutionInput.newExecutionInput()

.query(query)

.graphQLContext(["registry": registry])

.dataLoaderRegistry(registry)

.build()

def executionResult = graphQL.execute(executionInput)

then:

(executionResult.data["products"] as Map).size() == 5

personBatchLoadKeyCount == 26

personBatchLoadInvocationCount == 1

companyBatchLoadKeyCount == 26

companyBatchLoadInvocationCount == 1

}

}