DGS-22814 Allow function to be passed during Avro deserialization (#4004)

Author: rayokota

avro-serializer/src/main/java/io/confluent/kafka/serializers/AbstractKafkaAvroDeserializer.java

@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.ExecutionException;
+import java.util.function.Function;
 import org.apache.avro.Schema;
 import org.apache.avro.Schema.Type;
 import org.apache.avro.generic.GenericContainer;
@@ -253,12 +254,21 @@ protected GenericContainerWithVersion deserializeWithSchemaAndVersion(
   protected GenericContainerWithVersion deserializeWithSchemaAndVersion(
       String topic, boolean isKey, Headers headers, byte[] payload)
       throws SerializationException, InvalidConfigurationException {
-    return deserializeWithSchemaAndVersion(topic, isKey, headers, payload, null);
+    return deserializeWithSchemaAndVersion(topic, isKey, headers, payload, (Schema) null);
   }
 
   protected GenericContainerWithVersion deserializeWithSchemaAndVersion(
       String topic, boolean isKey, Headers headers, byte[] payload, Schema readerSchema)
       throws SerializationException, InvalidConfigurationException {
+    return deserializeWithSchemaAndVersion(
+        topic, isKey, headers, payload,
+        writerAvroSchema -> readerSchema != null ? new AvroSchema(readerSchema) : null);
+  }
+
+  protected GenericContainerWithVersion deserializeWithSchemaAndVersion(
+      String topic, boolean isKey, Headers headers, byte[] payload,
+      Function<AvroSchema, AvroSchema> writerToReaderSchemaFunc)
+      throws SerializationException, InvalidConfigurationException {
     // Even if the caller requests schema & version, if the payload is null we cannot include it.
     // The caller must handle this case.
     if (payload == null) {
@@ -276,8 +286,8 @@ protected GenericContainerWithVersion deserializeWithSchemaAndVersion(
     // explicit from the Connector).
     DeserializationContext context = new DeserializationContext(topic, isKey, headers, payload);
     AvroSchema schema = context.schemaForDeserialize();
-    AvroSchema readerAvroSchema = readerSchema != null
-        ? new AvroSchema(readerSchema)
+    AvroSchema readerAvroSchema = writerToReaderSchemaFunc != null
+        ? writerToReaderSchemaFunc.apply(schema)
         : specificAvroReaderSchema != null
             ? new AvroSchema(specificAvroReaderSchema)
             : null;

avro-serializer/src/main/java/io/confluent/kafka/serializers/KafkaAvroDeserializer.java

@@ -16,6 +16,8 @@
 
 package io.confluent.kafka.serializers;
 
+import io.confluent.kafka.schemaregistry.avro.AvroSchema;
+import java.util.function.Function;
 import org.apache.avro.Schema;
 import org.apache.kafka.common.header.Headers;
 
@@ -128,6 +130,12 @@ public GenericContainerWithVersion deserializeWithSchema(
     return deserializeWithSchemaAndVersion(topic, isKey, headers, bytes, readerSchema);
   }
 
+  public GenericContainerWithVersion deserializeWithSchema(
+      String topic, Headers headers, byte[] bytes,
+      Function<AvroSchema, AvroSchema> writerToReaderSchemaFunc) {
+    return deserializeWithSchemaAndVersion(topic, isKey, headers, bytes, writerToReaderSchemaFunc);
+  }
+
   @Override
   public void close() {
     try {

avro-serializer/src/test/java/io/confluent/kafka/serializers/KafkaAvroSerializerTest.java

@@ -1090,6 +1090,11 @@ public void testKafkaAvroSerializerWithProjection() {
         topic, headers, bytes, User.getClassSchema());
     assertEquals(new AvroSchema(ExtendedUser.SCHEMA$), schemaAndValue.getSchema());
     assertEquals(obj, schemaAndValue.getValue());
+
+    schemaAndValue = avroDeserializer.deserializeWithSchema(
+        topic, headers, bytes, x -> new AvroSchema(User.getClassSchema()));
+    assertEquals(new AvroSchema(ExtendedUser.SCHEMA$), schemaAndValue.getSchema());
+    assertEquals(obj, schemaAndValue.getValue());
   }
 
   @Test