SpringCloud/spring-cloud-stream
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

GH-2235 clean up PartitionAwareFunctionWrapper

Browse Source
This commit is contained in:
Oleg Zhurakousky
2021-10-18 15:42:35 +02:00
parent ce6c09dea9
commit a3414b77b7
2 changed files with 37 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
View File

@@ -16,9 +16,6 @@
package org.springframework.cloud.stream.function;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Function;
import java.util.function.Supplier;
@@ -48,7 +45,6 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
@SuppressWarnings("rawtypes")
private final Function function;
@SuppressWarnings("rawtypes")
private final Function<Object, Object> outputMessageEnricher;
PartitionAwareFunctionWrapper(Function<?, ?> function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
@@ -59,24 +55,13 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
PartitionHandler partitionHandler = new PartitionHandler(evaluationContext, producerProperties, context.getBeanFactory());
this.outputMessageEnricher = output -> {
if (ObjectUtils.isArray(output) && !(output instanceof byte[])) {
output = Arrays.asList(output);
}
if (output instanceof Iterable) {
Iterable elements = (Iterable) output;
List<Message> messages = new ArrayList<>();
for (Object element : elements) {
if (!(element instanceof Message)) {
element = MessageBuilder.withPayload(element).build();
}
messages.add(toMessageWithPartitionHeader((Message) element, partitionHandler));
}
return messages;
if ((ObjectUtils.isArray(output) && !(output instanceof byte[])) || output instanceof Iterable) {
return output;
}
else if (!(output instanceof Message)) {
output = MessageBuilder.withPayload(output).build();
}
return toMessageWithPartitionHeader((Message) output, partitionHandler);
return toMessageWithPartitionHeader((Message<?>) output, partitionHandler);
};
}
else {

34
spring-cloud-stream/src/test/java/org/springframework/cloud/stream/function/ImplicitFunctionBindingTests.java
View File

@@ -22,6 +22,7 @@ import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
@@ -763,6 +764,29 @@ public class ImplicitFunctionBindingTests {
}
}
@Test
public void partitionOnOutputPayloadAsListTest() {
System.clearProperty("spring.cloud.function.definition");
try (ConfigurableApplicationContext context = new SpringApplicationBuilder(TestChannelBinderConfiguration
.getCompleteConfiguration(PojoFunctionConfiguration.class))
.web(WebApplicationType.NONE).run("--spring.cloud.function.definition=persons",
"--spring.cloud.stream.bindings.persons-out-0.producer.partitionKeyExpression=payload.id",
"--spring.cloud.stream.bindings.persons-out-0.producer.partitionCount=5",
"--spring.jmx.enabled=false")) {
InputDestination inputDestination = context.getBean(InputDestination.class);
OutputDestination outputDestination = context.getBean(OutputDestination.class);
Message<byte[]> inputMessage = MessageBuilder.withPayload("Jim Lahey".getBytes()).build();
inputDestination.send(inputMessage, "persons-in-0");
assertThat(outputDestination.receive(100, "persons-out-0").getHeaders().get("scst_partition")).isEqualTo(3);
assertThat(outputDestination.receive(100)).isNull();
}
}
@Test
public void partitionOnOutputPayloadTestReactive() {
System.clearProperty("spring.cloud.function.definition");
@@ -1490,6 +1514,16 @@ public class ImplicitFunctionBindingTests {
};
}
@Bean
public Function<String, List<Message<Person>>> persons() {
return x -> {
Person person = new Person();
person.setName(x);
person.setId(3);
return Collections.singletonList(MessageBuilder.withPayload(person).build());
};
}
@Bean
public Function<Flux<String>, Flux<Person>> funcReactive() {
return flux -> flux.map(value -> {