feat: consumer schema validation for avro (#59)

ccm32004 · web-flow · commit 97d782f6991c · 2026-02-25T16:53:52.000-05:00
Signed-off-by: Cece Ma &lt;mayuqing131@gmail.com&gt;
diff --git a/.gitignore b/.gitignore
@@ -34,7 +34,7 @@ build/
 ### Local env ###
 .env
 
-### Register-schema scripts (dev/local) ###
+### Register-schema scripts (not fully tested; exclude from consumer-schema-validation) ###
 development/scripts/register-schema/
 
 ### Mac OS ###
diff --git a/docs/source/byte-array/manifests/api-key/byte-arr-consumer-config.yaml b/docs/source/byte-array/manifests/api-key/byte-arr-consumer-config.yaml
@@ -15,6 +15,7 @@ data:
             # Example: OAuth2 - authParams: "${PULSAR_OAUTH_CLIENT_SECRET}"
         consumer:  # see here for all configurations: https://pulsar.apache.org/reference/#/4.0.x/client/client-configuration-consumer
           enabled: true
+          useAutoConsumeSchema: true
           consumerConfig:
             # Single topic (string) or multiple topics (comma-separated string)
             topicNames: "persistent://public/default/test-topic"
diff --git a/pom.xml b/pom.xml
@@ -33,6 +33,12 @@
             <version>${lombok.version}</version>
             <optional>true</optional>
         </dependency>
+        <dependency>
+            <groupId>org.apache.avro</groupId>
+            <artifactId>avro</artifactId>
+            <version>1.11.3</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.springframework.boot</groupId>
             <artifactId>spring-boot-starter-test</artifactId>
diff --git a/src/main/java/io/numaproj/pulsar/config/consumer/PulsarConsumerProperties.java b/src/main/java/io/numaproj/pulsar/config/consumer/PulsarConsumerProperties.java
@@ -28,6 +28,15 @@ public class PulsarConsumerProperties {
 
     private Map<String, Object> consumerConfig = new HashMap<>(); // Default to an empty map
 
+    /**
+     * When true (default), the consumer uses Schema.AUTO_CONSUME so the client validates
+     * each message against the topic schema when decoding. When false, uses Schema.BYTES
+     * (no schema check; messages are read as raw bytes).
+     * When the topic has no schema (Pulsar treats it as BYTES), no validation and no decoding
+     * is performed on the message bytes; they are passed through as bytes.
+     */
+    private boolean useAutoConsumeSchema = true;
+
     @PostConstruct
     public void init() {
         // Pulsar expects topicNames to be type Set<String>. Config accepts a single string
diff --git a/src/main/java/io/numaproj/pulsar/consumer/PulsarConsumerManager.java b/src/main/java/io/numaproj/pulsar/consumer/PulsarConsumerManager.java
@@ -7,6 +7,7 @@
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.apache.pulsar.client.api.Schema;
 import org.apache.pulsar.client.api.SubscriptionType;
+import org.apache.pulsar.client.api.schema.GenericRecord;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.stereotype.Component;
@@ -17,9 +18,9 @@
 import java.util.concurrent.TimeUnit;
 
 /**
- * PulsarConsumerManager creates and maintains a single Consumer instance.
- * A new consumer is created based on the provided batch size and read timeout,
- * but once created it will be reused until explicitly removed.
+ * Creates and maintains Pulsar consumers: one for byte[] (Schema.BYTES) and one for
+ * schema-backed messages (Schema.AUTO_CONSUME / GenericRecord). Only the consumer
+ * matching the configured schema is created and used; the other remains null.
  */
 @Slf4j
 @Component
@@ -32,52 +33,72 @@ public class PulsarConsumerManager {
     @Autowired
     private PulsarClient pulsarClient;
 
-    // The current consumer instance.
-    private Consumer<byte[]> currentConsumer;
+    private Consumer<byte[]> bytesConsumer;
+    private Consumer<GenericRecord> genericRecordConsumer;
 
-    // Returns the current consumer if it exists. If not, creates a new one.
-    public Consumer<byte[]> getOrCreateConsumer(long count, long timeoutMillis)
-            throws PulsarClientException {
-        if (currentConsumer != null) {
-            return currentConsumer;
+    /** Returns the byte-array consumer, creating it if necessary. Use when not using AUTO_CONSUME. */
+    public Consumer<byte[]> getOrCreateBytesConsumer(long count, long timeoutMillis) throws PulsarClientException {
+        if (bytesConsumer != null) {
+            return bytesConsumer;
         }
-
         BatchReceivePolicy batchPolicy = BatchReceivePolicy.builder()
                 .maxNumMessages((int) count)
-                .timeout((int) timeoutMillis, TimeUnit.MILLISECONDS) // We do not expect user to specify a number larger
-                                                                     // than 2^63 - 1 which will cause an overflow
+                .timeout((int) timeoutMillis, TimeUnit.MILLISECONDS) // We do not expect user to specify a number larger than 2^63 - 1 which will cause an overflow
                 .build();
-
-        currentConsumer = pulsarClient.newConsumer(Schema.BYTES)
+        bytesConsumer = pulsarClient.newConsumer(Schema.BYTES)
                 .loadConf(pulsarConsumerProperties.getConsumerConfig())
                 .batchReceivePolicy(batchPolicy)
                 .subscriptionType(SubscriptionType.Shared) // Must be shared to support multiple pods
                 .subscribe();
+        log.info("Created byte-array consumer; batch receive: {}, timeoutMillis: {}", count, timeoutMillis);
+        return bytesConsumer;
+    }
 
-        log.info("Created new consumer with batch receive policy of: {} and timeoutMillis: {}", count, timeoutMillis);
-        return currentConsumer;
+    /** Returns the GenericRecord (AUTO_CONSUME) consumer, creating it if necessary. Use when using AUTO_CONSUME. */
+    public Consumer<GenericRecord> getOrCreateGenericRecordConsumer(long count, long timeoutMillis) throws PulsarClientException {
+        if (genericRecordConsumer != null) {
+            return genericRecordConsumer;
+        }
+        BatchReceivePolicy batchPolicy = BatchReceivePolicy.builder()
+                .maxNumMessages((int) count)
+                .timeout((int) timeoutMillis, TimeUnit.MILLISECONDS) // We do not expect user to specify a number larger than 2^63 - 1 which will cause an overflow
+                .build();
+        genericRecordConsumer = pulsarClient.newConsumer(Schema.AUTO_CONSUME())
+                .loadConf(pulsarConsumerProperties.getConsumerConfig())
+                .batchReceivePolicy(batchPolicy)
+                .subscriptionType(SubscriptionType.Shared) // Must be shared to support multiple pods
+                .subscribe();
+        log.info("Created AUTO_CONSUME (GenericRecord) consumer; batch receive: {}, timeoutMillis: {}", count, timeoutMillis);
+        return genericRecordConsumer;
     }
 
     @PreDestroy
     public void cleanup() {
-        if (currentConsumer != null) {
+        if (bytesConsumer != null) {
             try {
-                currentConsumer.close();
-                log.info("Consumer closed during cleanup.");
-            } catch (PulsarClientException e) {
-                log.error("Error while closing consumer in cleanup", e);
+                bytesConsumer.close();
+                log.info("Byte-array consumer closed.");
+            } catch (Exception e) {
+                log.error("Error closing byte-array consumer", e);
             }
+            bytesConsumer = null;
+        }
+        if (genericRecordConsumer != null) {
+            try {
+                genericRecordConsumer.close();
+                log.info("GenericRecord consumer closed.");
+            } catch (Exception e) {
+                log.error("Error closing GenericRecord consumer", e);
+            }
+            genericRecordConsumer = null;
         }
-
         if (pulsarClient != null) {
             try {
                 pulsarClient.close();
-                log.info("Pulsar client closed during cleanup.");
+                log.info("Pulsar client closed.");
             } catch (PulsarClientException e) {
-                log.error("Error while closing the Pulsar client in cleanup", e);
+                log.error("Error closing Pulsar client", e);
             }
-
         }
-
     }
-}
+}
diff --git a/src/main/java/io/numaproj/pulsar/consumer/PulsarSource.java b/src/main/java/io/numaproj/pulsar/consumer/PulsarSource.java
@@ -15,6 +15,9 @@
 import org.apache.pulsar.client.api.MessageId;
 import org.apache.pulsar.client.api.Messages;
 import org.apache.pulsar.client.api.PulsarClientException;
+import org.apache.pulsar.client.api.SchemaSerializationException;
+import org.apache.pulsar.client.api.schema.GenericRecord;
+import org.apache.pulsar.client.api.schema.Field;
 import org.apache.pulsar.client.admin.PulsarAdmin;
 import org.apache.pulsar.client.admin.PulsarAdminException;
 import org.apache.pulsar.common.policies.data.TopicStats;
@@ -24,6 +27,7 @@
 import org.springframework.stereotype.Component;
 
 import javax.annotation.PostConstruct;
+import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -37,8 +41,11 @@
 @ConditionalOnProperty(prefix = "spring.pulsar.consumer", name = "enabled", havingValue = "true")
 public class PulsarSource extends Sourcer {
 
-    // Map tracking received messages (keyed by topicName + messageId for multi-topic support)
-    private final Map<String, org.apache.pulsar.client.api.Message<byte[]>> messagesToAck = new HashMap<>();
+    // Map tracking received messages (keyed by topicName + messageId for multi-topic support).
+    // Value is Message<?> because we hold either Message<byte[]> or Message<GenericRecord> depending on
+    // useAutoConsumeSchema. only use the Message interface from the map (getMessageId, getTopicName, etc.
+    // for ack and buildHeaders), never the payload type, so a single map with Message<?> is appropriate;
+    private final Map<String, org.apache.pulsar.client.api.Message<?>> messagesToAck = new HashMap<>();
 
     private Server server;
 
@@ -60,49 +67,112 @@ public void startServer() throws Exception {
 
     @Override
     public void read(ReadRequest request, OutputObserver observer) {
-        // If there are messages not acknowledged, return
         if (!messagesToAck.isEmpty()) {
             log.trace("messagesToAck not empty: {}", messagesToAck);
             return;
         }
 
-        Consumer<byte[]> consumer = null;
-
         try {
-            // Obtain a consumer with the desired settings.
-            consumer = pulsarConsumerManager.getOrCreateConsumer(request.getCount(), request.getTimeout().toMillis());
+            if (pulsarConsumerProperties.isUseAutoConsumeSchema()) {
+                readWithAutoConsume(request, observer);
+            } else {
+                readWithBytes(request, observer);
+            }
+        } catch (Exception e) {
+            log.error("Failed to read from Pulsar", e);
+            throw new RuntimeException(e);
+        }
+    }
 
-            Messages<byte[]> batchMessages = consumer.batchReceive();
+    private void readWithBytes(ReadRequest request, OutputObserver observer) throws PulsarClientException {
+        Consumer<byte[]> consumer = pulsarConsumerManager.getOrCreateBytesConsumer(request.getCount(), request.getTimeout().toMillis());
+        Messages<byte[]> batchMessages = consumer.batchReceive();
 
-            if (batchMessages == null || batchMessages.size() == 0) {
-                log.trace("Received 0 messages, return early.");
-                return;
-            }
+        if (batchMessages == null || batchMessages.size() == 0) {
+            log.trace("Received 0 messages, return early.");
+            return;
+        }
 
-            // Process each message in the batch.
-            for (org.apache.pulsar.client.api.Message<byte[]> pMsg : batchMessages) {
-                String topicName = pMsg.getTopicName();
-                String msgId = pMsg.getMessageId().toString();
-                String topicMessageIdKey = topicName + msgId;
-                
-                // TODO : change to .debug or .trace to reduce log noise
-                log.info("Consumed Pulsar message [topic: {}, id: {}]: {}", topicName, pMsg.getMessageId(),
-                        new String(pMsg.getValue(), StandardCharsets.UTF_8));
+        for (org.apache.pulsar.client.api.Message<byte[]> pMsg : batchMessages) {
+
+            // TODO : change to .debug or .trace to reduce log noise
+            log.info("Consumed Pulsar message [topic: {}, id: {}]: {}", pMsg.getTopicName(), pMsg.getMessageId(),
+                    new String(pMsg.getValue(), StandardCharsets.UTF_8));
+            sendMessage(pMsg, pMsg.getValue(), observer);
+        }
+    }
+
+    private void readWithAutoConsume(ReadRequest request, OutputObserver observer) throws PulsarClientException {
+        Consumer<GenericRecord> consumer = pulsarConsumerManager.getOrCreateGenericRecordConsumer(request.getCount(), request.getTimeout().toMillis());
+        Messages<GenericRecord> batchMessages = consumer.batchReceive();
+
+        if (batchMessages == null || batchMessages.size() == 0) {
+            log.trace("Received 0 messages, return early.");
+            return;
+        }
+
+        for (org.apache.pulsar.client.api.Message<GenericRecord> pMsg : batchMessages) {
 
-                byte[] offsetBytes = topicMessageIdKey.getBytes(StandardCharsets.UTF_8);
-                Offset offset = new Offset(offsetBytes);
+            try {
+                GenericRecord record = pMsg.getValue(); // This will throw SchemaSerializationException if the message is not valid based on the topic schema 
+                byte[] payloadBytes = pMsg.getData();
 
-                Map<String, String> headers = buildHeaders(pMsg);
+                String decoded = recordToLogString(record);
+                // TODO : change to .debug or .trace to reduce log noise
+                log.info("Consumed Pulsar message (AUTO_CONSUME) [topic: {}, id: {}]: {} bytes, decoded={}", pMsg.getTopicName(), pMsg.getMessageId(), payloadBytes.length, decoded);
+                sendMessage(pMsg, payloadBytes, observer);
+            } catch (Exception e) {
+                if (isSchemaValidationFailure(e)) {
+                    throw new RuntimeException("Schema validation failure", e);
+                }
+                throw new RuntimeException(e);
+            }
+        }
+    }
 
-                Message message = new Message(pMsg.getValue(), offset, Instant.now(), headers);
-                observer.send(message);
+    /** Builds offset and headers, sends the message to the observer, and records it for ack. */
+    private void sendMessage(org.apache.pulsar.client.api.Message<?> pMsg, byte[] payloadBytes, OutputObserver observer) {
+        String topicMessageIdKey = pMsg.getTopicName() + pMsg.getMessageId().toString();
+        byte[] offsetBytes = topicMessageIdKey.getBytes(StandardCharsets.UTF_8);
+        Offset offset = new Offset(offsetBytes);
+        Map<String, String> headers = buildHeaders(pMsg);
+        observer.send(new Message(payloadBytes, offset, Instant.now(), headers));
+        messagesToAck.put(topicMessageIdKey, pMsg);
+    }
 
-                messagesToAck.put(topicMessageIdKey, pMsg);
+    /** True if the exception indicates schema deserialization failure. The Pulsar client throws
+     * SchemaSerializationException when decoding fails (e.g. schema mismatch, bad payload)
+     * and wraps underlying IOException from the decoder as the cause. */
+    private static boolean isSchemaValidationFailure(Throwable e) {
+        for (Throwable t = e; t != null; t = t.getCause()) {
+            if (t instanceof SchemaSerializationException) {
+                return true;
             }
-        } catch (PulsarClientException e) {
-            log.error("Failed to get consumer or receive messages from Pulsar", e);
-            throw new RuntimeException("Failed to get consumer or receive messages from Pulsar", e);
         }
+        return false;
+    }
+
+    // TODO : remove this logging later to reduce log noise
+    /** Builds a log-friendly string of the decoded record (field names and values). */
+    private static String recordToLogString(GenericRecord record) {
+        if (record == null) {
+            return "null";
+        }
+        List<Field> fields = record.getFields();
+        if (fields == null || fields.isEmpty()) {
+            return record.getSchemaType() + ":{}";
+        }
+        StringBuilder sb = new StringBuilder();
+        sb.append(record.getSchemaType()).append(":{");
+        for (int i = 0; i < fields.size(); i++) {
+            if (i > 0) sb.append(", ");
+            Field f = fields.get(i);
+            String name = f.getName();
+            Object value = record.getField(f);
+            sb.append(name).append("=").append(value);
+        }
+        sb.append("}");
+        return sb.toString();
     }
 
     @Override
@@ -128,8 +198,11 @@ public void ack(AckRequest request) {
                 .toList();
 
         try {
-            Consumer<byte[]> consumer = pulsarConsumerManager.getOrCreateConsumer(0, 0);
-            consumer.acknowledge(messageIds);
+            if (pulsarConsumerProperties.isUseAutoConsumeSchema()) {
+                pulsarConsumerManager.getOrCreateGenericRecordConsumer(0, 0).acknowledge(messageIds);
+            } else {
+                pulsarConsumerManager.getOrCreateBytesConsumer(0, 0).acknowledge(messageIds);
+            }
             log.info("Successfully acknowledged {} messages", messageIds.size());
         } catch (PulsarClientException e) {
             log.error("Failed to acknowledge Pulsar messages", e);
@@ -138,9 +211,9 @@ public void ack(AckRequest request) {
     }
 
     /**
-     * Builds headers from Pulsar message metadata
+     * Builds headers from Pulsar message metadata. Works for both Message<?> (byte[] or GenericRecord).
      */
-    private Map<String, String> buildHeaders(org.apache.pulsar.client.api.Message<byte[]> pulsarMessage) {
+    private Map<String, String> buildHeaders(org.apache.pulsar.client.api.Message<?> pulsarMessage) {
         Map<String, String> headers = new HashMap<>();
 
         headers.put(NumaHeaderKeys.PULSAR_PRODUCER_NAME, pulsarMessage.getProducerName());
@@ -150,7 +223,6 @@ private Map<String, String> buildHeaders(org.apache.pulsar.client.api.Message<by
         headers.put(NumaHeaderKeys.PULSAR_EVENT_TIME, String.valueOf(pulsarMessage.getEventTime()));
         headers.put(NumaHeaderKeys.PULSAR_REDELIVERY_COUNT, String.valueOf(pulsarMessage.getRedeliveryCount()));
 
-        // Add message properties as headers
         if (pulsarMessage.getProperties() != null && !pulsarMessage.getProperties().isEmpty()) {
             pulsarMessage.getProperties().forEach((key, value) -> {
                 if (key != null && value != null) {
diff --git a/src/test/java/io/numaproj/pulsar/config/consumer/PulsarConsumerPropertiesTest.java b/src/test/java/io/numaproj/pulsar/config/consumer/PulsarConsumerPropertiesTest.java
@@ -77,6 +77,14 @@ public void consumerProperties_DefaultSubscriptionName() {
         );
     }
 
+    @Test
+    public void consumerProperties_useAutoConsumeSchema_defaultTrue() {
+        PulsarConsumerProperties properties = new PulsarConsumerProperties();
+        assertTrue("useAutoConsumeSchema should default to true", properties.isUseAutoConsumeSchema());
+        properties.setUseAutoConsumeSchema(false);
+        assertFalse(properties.isUseAutoConsumeSchema());
+    }
+
     /**
      * Verify that if 'subscriptionName' is specified in the consumerConfig,
      * it is not overwritten by the default value.
diff --git a/src/test/java/io/numaproj/pulsar/consumer/PulsarConsumerManagerTest.java b/src/test/java/io/numaproj/pulsar/consumer/PulsarConsumerManagerTest.java
diff --git a/src/test/java/io/numaproj/pulsar/consumer/PulsarSourceTest.java b/src/test/java/io/numaproj/pulsar/consumer/PulsarSourceTest.java
