Merge pull request #2067 from Netflix/jjaco-debug-2058

Fix GraphQLContext propagation to data loaders with spring-graphql starter

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2024 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.netflix.graphql.dgs.example.datafetcher;
+
+import com.netflix.graphql.dgs.DgsQueryExecutor;
+import com.netflix.graphql.dgs.example.shared.datafetcher.RequestHeadersDataFetcher;
+import com.netflix.graphql.dgs.example.shared.datafetcher.MovieDataFetcher;
+import com.netflix.graphql.dgs.test.EnableDgsTest;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.mock.web.MockHttpServletRequest;
+import org.springframework.web.context.request.ServletWebRequest;
+
+import static com.netflix.graphql.dgs.example.shared.context.ExampleGraphQLContextContributor.CONTEXT_CONTRIBUTOR_HEADER_NAME;
+import static com.netflix.graphql.dgs.example.shared.context.ExampleGraphQLContextContributor.CONTEXT_CONTRIBUTOR_HEADER_VALUE;
+import static org.assertj.core.api.Assertions.assertThat;
+
+@EnableDgsTest
+@TestAppTestSlice
+@SpringBootTest(classes = {SpringGraphQLDataFetchers.class, com.netflix.graphql.dgs.example.datafetcher.HelloDataFetcher.class, WithHeader.class, WithCookie.class, MovieDataFetcher.class, RequestHeadersDataFetcher.class})
+public class GraphQLContextContributorTest {
+
+    @Autowired
+    DgsQueryExecutor queryExecutor;
+
+    @Test
+    void moviesExtensionShouldHaveContributedEnabledExtension() {
+        final MockHttpServletRequest mockServletRequest = new MockHttpServletRequest();
+        mockServletRequest.addHeader(CONTEXT_CONTRIBUTOR_HEADER_NAME, CONTEXT_CONTRIBUTOR_HEADER_VALUE);
+        ServletWebRequest servletWebRequest = new ServletWebRequest(mockServletRequest);
+        String contributorEnabled = queryExecutor.executeAndExtractJsonPath("{ movies { director } }", "extensions.contributorEnabled", servletWebRequest);
+        assertThat(contributorEnabled).isEqualTo("true");
+    }
+
+    @Test
+    void withDataloaderContext() {
+        String message = queryExecutor.executeAndExtractJsonPath("{withDataLoaderContext}", "data.withDataLoaderContext");
+        assertThat(message).isEqualTo("Custom state! A");
+    }
+
+    @Test
+    void withDataloaderGraphQLContext() {
+        final MockHttpServletRequest mockServletRequest = new MockHttpServletRequest();
+        mockServletRequest.addHeader(CONTEXT_CONTRIBUTOR_HEADER_NAME, CONTEXT_CONTRIBUTOR_HEADER_VALUE);
+        ServletWebRequest servletWebRequest = new ServletWebRequest(mockServletRequest);
+        String contributorEnabled = queryExecutor.executeAndExtractJsonPath("{ withDataLoaderGraphQLContext }", "data.withDataLoaderGraphQLContext", servletWebRequest);
+        assertThat(contributorEnabled).isEqualTo("true");
+    }
+}

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

@@ -20,6 +20,7 @@
 import com.netflix.graphql.dgs.example.context.MyContextBuilder;
 import com.netflix.graphql.dgs.example.shared.context.ExampleGraphQLContextContributor;
 import com.netflix.graphql.dgs.example.shared.dataLoader.ExampleLoaderWithContext;
+import com.netflix.graphql.dgs.example.shared.dataLoader.ExampleLoaderWithGraphQLContext;
 import com.netflix.graphql.dgs.example.shared.dataLoader.MessageDataLoader;
 import com.netflix.graphql.dgs.example.shared.instrumentation.ExampleInstrumentationDependingOnContextContributor;
 import com.netflix.graphql.dgs.pagination.DgsPaginationAutoConfiguration;
@@ -36,6 +37,7 @@
 @Import({MessageDataLoader.class,
         UploadScalar.class,
         ExampleLoaderWithContext.class,
+        ExampleLoaderWithGraphQLContext.class,
         ExampleGraphQLContextContributor.class,
         ExampleInstrumentationDependingOnContextContributor.class,
         MyContextBuilder.class

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

@@ -21,6 +21,7 @@ import com.jayway.jsonpath.JsonPath
 import com.jayway.jsonpath.TypeRef
 import com.jayway.jsonpath.spi.mapper.MappingException
 import com.netflix.graphql.dgs.DgsQueryExecutor
+import com.netflix.graphql.dgs.context.GraphQLContextContributor
 import com.netflix.graphql.dgs.exceptions.DgsQueryExecutionDataExtractionException
 import com.netflix.graphql.dgs.exceptions.QueryException
 import com.netflix.graphql.dgs.internal.BaseDgsQueryExecutor
@@ -42,6 +43,7 @@ class SpringGraphQLDgsQueryExecutor(
     private val dgsContextBuilder: DefaultDgsGraphQLContextBuilder,
     private val dgsDataLoaderProvider: DgsDataLoaderProvider,
     private val requestCustomizer: DgsQueryExecutorRequestCustomizer = DgsQueryExecutorRequestCustomizer.DEFAULT_REQUEST_CUSTOMIZER,
+    private val graphQLContextContributors: List<GraphQLContextContributor>,
 ) : DgsQueryExecutor {
     override fun execute(
         query: String,
@@ -63,8 +65,18 @@ class SpringGraphQLDgsQueryExecutor(
 
         val httpRequest = requestCustomizer.apply(webRequest ?: RequestContextHolder.getRequestAttributes() as? WebRequest, headers)
         val dgsContext = dgsContextBuilder.build(DgsWebMvcRequestData(request.extensions, headers, httpRequest))
-        lateinit var graphQLContext: GraphQLContext
-        val dataLoaderRegistry = dgsDataLoaderProvider.buildRegistryWithContextSupplier { graphQLContext }
+        val dataLoaderRegistry =
+            dgsDataLoaderProvider.buildRegistryWithContextSupplier {
+                val graphQLContext = request.toExecutionInput().graphQLContext
+                if (graphQLContextContributors.isNotEmpty()) {
+                    val requestData = dgsContext.requestData
+                    val builderForContributors = GraphQLContext.newContext()
+                    graphQLContextContributors.forEach { it.contribute(builderForContributors, extensions, requestData) }
+                    graphQLContext.putAll(builderForContributors)
+                }
+
+                graphQLContext
+            }
 
         request.configureExecutionInput { _, builder ->
             builder
@@ -74,8 +86,6 @@ class SpringGraphQLDgsQueryExecutor(
                 .build()
         }
 
-        graphQLContext = request.toExecutionInput().graphQLContext
-
         val response =
             executionService.execute(request).block() ?: throw IllegalStateException("Unexpected null response from Spring GraphQL client")
 

graphql-dgs-spring-graphql/src/main/kotlin/com/netflix/graphql/dgs/springgraphql/autoconfig/DgsSpringGraphQLAutoConfiguration.kt

@@ -159,12 +159,14 @@ open class DgsSpringGraphQLAutoConfiguration {
         dgsContextBuilder: DefaultDgsGraphQLContextBuilder,
         dgsDataLoaderProvider: DgsDataLoaderProvider,
         requestCustomizer: ObjectProvider<DgsQueryExecutorRequestCustomizer>,
+        graphQLContextContributors: List<GraphQLContextContributor>,
     ): DgsQueryExecutor =
         SpringGraphQLDgsQueryExecutor(
             executionService,
             dgsContextBuilder,
             dgsDataLoaderProvider,
             requestCustomizer = requestCustomizer.getIfAvailable(DgsQueryExecutorRequestCustomizer::DEFAULT_REQUEST_CUSTOMIZER),
+            graphQLContextContributors,
         )
 
     @Configuration(proxyBeanMethods = false)
@@ -176,11 +178,13 @@ open class DgsSpringGraphQLAutoConfiguration {
         open fun dgsGraphQlInterceptor(
             dgsDataLoaderProvider: DgsDataLoaderProvider,
             dgsDefaultContextBuilder: DefaultDgsGraphQLContextBuilder,
+            graphQLContextContributors: List<GraphQLContextContributor>,
         ): DgsWebMvcGraphQLInterceptor =
             DgsWebMvcGraphQLInterceptor(
                 dgsDataLoaderProvider,
                 dgsDefaultContextBuilder,
                 dgsSpringGraphQLConfigurationProperties,
+                graphQLContextContributors,
             )
     }
 

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

@@ -16,6 +16,7 @@
 
 package com.netflix.graphql.dgs.springgraphql.webmvc
 
+import com.netflix.graphql.dgs.context.GraphQLContextContributor
 import com.netflix.graphql.dgs.internal.DefaultDgsGraphQLContextBuilder
 import com.netflix.graphql.dgs.internal.DgsDataLoaderProvider
 import com.netflix.graphql.dgs.internal.DgsWebMvcRequestData
@@ -29,12 +30,12 @@ import org.springframework.web.context.request.ServletRequestAttributes
 import org.springframework.web.context.request.ServletWebRequest
 import org.springframework.web.context.request.WebRequest
 import reactor.core.publisher.Mono
-import java.util.concurrent.CompletableFuture
 
 class DgsWebMvcGraphQLInterceptor(
     private val dgsDataLoaderProvider: DgsDataLoaderProvider,
     private val dgsContextBuilder: DefaultDgsGraphQLContextBuilder,
     private val dgsSpringConfigurationProperties: DgsSpringGraphQLConfigurationProperties,
+    private val graphQLContextContributors: List<GraphQLContextContributor>,
 ) : WebGraphQlInterceptor {
     override fun intercept(
         request: WebGraphQlRequest,
@@ -55,8 +56,19 @@ class DgsWebMvcGraphQLInterceptor(
             } else {
                 dgsContextBuilder.build(DgsWebMvcRequestData(request.extensions, request.headers))
             }
-        val graphQLContextFuture = CompletableFuture<GraphQLContext>()
-        val dataLoaderRegistry = dgsDataLoaderProvider.buildRegistryWithContextSupplier { graphQLContextFuture.get() }
+        val dataLoaderRegistry =
+            dgsDataLoaderProvider.buildRegistryWithContextSupplier {
+                val graphQLContext = request.toExecutionInput().graphQLContext
+                if (graphQLContextContributors.isNotEmpty()) {
+                    val extensions = request.extensions
+                    val requestData = dgsContext.requestData
+                    val builderForContributors = GraphQLContext.newContext()
+                    graphQLContextContributors.forEach { it.contribute(builderForContributors, extensions, requestData) }
+                    graphQLContext.putAll(builderForContributors)
+                }
+
+                graphQLContext
+            }
 
         request.configureExecutionInput { _, builder ->
             builder
@@ -65,7 +77,6 @@ class DgsWebMvcGraphQLInterceptor(
                 .dataLoaderRegistry(dataLoaderRegistry)
                 .build()
         }
-        graphQLContextFuture.complete(request.toExecutionInput().graphQLContext)
 
         return if (dgsSpringConfigurationProperties.webmvc.asyncdispatch.enabled) {
             chain.next(request).doFinally {