Add interception for RSocket handler

See gh-339

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/support/DefaultExecutionGraphQlResponse.java

@@ -23,6 +23,7 @@
 import graphql.ErrorClassification;
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
+import graphql.ExecutionResultImpl;
 import graphql.GraphQLError;
 import graphql.language.SourceLocation;
 
@@ -153,4 +154,65 @@ public Map<String, Object> getExtensions() {
 
 	}
 
+
+	/**
+	 * Builder to transform the response's {@link ExecutionResult}.
+	 */
+	public static abstract class Builder<B extends Builder<B, R>, R extends ExecutionGraphQlResponse> {
+
+		private final R original;
+
+		private final ExecutionResultImpl.Builder executionResultBuilder;
+
+		protected Builder(R original) {
+			this.original = original;
+			this.executionResultBuilder = ExecutionResultImpl.newExecutionResult().from(original.getExecutionResult());
+		}
+
+		/**
+		 * Set the {@link ExecutionResult#getData() data} of the GraphQL execution result.
+		 * @param data the execution result data
+		 * @return the current builder
+		 */
+		public Builder<B, R> data(Object data) {
+			this.executionResultBuilder.data(data);
+			return this;
+		}
+
+		/**
+		 * Set the {@link ExecutionResult#getErrors() errors} of the GraphQL execution
+		 * result.
+		 * @param errors the execution result errors
+		 * @return the current builder
+		 */
+		public Builder<B, R> errors(@Nullable List<GraphQLError> errors) {
+			this.executionResultBuilder.errors(errors);
+			return this;
+		}
+
+		/**
+		 * Set the {@link ExecutionResult#getExtensions() extensions} of the GraphQL
+		 * execution result.
+		 * @param extensions the execution result extensions
+		 * @return the current builder
+		 */
+		public Builder<B, R> extensions(@Nullable Map<Object, Object> extensions) {
+			this.executionResultBuilder.extensions(extensions);
+			return this;
+		}
+
+		/**
+		 * Build the response with the transformed {@code ExecutionResult}.
+		 */
+		public R build() {
+			return build(this.original, this.executionResultBuilder.build());
+		}
+
+		/**
+		 * Subclasses to create the specific response instance.
+		 */
+		protected abstract R build(R original, ExecutionResult newResult);
+
+	}
+
 }

spring-graphql/src/main/java/org/springframework/graphql/web/DefaultWebGraphQlHandlerBuilder.java

@@ -25,6 +25,7 @@
 import org.springframework.graphql.ExecutionGraphQlService;
 import org.springframework.graphql.execution.ReactorContextManager;
 import org.springframework.graphql.execution.ThreadLocalAccessor;
+import org.springframework.graphql.web.WebGraphQlHandlerInterceptor.Chain;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -88,12 +89,11 @@ public WebGraphQlHandler.Builder threadLocalAccessors(List<ThreadLocalAccessor>
 	@Override
 	public WebGraphQlHandler build() {
 
-		WebGraphQlHandlerInterceptor.Chain endOfChain =
-				request -> this.service.execute(request).map(WebGraphQlResponse::new);
+		Chain endOfChain = request -> this.service.execute(request).map(WebGraphQlResponse::new);
 
-		WebGraphQlHandlerInterceptor.Chain chain = this.interceptors.stream()
+		Chain chain = this.interceptors.stream()
 				.reduce(WebGraphQlHandlerInterceptor::andThen)
-				.map(interceptor -> (WebGraphQlHandlerInterceptor.Chain) (request) -> interceptor.intercept(request, endOfChain))
+				.map(interceptor -> (Chain) (request) -> interceptor.intercept(request, endOfChain))
 				.orElse(endOfChain);
 
 		return new WebGraphQlHandler() {
@@ -112,7 +112,8 @@ public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
 
 			@Override
 			public WebSocketGraphQlHandlerInterceptor webSocketInterceptor() {
-				return (webSocketInterceptor != null ? webSocketInterceptor : new WebSocketGraphQlHandlerInterceptor() {});
+				return (webSocketInterceptor != null ?
+						webSocketInterceptor : new WebSocketGraphQlHandlerInterceptor() {});
 			}
 
 		};

spring-graphql/src/main/java/org/springframework/graphql/web/GraphQlRSocketHandler.java

@@ -17,6 +17,7 @@
 package org.springframework.graphql.web;
 
 
+import java.util.List;
 import java.util.Map;
 
 import graphql.ExecutionResult;
@@ -25,38 +26,39 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.ExecutionGraphQlResponse;
 import org.springframework.graphql.ExecutionGraphQlService;
-import org.springframework.graphql.support.DefaultExecutionGraphQlRequest;
+import org.springframework.graphql.web.RSocketGraphQlHandlerInterceptor.Chain;
+import org.springframework.util.AlternativeJdkIdGenerator;
+import org.springframework.util.IdGenerator;
 
 
 /**
  * Handler for GraphQL over RSocket requests.
  *
- * <p>This class can be extended from an {@code @Controller} that overrides
- * {@link #handle(Map)} and {@link #handleSubscription(Map)} in order to add
+ * <p>This class can be extended or wrapped from an {@code @Controller} in order
+ * to re-declare {@link #handle(Map)} and {@link #handleSubscription(Map)} with
  * {@link org.springframework.messaging.handler.annotation.MessageMapping @MessageMapping}
- * annotations with the route.
+ * annotations including the GraphQL endpoint route.
  *
  * <pre style="class">
  * &#064;Controller
- * private static class GraphQlRSocketController extends GraphQlRSocketHandler {
+ * private static class GraphQlRSocketController {
  *
- *    GraphQlRSocketController(ExecutionGraphQlService graphQlService) {
- *        super(graphQlService);
+ *    private final GraphQlRSocketHandler handler;
+ *
+ *    GraphQlRSocketController(GraphQlRSocketHandler handler) {
+ *        this.handler = handler;
  *    }
  *
- *    &#064;Override
  *    &#064;MessageMapping("graphql")
  *    public Mono<Map<String, Object>> handle(Map<String, Object> payload) {
- *        return super.handle(payload);
+ *        return this.handler.handle(payload);
  *    }
  *
- *    &#064;Override
  *    &#064;MessageMapping("graphql")
  *    public Flux<Map<String, Object>> handleSubscription(Map<String, Object> payload) {
- *        return super.handleSubscription(payload);
+ *        return this.handler.handleSubscription(payload);
  *    }
  * }
  * </pre>
@@ -66,26 +68,40 @@
  */
 public class GraphQlRSocketHandler {
 
-	private final ExecutionGraphQlService service;
+	private final Chain executionChain;
+
+	private final IdGenerator idGenerator = new AlternativeJdkIdGenerator();
+
+
+	/**
+	 * Create a new instance that handles requests through a chain of interceptors
+	 * followed by the given {@link ExecutionGraphQlService}.
+	 */
+	public GraphQlRSocketHandler(
+			ExecutionGraphQlService service, List<RSocketGraphQlHandlerInterceptor> interceptors) {
 
+		Chain endOfChain = request -> service.execute(request).map(RSocketGraphQlResponse::new);
 
-	public GraphQlRSocketHandler(ExecutionGraphQlService service) {
-		this.service = service;
+		this.executionChain = (interceptors.isEmpty() ? endOfChain :
+				interceptors.stream()
+						.reduce(RSocketGraphQlHandlerInterceptor::andThen)
+						.map(interceptor -> (Chain) request -> interceptor.intercept(request, endOfChain))
+						.orElse(endOfChain));
 	}
 
 
 	/**
 	 * Handle a {@code Request-Response} interaction. For queries and mutations.
 	 */
 	public Mono<Map<String, Object>> handle(Map<String, Object> payload) {
-		return this.service.execute(initRequest(payload)).map(ExecutionGraphQlResponse::toMap);
+		return handleInternal(payload).map(ExecutionGraphQlResponse::toMap);
 	}
 
 	/**
 	 * Handle a {@code Request-Stream} interaction. For subscriptions.
 	 */
 	public Flux<Map<String, Object>> handleSubscription(Map<String, Object> payload) {
-		return this.service.execute(initRequest(payload))
+		return handleInternal(payload)
 				.flatMapMany(response -> {
 					if (response.getData() instanceof Publisher) {
 						Publisher<ExecutionResult> publisher = response.getData();
@@ -100,12 +116,9 @@ public Flux<Map<String, Object>> handleSubscription(Map<String, Object> payload)
 				});
 	}
 
-	@SuppressWarnings("unchecked")
-	private ExecutionGraphQlRequest initRequest(Map<String, Object> payload) {
-		String query = (String) payload.get("query");
-		String operationName = (String) payload.get("operationName");
-		Map<String, Object> variables = (Map<String, Object>) payload.get("variables");
-		return new DefaultExecutionGraphQlRequest(query, operationName, variables, "1", null);
+	private Mono<RSocketGraphQlResponse> handleInternal(Map<String, Object> payload) {
+		String requestId = this.idGenerator.generateId().toString();
+		return this.executionChain.next(new RSocketGraphQlRequest(payload, requestId, null));
 	}
 
 }

spring-graphql/src/main/java/org/springframework/graphql/web/RSocketGraphQlHandlerInterceptor.java

@@ -0,0 +1,75 @@
+/*
+ * 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.web;
+
+import graphql.ExecutionInput;
+import graphql.ExecutionResult;
+import reactor.core.publisher.Mono;
+
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.graphql.ExecutionGraphQlService;
+
+
+/**
+ * Interceptor for server handling of GraphQL over RSocket requests,
+ * allowing customization of the {@link ExecutionInput} and
+ * the {@link ExecutionResult}.
+ *
+ * <p>Interceptors are typically declared as beans in Spring configuration and
+ * ordered as defined in {@link ObjectProvider#orderedStream()}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public interface RSocketGraphQlHandlerInterceptor {
+
+	/**
+	 * Intercept a request and delegate to the rest of the chain including other
+	 * interceptors and a {@link ExecutionGraphQlService}.
+	 * @param request the request to execute
+	 * @param chain the rest of the chain to execute the request
+	 * @return a {@link Mono} with the response
+	 */
+	Mono<RSocketGraphQlResponse> intercept(RSocketGraphQlRequest request, Chain chain);
+
+	/**
+	 * Return a new {@link RSocketGraphQlHandlerInterceptor} that invokes the current
+	 * interceptor first and then the one that is passed in.
+	 * @param nextInterceptor the interceptor to delegate to after the current
+	 * @return a new interceptor that chains the two
+	 */
+	default RSocketGraphQlHandlerInterceptor andThen(RSocketGraphQlHandlerInterceptor nextInterceptor) {
+		return (request, chain) -> intercept(request, nextRequest -> nextInterceptor.intercept(nextRequest, chain));
+	}
+
+
+	/**
+	 * Contract for delegation to the rest of the chain.
+	 */
+	interface Chain {
+
+		/**
+		 * Delegate to the rest of the chain to execute the request.
+		 * @param request the request to execute
+		 * the {@link ExecutionInput} for {@link graphql.GraphQL}.
+		 * @return {@code Mono} with the response
+		 */
+		Mono<RSocketGraphQlResponse> next(RSocketGraphQlRequest request);
+
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/web/RSocketGraphQlRequest.java

@@ -0,0 +1,57 @@
+/*
+ * 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.web;
+
+import java.util.Locale;
+import java.util.Map;
+
+import io.rsocket.exceptions.RejectedException;
+
+import org.springframework.graphql.ExecutionGraphQlRequest;
+import org.springframework.graphql.support.DefaultExecutionGraphQlRequest;
+import org.springframework.lang.Nullable;
+import org.springframework.util.StringUtils;
+
+/**
+ * {@link org.springframework.graphql.GraphQlRequest} implementation for server
+ * handling over RSocket.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public class RSocketGraphQlRequest extends DefaultExecutionGraphQlRequest implements ExecutionGraphQlRequest {
+
+
+	/**
+	 * Create an instance.
+	 * @param body the deserialized content of the GraphQL request
+	 * @param id an identifier for the GraphQL request
+	 * @param locale the locale from the HTTP request, if any
+	 */
+	public RSocketGraphQlRequest(Map<String, Object> body, String id, @Nullable Locale locale) {
+		super(getKey("query", body), getKey("operationName", body), getKey("variables", body), id, locale);
+	}
+
+	@SuppressWarnings("unchecked")
+	private static <T> T getKey(String key, Map<String, Object> body) {
+		if (key.equals("query") && !StringUtils.hasText((String) body.get(key))) {
+			throw new RejectedException("No \"query\" in the request document");
+		}
+		return (T) body.get(key);
+	}
+
+}