feat: pass graphQLContext to KotlinDataLoaderRegistryFactory

examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/dataloaders/BookDataLoader.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,20 +19,22 @@ package com.expediagroup.graphql.examples.server.ktor.schema.dataloaders
 import com.expediagroup.graphql.examples.server.ktor.schema.models.Book
 import com.expediagroup.graphql.dataloader.KotlinDataLoader
 import kotlinx.coroutines.runBlocking
+import graphql.GraphQLContext
 import org.dataloader.DataLoaderFactory
 import java.util.concurrent.CompletableFuture
 
 val BookDataLoader = object : KotlinDataLoader<List<Int>, List<Book>> {
     override val dataLoaderName = "BATCH_BOOK_LOADER"
-    override fun getDataLoader() = DataLoaderFactory.newDataLoader<List<Int>, List<Book>> { ids ->
-        CompletableFuture.supplyAsync {
-            val allBooks = runBlocking { Book.search(ids.flatten()).toMutableList() }
-            // produce lists of results from returned books
-            ids.fold(mutableListOf()) { acc: MutableList<List<Book>>, idSet ->
-                val matchingBooks = allBooks.filter { idSet.contains(it.id) }
-                acc.add(matchingBooks)
-                acc
+    override fun getDataLoader(graphQLContext: GraphQLContext) =
+        DataLoaderFactory.newDataLoader<List<Int>, List<Book>> { ids ->
+            CompletableFuture.supplyAsync {
+                val allBooks = runBlocking { Book.search(ids.flatten()).toMutableList() }
+                // produce lists of results from returned books
+                ids.fold(mutableListOf()) { acc: MutableList<List<Book>>, idSet ->
+                    val matchingBooks = allBooks.filter { idSet.contains(it.id) }
+                    acc.add(matchingBooks)
+                    acc
+                }
             }
         }
-    }
 }

examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/dataloaders/CourseDataLoader.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,14 +19,16 @@ package com.expediagroup.graphql.examples.server.ktor.schema.dataloaders
 import com.expediagroup.graphql.examples.server.ktor.schema.models.Course
 import com.expediagroup.graphql.dataloader.KotlinDataLoader
 import kotlinx.coroutines.runBlocking
+import graphql.GraphQLContext
 import org.dataloader.DataLoaderFactory
 import java.util.concurrent.CompletableFuture
 
 val CourseDataLoader = object : KotlinDataLoader<Int, Course?> {
     override val dataLoaderName = "COURSE_LOADER"
-    override fun getDataLoader() = DataLoaderFactory.newDataLoader<Int, Course?> { ids ->
-        CompletableFuture.supplyAsync {
-            runBlocking { Course.search(ids).toMutableList() }
+    override fun getDataLoader(graphQLContext: GraphQLContext) =
+        DataLoaderFactory.newDataLoader<Int, Course?> { ids ->
+            CompletableFuture.supplyAsync {
+                runBlocking { Course.search(ids).toMutableList() }
+            }
         }
-    }
 }

examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/dataloaders/UniversityDataLoader.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,14 +19,16 @@ package com.expediagroup.graphql.examples.server.ktor.schema.dataloaders
 import com.expediagroup.graphql.examples.server.ktor.schema.models.University
 import com.expediagroup.graphql.dataloader.KotlinDataLoader
 import kotlinx.coroutines.runBlocking
+import graphql.GraphQLContext
 import org.dataloader.DataLoaderFactory
 import java.util.concurrent.CompletableFuture
 
 val UniversityDataLoader = object : KotlinDataLoader<Int, University?> {
     override val dataLoaderName = "UNIVERSITY_LOADER"
-    override fun getDataLoader() = DataLoaderFactory.newDataLoader<Int, University?> { ids ->
-        CompletableFuture.supplyAsync {
-            runBlocking { University.search(ids).toMutableList() }
+    override fun getDataLoader(graphQLContext: GraphQLContext) =
+        DataLoaderFactory.newDataLoader<Int, University?> { ids ->
+            CompletableFuture.supplyAsync {
+                runBlocking { University.search(ids).toMutableList() }
+            }
         }
-    }
 }

examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/dataloaders/CompanyDataLoader.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,6 +20,7 @@ import com.expediagroup.graphql.examples.server.spring.model.Company
 import com.expediagroup.graphql.dataloader.KotlinDataLoader
 import org.dataloader.DataLoaderFactory
 import org.springframework.stereotype.Component
+import graphql.GraphQLContext
 import java.util.concurrent.CompletableFuture
 
 @Component
@@ -29,7 +30,8 @@ class CompanyDataLoader(private val service: CompanyService) : KotlinDataLoader<
     }
 
     override val dataLoaderName = name
-    override fun getDataLoader() = DataLoaderFactory.newDataLoader<Int, Company> { ids ->
-        CompletableFuture.supplyAsync { service.getCompanies(ids) }
-    }
+    override fun getDataLoader(graphQLContext: GraphQLContext) =
+        DataLoaderFactory.newDataLoader<Int, Company> { ids ->
+            CompletableFuture.supplyAsync { service.getCompanies(ids) }
+        }
 }

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt

@@ -42,6 +42,7 @@ import graphql.schema.idl.RuntimeWiring
 import graphql.schema.idl.SchemaGenerator
 import graphql.schema.idl.SchemaParser
 import graphql.schema.idl.TypeRuntimeWiring
+import io.mockk.mockk
 import io.mockk.spyk
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
@@ -206,7 +207,7 @@ object AstronautGraphQL {
                 AstronautDataLoader(),
                 MissionDataLoader(), MissionsByAstronautDataLoader(),
                 PlanetsByMissionDataLoader()
-            ).generate()
+            ).generate(mockk())
         )
 
         val graphQLContext = mapOf(

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/ProductGraphQL.kt

@@ -35,6 +35,7 @@ import graphql.schema.idl.RuntimeWiring
 import graphql.schema.idl.SchemaGenerator
 import graphql.schema.idl.SchemaParser
 import graphql.schema.idl.TypeRuntimeWiring
+import io.mockk.mockk
 import io.mockk.spyk
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
@@ -114,7 +115,7 @@ object ProductGraphQL {
         val kotlinDataLoaderRegistry = spyk(
             KotlinDataLoaderRegistryFactory(
                 ProductDataLoader()
-            ).generate()
+            ).generate(mockk())
         )
 
         val graphQLContext = mapOf(

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt

@@ -23,6 +23,7 @@ import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Missio
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Planet
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.extensions.toListOfNullables
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.repository.AstronautRepository
+import graphql.GraphQLContext
 import graphql.schema.DataFetchingEnvironment
 import org.dataloader.DataLoader
 import org.dataloader.DataLoaderFactory
@@ -38,7 +39,7 @@ data class CreateAstronautServiceRequest(val name: String)
 
 class AstronautDataLoader : KotlinDataLoader<AstronautServiceRequest, Astronaut?> {
     override val dataLoaderName: String = "AstronautDataLoader"
-    override fun getDataLoader(): DataLoader<AstronautServiceRequest, Astronaut?> =
+    override fun getDataLoader(graphQLContext: GraphQLContext): DataLoader<AstronautServiceRequest, Astronaut?> =
         DataLoaderFactory.newDataLoader(
             { keys ->
                 AstronautRepository

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/MissionService.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,6 +20,7 @@ import com.expediagroup.graphql.dataloader.KotlinDataLoader
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Mission
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.extensions.toListOfNullables
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.repository.MissionRepository
+import graphql.GraphQLContext
 import graphql.schema.DataFetchingEnvironment
 import org.dataloader.DataLoader
 import org.dataloader.DataLoaderFactory
@@ -32,28 +33,30 @@ data class MissionServiceRequest(val id: Int, val astronautId: Int = -1)
 
 class MissionDataLoader : KotlinDataLoader<MissionServiceRequest, Mission?> {
     override val dataLoaderName: String = "MissionDataLoader"
-    override fun getDataLoader(): DataLoader<MissionServiceRequest, Mission?> = DataLoaderFactory.newDataLoader(
-        { keys ->
-            MissionRepository
-                .getMissions(keys.map(MissionServiceRequest::id))
-                .collectList()
-                .map(List<Optional<Mission>>::toListOfNullables)
-                .toFuture()
-        },
-        DataLoaderOptions.newOptions().setStatisticsCollector(::SimpleStatisticsCollector)
-    )
+    override fun getDataLoader(graphQLContext: GraphQLContext): DataLoader<MissionServiceRequest, Mission?> =
+        DataLoaderFactory.newDataLoader(
+            { keys ->
+                MissionRepository
+                    .getMissions(keys.map(MissionServiceRequest::id))
+                    .collectList()
+                    .map(List<Optional<Mission>>::toListOfNullables)
+                    .toFuture()
+            },
+            DataLoaderOptions.newOptions().setStatisticsCollector(::SimpleStatisticsCollector)
+        )
 }
 
 class MissionsByAstronautDataLoader : KotlinDataLoader<MissionServiceRequest, List<Mission>> {
     override val dataLoaderName: String = "MissionsByAstronautDataLoader"
-    override fun getDataLoader(): DataLoader<MissionServiceRequest, List<Mission>> = DataLoaderFactory.newDataLoader(
-        { keys ->
-            MissionRepository
-                .getMissionsByAstronautIds(keys.map(MissionServiceRequest::astronautId))
-                .collectList().toFuture()
-        },
-        DataLoaderOptions.newOptions().setStatisticsCollector(::SimpleStatisticsCollector)
-    )
+    override fun getDataLoader(graphQLContext: GraphQLContext): DataLoader<MissionServiceRequest, List<Mission>> =
+        DataLoaderFactory.newDataLoader(
+            { keys ->
+                MissionRepository
+                    .getMissionsByAstronautIds(keys.map(MissionServiceRequest::astronautId))
+                    .collectList().toFuture()
+            },
+            DataLoaderOptions.newOptions().setStatisticsCollector(::SimpleStatisticsCollector)
+        )
 }
 
 class MissionService {

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/PlanetService.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,6 +19,7 @@ package com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher
 import com.expediagroup.graphql.dataloader.KotlinDataLoader
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Planet
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.repository.PlanetRepository
+import graphql.GraphQLContext
 import graphql.schema.DataFetchingEnvironment
 import org.dataloader.DataLoader
 import org.dataloader.DataLoaderFactory
@@ -30,15 +31,16 @@ data class PlanetServiceRequest(val id: Int, val missionId: Int = -1)
 
 class PlanetsByMissionDataLoader : KotlinDataLoader<PlanetServiceRequest, List<Planet>> {
     override val dataLoaderName: String = "PlanetsByMissionDataLoader"
-    override fun getDataLoader(): DataLoader<PlanetServiceRequest, List<Planet>> = DataLoaderFactory.newDataLoader(
-        { keys ->
-            PlanetRepository
-                .getPlanetsByMissionIds(keys.map(PlanetServiceRequest::missionId))
-                .collectList()
-                .toFuture()
-        },
-        DataLoaderOptions.newOptions().setStatisticsCollector(::SimpleStatisticsCollector)
-    )
+    override fun getDataLoader(graphQLContext: GraphQLContext): DataLoader<PlanetServiceRequest, List<Planet>> =
+        DataLoaderFactory.newDataLoader(
+            { keys ->
+                PlanetRepository
+                    .getPlanetsByMissionIds(keys.map(PlanetServiceRequest::missionId))
+                    .collectList()
+                    .toFuture()
+            },
+            DataLoaderOptions.newOptions().setStatisticsCollector(::SimpleStatisticsCollector)
+        )
 }
 
 class PlanetService {