From 60502aea5dff2f44128d74a6c751e87abc779005 Mon Sep 17 00:00:00 2001 From: Joe Barnett Date: Mon, 1 Aug 2022 16:27:36 -0700 Subject: [PATCH] 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. --- .../build.gradle.kts | 1 + .../graphql/dataloader/KotlinDataLoader.kt | 8 +++- .../KotlinDataLoaderOptionsFactory.kt | 7 ++++ .../KotlinDataLoaderRegistryFactory.kt | 5 ++- .../KotlinDataLoaderRegistryFactoryTest.kt | 38 ++++++++++++++++++- .../KotlinDataLoaderRegistryTest.kt | 5 ++- .../server/execution/GraphQLRequestHandler.kt | 6 ++- .../extensions/RequestExtensionsKtTest.kt | 3 +- .../SpringGraphQLSubscriptionHandler.kt | 6 ++- .../server/spring/SchemaConfigurationTest.kt | 4 +- 10 files changed, 70 insertions(+), 13 deletions(-) create mode 100644 executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderOptionsFactory.kt diff --git a/executions/graphql-kotlin-dataloader/build.gradle.kts b/executions/graphql-kotlin-dataloader/build.gradle.kts index 0dcad8938b..794f98e56e 100644 --- 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/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoader.kt b/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoader.kt index bfa9a6e1f5..08e42d75f1 100644 --- a/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoader.kt +++ b/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoader.kt @@ -17,6 +17,7 @@ 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 @@ -24,5 +25,10 @@ import org.dataloader.DataLoader */ interface KotlinDataLoader { val dataLoaderName: String - fun getDataLoader(): DataLoader + fun getDataLoader(options: DataLoaderOptions): DataLoader = getDataLoader() + @Deprecated( + "Should use getDataLoader(options) instead", + replaceWith = ReplaceWith("getDataLoader(DataLoaderOptions.newOptions)") + ) + fun getDataLoader(): DataLoader = TODO("${this::class} needs to implement getDataLoader") } diff --git a/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderOptionsFactory.kt b/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderOptionsFactory.kt new file mode 100644 index 0000000000..1fe9009398 --- /dev/null +++ b/executions/graphql-kotlin-dataloader/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/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactory.kt b/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactory.kt index da89e1d2fb..c3b6318f7b 100644 --- a/executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactory.kt +++ b/executions/graphql-kotlin-dataloader/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/executions/graphql-kotlin-dataloader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactoryTest.kt b/executions/graphql-kotlin-dataloader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactoryTest.kt index 9baa88b93a..090dcbebc0 100644 --- a/executions/graphql-kotlin-dataloader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactoryTest.kt +++ b/executions/graphql-kotlin-dataloader/src/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 = object : KotlinDataLoader { + override val dataLoaderName = "Unimplemented" + } + assertFailsWith(NotImplementedError::class) { + KotlinDataLoaderRegistryFactory(listOf(mockLoader)).generate() + } + } + + @Test + fun `generate registry with context in options`() = runBlocking { + val mockLoader = object : KotlinDataLoader { + override val dataLoaderName = "withGraphQLContext" + override fun getDataLoader(options: DataLoaderOptions): DataLoader { + return DataLoaderFactory.newDataLoader({ keys, environment -> + keys.map { (environment.getContext() as GraphQLContext).get() }.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(mockLoader.dataLoaderName).load("123") + registry.dispatchAll() + assertEquals(result.await(), "blah") + } } diff --git a/executions/graphql-kotlin-dataloader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryTest.kt b/executions/graphql-kotlin-dataloader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryTest.kt index 03bb1c44e4..b5ba8bab4f 100644 --- a/executions/graphql-kotlin-dataloader/src/test/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryTest.kt +++ b/executions/graphql-kotlin-dataloader/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 = object : KotlinDataLoader { override val dataLoaderName: String = "ToUppercaseDataLoader" - override fun getDataLoader(): DataLoader = DataLoaderFactory.newDataLoader { keys -> + override fun getDataLoader(options: DataLoaderOptions): DataLoader = DataLoaderFactory.newDataLoader { keys -> keys.toFlux().map(String::uppercase).collectList().delayElement(Duration.ofMillis(300)).toFuture() } } val stringToLowerCaseDataLoader: KotlinDataLoader = object : KotlinDataLoader { override val dataLoaderName: String = "ToLowercaseDataLoader" - override fun getDataLoader(): DataLoader = DataLoaderFactory.newDataLoader { keys -> + override fun getDataLoader(options: DataLoaderOptions): DataLoader = DataLoaderFactory.newDataLoader { keys -> keys.toFlux().map(String::lowercase).collectList().delayElement(Duration.ofMillis(300)).toFuture() } } diff --git a/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandler.kt b/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandler.kt index dc46546e14..f46c4cefcb 100644 --- a/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandler.kt +++ b/servers/graphql-kotlin-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? = @@ -71,7 +73,7 @@ open class GraphQLRequestHandler( context: GraphQLContext? = null, graphQLContext: Map<*, Any> = emptyMap() ): 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/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt b/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt index f01138cb45..9402e7698f 100644 --- a/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt +++ b/servers/graphql-kotlin-server/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 { override val dataLoaderName: String = "abc" - override fun getDataLoader(): DataLoader = mockk() + override fun getDataLoader(options: DataLoaderOptions): DataLoader = mockk() } ).generate() diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt index 57747d8a5d..e207bea95e 100644 --- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt +++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/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,7 +36,8 @@ 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( @@ -43,7 +45,7 @@ open class SpringGraphQLSubscriptionHandler( context: GraphQLContext?, graphQLContext: Map<*, Any> = emptyMap() ): Flow> { - 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/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/SchemaConfigurationTest.kt b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/SchemaConfigurationTest.kt index 88639af5d3..9bd2063000 100644 --- a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/SchemaConfigurationTest.kt +++ b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/SchemaConfigurationTest.kt @@ -16,6 +16,8 @@ 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 @@ -23,8 +25,6 @@ import com.expediagroup.graphql.generator.execution.KotlinDataFetcherFactoryProv 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