Spring Boot - update response of every API using ResponseBodyAdvice

Not sure why we need it but I've seen several codebase where developers were wrapping the response body in some structure like below - where the actual content is put under an element eg: body and the root object has several other values.

{
"body": { //the actual body
"greeting": "Hello World!"
},

//additional elements
"timestamp": 1741653577468,
"traceId": "67cf8649f3c848a71d8171473d3b4955",
"spanId": "1d8171473d3b4955",
"durationMs": 0,
"status": 200
}

The corresponding controller method would look something like below where they would explicitly create the ResponseObject and map other parameters in every single endpoint definition:

@GetMapping("/object-bad-way/hi")
public ResponseObject<Hi> oldWayHi() {
log.info("Got request - object - old way - don't do this, use ResponseBodyAdvice");
var resp = new ResponseObject<Hi>();
var span = tracer.currentSpan();
if (span != null) {
resp.setTraceId(span.context().traceId());
resp.setSpanId(span.context().spanId());
}

resp.setTimestamp(System.currentTimeMillis());
try {
resp.setBody(myService.getHi()); //service call
resp.setStatus(HttpStatus.OK.value());
} catch (Exception e) {
resp.setStatus(HttpStatus.INTERNAL_SERVER_ERROR.value());
}

return resp;
}

The ResponseObject:

@Getter
@Setter
class ResponseObject<T> {
T body;
long timestamp;

String traceId;
String spanId;

long durationMs;
int status;
}


But wait, there's a better way to do this if you really want this. The idea is to use Spring's @ControllerAdvice that implements ResponseBodyAdvice to map the additional params so that you don't need to create the new ResponseObject instance on all endpoints.

After the change, your controller method would look simple as below: wouldn't that be nice?

@GetMapping("/better-way/hi")
public Hi objectHi() {
log.info("Got request - object");
return service.getHi();
}

 

ResponseBodyAdvice

From the docs:

ResponseBodyAdvice allows customizing the response after the execution of an @ResponseBody or a ResponseEntity controller method but before the body is written with an HttpMessageConverter.
Implementations may be registered directly with RequestMappingHandlerAdapter and ExceptionHandlerExceptionResolver or more likely annotated with @ControllerAdvice in which case they will be auto-detected by both.

 

A simple (empty) usage would look like this. On the supports method, we can decide if the response body needs to be modified. The beforeBodyWrite is where we can modify the request body or create new.

 

@ControllerAdvice
class ObjectResponseAdvice implements ResponseBodyAdvice<Object> {

//Whether this component supports the given controller method return type and the selected HttpMessageConverter type.
public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
return true;
}

//Invoked after an HttpMessageConverter is selected and just before its write method is invoked.
public Object beforeBodyWrite(Object body, MethodParameter returnType,
MediaType selectedContentType, Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
//modify body
return body;
}
}

 

A simple implementation of beforeBodyWrite that would work with both String response type and any other response type. Here we are handling the String response specially by checking if the original body(response body) is String. Otherwise, we are creating a ObjectNode and mapping timestamp and body elements under it.

 

supports() method: 
    return true;


beforeBodyWrit() method:
  ObjectNode root = objectMapper.createObjectNode();
root.put("timestamp", System.currentTimeMillis());

if (body instanceof String) {
root.set("body", objectMapper.convertValue(body, JsonNode.class));
return objectMapper.writeValueAsString(root);
} else {
root.put("body", String.valueOf(body));
return root;
}

 

Handle more stuff and map more stuff:

Here we are mapping the traceId/spanId from Micrometer Tracing , mapping 'status' from ServletServerHttpResponse. Also for backward compatibility with existing RestController methods that already return ResponseObject (with all elements we want on the response), we are skipping the additional mapping and simply returning body.

if (response instanceof ServletServerHttpResponse resp
//if body is already ResponseObject - do not convert
&& !(body instanceof ResponseObject)) {

ObjectNode root = objectMapper.createObjectNode();
var span = tracer.currentSpan();
if (span != null) {
root.put("traceId", span.context().traceId());
root.put("spanId", span.context().spanId());
}
root.put("status", resp.getServletResponse().getStatus());
if (BeanUtils.isSimpleValueType(body.getClass())) {
root.put("body", String.valueOf(body));
} else {
root.set("body", objectMapper.convertValue(body, JsonNode.class));
}
root.put("timestamp", System.currentTimeMillis());

if (body instanceof String) {
// String are handled special case with StringHttpMessageConverter
// we need to return String from this method
return objectMapper.writeValueAsString(root);
} else {
return root;
}
}

 

 

Full code in action: Note that it supports, string type, int type, map, and any object response.

 

Test it out with following endpoints:

###
GET http://localhost:8080/string

###
GET http://localhost:8080/map

###
GET http://localhost:8080/object/hi

###
GET http://localhost:8080/object/hello

###
GET http://localhost:8080/object-bad-way/hi

###
GET http://localhost:8080/object-bad-way/hello

 

 Full Code:

package responsebody;

import com.fasterxml.jackson.databind.*;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.micrometer.tracing.Tracer;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.core.MethodParameter;
import org.springframework.http.*;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.*;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.net.*;
import java.util.Map;

@SpringBootApplication
public class DemoApplication {

public static void main(String[] args) {
SpringApplication.run(DemoApplication.class, args);
}
}

@RestController
@Slf4j
@RequiredArgsConstructor
class GreetingController {

final Tracer tracer;

@GetMapping("/string")
public String string() {
log.info("Got request - string");
return "Hello World!";
}

@GetMapping("/int")
public int intVal() {
log.info("Got request - int");
return 1;
}

@GetMapping("/url")
public URL url() throws MalformedURLException {
log.info("Got request - url");
return URI.create("https://localhost:8080/").toURL();
}

@GetMapping("/map")
public Map<String, String> map() {
log.info("Got request - map");
return Map.of("greeting", "Hello World!");
}

@GetMapping("/object/hello")
public Hello objectHello() {
log.info("Got request - object");
return new Hello("Hello World!");
}

@GetMapping("/better-way/hi")
public Hi objectHi() {
log.info("Got request - object");
return new Hi("Hi World!");
}

@GetMapping("/object-bad-way/hello")
public ResponseObject<Hello> oldWay() {
log.info("Got request - object - old way - don't do this, use ResponseBodyAdvice");
var resp = new ResponseObject<Hello>();
var span = tracer.currentSpan();
if (span != null) {
resp.setTraceId(span.context().traceId());
resp.setSpanId(span.context().spanId());
}

resp.setTimestamp(System.currentTimeMillis());
try {
resp.setBody(new Hello("Hello World!")); //service call
resp.setStatus(HttpStatus.OK.value());
} catch (Exception e) {
resp.setStatus(HttpStatus.INTERNAL_SERVER_ERROR.value());
}

return resp;
}

@GetMapping("/object-bad-way/hi")
public ResponseObject<Hi> oldWayHi() {
log.info("Got request - object - old way - don't do this, use ResponseBodyAdvice");
var resp = new ResponseObject<Hi>();
var span = tracer.currentSpan();
if (span != null) {
resp.setTraceId(span.context().traceId());
resp.setSpanId(span.context().spanId());
}

resp.setTimestamp(System.currentTimeMillis());
try {
resp.setBody(new Hi("Hello World!")); //service call
resp.setStatus(HttpStatus.OK.value());
} catch (Exception e) {
resp.setStatus(HttpStatus.INTERNAL_SERVER_ERROR.value());
}

return resp;
}


}


@Slf4j
@ControllerAdvice
@RequiredArgsConstructor
class ObjectResponseAdvice implements ResponseBodyAdvice<Object> {

final Tracer tracer;
final ObjectMapper objectMapper;

@Override
public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
return true;
}

@SneakyThrows
@Override
public Object beforeBodyWrite(Object body, MethodParameter returnType,
MediaType selectedContentType, Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {


if (response instanceof ServletServerHttpResponse resp
//if body is already ResponseObject - do not convert
&& !(body instanceof ResponseObject)) {

ObjectNode root = objectMapper.createObjectNode();
var span = tracer.currentSpan();
if (span != null) {
root.put("traceId", span.context().traceId());
root.put("spanId", span.context().spanId());
}
root.put("status", resp.getServletResponse().getStatus());
if (BeanUtils.isSimpleValueType(body.getClass())) {
root.put("body", String.valueOf(body));
} else {
root.set("body", objectMapper.convertValue(body, JsonNode.class));
}
root.put("timestamp", System.currentTimeMillis());

if (body instanceof String) {
// String are handled special case with StringHttpMessageConverter
// we need to return String from this method
return objectMapper.writeValueAsString(root);
} else {
return root;
}
}

return body;
}
}


@Getter
@NoArgsConstructor
@AllArgsConstructor
class Hello {
String greeting;
}

@Getter
@NoArgsConstructor
@AllArgsConstructor
class Hi {
String hi;
}

@Getter
@Setter
class ResponseObject<T> {
T body;
long timestamp;

String traceId;
String spanId;

long durationMs;
int status;
}