Merge branch '1.1.x'

2023-05-10 15:47:22 +02:00
parent 577014f4f1 0a41104226
commit c45dddc9d5
3 changed files with 173 additions and 72 deletions
--- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java
@@ -19,6 +19,7 @@ package org.springframework.graphql.execution;
 import java.util.List;

 import graphql.ExecutionInput;
+import graphql.GraphQLContext;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.GraphQLCodeRegistry;
@@ -72,7 +73,18 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
 	@Override
 	public Object get(DataFetchingEnvironment environment) throws Exception {

-		ContextSnapshot snapshot = ContextSnapshot.captureFrom(environment.getGraphQlContext());
+		GraphQLContext context;
+		// temporarily merge global and local graphql context until https://github.com/micrometer-metrics/context-propagation/pull/98
+		if (environment.getLocalContext() instanceof GraphQLContext localContext) {
+			context = GraphQLContext.newContext()
+					.of(environment.getGraphQlContext())
+					.of(localContext)
+					.build();
+		}
+		else {
+			context = environment.getGraphQlContext();
+		}
+		ContextSnapshot snapshot = ContextSnapshot.captureFrom(context);
 		Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();

 		if (this.subscription) {
--- a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java
@@ -17,8 +17,8 @@
 package org.springframework.graphql.observation;

 import graphql.ExecutionResult;
-import graphql.execution.ExecutionStepInfo;
-import graphql.execution.ResultPath;
+import graphql.GraphQLContext;
+import graphql.execution.DataFetcherResult;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.SimpleInstrumentationContext;
@@ -27,12 +27,16 @@ import graphql.execution.instrumentation.parameters.InstrumentationCreateStatePa
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
 import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.lang.Nullable;

+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
-import java.util.concurrent.ConcurrentHashMap;

 /**
 * {@link graphql.execution.instrumentation.Instrumentation} that creates
@@ -53,6 +57,8 @@ import java.util.concurrent.ConcurrentHashMap;
 */
 public class GraphQlObservationInstrumentation extends SimplePerformantInstrumentation {

+	private static final Log logger = LogFactory.getLog(GraphQlObservationInstrumentation.class);
+
 	private static final ExecutionRequestObservationConvention DEFAULT_REQUEST_CONVENTION =
 			new DefaultExecutionRequestObservationConvention();

@@ -94,7 +100,7 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen

 	@Override
 	public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
-		return new RequestObservationInstrumentationState();
+		return RequestObservationInstrumentationState.INSTANCE;
 	}

 	@Override
@@ -102,10 +108,11 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
 			InstrumentationState state) {
 		if (state instanceof RequestObservationInstrumentationState instrumentationState) {
 			ExecutionRequestObservationContext observationContext = new ExecutionRequestObservationContext(parameters.getExecutionInput());
-			Observation parentObservation = parameters.getGraphQLContext().get(ObservationThreadLocalAccessor.KEY);
-			Observation requestObservation = instrumentationState.createRequestObservation(this.requestObservationConvention,
-					observationContext, this.observationRegistry);
-			requestObservation.parentObservation(parentObservation);
+			Observation requestObservation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(this.requestObservationConvention,
+					DEFAULT_REQUEST_CONVENTION, () -> observationContext, this.observationRegistry);
+			requestObservation.parentObservation(getCurrentObservation(parameters.getGraphQLContext()));
+			GraphQLContext graphQLContext = parameters.getGraphQLContext();
+			graphQLContext.put(ObservationThreadLocalAccessor.KEY, requestObservation);
 			requestObservation.start();
 			return new SimpleInstrumentationContext<>() {
 				@Override
@@ -121,6 +128,11 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
 		return super.beginExecution(parameters, state);
 	}

+	@Nullable
+	private static Observation getCurrentObservation(GraphQLContext graphQLContext) {
+		return graphQLContext.get(ObservationThreadLocalAccessor.KEY);
+	}
+
 	@Override
 	public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
 			InstrumentationFieldFetchParameters parameters, InstrumentationState state) {
@@ -128,23 +140,29 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
 				&& state instanceof RequestObservationInstrumentationState instrumentationState) {
 			return (environment) -> {
 				DataFetcherObservationContext observationContext = new DataFetcherObservationContext(parameters.getEnvironment());
-				Observation dataFetcherObservation = instrumentationState.createDataFetcherObservation(this.dataFetcherObservationConvention, observationContext, this.observationRegistry);
+				Observation dataFetcherObservation = GraphQlObservationDocumentation.DATA_FETCHER.observation(this.dataFetcherObservationConvention,
+						DEFAULT_DATA_FETCHER_CONVENTION, () -> observationContext, this.observationRegistry);
+				dataFetcherObservation.parentObservation(getCurrentObservation(environment));
 				dataFetcherObservation.start();
 				try {
 					Object value = dataFetcher.get(environment);
 					if (value instanceof CompletionStage<?> completion) {
-						return completion.whenComplete((result, error) -> {
+						return completion.handle((result, error) -> {
 							observationContext.setValue(result);
 							if (error != null) {
 								dataFetcherObservation.error(error);
+								dataFetcherObservation.stop();
+								return CompletableFuture.failedStage(error);
 							}
 							dataFetcherObservation.stop();
+							return CompletableFuture.completedStage(wrapAsDataFetcherResult(result, dataFetcherObservation));
+
 						});
 					}
 					else {
 						observationContext.setValue(value);
 						dataFetcherObservation.stop();
-						return value;
+						return wrapAsDataFetcherResult(value, dataFetcherObservation);
 					}
 				}
 				catch (Throwable throwable) {
@@ -157,33 +175,46 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
 		return dataFetcher;
 	}

-
-	static class RequestObservationInstrumentationState implements InstrumentationState {
-
-		private Observation requestObservation;
-
-		private final ConcurrentHashMap<ResultPath, Observation> activeObservations = new ConcurrentHashMap<>(8);
-
-		Observation createRequestObservation(ExecutionRequestObservationConvention convention,
-											 ExecutionRequestObservationContext context, ObservationRegistry registry) {
-			Observation observation = GraphQlObservationDocumentation.EXECUTION_REQUEST.observation(convention,
-					DEFAULT_REQUEST_CONVENTION, () -> context, registry);
-			this.requestObservation = observation;
-			return observation;
+	@Nullable
+	private static Observation getCurrentObservation(DataFetchingEnvironment environment) {
+		Observation currentObservation = null;
+		if (environment.getLocalContext() != null && environment.getLocalContext() instanceof GraphQLContext) {
+			GraphQLContext localContext = environment.getLocalContext();
+			currentObservation = localContext.get(ObservationThreadLocalAccessor.KEY);
 		}
+		if (currentObservation == null) {
+			return environment.getGraphQlContext().get(ObservationThreadLocalAccessor.KEY);
+		}
+		return currentObservation;
+	}

-		Observation createDataFetcherObservation(DataFetcherObservationConvention convention,
-												 DataFetcherObservationContext context, ObservationRegistry registry) {
-			ExecutionStepInfo executionStepInfo = context.getEnvironment().getExecutionStepInfo();
-			Observation observation = GraphQlObservationDocumentation.DATA_FETCHER.observation(convention,
-					DEFAULT_DATA_FETCHER_CONVENTION, () -> context, registry);
-			Observation parentObservation = (executionStepInfo.hasParent()) ? this.activeObservations.getOrDefault(executionStepInfo.getParent().getPath(), this.requestObservation)
-					: this.requestObservation;
-			observation.parentObservation(parentObservation);
-			this.activeObservations.put(executionStepInfo.getPath(), observation);
-			return observation;
+	private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation) {
+		if (value instanceof DataFetcherResult<?> result) {
+			if (result.getLocalContext() == null) {
+				return result.transform(builder -> builder.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build()));
+			}
+			else if (result.getLocalContext() instanceof GraphQLContext) {
+				((GraphQLContext) result.getLocalContext()).put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation);
+			} else {
+				logger.debug("Cannot add observation to localContext as it is not a GraphQLContext but a "
+						+ result.getLocalContext().getClass().toString());
+			}
+			return result;
+		}
+		else {
+			return DataFetcherResult.newResult()
+					.data(value)
+					.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build())
+					.build();
 		}

 	}

+
+	static class RequestObservationInstrumentationState implements InstrumentationState {
+
+		static final RequestObservationInstrumentationState INSTANCE = new RequestObservationInstrumentationState();
+
+	}
+
 }
--- a/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java
@@ -17,24 +17,27 @@
 package org.springframework.graphql.observation;

 import graphql.GraphqlErrorBuilder;
+import graphql.execution.DataFetcherResult;
+import graphql.schema.AsyncDataFetcher;
+import graphql.schema.DataFetcher;
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 import io.micrometer.observation.tck.TestObservationRegistry;
 import io.micrometer.observation.tck.TestObservationRegistryAssert;
 import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Mono;
-
-import org.springframework.graphql.BookSource;
-import org.springframework.graphql.ExecutionGraphQlRequest;
-import org.springframework.graphql.ExecutionGraphQlResponse;
-import org.springframework.graphql.GraphQlSetup;
-import org.springframework.graphql.ResponseHelper;
-import org.springframework.graphql.TestExecutionRequest;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.springframework.graphql.*;
 import org.springframework.graphql.execution.DataFetcherExceptionResolver;
 import org.springframework.graphql.execution.ErrorType;
+import reactor.core.publisher.Mono;

 import java.util.concurrent.CompletableFuture;
+import java.util.stream.Stream;
+
+import static org.assertj.core.api.Assertions.assertThat;

 /**
 * Tests for {@link GraphQlObservationInstrumentation}.
@@ -50,8 +53,9 @@ class GraphQlObservationInstrumentationTests {
 	private final GraphQlSetup graphQlSetup = GraphQlSetup.schemaResource(BookSource.schema).instrumentation(this.instrumentation);


-	@Test
-	void instrumentGraphQlRequestWhenSuccess() {
+	@ParameterizedTest
+	@MethodSource("successDataFetchers")
+	void instrumentGraphQlRequestWhenSuccess(DataFetcher<?> dataFetcher) {
 		String document = """
 				{
 					bookById(id: 1) {
@@ -60,7 +64,7 @@ class GraphQlObservationInstrumentationTests {
 				}
 				""";
 		Mono<ExecutionGraphQlResponse> responseMono = graphQlSetup
-				.queryFetcher("bookById", env -> BookSource.getBookWithoutAuthor(1L))
+				.queryFetcher("bookById", dataFetcher)
 				.toGraphQlService()
 				.execute(TestExecutionRequest.forDocument(document));
 		ResponseHelper response = ResponseHelper.forResponse(responseMono);
@@ -77,6 +81,16 @@ class GraphQlObservationInstrumentationTests {
 				.hasHighCardinalityKeyValue("graphql.field.path", "/bookById");
 	}

+	static Stream<Arguments> successDataFetchers() {
+		DataFetcher<Book> bookDataFetcher = environment -> BookSource.getBookWithoutAuthor(1L);
+		return Stream.of(
+				Arguments.of(bookDataFetcher),
+				Arguments.of(new AsyncDataFetcher<>(bookDataFetcher)),
+				Arguments.of((DataFetcher<DataFetcherResult<?>>) environment ->
+						DataFetcherResult.newResult().data(BookSource.getBookWithoutAuthor(1L)).build())
+		);
+	}
+
 	@Test
 	void instrumentGraphQlRequestWhenInvalidRequest() {
 		String document = "invalid";
@@ -131,8 +145,9 @@ class GraphQlObservationInstrumentationTests {
 				.hasAnObservationWithAKeyValue("graphql.field.path", "/bookById/author");
 	}

-	@Test
-	void instrumentGraphQlRequestWhenDataFetchingFailure() {
+	@ParameterizedTest
+	@MethodSource("failureDataFetchers")
+	void instrumentGraphQlRequestWhenDataFetchingFailure(DataFetcher<?> dataFetcher) {
 		String document = """
 				{
 					bookById(id: 1) {
@@ -153,7 +168,7 @@ class GraphQlObservationInstrumentationTests {
 				.execute(TestExecutionRequest.forDocument(document));
 		ResponseHelper response = ResponseHelper.forResponse(responseMono);
 		TestObservationRegistryAssert.assertThat(this.observationRegistry).hasObservationWithNameEqualTo("graphql.request")
-				.that().hasLowCardinalityKeyValue("graphql.outcome", "REQUEST_ERROR")
+				.that().hasLowCardinalityKeyValue("graphql.outcome", "SUCCESS")
 				.hasHighCardinalityKeyValueWithKey("graphql.execution.id");

 		TestObservationRegistryAssert.assertThat(this.observationRegistry)
@@ -166,29 +181,15 @@ class GraphQlObservationInstrumentationTests {
 				.hasHighCardinalityKeyValue("graphql.field.path", "/bookById");
 	}

-	@Test
-	void propagatesContextBetweenObservations() {
-		String document = """
-				{
-					bookById(id: 1) {
-						name
-					}
-				}
-				""";
-		Mono<ExecutionGraphQlResponse> responseMono = graphQlSetup
-				.queryFetcher("bookById", env -> BookSource.getBookWithoutAuthor(1L))
-				.toGraphQlService()
-				.execute(TestExecutionRequest.forDocument(document));
-		ResponseHelper response = ResponseHelper.forResponse(responseMono);
-
-		TestObservationRegistryAssert.assertThat(this.observationRegistry).hasObservationWithNameEqualTo("graphql.request")
-				.that().hasLowCardinalityKeyValue("graphql.outcome", "SUCCESS")
-				.hasHighCardinalityKeyValueWithKey("graphql.execution.id");
-
-		TestObservationRegistryAssert.assertThat(this.observationRegistry)
-				.hasObservationWithNameEqualTo("graphql.datafetcher")
-				.that()
-				.hasParentObservationContextMatching(context -> context instanceof ExecutionRequestObservationContext);
+	static Stream<Arguments> failureDataFetchers() {
+		DataFetcher<Book> bookDataFetcher = environment -> {
+			throw new IllegalStateException("book fetching failure");
+		};
+		return Stream.of(
+				Arguments.of(bookDataFetcher),
+				Arguments.of((DataFetcher<?>) environment ->
+						CompletableFuture.failedStage(new IllegalStateException("book fetching failure")))
+		);
 	}

 	@Test
@@ -215,4 +216,61 @@ class GraphQlObservationInstrumentationTests {
 		incoming.stop();
 	}

+	@ParameterizedTest
+	@MethodSource("successDataFetchers")
+	void propagatesContextBetweenObservations(DataFetcher<?> dataFetcher) {
+		String document = """
+				{
+					bookById(id: 1) {
+						name
+					}
+				}
+				""";
+		Mono<ExecutionGraphQlResponse> responseMono = graphQlSetup
+				.queryFetcher("bookById", dataFetcher)
+				.toGraphQlService()
+				.execute(TestExecutionRequest.forDocument(document));
+		ResponseHelper response = ResponseHelper.forResponse(responseMono);
+
+		TestObservationRegistryAssert.assertThat(this.observationRegistry).hasObservationWithNameEqualTo("graphql.request")
+				.that().hasLowCardinalityKeyValue("graphql.outcome", "SUCCESS")
+				.hasHighCardinalityKeyValueWithKey("graphql.execution.id");
+
+		TestObservationRegistryAssert.assertThat(this.observationRegistry)
+				.hasObservationWithNameEqualTo("graphql.datafetcher")
+				.that()
+				.hasParentObservationContextMatching(context -> context instanceof ExecutionRequestObservationContext);
+	}
+
+	@Test
+	void currentObservationSetInDataFetcherContext() {
+		String document = """
+				{
+					bookById(id: 1) {
+						author {
+							firstName,
+							lastName
+						}
+					}
+				}
+				""";
+		DataFetcher<Book> bookDataFetcher = environment -> {
+			assertThat(observationRegistry.getCurrentObservation().getContext())
+					.isInstanceOf(ExecutionRequestObservationContext.class);
+			return BookSource.getBookWithoutAuthor(1L);
+		};
+		DataFetcher<Author> authorDataFetcher = environment -> {
+			assertThat(observationRegistry.getCurrentObservation().getContext())
+					.isInstanceOf(DataFetcherObservationContext.class);
+			return BookSource.getAuthor(101L);
+		};
+
+		Mono<ExecutionGraphQlResponse> responseMono = graphQlSetup
+				.queryFetcher("bookById", bookDataFetcher)
+				.dataFetcher("Book", "author", authorDataFetcher)
+				.toGraphQlService()
+				.execute(TestExecutionRequest.forDocument(document));
+		ResponseHelper response = ResponseHelper.forResponse(responseMono);
+	}
+
 }