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 475675534..6faa8e6d1 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 @@ else if (result.getLocalContext() instanceof GraphQLContext) {
 		}
 		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 b664ffc9a..9a652fe80 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 @@ void currentObservationSetInDataFetcherContext() {
 		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);
+	}
+
 }