Response builder which wraps around JAX-RS responses, provides built in error handling.

ApiResponseBuilder.java

import javax.ws.rs.core.Response;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

import static javax.ws.rs.core.MediaType.TEXT_PLAIN_TYPE;
import static javax.ws.rs.core.Response.Status.INTERNAL_SERVER_ERROR;
import static javax.ws.rs.core.Response.Status.OK;

/**
 * Examples:
 *
 *    // If getResponseFromSomewhereThatMightThrowAnError runs without error: object returned by response supplier is used as the response entity with a status of OK
 *    // If getResponseFromSomewhereThatMightThrowAnError throws an error: error message as plain text is used as the response entity with a status of INTERNAL_SERVER_ERROR
 *    ApiResponseBuilder.apiResponse()
 *      .response(() -> getResponseFromSomewhereThatMightThrowAnError())
 *      .successStatus(Status.OK) // successStatus is optional here, default is OK
 *      .build()
 *
 *    // If the below code runs without error: object returned by response lambda is used as the response entity with a status of OK
 *    // If getResponseFromSomewhereThatMightThrowAnError throws an error: generate error response using buildErrorResponseForException and set status using errorStatus supplier
 *    ApiResponseBuilder.apiResponse()
 *      .response(() -> getResponseFromSomewhereThatMightThrowAnError())
 *      .successStatus(Status.OK)
 *      .errorResponse((e) -> buildErrorResponseForException(e))
 *      .errorStatus((e) -> e instanceof InvalidUserInputException ? Status.BAD_REQUEST : Status.INTERNAL_SERVER_ERROR)
 *      .build()
 *
 *    // If the below code runs without error: no response entity set and a status of OK returned
 *    // If requestProcessor runnable throws an error: set status using errorStatus supplier
 *    ApiResponseBuilder.apiResponse()
 *      .requestProcessor(() -> processRequest()) // used to handle errors from logic that does not return an exception, executed even if response is specified
 *      .errorStatus((e) -> e instanceof InvalidUserInputException ? Status.BAD_REQUEST : Status.INTERNAL_SERVER_ERROR)
 *      .build()
 *
 *    ApiResponseBuilder.apiResponse()
 *      .customResponse(() -> Response.ok().build()) // return a JAX-RS Response object directly (with same error handling as above)
 *      .build()
 */
public class ApiResponseBuilder {
    private Optional<Runnable> requestProcessor;
    private Optional<Supplier<Object>> responseObjectSupplier;
    private Optional<Supplier<Response>> responseSupplier;
    private Optional<Supplier<Response.Status>> successSupplier;
    private Optional<Function<Throwable, Object>> errorResponseFunction;
    private Optional<Function<Throwable, Response.Status>> errorStatusFunction;
    private Optional<Consumer<Response.ResponseBuilder>> responseCustomiser;

    private ApiResponseBuilder() {
        requestProcessor = Optional.empty();
        responseObjectSupplier = Optional.empty();
        responseSupplier = Optional.empty();
        successSupplier = Optional.empty();
        errorResponseFunction = Optional.empty();
        errorStatusFunction = Optional.empty();
        responseCustomiser = Optional.empty();
    }

    public static ApiResponseBuilder apiResponse() {
        return new ApiResponseBuilder();
    }

    public ApiResponseBuilder requestProcessor(Runnable requestProcessor) {
        this.requestProcessor = Optional.ofNullable(requestProcessor);
        return this;
    }

    public ApiResponseBuilder response(Object response) {
        this.responseObjectSupplier = Optional.of(() -> responseObjectSupplier);
        return this;
    }

    public ApiResponseBuilder response(Supplier<Object> responseSupplier) {
        this.responseObjectSupplier = Optional.ofNullable(responseSupplier);
        return this;
    }

    public ApiResponseBuilder customResponse(Supplier<Response> responseSupplier) {
        this.responseSupplier = Optional.ofNullable(responseSupplier);
        return this;
    }

    public ApiResponseBuilder successStatus(Response.Status status) {
        this.successSupplier = Optional.of(() -> status);
        return this;
    }

    public ApiResponseBuilder successStatus(Supplier<Response.Status> successStatusSupplier) {
        this.successSupplier = Optional.ofNullable(successStatusSupplier);
        return this;
    }


    public ApiResponseBuilder errorResponse(Object errorResponse) {
        this.errorResponseFunction = Optional.of((ex) -> errorResponse);
        return this;
    }

    public ApiResponseBuilder errorResponse(Function<Throwable, Object> errorResponseSupplier) {
        this.errorResponseFunction = Optional.ofNullable(errorResponseSupplier);
        return this;
    }

    public ApiResponseBuilder errorStatus(Response.Status errorStatus) {
        this.errorStatusFunction = Optional.of((ex) -> errorStatus);
        return this;
    }

    public ApiResponseBuilder errorStatus(Function<Throwable, Response.Status> errorStatusSupplier) {
        this.errorStatusFunction = Optional.ofNullable(errorStatusSupplier);
        return this;
    }

    public ApiResponseBuilder customiseResponse(Consumer<Response.ResponseBuilder> responseConsumer) {
        this.responseCustomiser = Optional.ofNullable(responseConsumer);
        return this;
    }

    public Response build() {
        Optional<Object> apiResponse = Optional.empty();
        Optional<Throwable> error = Optional.empty();

        try {
            requestProcessor.ifPresent(Runnable::run);

            if (responseSupplier.isPresent()) {
                // get custom response
                return responseSupplier.get().get();
            }

            if (responseObjectSupplier.isPresent()) {
                // get normal response
                Supplier<Object> responseObjectSupplier = this.responseObjectSupplier.get();

                apiResponse = Optional.ofNullable(responseObjectSupplier.get());
            }
        } catch (Exception ex) {
            error = Optional.of(ex);

            if (errorResponseFunction.isPresent()) {
                // get error response
                Function<Throwable, Object> errorResponseObjectFunction = errorResponseFunction.get();

                apiResponse = Optional.ofNullable(errorResponseObjectFunction.apply(ex));
            }
        }

        Response.ResponseBuilder responseBuilder;

        if (!error.isPresent()) {
            // return normal status code
            Supplier<Response.Status> statusSupplier = successSupplier.orElse(() -> OK);

            responseBuilder = Response.status(statusSupplier.get());
        } else {
            // return error status code
            Function<Throwable, Response.Status> statusFunction = errorStatusFunction.orElse((ex) -> INTERNAL_SERVER_ERROR);

            responseBuilder = Response.status(statusFunction.apply(error.get()));
        }

        if (apiResponse.isPresent()) {
            // If return value from responseObjectSupplier/errorResponseObjectFunction is not
            // null, use it as the response entity.
            responseBuilder = responseBuilder.entity(apiResponse.get());
        } else if (error.isPresent()) {
            // When responseSupplier throws an exception or errorResponseObjectFunction is not
            // present/returns null, set response body to the exception message as plain text.
            responseBuilder = responseBuilder.entity(error.get().getMessage()).type(TEXT_PLAIN_TYPE);
        }

        if (responseCustomiser.isPresent()) {
            // allow user to customise response before building
            Consumer<Response.ResponseBuilder> responseConsumer = responseCustomiser.get();

            responseConsumer.accept(responseBuilder);
        }

        return responseBuilder.build();
    }
}