Add annotation for id class mapping

[markbanierink]

Feature request

PROBLEM
Currently, mapping id's to classes is, for example, done through configuration beans. There a TypePrecedence, mapping and TypeMapper can be configured. For any additional message, you also have to add changes to your configuration.

SUGGESTED SOLUTION
It would be nice if setting this configuration could be simplified by annotating a class with, for example:

@KafkaMessage(typeId = "foo")
public class Bar {}

Upon setting this annotation, the TypePrecedence and a default TypeMapper could be set as well.

[garyrussell]

Does the new Using Methods to Determine Types feature not meet your needs?

private static final JavaType barType = TypeFactory.defaultInstance()
				.constructType(Bar.class);

public static JavaType whichType(byte[] data, Headers headers) {
	Header header = headers.lastHeader("foo");
	if (header != null) {
		return this.barType;
	}
	else {
		return TypeFactory.defaultInstance().constructType(Object.class);
	}
}

spring.kafka.consumer.properties.spring.json.value.type.method=com.example.demo.MyClass.returnType

[markbanierink]

But that still requires you to define the mapping in a centralized place, right?

[markbanierink]

I have now created the following for our project, which allows you to just annotate a class with @KafkaMessage and a type id:

@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface KafkaMessage {

    String[] typeId() default "";

}

public class KafkaMessageAnnotationPostProcessor implements BeanPostProcessor {

    private static final String SCAN_PACKAGE_PROPERTY_KEY = "spring.kafka.message.package";
    private static final String DEFAULT_SCAN_PACKAGE = "*";

    private final Environment environment;

    public KafkaMessageAnnotationPostProcessor(Environment environment) {
        this.environment = environment;
    }

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (JsonMessageConverter.class.isAssignableFrom(bean.getClass())) {
            setTypeMapper((JsonMessageConverter)bean, getMapping());
        }
        return bean;
    }

    private Map<String, Class<?>> getMapping() {
        Map<String, Class<?>> mapping = new HashMap<>();
        for (BeanDefinition beanDefinition : findAnnotatedClasses(getScanPackage())) {
            addToMapping(mapping, getClass(beanDefinition));
        }
        return mapping;
    }

    private String getScanPackage() {
        return environment.getProperty(SCAN_PACKAGE_PROPERTY_KEY, DEFAULT_SCAN_PACKAGE);
    }

    private Set<BeanDefinition> findAnnotatedClasses(String scanPackage) {
        ClassPathScanningCandidateComponentProvider provider = createComponentProvider();
        return provider.findCandidateComponents(scanPackage);
    }

    private ClassPathScanningCandidateComponentProvider createComponentProvider() {
        ClassPathScanningCandidateComponentProvider provider = new ClassPathScanningCandidateComponentProvider(false);
        provider.addIncludeFilter(new AnnotationTypeFilter(KafkaMessage.class));
        return provider;
    }

    private Class<?> getClass(BeanDefinition beanDefinition) {
        try {
            return Class.forName(beanDefinition.getBeanClassName());
        }
        catch (ClassNotFoundException e) {
            throw new IllegalStateException("Cannot find " + beanDefinition.getBeanClassName(), e);
        }
    }

    private void addToMapping(Map<String, Class<?>> mapping, Class<?> clazz) {
        for (String i : getTypeIds(clazz)) {
            mapping.put(i, clazz);
        }
    }

    private List<String> getTypeIds(Class<?> clazz) {
        KafkaMessage annotation = clazz.getAnnotation(KafkaMessage.class);
        // Use simple name as default
        return annotation.typeId().length > 0 ? Arrays.asList(annotation.typeId()) : Collections.singletonList(clazz.getSimpleName());
    }

    private void setTypeMapper(JsonMessageConverter converter, Map<String, Class<?>> mapping) {
        // Only set a mapping if an @KafkaMessage annotation is set on any class
        if (!mapping.isEmpty()) {
            Jackson2JavaTypeMapper typeMapper = converter.getTypeMapper();
            if (!AbstractJavaTypeMapper.class.isAssignableFrom(typeMapper.getClass())) {
                throw new IllegalStateException("TypeMapper must be of type AbstractJavaTypeMapper");
            }
            typeMapper.setTypePrecedence(Jackson2JavaTypeMapper.TypePrecedence.TYPE_ID);
            ((AbstractJavaTypeMapper)typeMapper).setIdClassMapping(mapping);
        }
    }

}

Maybe something like the above could be adopted by the spring-kafka project?

[garyrussell]

Sure; will take a look.