Ensure reactive context is propagated

gradle/libs.versions.toml

@@ -5,7 +5,6 @@ micronaut-platform = "4.1.2"
 micronaut-test = "4.0.0"
 groovy = "4.0.13"
 spock = "2.3-groovy-4.0"
-
 graal-svm = "23.1.2"
 kotlin = '1.9.22'
 shadow = '8.0.0'
@@ -14,6 +13,7 @@ micronaut-kotlin = "4.2.0"
 micronaut-security = "4.6.6"
 micronaut-serde = "2.8.1"
 micronaut-logging = "1.1.2"
+micronaut-reactor = "3.2.1"
 
 managed-graphql-java = "21.3"
 managed-graphql-java-extended-scalars = "2023-01-24T02-11-56-babda5f"
@@ -26,6 +26,7 @@ micronaut-core = { module = 'io.micronaut:micronaut-core-bom', version.ref = 'mi
 micronaut-kotlin = { module = "io.micronaut.kotlin:micronaut-kotlin-bom", version.ref = "micronaut-kotlin" }
 micronaut-security = { module = "io.micronaut.security:micronaut-security-bom", version.ref = "micronaut-security" }
 micronaut-serde = { module = "io.micronaut.serde:micronaut-serde-bom", version.ref = "micronaut-serde" }
+micronaut-reactor = { module = "io.micronaut.reactor:micronaut-reactor-bom", version.ref = "micronaut-reactor" }
 
 managed-graphql-java = { module = 'com.graphql-java:graphql-java', version.ref = 'managed-graphql-java' }
 managed-graphql-java-extended-scalars = { module = 'com.graphql-java:graphql-java-extended-scalars', version.ref = 'managed-graphql-java-extended-scalars' }

graphql/build.gradle

@@ -17,6 +17,8 @@ dependencies {
     testImplementation(mn.micronaut.http.server.netty)
     testImplementation(mn.micronaut.http.client)
     testRuntimeOnly(mn.snakeyaml)
+    testRuntimeOnly(mnReactor.micronaut.reactor)
+    testRuntimeOnly(mnReactor.micrometer.context.propagation)
 }
 
 micronautBuild {

graphql/src/main/java/io/micronaut/configuration/graphql/DefaultGraphQLExecutionInputCustomizer.java

@@ -18,11 +18,11 @@
 import graphql.ExecutionInput;
 import io.micronaut.context.annotation.Requires;
 import io.micronaut.core.annotation.Nullable;
-import io.micronaut.core.async.publisher.Publishers;
 import io.micronaut.http.HttpRequest;
 import io.micronaut.http.MutableHttpResponse;
 import jakarta.inject.Singleton;
 import org.reactivestreams.Publisher;
+import reactor.core.publisher.Mono;
 
 /**
  * The default implementation for customizing GraphQL execution inputs.
@@ -42,6 +42,6 @@ public Publisher<ExecutionInput> customize(
             ExecutionInput executionInput,
             HttpRequest httpRequest,
             @Nullable MutableHttpResponse<String> httpResponse) {
-        return Publishers.just(executionInput);
+        return Mono.just(executionInput);
     }
 }

graphql/src/main/java/io/micronaut/configuration/graphql/DefaultGraphQLInvocation.java

@@ -20,13 +20,13 @@
 import graphql.GraphQL;
 import io.micronaut.context.BeanProvider;
 import io.micronaut.core.annotation.Nullable;
-import io.micronaut.core.async.publisher.Publishers;
 import io.micronaut.http.HttpRequest;
 import io.micronaut.http.MutableHttpResponse;
 import jakarta.inject.Singleton;
 import org.dataloader.DataLoaderRegistry;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import java.util.concurrent.CompletableFuture;
 
@@ -81,11 +81,11 @@ public Publisher<ExecutionResult> invoke(
         ExecutionInput executionInput = executionInputBuilder.build();
         return Flux
                 .from(graphQLExecutionInputCustomizer.customize(executionInput, httpRequest, httpResponse))
-                .flatMap(customizedExecutionInput -> Publishers.fromCompletableFuture(() -> {
+                .flatMap(customizedExecutionInput -> Mono.fromFuture(() -> {
                     try {
                         return graphQL.executeAsync(customizedExecutionInput);
                     } catch (Throwable e) {
-                        CompletableFuture future = new CompletableFuture();
+                        CompletableFuture<ExecutionResult> future = new CompletableFuture<>();
                         future.completeExceptionally(e);
                         return future;
                     }

graphql/src/test/groovy/io/micronaut/configuration/graphql/GraphQLConfigurationSpec.groovy

@@ -17,11 +17,15 @@
 package io.micronaut.configuration.graphql
 
 import graphql.GraphQL
+import graphql.schema.GraphQLSchema
 import io.micronaut.context.ApplicationContext
-import io.micronaut.context.DefaultApplicationContext
+import io.micronaut.context.annotation.Bean
+import io.micronaut.context.annotation.Factory
+import io.micronaut.context.annotation.Requires
 import io.micronaut.context.env.Environment
-import io.micronaut.context.env.PropertySource
+import io.micronaut.core.util.StringUtils
 import io.micronaut.http.annotation.Controller
+import jakarta.inject.Singleton
 import spock.lang.Specification
 
 /**
@@ -33,11 +37,8 @@ class GraphQLConfigurationSpec extends Specification {
 
     void "test no graphql bean provided"() {
         given:
-        ApplicationContext context = new DefaultApplicationContext(Environment.TEST)
-        context.environment.addPropertySource(PropertySource.of(
-                ["graphql.factory": false]
-        ))
-        context.start()
+        ApplicationContext context = ApplicationContext.run(["spec.name": GraphQLConfigurationSpec.simpleName,
+                                                             "graphql.factory": false], Environment.TEST)
 
         expect:
         !context.containsBean(GraphQLExecutionResultHandler)
@@ -50,9 +51,7 @@ class GraphQLConfigurationSpec extends Specification {
 
     void "test graphql bean provided"() {
         given:
-        ApplicationContext context = new DefaultApplicationContext(Environment.TEST)
-        context.registerSingleton(Mock(GraphQL))
-        context.start()
+        ApplicationContext context = ApplicationContext.run(["spec.name": GraphQLConfigurationSpec.simpleName], Environment.TEST)
 
         expect:
         context.containsBean(GraphQLExecutionResultHandler)
@@ -66,11 +65,8 @@ class GraphQLConfigurationSpec extends Specification {
 
     void "test custom graphql path"() {
         given:
-        ApplicationContext context = new DefaultApplicationContext(Environment.TEST)
-        context.environment.addPropertySource(PropertySource.of(
-                ["graphql.path": "/custom-graphql"]
-        ))
-        context.start()
+        ApplicationContext context = ApplicationContext.run(["spec.name": GraphQLConfigurationSpec.simpleName,
+                                                             "graphql.path": "/custom-graphql"], Environment.TEST)
 
         expect:
         context.containsBean(GraphQLExecutionResultHandler)
@@ -85,11 +81,8 @@ class GraphQLConfigurationSpec extends Specification {
 
     void "test graphql disabled"() {
         given:
-        ApplicationContext context = new DefaultApplicationContext(Environment.TEST)
-        context.environment.addPropertySource(PropertySource.of(
-                ["graphql.enabled": false]
-        ))
-        context.start()
+        ApplicationContext context = ApplicationContext.run(["spec.name": GraphQLConfigurationSpec.simpleName,
+                                                             "graphql.enabled": false], Environment.TEST)
 
         expect:
         !context.containsBean(GraphQLExecutionResultHandler)
@@ -99,4 +92,17 @@ class GraphQLConfigurationSpec extends Specification {
         cleanup:
         context.close()
     }
+
+    @Factory
+    static class GraphQLFactory {
+
+        @Bean
+        @Singleton
+        @Requires(property = "graphql.factory", notEquals = StringUtils.FALSE)
+        @Requires(property = "spec.name", value = "GraphQLConfigurationSpec")
+        GraphQL graphQL() {
+            def schema = GraphQLSchema.newSchema().build()
+            GraphQL.newGraphQL(schema).build()
+        }
+    }
 }

graphql/src/test/groovy/io/micronaut/configuration/graphql/GraphQLContextPropagationSpec.groovy

@@ -0,0 +1,104 @@
+package io.micronaut.configuration.graphql
+
+import graphql.GraphQL
+import graphql.schema.DataFetcher
+import graphql.schema.DataFetchingEnvironment
+import graphql.schema.GraphQLSchema
+import graphql.schema.idl.RuntimeWiring
+import graphql.schema.idl.SchemaGenerator
+import graphql.schema.idl.SchemaParser
+import graphql.schema.idl.TypeDefinitionRegistry
+import io.micronaut.context.ApplicationContext
+import io.micronaut.context.annotation.Bean
+import io.micronaut.context.annotation.Factory
+import io.micronaut.context.annotation.Requires
+import io.micronaut.context.env.Environment
+import io.micronaut.core.annotation.Nullable
+import io.micronaut.core.io.ResourceResolver
+import io.micronaut.http.HttpResponse
+import io.micronaut.http.annotation.Body
+import io.micronaut.http.annotation.Get
+import io.micronaut.http.annotation.Post
+import io.micronaut.http.annotation.QueryValue
+import io.micronaut.http.client.annotation.Client
+import io.micronaut.http.context.ServerRequestContext
+import io.micronaut.runtime.server.EmbeddedServer
+import jakarta.inject.Singleton
+import spock.lang.AutoCleanup
+import spock.lang.Specification
+
+import java.net.http.HttpRequest
+
+class GraphQLContextPropagationSpec extends Specification {
+
+    @AutoCleanup
+    EmbeddedServer embeddedServer
+
+    GraphQLClient graphQLClient
+
+    def setup() {
+        embeddedServer = ApplicationContext.run(
+                EmbeddedServer,
+                ["spec.name": GraphQLContextPropagationSpec.simpleName])
+        graphQLClient = embeddedServer.applicationContext.getBean(GraphQLClient)
+    }
+
+    void "server request context is propagated to data fetcher"() {
+        when:
+        GraphQLResponseBody response = graphQLClient.hello("query { hello }")
+
+        then:
+        response
+        response.getSpecification()["data"]["hello"] == "Hello World!"
+    }
+
+    @Client("/graphql")
+    static interface GraphQLClient {
+
+        @Get("{?query}")
+        GraphQLResponseBody hello(@QueryValue String query)
+
+    }
+
+    @Factory
+    static class GraphQLFactory {
+
+        @Bean
+        @Singleton
+        @Requires(property = "spec.name", value = "GraphQLContextPropagationSpec")
+        GraphQL graphQL(ResourceResolver resourceResolver, HelloDataFetcher helloDataFetcher) {
+
+            SchemaParser schemaParser = new SchemaParser()
+            SchemaGenerator schemaGenerator = new SchemaGenerator()
+
+            TypeDefinitionRegistry typeRegistry = new TypeDefinitionRegistry()
+            typeRegistry.merge(schemaParser.parse(new BufferedReader(new InputStreamReader(
+                    resourceResolver.getResourceAsStream("classpath:schema.graphqls").get()))))
+
+            RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring()
+                    .type("Query", typeWiring -> typeWiring
+                            .dataFetcher("hello", helloDataFetcher))
+                    .build()
+
+            GraphQLSchema graphQLSchema = schemaGenerator.makeExecutableSchema(typeRegistry, runtimeWiring)
+
+            return GraphQL.newGraphQL(graphQLSchema).build()
+        }
+    }
+
+    @Singleton
+    @Requires(property = "spec.name", value = "GraphQLContextPropagationSpec")
+    static class HelloDataFetcher implements DataFetcher<String> {
+
+        @Override
+        String get(DataFetchingEnvironment env) {
+            Optional<HttpRequest> request = ServerRequestContext.currentRequest()
+            assert request.isPresent()
+            String name = env.getArgument("name")
+            if (name == null || name.trim().length() == 0) {
+                name = "World"
+            }
+            return String.format("Hello %s!", name)
+        }
+    }
+}

graphql/src/test/groovy/io/micronaut/configuration/graphql/GraphQLControllerSpec.groovy

@@ -301,7 +301,7 @@ class GraphQLControllerSpec extends Specification {
     }
 
     @Factory
-    class GraphQLFactory {
+    static class GraphQLFactory {
 
         @Bean
         @Singleton