Allow setting custom ExecutionId on RequestInput

Prior to this commit, we would use the request id as an execution id,
regardless of the overall GraphQL configuration.

This commit:

* adds a new executionId field on `RequestInput` and uses it in the
  `ExecutionInput`, if provided
* only uses the request id as a fallback if it's been asked to
* auto-detect in the `GraphQlService` whether a custom
  `ExecutionIdProvider` implementation was configured on `GraphQL`.

Closes gh-243

Author: bclozel

spring-graphql/src/main/java/org/springframework/graphql/RequestInput.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,7 +33,7 @@
 
 /**
  * Common representation for GraphQL request input. This can be converted to
- * {@link ExecutionInput} via {@link #toExecutionInput()} and the
+ * {@link ExecutionInput} via {@link #toExecutionInput(boolean)} and the
  * {@code ExecutionInput} further customized via
  * {@link #configureExecutionInput(BiFunction)}.
  *
@@ -57,6 +57,8 @@ public class RequestInput {
 
 	private final List<BiFunction<ExecutionInput, ExecutionInput.Builder, ExecutionInput>> executionInputConfigurers = new ArrayList<>();
 
+	@Nullable
+	private ExecutionId executionId;
 
 	/**
 	 * Create an instance.
@@ -80,14 +82,9 @@ public RequestInput(
 	}
 
 
-	@SuppressWarnings("unchecked")
-	private static <T> T getKey(String key, Map<String, Object> body) {
-		return (T) body.get(key);
-	}
-
 	/**
-	 * Return an identifier for the request. This id is later propagated
-	 * as the {@link ExecutionId} of the execution input.
+	 * Return an identifier for the request. This id can be later propagated
+	 * as the {@link ExecutionId} if {@link #executionId(ExecutionId) none has been set}.
 	 * <p>For web transports, this identifier can be used to correlate
 	 * request and response messages on a multiplexed connection.
 	 * @return the request id.
@@ -131,6 +128,15 @@ public Locale getLocale() {
 		return this.locale;
 	}
 
+	/**
+	 * Set an {@link ExecutionId} to be used for the {@link ExecutionInput}
+	 * @param executionId the execution id to use with the {@link ExecutionInput}.
+	 */
+	public void executionId(ExecutionId executionId) {
+		Assert.notNull(executionId, "executionId should not be null");
+		this.executionId = executionId;
+	}
+
 	/**
 	 * Provide a consumer to configure the {@link ExecutionInput} used for input to
 	 * {@link graphql.GraphQL#executeAsync(ExecutionInput)}. The builder is initially
@@ -148,16 +154,22 @@ public void configureExecutionInput(BiFunction<ExecutionInput, ExecutionInput.Bu
 	 * populated from {@link #getQuery()}, {@link #getOperationName()}, and
 	 * {@link #getVariables()}, and is then further customized through
 	 * {@link #configureExecutionInput(BiFunction)}.
+	 * @param useRequestId whether the {@link #getId()} should be used as a fallback for {@link ExecutionId}.
 	 * @return the execution input
 	 */
-	public ExecutionInput toExecutionInput() {
-		ExecutionInput executionInput = ExecutionInput.newExecutionInput()
+	public ExecutionInput toExecutionInput(boolean useRequestId) {
+		ExecutionInput.Builder inputBuilder = ExecutionInput.newExecutionInput()
 				.query(this.query)
 				.operationName(this.operationName)
 				.variables(this.variables)
-				.locale(this.locale)
-				.executionId(ExecutionId.from(this.id))
-				.build();
+				.locale(this.locale);
+		if (this.executionId != null) {
+			inputBuilder.executionId(this.executionId);
+		}
+		else if (useRequestId) {
+			inputBuilder.executionId(ExecutionId.from(this.id));
+		}
+		ExecutionInput executionInput = inputBuilder.build();
 
 		for (BiFunction<ExecutionInput, ExecutionInput.Builder, ExecutionInput> configurer : this.executionInputConfigurers) {
 			ExecutionInput current = executionInput;

spring-graphql/src/main/java/org/springframework/graphql/execution/ExecutionGraphQlService.java

@@ -22,6 +22,7 @@
 import graphql.ExecutionInput;
 import graphql.GraphQL;
 import graphql.GraphQLContext;
+import graphql.execution.ExecutionIdProvider;
 import org.dataloader.DataLoaderRegistry;
 import reactor.core.publisher.Mono;
 
@@ -42,9 +43,12 @@ public class ExecutionGraphQlService implements GraphQlService {
 
 	private final List<DataLoaderRegistrar> dataLoaderRegistrars = new ArrayList<>();
 
+	private final boolean hasDefaultExecutionIdProvider;
+
 
 	public ExecutionGraphQlService(GraphQlSource graphQlSource) {
 		this.graphQlSource = graphQlSource;
+		this.hasDefaultExecutionIdProvider = ExecutionIdProvider.DEFAULT_EXECUTION_ID_PROVIDER == graphQlSource.graphQl().getIdProvider();
 	}
 
 
@@ -61,7 +65,7 @@ public void addDataLoaderRegistrar(DataLoaderRegistrar registrar) {
 	@Override
 	public final Mono<RequestOutput> execute(RequestInput requestInput) {
 		return Mono.deferContextual((contextView) -> {
-			ExecutionInput executionInput = requestInput.toExecutionInput();
+			ExecutionInput executionInput = requestInput.toExecutionInput(this.hasDefaultExecutionIdProvider);
 			ReactorContextManager.setReactorContext(contextView, executionInput);
 			ExecutionInput updatedExecutionInput = registerDataLoaders(executionInput);
 			return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(updatedExecutionInput))

spring-graphql/src/test/java/org/springframework/graphql/RequestInputTests.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2020-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql;
+
+import graphql.execution.ExecutionId;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link RequestInput}.
+ *
+ * @author Brian Clozel
+ */
+class RequestInputTests {
+
+	private RequestInput requestInput = new RequestInput("greeting", "Greeting", null, null, "id");
+
+	@Test
+	void shouldUseCustomExecutionIdIfPresent() {
+		ExecutionId customId = ExecutionId.from("customId");
+		this.requestInput.executionId(customId);
+		assertThat(this.requestInput.toExecutionInput(true).getExecutionId()).isEqualTo(customId);
+		assertThat(this.requestInput.toExecutionInput(false).getExecutionId()).isEqualTo(customId);
+	}
+
+	@Test
+	void executionIdShouldFallBackToRequestId() {
+		assertThat(this.requestInput.toExecutionInput(true).getExecutionId()).isEqualTo(ExecutionId.from("id"));
+	}
+
+	@Test
+	void executionIdShouldFallBackToProvider() {
+		assertThat(this.requestInput.toExecutionInput(false).getExecutionId()).isNull();
+	}
+}

spring-graphql/src/test/java/org/springframework/graphql/web/WebInterceptorTests.java

@@ -79,7 +79,7 @@ void executionInputCustomization() {
 
 		WebGraphQlHandler handler = WebGraphQlHandler
 				.builder((input) -> {
-					actualName.set(input.toExecutionInput().getOperationName());
+					actualName.set(input.toExecutionInput(true).getOperationName());
 					return emptyExecutionResult(input);
 				})
 				.interceptor((webInput, next) -> {