Add DataLoaderOptions to dataloader creation

To allow setting eg. the GraphQLContext as the batch context in a dataloader, pass it through the KotlinDataLoaderRegistryFactory.generate() call.
ExpediaGroup · Aug 2, 2022 · 60502ae · 60502ae
1 parent 7711942
commit 60502ae
Show file tree

Hide file tree

Showing 10 changed files with 70 additions and 13 deletions.
diff --git a/executions/graphql-kotlin-dataloader/build.gradle.kts b/executions/graphql-kotlin-dataloader/build.gradle.kts
@@ -7,6 +7,7 @@ val reactorExtensionsVersion: String by project
 
 dependencies {
     api("com.graphql-java:java-dataloader:$graphQLJavaDataLoaderVersion")
+    testImplementation(project(path = ":graphql-kotlin-schema-generator"))
     testImplementation("io.projectreactor.kotlin:reactor-kotlin-extensions:$reactorExtensionsVersion")
     testImplementation("io.projectreactor:reactor-core:$reactorVersion")
     testImplementation("org.junit.jupiter:junit-jupiter-api:$junitVersion")

diff --git a/...kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoader.kt b/...kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoader.kt
@@ -17,12 +17,18 @@
 package com.expediagroup.graphql.dataloader
 
 import org.dataloader.DataLoader
+import org.dataloader.DataLoaderOptions
 
 /**
  * Wrapper around the [DataLoader] class so we can have common logic around registering the loaders
  * by return type and loading values in the data fetchers.
  */
 interface KotlinDataLoader<K, V> {
     val dataLoaderName: String
-    fun getDataLoader(): DataLoader<K, V>
+    fun getDataLoader(options: DataLoaderOptions): DataLoader<K, V> = getDataLoader()
+    @Deprecated(
+        "Should use getDataLoader(options) instead",
+        replaceWith = ReplaceWith("getDataLoader(DataLoaderOptions.newOptions)")
+    )
+    fun getDataLoader(): DataLoader<K, V> = TODO("${this::class} needs to implement getDataLoader")
 }
diff --git a/...der/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderOptionsFactory.kt b/...der/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderOptionsFactory.kt
@@ -0,0 +1,7 @@
+package com.expediagroup.graphql.dataloader
+
+import org.dataloader.DataLoaderOptions
+
+open class KotlinDataLoaderOptionsFactory {
+    open fun generate(contextMap: Map<*, Any>): DataLoaderOptions = DataLoaderOptions.newOptions()
+}
diff --git a/...er/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactory.kt b/...er/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactory.kt
@@ -16,6 +16,7 @@
 
 package com.expediagroup.graphql.dataloader
 
+import org.dataloader.DataLoaderOptions
 import org.dataloader.DataLoaderRegistry
 
 /**
@@ -30,12 +31,12 @@ class KotlinDataLoaderRegistryFactory(
     /**
      * Generate [KotlinDataLoaderRegistry] to be used for GraphQL request execution.
      */
-    fun generate(): KotlinDataLoaderRegistry {
+    fun generate(options: DataLoaderOptions = DataLoaderOptions.newOptions()): KotlinDataLoaderRegistry {
         val registry = DataLoaderRegistry()
         dataLoaders.forEach { dataLoader ->
             registry.register(
                 dataLoader.dataLoaderName,
-                dataLoader.getDataLoader()
+                dataLoader.getDataLoader(options)
             )
         }
         return KotlinDataLoaderRegistry(registry)

diff --git a/...rc/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactoryTest.kt b/...rc/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactoryTest.kt
@@ -15,11 +15,18 @@
  */
 
 package com.expediagroup.graphql.dataloader
-
+import com.expediagroup.graphql.generator.extensions.get
+import graphql.GraphQLContext
 import io.mockk.mockk
+import kotlinx.coroutines.future.await
+import kotlinx.coroutines.runBlocking
 import org.dataloader.DataLoader
+import org.dataloader.DataLoaderFactory
+import org.dataloader.DataLoaderOptions
 import org.junit.jupiter.api.Test
+import reactor.kotlin.core.publisher.toMono
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 import kotlin.test.assertTrue
 
 class KotlinDataLoaderRegistryFactoryTest {
@@ -45,4 +52,33 @@ class KotlinDataLoaderRegistryFactoryTest {
         val registry = KotlinDataLoaderRegistryFactory(listOf(mockLoader)).generate()
         assertEquals(1, registry.dataLoaders.size)
     }
+
+    @Test
+    fun `generate registry with minimal compilable loader throws TODO`() {
+        val mockLoader: KotlinDataLoader<String, String> = object : KotlinDataLoader<String, String> {
+            override val dataLoaderName = "Unimplemented"
+        }
+        assertFailsWith(NotImplementedError::class) {
+            KotlinDataLoaderRegistryFactory(listOf(mockLoader)).generate()
+        }
+    }
+
+    @Test
+    fun `generate registry with context in options`() = runBlocking {
+        val mockLoader = object : KotlinDataLoader<String, String> {
+            override val dataLoaderName = "withGraphQLContext"
+            override fun getDataLoader(options: DataLoaderOptions): DataLoader<String, String> {
+                return DataLoaderFactory.newDataLoader({ keys, environment ->
+                    keys.map { (environment.getContext() as GraphQLContext).get<String>() }.toMono().toFuture()
+                }, options)
+            }
+        }
+        val options = DataLoaderOptions.newOptions().setBatchLoaderContextProvider {
+            GraphQLContext.of(mapOf(String::class to "blah"))
+        }
+        val registry = KotlinDataLoaderRegistryFactory(mockLoader).generate(options)
+        val result = registry.getDataLoader<String, String>(mockLoader.dataLoaderName).load("123")
+        registry.dispatchAll()
+        assertEquals(result.await(), "blah")
+    }
 }
diff --git a/...oader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryTest.kt b/...oader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryTest.kt
@@ -18,6 +18,7 @@ package com.expediagroup.graphql.dataloader
 
 import org.dataloader.DataLoader
 import org.dataloader.DataLoaderFactory
+import org.dataloader.DataLoaderOptions
 import org.junit.jupiter.api.Test
 import reactor.kotlin.core.publisher.toFlux
 import java.time.Duration
@@ -30,14 +31,14 @@ class KotlinDataLoaderRegistryTest {
     fun `Decorator will keep track of DataLoaders futures`() {
         val stringToUpperCaseDataLoader: KotlinDataLoader<String, String> = object : KotlinDataLoader<String, String> {
             override val dataLoaderName: String = "ToUppercaseDataLoader"
-            override fun getDataLoader(): DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys ->
+            override fun getDataLoader(options: DataLoaderOptions): DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys ->
                 keys.toFlux().map(String::uppercase).collectList().delayElement(Duration.ofMillis(300)).toFuture()
             }
         }
 
         val stringToLowerCaseDataLoader: KotlinDataLoader<String, String> = object : KotlinDataLoader<String, String> {
             override val dataLoaderName: String = "ToLowercaseDataLoader"
-            override fun getDataLoader(): DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys ->
+            override fun getDataLoader(options: DataLoaderOptions): DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys ->
                 keys.toFlux().map(String::lowercase).collectList().delayElement(Duration.ofMillis(300)).toFuture()
             }
         }

diff --git a/...server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandler.kt b/...server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandler.kt
@@ -16,6 +16,7 @@
 
 package com.expediagroup.graphql.server.execution
 
+import com.expediagroup.graphql.dataloader.KotlinDataLoaderOptionsFactory
 import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistry
 import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
 import com.expediagroup.graphql.dataloader.instrumentation.level.DataLoaderLevelDispatchedInstrumentation
@@ -45,7 +46,8 @@ import kotlinx.coroutines.supervisorScope
 
 open class GraphQLRequestHandler(
     private val graphQL: GraphQL,
-    private val dataLoaderRegistryFactory: KotlinDataLoaderRegistryFactory? = null
+    private val dataLoaderRegistryFactory: KotlinDataLoaderRegistryFactory? = null,
+    private val dataLoaderOptions: KotlinDataLoaderOptionsFactory = KotlinDataLoaderOptionsFactory()
 ) {
 
     private val batchDataLoaderInstrumentationType: Class<Instrumentation>? =
@@ -71,7 +73,7 @@ open class GraphQLRequestHandler(
         context: GraphQLContext? = null,
         graphQLContext: Map<*, Any> = emptyMap<Any, Any>()
     ): GraphQLServerResponse {
-        val dataLoaderRegistry = dataLoaderRegistryFactory?.generate()
+        val dataLoaderRegistry = dataLoaderRegistryFactory?.generate(dataLoaderOptions.generate(graphQLContext))
         return when (graphQLRequest) {
             is GraphQLRequest -> {
                 val batchGraphQLContext = graphQLContext + getBatchContext(1, dataLoaderRegistry)

diff --git a/...ver/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt b/...ver/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt
@@ -21,6 +21,7 @@ import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
 import com.expediagroup.graphql.server.types.GraphQLRequest
 import io.mockk.mockk
 import org.dataloader.DataLoader
+import org.dataloader.DataLoaderOptions
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
@@ -64,7 +65,7 @@ class RequestExtensionsKtTest {
         val dataLoaderRegistry = KotlinDataLoaderRegistryFactory(
             object : KotlinDataLoader<String, String> {
                 override val dataLoaderName: String = "abc"
-                override fun getDataLoader(): DataLoader<String, String> = mockk()
+                override fun getDataLoader(options: DataLoaderOptions): DataLoader<String, String> = mockk()
             }
         ).generate()
 

diff --git a/.../com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt b/.../com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt
@@ -16,6 +16,7 @@
 
 package com.expediagroup.graphql.server.spring.subscriptions
 
+import com.expediagroup.graphql.dataloader.KotlinDataLoaderOptionsFactory
 import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.extensions.toExecutionInput
@@ -35,15 +36,16 @@ import kotlinx.coroutines.flow.map
  */
 open class SpringGraphQLSubscriptionHandler(
     private val graphQL: GraphQL,
-    private val dataLoaderRegistryFactory: KotlinDataLoaderRegistryFactory? = null
+    private val dataLoaderRegistryFactory: KotlinDataLoaderRegistryFactory? = null,
+    private val dataLoaderOptions: KotlinDataLoaderOptionsFactory = KotlinDataLoaderOptionsFactory()
 ) {
 
     fun executeSubscription(
         graphQLRequest: GraphQLRequest,
         context: GraphQLContext?,
         graphQLContext: Map<*, Any> = emptyMap<Any, Any>()
     ): Flow<GraphQLResponse<*>> {
-        val dataLoaderRegistry = dataLoaderRegistryFactory?.generate()
+        val dataLoaderRegistry = dataLoaderRegistryFactory?.generate(dataLoaderOptions.generate(graphQLContext))
         val input = graphQLRequest.toExecutionInput(dataLoaderRegistry, context, graphQLContext)
 
         return graphQL.execute(input)

diff --git a/...-server/src/test/kotlin/com/expediagroup/graphql/server/spring/SchemaConfigurationTest.kt b/...-server/src/test/kotlin/com/expediagroup/graphql/server/spring/SchemaConfigurationTest.kt
@@ -16,15 +16,15 @@
 
 package com.expediagroup.graphql.server.spring
 
+import com.expediagroup.graphql.dataloader.KotlinDataLoader
+import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
 import com.expediagroup.graphql.generator.SchemaGeneratorConfig
 import com.expediagroup.graphql.generator.TopLevelObject
 import com.expediagroup.graphql.generator.annotations.GraphQLDescription
 import com.expediagroup.graphql.generator.execution.KotlinDataFetcherFactoryProvider
 import com.expediagroup.graphql.generator.toSchema
 import com.expediagroup.graphql.server.execution.GraphQLContextFactory
 import com.expediagroup.graphql.server.execution.GraphQLRequestHandler
-import com.expediagroup.graphql.dataloader.KotlinDataLoader
-import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
 import com.expediagroup.graphql.server.extensions.getValueFromDataLoader
 import com.expediagroup.graphql.server.operations.Query
 import com.expediagroup.graphql.server.spring.execution.SpringGraphQLContextFactory