From 6fa721cc7f954acb77d02be220093c25007bef70 Mon Sep 17 00:00:00 2001 From: kjg Date: Sat, 8 Nov 2025 10:13:52 +0900 Subject: [PATCH] Fix POJO functions to return Message consistently POJO functions now return Message when input is Message, maintaining consistency with regular Function implementations. This ensures headers are preserved in POJO functions just like regular functions. Changes: - Add isPojoFunction flag to identify POJO functions - Mark POJO functions during discovery - Re-wrap output in Message for POJO functions when input is Message - Add test verifying Message return with header preservation - Add test verifying plain String input returns plain String output Fixes gh-1307 Signed-off-by: kjg --- .../BeanFactoryAwareFunctionRegistry.java | 5 ++ .../catalog/SimpleFunctionRegistry.java | 29 +++++++++++ ...FactoryAwarePojoFunctionRegistryTests.java | 52 ++++++++++++++++++- 3 files changed, 84 insertions(+), 2 deletions(-) diff --git a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java index 062a25dac..732a08e70 100644 --- a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java +++ b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java @@ -21,6 +21,7 @@ import java.lang.reflect.Type; import java.util.Arrays; import java.util.Collection; +import java.util.Collections; import java.util.Set; import java.util.function.BiConsumer; import java.util.function.BiFunction; @@ -171,6 +172,10 @@ else if (this.isFunctionPojo(functionCandidate, functionName)) { Method functionalMethod = FunctionTypeUtils.discoverFunctionalMethod(functionCandidate.getClass()); functionCandidate = this.proxyTarget(functionCandidate, functionalMethod); functionType = FunctionTypeUtils.fromFunctionMethod(functionalMethod); + // GH-1307: Mark this as a POJO function for special handling + functionRegistration = new FunctionRegistration(functionCandidate, functionName) + .type(functionType) + .properties(Collections.singletonMap("isPojoFunction", "true")); } else if (this.isSpecialFunctionRegistration(functionNames, functionName)) { functionRegistration = this.applicationContext diff --git a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java index debe5653c..1686607cc 100644 --- a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java +++ b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java @@ -287,6 +287,17 @@ private FunctionInvocationWrapper findFunctionInFunctionRegistrations(String fun // ignore } } + // GH-1307: Mark POJO functions for special Message wrapping behavior + if (functionRegistration != null && + functionRegistration.getProperties().containsKey("isPojoFunction")) { + try { + String isPojoValue = functionRegistration.getProperties().get("isPojoFunction"); + function.setPojoFunction(Boolean.parseBoolean(isPojoValue)); + } + catch (Exception e) { + // ignore + } + } return function; } @@ -439,6 +450,8 @@ public class FunctionInvocationWrapper implements Function, Cons private boolean wrappedBiConsumer; + private boolean isPojoFunction; + FunctionInvocationWrapper(String functionDefinition, Object target, Type inputType, Type outputType) { if (target instanceof PostProcessingFunction) { this.postProcessor = (PostProcessingFunction) target; @@ -489,6 +502,14 @@ public void setWrappedBiConsumer(boolean wrappedBiConsumer) { this.wrappedBiConsumer = wrappedBiConsumer; } + public void setPojoFunction(boolean isPojoFunction) { + this.isPojoFunction = isPojoFunction; + } + + public boolean isPojoFunction() { + return this.isPojoFunction; + } + public boolean isSkipOutputConversion() { return skipOutputConversion; } @@ -1245,6 +1266,14 @@ else if (isExtractPayload((Message) convertedOutput, type)) { } if (ObjectUtils.isEmpty(contentType)) { + // GH-1307: For POJO functions, wrap output in Message to maintain + // consistency with regular functions + if (this.isPojoFunction && output instanceof Message + && !(convertedOutput instanceof Message)) { + convertedOutput = MessageBuilder.withPayload(convertedOutput) + .copyHeaders(((Message) output).getHeaders()) + .build(); + } return convertedOutput; } diff --git a/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwarePojoFunctionRegistryTests.java b/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwarePojoFunctionRegistryTests.java index 81a58f92f..b13790a37 100644 --- a/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwarePojoFunctionRegistryTests.java +++ b/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwarePojoFunctionRegistryTests.java @@ -83,8 +83,10 @@ public void testWithPojoFunction() { Function f2conversion = catalog.lookup("myFunctionLike"); assertThat(f2conversion.apply(123)).isEqualTo("123"); - Function, String> f2message = catalog.lookup("myFunctionLike"); - assertThat(f2message.apply(MessageBuilder.withPayload("message").build())).isEqualTo("MESSAGE"); + // GH-1307: POJO functions now return Message for consistency + Function, Message> f2message = catalog.lookup("myFunctionLike"); + Message messageResult = f2message.apply(MessageBuilder.withPayload("message").build()); + assertThat(messageResult.getPayload()).isEqualTo("MESSAGE"); Function, Flux> f3 = catalog.lookup("myFunctionLike"); assertThat(f3.apply(Flux.just("foo")).blockFirst()).isEqualTo("FOO"); @@ -100,6 +102,52 @@ public void testWithPojoFunctionComposition() { assertThat(f1.apply("foo")).isEqualTo("FOO"); } + /** + * GH-1307: POJO function should return Message consistently with regular functions + * when no contentType is specified. + */ + @Test + public void testPojoFunctionReturnsMessageWithoutContentType() { + FunctionCatalog catalog = this.configureCatalog(); + + // Test POJO function without contentType + Function, Object> pojoFunction = catalog.lookup("myFunctionLike"); + Message input = MessageBuilder.withPayload("test") + .setHeader("correlationId", "123") + .build(); + + Object result = pojoFunction.apply(input); + + // GH-1307: Verify POJO functions return Message for consistency + assertThat(result) + .as("POJO function should return Message, not plain value when input is Message") + .isInstanceOf(Message.class); + + Message messageResult = (Message) result; + assertThat(messageResult.getPayload()).isEqualTo("TEST"); + assertThat(messageResult.getHeaders().get("correlationId")) + .as("Headers should be preserved") + .isEqualTo("123"); + } + + /** + * GH-1307: POJO function should NOT wrap output when input is plain String + */ + @Test + public void testPojoFunctionDoesNotWrapPlainStringInput() { + FunctionCatalog catalog = this.configureCatalog(); + + // GH-1307: POJO function with plain String input should return plain String + Function pojoFunction = catalog.lookup("myFunctionLike"); + Object result = pojoFunction.apply("plainInput"); + + // Should return String, not Message + assertThat(result) + .as("POJO function should return plain String when input is plain String, not wrap in Message") + .isInstanceOf(String.class) + .isEqualTo("PLAININPUT"); + } + @EnableAutoConfiguration @Configuration(proxyBeanMethods = false)