Spring Boot + ControllerAdvice + ResponseBodyAdvice или как обернуть ответ контроллеров

Введение

Всем привет, друзья! Сегодня хочу рассказать про способ использования ControllerAdvice для оборачивания объекта, возвращаемого контроллерами, в новый класс на уровне DispatcherServlet.

Пример:

Допустим, некоторый метод отдавал информацию о пользователе

{
        "name": "Ivan",
        "surname": "Ivanov"
}

И есть еще десяток методов, которые отдают некоторую информацию о пользователе

Но теперь мы хотим, чтобы каждый метод отдавал дополнительно еще несколько общих полей (например, серию и номер паспорта)

{
        "name": "Ivan",
        "surname": "Ivanov",
        "passport": "1111 111111"
}

Я расскажу про интересное применение ControllerAdvice и покажу один из способов, которым можно решить такую задачу

P.S. В жизни такой подход был удобен, когда в проекте существовала отдельная библиотека, импортирующая в микросервисы модель данных API, а по бизнес требованиям стало необходимо добавить в некоторых микросервисах к моделям дополнительные данные. Возможны и иные кейсы

Решение задачи

Решение удобно при его многократном использовании. В одном сервисе будет удобнее, скорее всего, использовать иной подход. Поэтому будем писать стартер

Наши задачи:

1. Создать удобный способ использования стартера в коде - через аннотации

2. Создать возможность гибкой настройки добавляемых данных в класс-обертку

3. Создать класс с @ControllerAdvicе, обрабатывающий методы контроллеров

4. Собрать все в стартер Spring Boot

Аннотации

Для начала создадим 2 аннотации, который будут включать оборачивание для контроллера и выключать его для конкретного метода

Аннотация, включающая обработку методов

@RestController
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
public @interface EnableResponseWrapper {
    Class<? extends IWrapperModel> wrapperClass();
}

Так как наша аннотация всегда висит над контроллером, а у аннотаций в java не существует понятия наследования - то вешаем над нашей аннотацией @RestController - это позволит использовать @EnableResponseWrapper вместо @RestController

@Target(ElementType.TYPE) - указывает на то, что наша аннотация может висеть над классом, интерфейсом или enum-ом

@Retention(RetentionPolicy.RUNTIME) - указывает область видимости аннотации - во время выполнения кода

Аннотация, в качестве аргумента, принимает класс, который описывает оболочку Class<? extends IWrapperModel> wrapperClass();

Аннотация, отключающая обработку метода

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface DisableResponseWrapper {
}

Сами по себе аннотации, безусловно, никакой функциональности не добавляют. Обработкой аннотаций займемся позднее

Интерфейсы

Для работы нам понадобится сущность, возвращающая дополнительные данные и сущность, описывающая обертку. Создадим интерфейсы

Интерфейс сервиса, через который будем получать данные для наполнения обертки

@Service
public interface IWrapperService {
    Object getData(Object body);
}

Через метод getData(Object body) будем получать данные, затем кладем в класс-обертку.

Интерфейс класса-обертки

public interface IWrapperModel {
    void setData(Object object);
    void setBody(Object object);
}

Через метод setData(Object object) устанавливаем те данные, которые получили в методе getData(Object object).
Через setBody(Object object) устанавливаем объект-ответ, который вернул обрабатываемый метод

Эдвайс

Создадим основной класс стартера, обрабатывающий методы контроллеров

UPD: Инжект бина сервиса происходит через внутренние функции спринга. В случае, если интерфейс реализуют несколько классов - при инжекте произойдет ошибка. В проекте возможно решить эту проблему с помощью @Primary . Если вам необходима различная реализация сервиса - возможна доработка проекта с добавлением передачи сервиса вместе с классом-обертки, инжектом мапы сервисов либо иными более простыми способами решения.
Данный pet проект я не стал дорабатывать из-за нежелания усложнять код - в первую очередь цели проекта и статьи - рассказать об интересном решении, а не создание библиотеки. Выражаю благодарность пользователю, нашедшему данное неудобство @maxzh83

@AllArgsConstructor
@ControllerAdvice(annotations = EnableResponseWrapper.class)
public class ResponseWrapperAdvice implements ResponseBodyAdvice<Object> {
	...
}

Аннотация
@ControllerAdvice(annotations = EnableResponseWrapper.class) указывает, что методы данного компонента будут использоваться сразу несколькими контроллерами. Также указываем, что наши методы будут обрабатывать только те контроллеры, которые помечены @EnableResponseWrapper

Класс реализует интерфейс ResponseBodyAdvice<>, который позволяет настраивать ответ, после его возвращения методом @ResponseBody или контроллером ResponseEntity, но до того, как тело будет записано с помощью HttpMessageConverter

В классе необходимо реализовать 2 метода

Первый позволяет отсеить методы на обрабатываемые и не обрабатываемые
Именно в нем мы проверим, не аннотирован ли случайно наш метод @DisableResponseWrapper. Получим все аннотации, которые висят над методом, и поищем среди них нужную нам аннотацию.

...
@Override
public boolean supports(MethodParameter returnType, @NonNull Class converterType) {
		for (Annotation a : returnType.getMethodAnnotations()) {
				if (a.annotationType() == DisableResponseWrapper.class) {
						return false;
				}
		}

		return true;
}
...

Второй метод класса вызывается только для тех методов, для которых метод supports возвращает true

...
@SneakyThrows
@Override
public Object beforeBodyWrite(
		@Nullable Object body,
		@NonNull MethodParameter returnType,
		@NonNull MediaType selectedContentType,
		@NonNull Class selectedConverterType,
		@NonNull ServerHttpRequest request,
		@NonNull ServerHttpResponse response
		) {
				if (body == null) {
						return null;
        }

        // получаем wrapperClass из аннотации
        Class<? extends IWrapperModel> wrapperClass = null;
        for (Annotation annotation : returnType.getContainingClass().getAnnotations()) {
            if (annotation.annotationType() == EnableResponseWrapper.class) {
                wrapperClass = ((EnableResponseWrapper) annotation).wrapperClass();
                break;
            }
        }

				if (wrapperClass == null) {
            return body;
        }
        ...

Достаем класс-обертку из аннотации

Далее будем работать с объектом, который возвращает наш метод (в методе beforeBodyWrite() он передается первым параметром Object body)

Рассмотрим две ситуации: когда метод возвращает коллекцию и когда возвращает единичный объект. В случае коллекции мы хотим, чтобы был обернут каждый объект коллекции:

...
// проверяем, был ли передан Collection или наследник Collection
if (Collection.class.isAssignableFrom(body.getClass())) {
		try {
				Collection<?> bodyCollection = (Collection<?>) body;

				// проверяем, что collection не пустой
				if (bodyCollection.isEmpty()) {
						return body;
				}
				// оборачиваем каждый элемент коллекции
				return generateListOfResponseWrapper(bodyCollection, wrapperClass);
		} catch (Exception e) {
				return body;
		}
}
...

И если обрабатываемый метод отдает не коллекцию:

...
return generateResponseWrapper(body, wrapperClass);
...

Функции generateListOfResponseWrapper и generateResponseWrapper генерируют обертку для коллекции и для единичного элемента:

...
private List<IWrapperModel> generateListOfResponseWrapper(Collection<?> bodyCollection, Class<? extends IWrapperModel> wrapperClass) {
		return bodyCollection.stream()
				.map((t) -> t == null ?
						null :
						generateResponseWrapper(t, wrapperClass)
				)
						.collect(Collectors.toList());
}
...

...
@SneakyThrows
private IWrapperModel generateResponseWrapper(Object body, Class<? extends IWrapperModel> wrapperClass) {
		// wrapperClass должен иметь конструктор без параметров - получаем объект класса, реализующего IWrapperModel
		IWrapperModel wrapper = wrapperClass.getDeclaredConstructor().newInstance();
		wrapper.setBody(body);
		wrapper.setData(wrapperService.getData(body));
		return wrapper;
}
...

Обратим внимание, что из класса нам необходимо получить объект
IWrapperModel wrapper = wrapperClass.getDeclaredConstructor().newInstance(); , но такой подход требует наличия в классе конструктора без параметров. Используем @SneakyThrows библиотеки Lombok для того, чтобы обработать это исключение

@AllArgsConstructor
@ControllerAdvice(annotations = EnableResponseWrapper.class)
public class ResponseWrapperAdvice implements ResponseBodyAdvice<Object> {
    private final IWrapperService wrapperService;

    /**
     * Метод не будет обработан, если помечен аннотацией {@link DisableResponseWrapper} <br/> <br/>
     *
     * @param returnType    the return type
     * @param converterType the selected converter type
     * @return {@code true} if {@link #beforeBodyWrite} should be invoked;
     * {@code false} otherwise
     */
    @Override
    public boolean supports(MethodParameter returnType, @NonNull Class converterType) {
        for (Annotation a : returnType.getMethodAnnotations()) {
            if (a.annotationType() == DisableResponseWrapper.class) {
                return false;
            }
        }

        return true;
    }

    /**
     * Оборачиваем ответ
     *
     * @param body                  the body to be written
     * @param returnType            the return type of the controller method
     * @param selectedContentType   the content type selected through content negotiation
     * @param selectedConverterType the converter type selected to write to the response
     * @param request               the current request
     * @param response              the current response
     * @return the body that was passed in or a modified (possibly new) instance
     */
    @SneakyThrows
    @Override
    public Object beforeBodyWrite(
            @Nullable Object body,
            @NonNull MethodParameter returnType,
            @NonNull MediaType selectedContentType,
            @NonNull Class selectedConverterType,
            @NonNull ServerHttpRequest request,
            @NonNull ServerHttpResponse response
    ) {
        if (body == null) {
            return null;
        }

        // получаем wrapperClass из аннотации
        Class<? extends IWrapperModel> wrapperClass = null;
        for (Annotation annotation : returnType.getContainingClass().getAnnotations()) {
            if (annotation.annotationType() == EnableResponseWrapper.class) {
                wrapperClass = ((EnableResponseWrapper) annotation).wrapperClass();
                break;
            }
        }

        if (wrapperClass == null) {
            return body;
        }

        // проверяем, был ли передан Collection или наследник Collection
        if (Collection.class.isAssignableFrom(body.getClass())) {
            try {
                Collection<?> bodyCollection = (Collection<?>) body;

                // проверяем, что collection не пустой
                if (bodyCollection.isEmpty()) {
                    return body;
                }
                // оборачиваем каждый элемент коллекции
                return generateListOfResponseWrapper(bodyCollection, wrapperClass);
            } catch (Exception e) {
                return body;
            }
        }

        // если не collection
        return generateResponseWrapper(body, wrapperClass);
    }

    /**
     * Генерируем список оберток для коллекции (те информация добавляется внутрь списка)
     *
     * @param bodyCollection список объектов, которые необходимо обернуть
     * @param wrapperClass   объект обертки
     * @return список оберток
     */
    private List<IWrapperModel> generateListOfResponseWrapper(Collection<?> bodyCollection, Class<? extends IWrapperModel> wrapperClass) {
        return bodyCollection.stream()
                .map((t) -> t == null ?
                        null :
                        generateResponseWrapper(t, wrapperClass)
                )
                .collect(Collectors.toList());
    }

    /**
     * Генерируем обертку вокруг объекта
     *
     * @param body         объект который необходимо поместить в обертку
     * @param wrapperClass объект обертки
     * @return обертка
     */
    @SneakyThrows
    private IWrapperModel generateResponseWrapper(Object body, Class<? extends IWrapperModel> wrapperClass) {
        // wrapperClass должен иметь конструктор без параметров - получаем объект IWrapperModel
        IWrapperModel wrapper = wrapperClass.getDeclaredConstructor().newInstance();
        wrapper.setBody(body);
        wrapper.setData(wrapperService.getData(body));
        return wrapper;
    }

}

Стартер

Теперь нам необходимо превратить наш проект в стартер. Для этого создадим класс автоконфигурации, в котором будем создавать бин из класса ResponseWrapperAdvice

@Configuration
@AutoConfigureAfter(WebMvcAutoConfiguration.class)
@AllArgsConstructor
public class ResponseWrapperAutoConfiguration {
    private final IWrapperService wrapperService;

    @Bean
    @ConditionalOnMissingBean
    public ResponseWrapperAdvice responseWrapperAdvice() {
        return new ResponseWrapperAdvice(wrapperService);
    }
}

@AutoConfigureAfter(WebMvcAutoConfiguration.class) - говорит о том, что наши бины подключатся после того, как сконфигурируются и подключатся бины web mvc

А также в resources/META-INF/ создадим файл spring.factories, в котором укажем, где Spring Boot-у искать наши настроенные бины для добавления в контекст

org.springframework.boot.autoconfigure.EnableAutoConfiguration=ru.emilnasyrov.lib.response.wrapper.config.ResponseWrapperAutoConfiguration

Сборка

Соберем стартер в jar файл с помощью команды

gradle jar

Наш jar-ник появится в build/libs/response-wrapper-starter-0.0.1-SNAPSHOT-plain.jar - для удобства обрезаем -plain. Стартер готов и нам остается только подключить его к проекту

Демо

Для демо создадим отдельный проект, в котором будем использовать стартер

В проекте создадим папку libs, в которую положим jar стартера

В build.gradle стартер подключаем следующим образом:

repositories {
    ...
    flatDir {
        dirs 'libs'
    }
}

dependencies {
		...
    implementation 'ru.emilnasyrov.lib:response-wrapper-starter:0.0.1-SNAPSHOT'
		...
}

Модели данных

@Data
@AllArgsConstructor
public class MainModel {
    private String name;
    private String surname;

    // переопределяем toString, hashCode, equals
}

И Wrapper - класс-обертка
@Data
@AllArgsConstructor
@NoArgsConstructor
public class Wrapper implements IWrapperModel {
    @JsonUnwrapped
    Object main;

    String someInfo;

    @Override
    public void setData(Object object) {
        someInfo = object.toString();
    }

    @Override
    public void setBody(Object object) {
        main = object;
    }

    // переопределяем toString, hashCode, equals
}

@Service
public class WrapperServiceImpl implements IWrapperService {
    @Override
    public Object getData(Object body) {
        return "Additional Information";
    }
}

Контроллер

@EnableResponseWrapper(wrapperClass = Wrapper.class)
@RequestMapping("/test")
public class Controller {

    @GetMapping
    public MainModel test() {
        return new MainModel("Name", "Surname");
    }

    @GetMapping("/collection")
    public Collection<MainModel> testList() {
        Collection<MainModel> mainModels = new ArrayList<>();
        mainModels.add(new MainModel("Name1", "Surname1"));
        mainModels.add(new MainModel("Name2", "Surname2"));

        return mainModels;
    }

    @DisableResponseWrapper
    @GetMapping("/unwrapped")
    public MainModel unwrapped() {
        return new MainModel("Name", "Surname");
    }
}

Запустим проект и проверим запросы

Postman

Резюме

Мы рассмотрели интересный способ использования ControllerAdvice для работы с ответом точек входа контроллера. Также создали pet-библиотеку с реализацией в виде стартера Spring Boot

Почему не будет работать обычное AOP? Потому что AOP создает прокси класса с помощью CGLib или JDK Dynamic Proxy. Но когда DispatcherServlet, сканируя наши контроллеры, увидит, что контроллер должен возвращать один класс, а возвращает в итоге иной  (обертку), то он отдаст ошибку. Однако подход, описанный в статье, позволяет провернуть такую интересную штуку

Ссылка на на полный код проекта: GitHub