Skip to content

Commit

Permalink
spring-projectsGH-3067: Draft of mapping multiple headers with same k…
Browse files Browse the repository at this point in the history 
…ey with SimpleKafkaHeaderMapper
  • Loading branch information
@poznachowski
poznachowski committed Mar 12, 2024
1 parent 265e55f commit 0f20c3b
Show file tree
Hide file tree
Showing 4 changed files with 197 additions and 74 deletions.
154 changes: 104 additions & 50 deletions spring-kafka/src/main/java/org/springframework/kafka/support/DefaultKafkaHeaderMapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,21 @@
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeader;
import org.assertj.core.util.Streams;

import org.springframework.messaging.MessageHeaders;
import org.springframework.util.Assert;
Expand All @@ -48,12 +52,14 @@
*
* @author Gary Russell
* @author Artem Bilan
* @author Grzegorz Poznachowski
*
* @since 1.3
*
*/
public class DefaultKafkaHeaderMapper extends AbstractKafkaHeaderMapper {

private static final String ITERABLE_HEADER_TYPE_PATTERN = "%s#%s";

private static final String JAVA_LANG_STRING = "java.lang.String";

private static final Set<String> TRUSTED_ARRAY_TYPES = Set.of(
Expand Down Expand Up @@ -96,6 +102,7 @@ public class DefaultKafkaHeaderMapper extends AbstractKafkaHeaderMapper {
* {@code "!id", "!timestamp" and "*"}. In addition, most of the headers in
* {@link KafkaHeaders} are never mapped as headers since they represent data in
* consumer/producer records.
*
* @see #DefaultKafkaHeaderMapper(ObjectMapper)
*/
public DefaultKafkaHeaderMapper() {
Expand All @@ -110,6 +117,7 @@ public DefaultKafkaHeaderMapper() {
* {@code "!id", "!timestamp" and "*"}. In addition, most of the headers in
* {@link KafkaHeaders} are never mapped as headers since they represent data in
* consumer/producer records.
*
* @param objectMapper the object mapper.
* @see org.springframework.util.PatternMatchUtils#simpleMatch(String, String)
*/
Expand All @@ -128,6 +136,7 @@ public DefaultKafkaHeaderMapper(ObjectMapper objectMapper) {
* generally should not map the {@code "id" and "timestamp"} headers. Note:
* most of the headers in {@link KafkaHeaders} are ever mapped as headers since they
* represent data in consumer/producer records.
*
* @param patterns the patterns.
* @see org.springframework.util.PatternMatchUtils#simpleMatch(String, String)
*/
Expand All @@ -143,8 +152,9 @@ public DefaultKafkaHeaderMapper(String... patterns) {
* you generally should not map the {@code "id" and "timestamp"} headers. Note: most
* of the headers in {@link KafkaHeaders} are never mapped as headers since they
* represent data in consumer/producer records.
*
* @param objectMapper the object mapper.
* @param patterns the patterns.
* @param patterns the patterns.
* @see org.springframework.util.PatternMatchUtils#simpleMatch(String, String)
*/
public DefaultKafkaHeaderMapper(ObjectMapper objectMapper, String... patterns) {
Expand All @@ -160,6 +170,7 @@ private DefaultKafkaHeaderMapper(boolean outbound, ObjectMapper objectMapper, St

/**
* Create an instance for inbound mapping only with pattern matching.
*
* @param patterns the patterns to match.
* @return the header mapper.
* @since 2.8.8
Expand All @@ -170,8 +181,9 @@ public static DefaultKafkaHeaderMapper forInboundOnlyWithMatchers(String... patt

/**
* Create an instance for inbound mapping only with pattern matching.
*
* @param objectMapper the object mapper.
* @param patterns the patterns to match.
* @param patterns the patterns to match.
* @return the header mapper.
* @since 2.8.8
*/
Expand All @@ -181,6 +193,7 @@ public static DefaultKafkaHeaderMapper forInboundOnlyWithMatchers(ObjectMapper o

/**
* Return the object mapper.
*
* @return the mapper.
*/
protected ObjectMapper getObjectMapper() {
Expand All @@ -189,6 +202,7 @@ protected ObjectMapper getObjectMapper() {

/**
* Provide direct access to the trusted packages set for subclasses.
*
* @return the trusted packages.
* @since 2.2
*/
Expand All @@ -198,6 +212,7 @@ protected Set<String> getTrustedPackages() {

/**
* Provide direct access to the toString() classes by subclasses.
*
* @return the toString() classes.
* @since 2.2
*/
Expand All @@ -214,6 +229,7 @@ protected boolean isEncodeStrings() {
* raw String value is converted to a byte array using the configured charset. Set to
* true if a consumer of the outbound record is using Spring for Apache Kafka version
* less than 2.3
*
* @param encodeStrings true to encode (default false).
* @since 2.3
*/
Expand All @@ -234,6 +250,7 @@ public void setEncodeStrings(boolean encodeStrings) {
* If any of the supplied packages is {@code "*"}, all packages are trusted.
* If a class for a non-trusted package is encountered, the header is returned to the
* application with value of type {@link NonTrustedHeaderType}.
*
* @param packagesToTrust the packages to trust.
*/
public void addTrustedPackages(String... packagesToTrust) {
Expand All @@ -253,6 +270,7 @@ public void addTrustedPackages(String... packagesToTrust) {
/**
* Add class names that the outbound mapper should perform toString() operations on
* before mapping.
*
* @param classNames the class names.
* @since 2.2
*/
Expand All @@ -264,32 +282,17 @@ public void addToStringClasses(String... classNames) {
public void fromHeaders(MessageHeaders headers, Headers target) {
final Map<String, String> jsonHeaders = new HashMap<>();
final ObjectMapper headerObjectMapper = getObjectMapper();
headers.forEach((key, rawValue) -> {
if (matches(key, rawValue)) {
Object valueToAdd = headerValueToAddOut(key, rawValue);
if (valueToAdd instanceof byte[]) {
target.add(new RecordHeader(key, (byte[]) valueToAdd));
headers.forEach((key, value) -> {
if (matches(key, value)) {
if (value instanceof Collection<?> values) {
int i = 0;
for (Object element : values) {
resolveSingleHeader(key, element, target, jsonHeaders, i);
i++;
}
}
else {
try {
String className = valueToAdd.getClass().getName();
boolean encodeToJson = this.encodeStrings;
if (this.toStringClasses.contains(className)) {
valueToAdd = valueToAdd.toString();
className = JAVA_LANG_STRING;
encodeToJson = true;
}
if (!encodeToJson && valueToAdd instanceof String) {
target.add(new RecordHeader(key, ((String) valueToAdd).getBytes(getCharset())));
}
else {
target.add(new RecordHeader(key, headerObjectMapper.writeValueAsBytes(valueToAdd)));
}
jsonHeaders.put(key, className);
}
catch (Exception e) {
logger.error(e, () -> "Could not map " + key + " with type " + rawValue.getClass().getName());
}
resolveSingleHeader(key, value, target, jsonHeaders);
}
}
});
Expand All @@ -303,30 +306,82 @@ public void fromHeaders(MessageHeaders headers, Headers target) {
}
}

@Override
public void toHeaders(Headers source, final Map<String, Object> headers) {
final Map<String, String> jsonTypes = decodeJsonTypes(source);
source.forEach(header -> {
String headerName = header.key();
if (headerName.equals(KafkaHeaders.DELIVERY_ATTEMPT) && matchesForInbound(headerName)) {
headers.put(headerName, ByteBuffer.wrap(header.value()).getInt());
}
else if (headerName.equals(KafkaHeaders.LISTENER_INFO) && matchesForInbound(headerName)) {
headers.put(headerName, new String(header.value(), getCharset()));
}
else if (!(headerName.equals(JSON_TYPES)) && matchesForInbound(headerName)) {
if (jsonTypes.containsKey(headerName)) {
String requestedType = jsonTypes.get(headerName);
populateJsonValueHeader(header, requestedType, headers);
private void resolveSingleHeader(String headerName, Object value, Headers target, Map<String, String> jsonHeaders) {
resolveSingleHeader(headerName, value, target, jsonHeaders, null);
}

private void resolveSingleHeader(String headerName, Object value, Headers target, Map<String, String> jsonHeaders, Integer headerIndex) {
Object valueToAdd = headerValueToAddOut(headerName, value);
if (valueToAdd instanceof byte[] byteArray) {
target.add(new RecordHeader(headerName, byteArray));
}
else {
try {
String className = valueToAdd.getClass().getName();
boolean encodeToJson = this.encodeStrings;
if (this.toStringClasses.contains(className)) {
valueToAdd = valueToAdd.toString();
className = JAVA_LANG_STRING;
encodeToJson = true;
}
if (!encodeToJson && valueToAdd instanceof String stringValue) {
target.add(new RecordHeader(headerName, stringValue.getBytes(getCharset())));
}
else {
headers.put(headerName, headerValueToAddIn(header));
target.add(new RecordHeader(headerName, this.objectMapper.writeValueAsBytes(valueToAdd)));
}
jsonHeaders.put(headerIndex == null ?
headerName :
ITERABLE_HEADER_TYPE_PATTERN.formatted(headerName, headerIndex), className);
}
});
catch (Exception e) {
logger.error(e, () -> "Could not map " + headerName + " with type " + value.getClass().getName());
}
}
}

@Override
public void toHeaders(Headers source, final Map<String, Object> target) {
final Map<String, String> jsonTypes = decodeJsonTypes(source);

Streams.stream(source)
.collect(Collectors.groupingBy(Header::key))
.forEach((headerName, headers) -> {
if (headerName.equals(KafkaHeaders.DELIVERY_ATTEMPT) && matchesForInbound(headerName)) {
target.put(headerName, ByteBuffer.wrap(headers.get(headers.size() - 1).value()).getInt());
}
else if (headerName.equals(KafkaHeaders.LISTENER_INFO) && matchesForInbound(headerName)) {
target.put(headerName, new String(headers.get(headers.size() - 1).value(), getCharset()));
}
else if (!(headerName.equals(JSON_TYPES)) && matchesForInbound(headerName)) {
if (headers.size() == 1) {
if (jsonTypes.containsKey(headerName)) {
String requestedType = jsonTypes.get(headerName);
target.put(headerName, resolveJsonValueHeader(headers.get(0), requestedType));
}
else {
target.put(headerName, headerValueToAddIn(headers.get(0)));
}
}
else {
List<Object> valueList = new ArrayList<>();
for (int i = 0; i < headers.size(); i++) {
var jsonTypeIterableHeader = ITERABLE_HEADER_TYPE_PATTERN.formatted(headerName, i);
if (jsonTypes.containsKey(jsonTypeIterableHeader)) {
String requestedType = jsonTypes.get(jsonTypeIterableHeader);
valueList.add(resolveJsonValueHeader(headers.get(i), requestedType));
}
else {
valueList.add(headerValueToAddIn(headers.get(i)));
}
}
target.put(headerName, valueList);
}
}
});
}

private void populateJsonValueHeader(Header header, String requestedType, Map<String, Object> headers) {
private Object resolveJsonValueHeader(Header header, String requestedType) {
Class<?> type = Object.class;
boolean trusted = false;
try {
Expand All @@ -339,22 +394,21 @@ private void populateJsonValueHeader(Header header, String requestedType, Map<St
logger.error(e, () -> "Could not load class for header: " + header.key());
}
if (String.class.equals(type) && (header.value().length == 0 || header.value()[0] != '"')) {
headers.put(header.key(), new String(header.value(), getCharset()));
return new String(header.value(), getCharset());
}
else {
if (trusted) {
try {
Object value = decodeValue(header, type);
headers.put(header.key(), value);
return decodeValue(header, type);
}
catch (IOException e) {
logger.error(e, () ->
"Could not decode json type: " + requestedType + " for key: " + header.key());
headers.put(header.key(), header.value());
return header.value();
}
}
else {
headers.put(header.key(), new NonTrustedHeaderType(header.value(), requestedType));
return new NonTrustedHeaderType(header.value(), requestedType);
}
}
}
Expand Down
Loading

0 comments on commit 0f20c3b

Please sign in to comment.