Skip to content

Commit 01710bb

Browse files
paulbakkerpaulbakker
authored
Merge pull request #2131 from Netflix/fix/graphqlcontext
Move initialization of dataloader registry to the callback where we have GraphQLContext constructed
2 parents 0317b08 + 4e8472c commit 01710bb

File tree

6 files changed

+35
-40
lines changed

6 files changed

+35
-40
lines changed

‎graphql-dgs-example-shared/src/main/java/com/netflix/graphql/dgs/example/shared/datafetcher/HelloDataFetcher.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,14 @@ public CompletableFuture<String> withDataLoaderGraphQLContext(DataFetchingEnviro
106106
return exampleLoaderWithContext.load(CONTRIBUTOR_ENABLED_CONTEXT_KEY);
107107
}
108108

109+
@DgsData(parentType = "Query", field = "withDataLoaderGraphQLContextWithFromDfe")
110+
@DgsEnableDataFetcherInstrumentation
111+
public CompletableFuture<String> withDataLoaderGraphQLContextWithFromDfe(DataFetchingEnvironment dfe) {
112+
dfe.getGraphQlContext().put(CONTRIBUTOR_ENABLED_CONTEXT_KEY, "override");
113+
DataLoader<String, String> exampleLoaderWithContext = dfe.getDataLoader("exampleLoaderWithGraphQLContext");
114+
return exampleLoaderWithContext.load(CONTRIBUTOR_ENABLED_CONTEXT_KEY);
115+
}
116+
109117
@DgsData(parentType = "Query", field = "withGraphqlException")
110118
public String withGraphqlException() {
111119
throw new GraphQLException("that's not going to work!");

‎graphql-dgs-example-shared/src/main/resources/schema/schema.graphqls‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ type Query {
44
withContext: String
55
withDataLoaderContext: String
66
withDataLoaderGraphQLContext: String
7+
withDataLoaderGraphQLContextWithFromDfe: String
78
movies: [Movie]
89
messageFromBatchLoader: String
910
messageFromBatchLoaderWithGreetings: String

‎graphql-dgs-spring-graphql-example-java/src/test/java/com/netflix/graphql/dgs/example/datafetcher/GraphQLContextContributorTest.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,4 +61,13 @@ void withDataloaderGraphQLContext() {
6161
String contributorEnabled = queryExecutor.executeAndExtractJsonPath("{ withDataLoaderGraphQLContext }", "data.withDataLoaderGraphQLContext", servletWebRequest);
6262
assertThat(contributorEnabled).isEqualTo("true");
6363
}
64+
65+
@Test
66+
void withDataloaderGraphQLContextOverride() {
67+
final MockHttpServletRequest mockServletRequest = new MockHttpServletRequest();
68+
mockServletRequest.addHeader(CONTEXT_CONTRIBUTOR_HEADER_NAME, CONTEXT_CONTRIBUTOR_HEADER_VALUE);
69+
ServletWebRequest servletWebRequest = new ServletWebRequest(mockServletRequest);
70+
String contributorEnabled = queryExecutor.executeAndExtractJsonPath("{ withDataLoaderGraphQLContextWithFromDfe }", "data.withDataLoaderGraphQLContextWithFromDfe", servletWebRequest);
71+
assertThat(contributorEnabled).isEqualTo("override");
72+
}
6473
}

‎graphql-dgs-spring-graphql/src/main/kotlin/com/netflix/graphql/dgs/springgraphql/SpringGraphQLDgsQueryExecutor.kt‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ import com.netflix.graphql.dgs.internal.DgsDataLoaderProvider
3030
import com.netflix.graphql.dgs.internal.DgsQueryExecutorRequestCustomizer
3131
import com.netflix.graphql.dgs.internal.DgsWebMvcRequestData
3232
import graphql.ExecutionResult
33-
import graphql.GraphQLContext
3433
import org.springframework.graphql.ExecutionGraphQlService
3534
import org.springframework.graphql.support.DefaultExecutionGraphQlRequest
3635
import org.springframework.http.HttpHeaders
@@ -65,20 +64,9 @@ class SpringGraphQLDgsQueryExecutor(
6564

6665
val httpRequest = requestCustomizer.apply(webRequest ?: RequestContextHolder.getRequestAttributes() as? WebRequest, headers)
6766
val dgsContext = dgsContextBuilder.build(DgsWebMvcRequestData(request.extensions, headers, httpRequest))
68-
val dataLoaderRegistry =
69-
dgsDataLoaderProvider.buildRegistryWithContextSupplier {
70-
val graphQLContext = request.toExecutionInput().graphQLContext
71-
if (graphQLContextContributors.isNotEmpty()) {
72-
val requestData = dgsContext.requestData
73-
val builderForContributors = GraphQLContext.newContext()
74-
graphQLContextContributors.forEach { it.contribute(builderForContributors, extensions, requestData) }
75-
graphQLContext.putAll(builderForContributors)
76-
}
77-
78-
graphQLContext
79-
}
80-
81-
request.configureExecutionInput { _, builder ->
67+
68+
request.configureExecutionInput { e, builder ->
69+
val dataLoaderRegistry = dgsDataLoaderProvider.buildRegistryWithContextSupplier { e.graphQLContext }
8270
builder
8371
.context(dgsContext)
8472
.graphQLContext(dgsContext)

‎graphql-dgs-spring-graphql/src/main/kotlin/com/netflix/graphql/dgs/springgraphql/webflux/DgsWebFluxGraphQLInterceptor.kt‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,13 @@ package com.netflix.graphql.dgs.springgraphql.webflux
1919
import com.netflix.graphql.dgs.internal.DgsDataLoaderProvider
2020
import com.netflix.graphql.dgs.reactive.internal.DefaultDgsReactiveGraphQLContextBuilder
2121
import com.netflix.graphql.dgs.reactive.internal.DgsReactiveRequestData
22-
import graphql.GraphQLContext
22+
import org.dataloader.DataLoaderRegistry
2323
import org.springframework.graphql.server.WebGraphQlInterceptor
2424
import org.springframework.graphql.server.WebGraphQlRequest
2525
import org.springframework.graphql.server.WebGraphQlResponse
2626
import org.springframework.web.filter.reactive.ServerWebExchangeContextFilter
2727
import org.springframework.web.reactive.function.server.ServerRequest
2828
import reactor.core.publisher.Mono
29-
import java.util.concurrent.CompletableFuture
3029

3130
class DgsWebFluxGraphQLInterceptor(
3231
private val dgsDataLoaderProvider: DgsDataLoaderProvider,
@@ -48,21 +47,19 @@ class DgsWebFluxGraphQLInterceptor(
4847
),
4948
)
5049
}.flatMap { dgsContext ->
51-
val graphQLContextFuture = CompletableFuture<GraphQLContext>()
52-
val dataLoaderRegistry =
53-
dgsDataLoaderProvider.buildRegistryWithContextSupplier { graphQLContextFuture.get() }
54-
55-
request.configureExecutionInput { _, builder ->
50+
var dataLoaderRegistry: DataLoaderRegistry? = null
51+
request.configureExecutionInput { e, builder ->
52+
dataLoaderRegistry = dgsDataLoaderProvider.buildRegistryWithContextSupplier { e.graphQLContext }
5653
builder
5754
.context(dgsContext)
5855
.graphQLContext(dgsContext)
5956
.dataLoaderRegistry(dataLoaderRegistry)
6057
.build()
6158
}
62-
graphQLContextFuture.complete(request.toExecutionInput().graphQLContext)
59+
6360
chain.next(request).doFinally {
6461
if (dataLoaderRegistry is AutoCloseable) {
65-
dataLoaderRegistry.close()
62+
(dataLoaderRegistry as AutoCloseable).close()
6663
}
6764
}
6865
}

‎graphql-dgs-spring-graphql/src/main/kotlin/com/netflix/graphql/dgs/springgraphql/webmvc/DgsWebMvcGraphQLInterceptor.kt‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import com.netflix.graphql.dgs.internal.DefaultDgsGraphQLContextBuilder
2121
import com.netflix.graphql.dgs.internal.DgsDataLoaderProvider
2222
import com.netflix.graphql.dgs.internal.DgsWebMvcRequestData
2323
import com.netflix.graphql.dgs.springgraphql.autoconfig.DgsSpringGraphQLConfigurationProperties
24-
import graphql.GraphQLContext
24+
import org.dataloader.DataLoaderRegistry
2525
import org.springframework.graphql.server.WebGraphQlInterceptor
2626
import org.springframework.graphql.server.WebGraphQlRequest
2727
import org.springframework.graphql.server.WebGraphQlResponse
@@ -56,21 +56,13 @@ class DgsWebMvcGraphQLInterceptor(
5656
} else {
5757
dgsContextBuilder.build(DgsWebMvcRequestData(request.extensions, request.headers))
5858
}
59-
val dataLoaderRegistry =
60-
dgsDataLoaderProvider.buildRegistryWithContextSupplier {
61-
val graphQLContext = request.toExecutionInput().graphQLContext
62-
if (graphQLContextContributors.isNotEmpty()) {
63-
val extensions = request.extensions
64-
val requestData = dgsContext.requestData
65-
val builderForContributors = GraphQLContext.newContext()
66-
graphQLContextContributors.forEach { it.contribute(builderForContributors, extensions, requestData) }
67-
graphQLContext.putAll(builderForContributors)
68-
}
6959

70-
graphQLContext
71-
}
60+
var dataLoaderRegistry: DataLoaderRegistry? = null
61+
request.configureExecutionInput { e, builder ->
62+
63+
dataLoaderRegistry =
64+
dgsDataLoaderProvider.buildRegistryWithContextSupplier { e.graphQLContext }
7265

73-
request.configureExecutionInput { _, builder ->
7466
builder
7567
.context(dgsContext)
7668
.graphQLContext(dgsContext)
@@ -81,14 +73,14 @@ class DgsWebMvcGraphQLInterceptor(
8173
return if (dgsSpringConfigurationProperties.webmvc.asyncdispatch.enabled) {
8274
chain.next(request).doFinally {
8375
if (dataLoaderRegistry is AutoCloseable) {
84-
dataLoaderRegistry.close()
76+
(dataLoaderRegistry as AutoCloseable).close()
8577
}
8678
}
8779
} else {
8880
@Suppress("BlockingMethodInNonBlockingContext")
8981
val response = chain.next(request).block()!!
9082
if (dataLoaderRegistry is AutoCloseable) {
91-
dataLoaderRegistry.close()
83+
(dataLoaderRegistry as AutoCloseable).close()
9284
}
9385
return Mono.just(response)
9486
}

0 commit comments

Comments
 (0)