From 94ba9fed8fc9e0fb8a81e43e65ae30feb80f61c9 Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@vmware.com>
Date: Wed, 16 Aug 2023 18:25:15 +0200
Subject: [PATCH] Add test case for gh-761

This commit also ensures that a new local context is created, copying
the existing values. This avoids mutating the parent local context and
polluting it with local values.
This could cause unintended side effects on other child datafetchers.

Fixes gh-761
---
 .../GraphQlObservationInstrumentation.java    |  2 +-
 ...raphQlObservationInstrumentationTests.java | 36 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java
index 47567553..6faa8e6d 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java
@@ -203,7 +203,7 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen
 		}
 		else {
 			GraphQLContext localContext = dataFetcherLocalContext == null ?
-					GraphQLContext.getDefault() : dataFetcherLocalContext;
+					GraphQLContext.getDefault() : GraphQLContext.newContext().of(dataFetcherLocalContext).build();
 			return DataFetcherResult.newResult()
 					.data(value)
 					.localContext(localContext.put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation))
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java
index b664ffc9..9a652fe8 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/observation/GraphQlObservationInstrumentationTests.java
@@ -16,6 +16,7 @@
 
 package org.springframework.graphql.observation;
 
+import graphql.GraphQLContext;
 import graphql.GraphqlErrorBuilder;
 import graphql.execution.DataFetcherResult;
 import graphql.schema.AsyncDataFetcher;
@@ -279,4 +280,39 @@ class GraphQlObservationInstrumentationTests {
 		ResponseHelper.forResponse(responseMono);
 	}
 
+	@Test
+	void shouldNotOverrideExistingLocalContext() {
+		
+		String document = """
+				{
+					bookById(id: 1) {
+						author {
+							firstName,
+							lastName
+						}
+					}
+				}
+				""";
+		DataFetcher<DataFetcherResult<Object>> bookDataFetcher = environment -> DataFetcherResult.newResult()
+				.data(BookSource.getBook(1L))
+				.localContext(GraphQLContext.newContext().of("test", "value").build())
+				.build();
+		DataFetcher<Author> authorDataFetcher = environment -> BookSource.getAuthor(101L);
+		DataFetcher<String> authorFirstNameDataFetcher = environment -> {
+			GraphQLContext context = environment.getLocalContext();
+			String value = context.get("test");
+			assertThat(value).isEqualTo("value");
+			return BookSource.getAuthor(101L).getFirstName();
+		};
+
+        ExecutionGraphQlRequest request = TestExecutionRequest.forDocument(document);
+		Mono<ExecutionGraphQlResponse> responseMono = graphQlSetup
+				.queryFetcher("bookById", bookDataFetcher)
+				.dataFetcher("Book", "author", authorDataFetcher)
+				.dataFetcher("Author", "firstName", authorFirstNameDataFetcher)
+				.toGraphQlService()
+				.execute(request);
+		ResponseHelper.forResponse(responseMono);
+	}
+
 }