2,875 reads

Custom Controller Annotation with Spring Boot

by Alexander KozhenkovOctober 24th, 2022

Too Long; Didn't Read

This article will describe how to create your custom annotation for request routing. To implement this, we need to create an annotation. Then a custom BeanPostProcessor should register all methods with the annotation. QueryRouter implementation helps with object mapping and routing requests to target controllers. The usage of the annotation looks like casual Spring MVC.

Company Mentioned

featured image - Custom Controller Annotation with Spring Boot

Sometimes it becomes necessary to write your implementation of controllers using Spring Boot. For example, if you are not using REST, requests come to you from a non-standard communication channel.

This article will describe how to create your custom annotation for request routing. I will be using Kotlin to demonstrate the use of

suspend

functions in controllers. But nothing prevents you from using exactly the same approach in Java.

Creating an annotation

@Target(AnnotationTarget.FUNCTION)
@Retention(AnnotationRetention.RUNTIME)
annotation class Route(
    val value: RouteCommand,
)

You can use an Enum or a REST-like string as a value, then use

AntPathMatcher

for parsing.

Methods registration

To register methods, we need to implement

BeanPostProcessor

and scan all beans containing our annotation. Then put the parsed methods into a map.

@Component
class RouteMethodRegistry : BeanPostProcessor {

    private val registry: MutableMap<RouteCommand, RouteMethod> = ConcurrentHashMap()

    fun resolveMethod(command: RouteCommand): RouteMethod =
        registry[command] ?: throw IllegalStateException("No registered method for $command command")

    override fun postProcessAfterInitialization(bean: Any, beanName: String): Any {
        val targetClass = AopProxyUtils.ultimateTargetClass(bean)

        if (AnnotationUtils.isCandidateClass(targetClass, Route::class.java)) {
            val methodAnnotations: Map<Method, Route> = getMethodsWithRouteAnnotation(targetClass)

            for ((method, annotation) in methodAnnotations) {
                validateNoRegisteredMethods(annotation.value)

                val function = method.kotlinFunction
                    ?: throw IllegalStateException("Kotlin function for ${method.name} method is null")
                val (inputType, returnType) = parseTypes(function)

                registry[annotation.value] = RouteMethod(
                    instance = bean,
                    function = function,
                    inputType = inputType,
                    returnType = returnType,
                )
            }
        }

        return bean
    }

    private fun parseTypes(function: KFunction<*>): Pair<KClass<*>?, KClass<*>> {
        // The first element is the instance itself
        val inputParams = function.parameters.drop(1)

        if (inputParams.size > 1)
            throw IllegalStateException("Found ${inputParams.size} parameters for function $function. Please use 1 parameter")

        val inputType = inputParams.map { it.type.classifier as KClass<*>}.firstOrNull()
        val returnType = function.returnType.classifier!! as KClass<*>

        return inputType to returnType
    }

    private fun getMethodsWithRouteAnnotation(targetType: Class<*>): Map<Method, Route> =
        MethodIntrospector.selectMethods<Route>(targetType) {
                method -> AnnotatedElementUtils.getMergedAnnotation(method, Route::class.java)
        }

    private fun validateNoRegisteredMethods(command: RouteCommand) {
        if (registry.containsKey(command))
            throw IllegalStateException("More than 1 methods use the same route ${command.name}")
    }
}

data class RouteMethod(
    val instance: Any,
    val function: KFunction<*>,
    val inputType: KClass<*>?,
    val returnType: KClass<*>,
)

In the

parseTypes()

method, the input and output parameters are extracted to then convert the data into them using the

objectMapper

The

resolveMethod()

method will be used later for routing requests.

Requests routing

Then we need to implement the routing itself. In this example, when the method completes, we invoke the callback function.

@Component
class QueryRouter(
    private val objectMapper: ObjectMapper,
    private val routeMethodRegistry: RouteMethodRegistry,
) {

    fun route(query: String, callback: Callback) = CoroutineScope(Dispatchers.Default).launch {
        try {
            val output = route(query)

            callback.success(output)
        } catch (e: Exception) {
            callback.failure(500, e.message)
        }
    }

    private suspend fun route(query: String): String {
        val (command, input) = parseRouteParams(query)
        val method = routeMethodRegistry.resolveMethod(command)

        val output = if (method.inputType == null)
            method.function.callSuspend(method.instance)
        else
            method.function.callSuspend(method.instance, input)

        return objectMapper.writeValueAsString(output)
    }

    private fun parseRouteParams(query: String): Pair<RouteCommand, Any?> {
        TODO("Implement according to your data format")
    }
}

In order for controller methods to be marked with the suspend keyword, we call functions via

callSuspend()

You can use the standard

objectMapper

to convert input and output data. The implementation of the

parseRouteParams()

method depends on the data format used.

Using annotations in controllers

@Controller
class TestController {

    @Route(RouteCommand.TEST_CONTROLLER)
    suspend fun someMethod(request: MyRequest): MyResponse {
        delay(1000)
        return MyResponse(request.data)
    }
}