Spring Webflux — wrap response

Spring webflux — the reactive universe of the Spring framework is very interesting place.

This post is dedicated to wrapping around the response to a common response that solves the purpose of having common fields which are always required for all API endpoints.

The code repository is here — https://github.com/krnbr/spring-webflux-wrap-response

Here I make use of the ResponseBodyResultHandler to wrap the request inside a common structure of the response if the api mapping method or the controller class is mapped with a specific @ApiResponse annotation.

The example code in the repository shared above works well in the case of the Mono and non Mono wrapped objects, for the Flux based objects there will be required some additional changes.

The main code is in the class — CustomResponseBodyResultHandler

@Slf4j
public class CustomResponseBodyResultHandler extends ResponseBodyResultHandler {

    public CustomResponseBodyResultHandler(List<HttpMessageWriter<?>> writers, RequestedContentTypeResolver resolver) {
        super(writers, resolver);
    }

    @Override
    public boolean supports(HandlerResult result) {
        var className = result.getReturnTypeSource().getDeclaringClass().getName();
        var methodName = result.getReturnTypeSource().getMethod().getName();
        var classAnnotations = result.getReturnTypeSource().getDeclaringClass().getAnnotations();
        var methodAnnotations = result.getReturnTypeSource().getMethodAnnotations();
        var annotations = result.getReturnTypeSource().getDeclaringClass().getAnnotations();

        if (Arrays.stream(classAnnotations).anyMatch(a -> a.annotationType() == ApiResponse.class)) {
            log.info("{} is marked with ApiResponse annotation", className);
            return true;
        } else if (Arrays.stream(methodAnnotations).anyMatch(a -> a.annotationType() == ApiResponse.class)) {
            log.info("{} inside {} is marked with ApiResponse annotation", methodName, className);
            return true;
        }

        return false;
    }

    @Override
    public Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result) {
        ServiceResponse s = new ServiceResponse();
        s.setMethod(exchange.getRequest().getMethod().name());
        s.setStatus(exchange.getResponse().getStatusCode().value());
        s.setCorrelationId(exchange.getAttribute("correlation-id"));
        var adapter = getAdapter(result);
        // modify the result as you want
        if (adapter != null) { // if the response was wrapped inside Mono?
            Mono<ServiceResponse> body = ((Mono<Object>) result.getReturnValue()).map(o -> {
                s.setData(o);
                return s;
            });
            return writeBody(body, result.getReturnTypeSource().nested(), exchange);
        } else { // if the response was not wrapped inside Mono
            s.setData(result.getReturnValue());
            Mono<ServiceResponse> body = Mono.just(s);
            return writeBody(body, result.getReturnTypeSource().nested(), exchange);
        }
    }
}

It is implementation of the ResponseBodyResultHandler

• it runs after the controller has finished processing
• there is a method — supports, to make logical evaluation if the handler is to be used or not
• other method is the main one handleResult — where we construct the wrapper response, currently the code only supports Mono or normal responses from the controller, but support for Flux can also be added for sure.

For any questions, queries, feedback please comment below or contact me over LinkedIn here

Karanbir Singh — Software Engineer — Telstra | LinkedIn