Upgrade to Micrometer context-propagation snapshots

rstoyanchev · rstoyanchev · commit 973fc567f301 · 2022-09-06T12:33:56.000+01:00
Closes gh-477
diff --git a/platform/build.gradle b/platform/build.gradle
@@ -23,7 +23,7 @@ dependencies {
 
 	constraints {
 		api("com.graphql-java:graphql-java:${graphQlJavaVersion}")
-		api("io.micrometer:context-propagation:1.0.0-M4")
+		api("io.micrometer:context-propagation:1.0.0-SNAPSHOT")
 
 		api("jakarta.annotation:jakarta.annotation-api:2.0.0")
 		api("jakarta.servlet:jakarta.servlet-api:5.0.0")
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java
@@ -114,7 +114,7 @@ private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object
 			return CompletableFuture.supplyAsync(
 					() -> {
 						try {
-							return ContextSnapshot.capture(graphQLContext).wrap((Callable<?>) result).call();
+							return ContextSnapshot.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
 						}
 						catch (Exception ex) {
 							throw new IllegalStateException(
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java
@@ -70,7 +70,7 @@ private ContextDataFetcherDecorator(
 	@Override
 	public Object get(DataFetchingEnvironment environment) throws Exception {
 
-		ContextSnapshot snapshot = ContextSnapshot.capture(environment.getGraphQlContext());
+		ContextSnapshot snapshot = ContextSnapshot.captureFrom(environment.getGraphQlContext());
 		Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();
 
 		if (this.subscription) {
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java
@@ -98,7 +98,7 @@ private List<GraphQLError> resolveInternal(Throwable exception, DataFetchingEnvi
 			return resolveToMultipleErrors(exception, env);
 		}
 		try {
-			return ContextSnapshot.capture(env.getGraphQlContext())
+			return ContextSnapshot.captureFrom(env.getGraphQlContext())
 					.wrap(() -> resolveToMultipleErrors(exception, env))
 					.call();
 		}
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java
@@ -190,7 +190,7 @@ public DataLoaderOptions getOptionsOrDefault(
 		@Override
 		public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
 			GraphQLContext graphQLContext = environment.getContext();
-			ContextSnapshot snapshot = ContextSnapshot.capture(graphQLContext);
+			ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
 			try {
 				return snapshot.wrap(() ->
 								this.loader.apply(keys, environment)
@@ -245,7 +245,7 @@ public DataLoaderOptions getOptionsOrDefault(
 		@Override
 		public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
 			GraphQLContext graphQLContext = environment.getContext();
-			ContextSnapshot snapshot = ContextSnapshot.capture(graphQLContext);
+			ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
 			try {
 				return snapshot.wrap(() ->
 								this.loader.apply(keys, environment)
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java
@@ -77,7 +77,7 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
 				request.configureExecutionInput(RESET_EXECUTION_ID_CONFIGURER);
 			}
 			ExecutionInput executionInput = request.toExecutionInput();
-			ContextSnapshot.capture(contextView).updateContext(executionInput.getGraphQLContext());
+			ContextSnapshot.captureFrom(contextView).updateContext(executionInput.getGraphQLContext());
 			ExecutionInput updatedExecutionInput = registerDataLoaders(executionInput);
 			return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(updatedExecutionInput))
 					.map(result -> new DefaultExecutionGraphQlResponse(updatedExecutionInput, result));
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java
@@ -70,7 +70,7 @@ public DataFetcherExceptionHandlerResult onException(DataFetcherExceptionHandler
 	public CompletableFuture<DataFetcherExceptionHandlerResult> handleException(DataFetcherExceptionHandlerParameters params) {
 		Throwable exception = unwrapException(params);
 		DataFetchingEnvironment env = params.getDataFetchingEnvironment();
-		ContextSnapshot snapshot = ContextSnapshot.capture(env.getGraphQlContext());
+		ContextSnapshot snapshot = ContextSnapshot.captureFrom(env.getGraphQlContext());
 		try {
 			return Flux.fromIterable(this.resolvers)
 					.flatMap(resolver -> resolver.resolveException(exception, env))
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java
@@ -83,7 +83,7 @@ public boolean isThreadLocalContextAware() {
     public final Mono<List<GraphQLError>> resolveException(Throwable exception) {
         if (this.threadLocalContextAware) {
             return Mono.deferContextual(contextView -> {
-                ContextSnapshot snapshot = ContextSnapshot.capture(contextView);
+                ContextSnapshot snapshot = ContextSnapshot.captureFrom(contextView);
                 try {
                     List<GraphQLError> errors = snapshot.wrap(() -> resolveToMultipleErrors(exception)).call();
                     return Mono.justOrEmpty(errors);
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java b/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java
@@ -87,7 +87,7 @@ public WebSocketGraphQlInterceptor getWebSocketInterceptor() {
 
 			@Override
 			public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
-				ContextSnapshot snapshot = ContextSnapshot.capture();
+				ContextSnapshot snapshot = ContextSnapshot.captureAll();
 				return executionChain.next(request).contextWrite(snapshot::updateContext);
 			}
 		};
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java
@@ -360,7 +360,7 @@ public boolean beforeHandshake(
 				ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
 				Map<String, Object> attributes) {
 
-			attributes.put(KEY, ContextSnapshot.capture());
+			attributes.put(KEY, ContextSnapshot.captureAll());
 			return true;
 		}
 
@@ -373,7 +373,7 @@ public void afterHandshake(
 		public static AutoCloseable setThreadLocals(WebSocketSession session) {
 			ContextSnapshot snapshot = (ContextSnapshot) session.getAttributes().get(KEY);
 			Assert.notNull(snapshot, "Expected ContextSnapshot in WebSocketSession attributes");
-			return snapshot.setThreadLocalValues();
+			return snapshot.setThreadLocals();
 		}
 	}
 
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java
@@ -60,7 +60,7 @@ public class BatchMappingPrincipalMethodArgumentResolverTests extends BatchMappi
 			ReactiveSecurityContextHolder.withAuthentication(this.authentication);
 
 	private final Function<Context, Context> threadLocalContextWriter = context ->
-			ContextSnapshot.capture().updateContext(context);
+			ContextSnapshot.captureAll().updateContext(context);
 
 
 	private static Stream<Arguments> controllers() {
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java
@@ -65,7 +65,7 @@ public class SchemaMappingPrincipalMethodArgumentResolverTests {
 			ReactiveSecurityContextHolder.withAuthentication(this.authentication);
 
 	private final Function<Context, Context> threadLocalContextWriter = context ->
-			ContextSnapshot.capture().updateContext(context);
+			ContextSnapshot.captureAll().updateContext(context);
 
 	private final GreetingController greetingController = new GreetingController();
 
@@ -206,4 +206,4 @@ Flux<String> greetingSubscription(Principal principal) {
 
 	}
 
-}
+}
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/CompositeSubscriptionExceptionResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/CompositeSubscriptionExceptionResolverTests.java
@@ -108,7 +108,7 @@ void resolveExceptionWithThreadLocal() {
 					.toGraphQl();
 
 			ExecutionInput input = ExecutionInput.newExecutionInput(query).build();
-			ContextSnapshot.capture().updateContext(input.getGraphQLContext());
+			ContextSnapshot.captureAll().updateContext(input.getGraphQLContext());
 
 			Flux<ResponseHelper> flux = Mono.defer(() -> Mono.fromFuture(graphQL.executeAsync(input)))
 					.map(ResponseHelper::forSubscription)
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java
@@ -163,7 +163,7 @@ void dataFetcherWithThreadLocalContext() {
 					.toGraphQl();
 
 			ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-			ContextSnapshot.capture().updateContext(input.getGraphQLContext());
+			ContextSnapshot.captureAll().updateContext(input.getGraphQLContext());
 
 			Mono<ExecutionResult> resultMono = Mono.delay(Duration.ofMillis(10))
 					.flatMap((aLong) -> Mono.fromFuture(graphQl.executeAsync(input)));
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java
@@ -101,7 +101,7 @@ void resolveExceptionWithThreadLocal() {
 									.build());
 
 			resolver.setThreadLocalContextAware(true);
-			ContextSnapshot.capture().updateContext(this.input.getGraphQLContext());
+			ContextSnapshot.captureAll().updateContext(this.input.getGraphQLContext());
 
 			Mono<ExecutionResult> result = Mono.delay(Duration.ofMillis(10)).flatMap((aLong) ->
 					Mono.fromFuture(this.graphQlSetup.exceptionResolver(resolver).toGraphQl().executeAsync(this.input)));
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java
@@ -379,7 +379,7 @@ void contextPropagation() throws Exception {
 			GraphQlWebSocketHandler handler = initWebSocketHandler(threadLocalInterceptor);
 
 			// Ensure ContextSnapshot is present in WebSocketSession attributes
-			this.session.getAttributes().put(ContextSnapshot.class.getName(), ContextSnapshot.capture());
+			this.session.getAttributes().put(ContextSnapshot.class.getName(), ContextSnapshot.captureAll());
 
 			// Context should propagate, if message is handled on different thread
 			Thread thread = new Thread(() -> {
@@ -411,7 +411,7 @@ private void handle(GraphQlWebSocketHandler handler, TextMessage... textMessages
 
 		if (!this.session.getAttributes().containsKey(ContextSnapshot.class.getName())) {
 			// Ensure ContextSnapshot is present in WebSocketSession attributes
-			this.session.getAttributes().put(ContextSnapshot.class.getName(), ContextSnapshot.capture());
+			this.session.getAttributes().put(ContextSnapshot.class.getName(), ContextSnapshot.captureAll());
 		}
 
 		for (TextMessage message : textMessages) {

Original file line number	Diff line number	Diff line change
`@@ -114,7 +114,7 @@ private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object`
`114`	`114`	`return CompletableFuture.supplyAsync(`
`115`	`115`	`() -> {`
`116`	`116`	`try {`
`117`		`- return ContextSnapshot.capture(graphQLContext).wrap((Callable<?>) result).call();`
	`117`	`+ return ContextSnapshot.captureFrom(graphQLContext).wrap((Callable<?>) result).call();`
`118`	`118`	`}`
`119`	`119`	`catch (Exception ex) {`
`120`	`120`	`throw new IllegalStateException(`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ private List<GraphQLError> resolveInternal(Throwable exception, DataFetchingEnvi`
`98`	`98`	`return resolveToMultipleErrors(exception, env);`
`99`	`99`	`}`
`100`	`100`	`try {`
`101`		`- return ContextSnapshot.capture(env.getGraphQlContext())`
	`101`	`+ return ContextSnapshot.captureFrom(env.getGraphQlContext())`
`102`	`102`	`.wrap(() -> resolveToMultipleErrors(exception, env))`
`103`	`103`	`.call();`
`104`	`104`	`}`
Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,7 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ`
`77`	`77`	`request.configureExecutionInput(RESET_EXECUTION_ID_CONFIGURER);`
`78`	`78`	`}`
`79`	`79`	`ExecutionInput executionInput = request.toExecutionInput();`
`80`		`- ContextSnapshot.capture(contextView).updateContext(executionInput.getGraphQLContext());`
	`80`	`+ ContextSnapshot.captureFrom(contextView).updateContext(executionInput.getGraphQLContext());`
`81`	`81`	`ExecutionInput updatedExecutionInput = registerDataLoaders(executionInput);`
`82`	`82`	`return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(updatedExecutionInput))`
`83`	`83`	`.map(result -> new DefaultExecutionGraphQlResponse(updatedExecutionInput, result));`
Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ public WebSocketGraphQlInterceptor getWebSocketInterceptor() {`
`87`	`87`
`88`	`88`	`@Override`
`89`	`89`	`public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {`
`90`		`- ContextSnapshot snapshot = ContextSnapshot.capture();`
	`90`	`+ ContextSnapshot snapshot = ContextSnapshot.captureAll();`
`91`	`91`	`return executionChain.next(request).contextWrite(snapshot::updateContext);`
`92`	`92`	`}`
`93`	`93`	`};`
Original file line number	Diff line number	Diff line change
`@@ -360,7 +360,7 @@ public boolean beforeHandshake(`
`360`	`360`	`ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,`
`361`	`361`	`Map<String, Object> attributes) {`
`362`	`362`
`363`		`- attributes.put(KEY, ContextSnapshot.capture());`
	`363`	`+ attributes.put(KEY, ContextSnapshot.captureAll());`
`364`	`364`	`return true;`
`365`	`365`	`}`
`366`	`366`
`@@ -373,7 +373,7 @@ public void afterHandshake(`
`373`	`373`	`public static AutoCloseable setThreadLocals(WebSocketSession session) {`
`374`	`374`	`ContextSnapshot snapshot = (ContextSnapshot) session.getAttributes().get(KEY);`
`375`	`375`	`Assert.notNull(snapshot, "Expected ContextSnapshot in WebSocketSession attributes");`
`376`		`- return snapshot.setThreadLocalValues();`
	`376`	`+ return snapshot.setThreadLocals();`
`377`	`377`	`}`
`378`	`378`	`}`
`379`	`379`