[HUDI-1648] Added custom kafka meta fields and custom kafka avro decoder.

hudi-utilities/src/main/java/org/apache/hudi/utilities/decoders/KafkaAvroSchemaDeserializer.java

@@ -0,0 +1,175 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.utilities.decoders;
+
+import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException;
+import io.confluent.kafka.serializers.AbstractKafkaAvroDeserializer;
+import io.confluent.kafka.serializers.KafkaAvroDeserializer;
+import io.confluent.kafka.serializers.NonRecordContainer;
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.io.DatumReader;
+import org.apache.avro.io.DecoderFactory;
+import org.apache.hudi.common.config.TypedProperties;
+import org.apache.hudi.utilities.UtilHelpers;
+import org.apache.hudi.utilities.schema.SchemaProvider;
+import org.apache.kafka.common.errors.SerializationException;
+import org.codehaus.jackson.node.JsonNodeFactory;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Objects;
+
+/**
+ * Extending {@link KafkaAvroSchemaDeserializer} as we need to be able to inject reader schema during deserialization.
+ */
+public class KafkaAvroSchemaDeserializer extends KafkaAvroDeserializer {
+
+  private static final String SCHEMA_PROVIDER_CLASS_PROP = "hoodie.deltastreamer.schemaprovider.class";
+
+  private final DecoderFactory decoderFactory = DecoderFactory.get();
+  private Schema sourceSchema;
+
+  public KafkaAvroSchemaDeserializer() {}
+
+  @Override
+  public void configure(Map<String, ?> configs, boolean isKey) {
+    super.configure(configs, isKey);
+    try {
+      TypedProperties props = getConvertToTypedProperties(configs);
+      SchemaProvider schemaProvider = UtilHelpers.createSchemaProvider(
+          props.getString(SCHEMA_PROVIDER_CLASS_PROP), props, null);
+      sourceSchema = Objects.requireNonNull(schemaProvider).getSourceSchema();
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  /**
+   * Pretty much copy-paste from the {@link AbstractKafkaAvroDeserializer} except line 87:
+   * DatumReader reader = new GenericDatumReader(schema, sourceSchema);
+   * <p>
+   * We need to inject reader schema during deserialization or later stages of the pipeline break.
+   *
+   * @param includeSchemaAndVersion
+   * @param topic
+   * @param isKey
+   * @param payload
+   * @param readerSchema
+   * @return
+   * @throws SerializationException
+   */
+  @Override
+  protected Object deserialize(
+      boolean includeSchemaAndVersion,
+      String topic,
+      Boolean isKey,
+      byte[] payload,
+      Schema readerSchema)
+      throws SerializationException {
+    // 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) {
+      return null;
+    }
+    int id = -1;
+    try {
+      ByteBuffer buffer = getByteBuffer(payload);
+      id = buffer.getInt();
+      Schema schema = schemaRegistry.getByID(id);
+
+      int length = buffer.limit() - 1 - idSize;
+      final Object result;
+      if (schema.getType().equals(Schema.Type.BYTES)) {
+        byte[] bytes = new byte[length];
+        buffer.get(bytes, 0, length);
+        result = bytes;
+      } else {
+        int start = buffer.position() + buffer.arrayOffset();
+        DatumReader reader = new GenericDatumReader(schema, sourceSchema);
+        Object object =
+            reader.read(null, decoderFactory.binaryDecoder(buffer.array(), start, length, null));
+
+        if (schema.getType().equals(Schema.Type.STRING)) {
+          object = object.toString(); // Utf8 -> String
+        }
+        result = object;
+      }
+
+      if (includeSchemaAndVersion) {
+        // Annotate the schema with the version. Note that we only do this if the schema +
+        // version are requested, i.e. in Kafka Connect converters. This is critical because that
+        // code *will not* rely on exact schema equality. Regular deserializers *must not* include
+        // this information because it would return schemas which are not equivalent.
+        //
+        // Note, however, that we also do not fill in the connect.version field. This allows the
+        // Converter to let a version provided by a Kafka Connect source take priority over the
+        // schema registry's ordering (which is implicit by auto-registration time rather than
+        // explicit from the Connector).
+        Integer version = schemaRegistry.getVersion(getSubjectName(topic, isKey), schema);
+        if (schema.getType() == Schema.Type.UNION) {
+          // Can't set additional properties on a union schema since it's just a list, so set it
+          // on the first non-null entry
+          for (Schema memberSchema : schema.getTypes()) {
+            if (memberSchema.getType() != Schema.Type.NULL) {
+              memberSchema.addProp(
+                  SCHEMA_REGISTRY_SCHEMA_VERSION_PROP,
+                  JsonNodeFactory.instance.numberNode(version));
+              break;
+            }
+          }
+        } else {
+          schema.addProp(
+              SCHEMA_REGISTRY_SCHEMA_VERSION_PROP, JsonNodeFactory.instance.numberNode(version));
+        }
+        if (schema.getType().equals(Schema.Type.RECORD)) {
+          return result;
+        } else {
+          return new NonRecordContainer(schema, result);
+        }
+      } else {
+        return result;
+      }
+    } catch (IOException | RuntimeException e) {
+      // avro deserialization may throw AvroRuntimeException, NullPointerException, etc
+      throw new SerializationException("Error deserializing Avro message for id " + id, e);
+    } catch (RestClientException e) {
+      throw new SerializationException("Error retrieving Avro schema for id " + id, e);
+    }
+  }
+
+  private ByteBuffer getByteBuffer(byte[] payload) {
+    ByteBuffer buffer = ByteBuffer.wrap(payload);
+    if (buffer.get() != MAGIC_BYTE) {
+      throw new SerializationException("Unknown magic byte!");
+    }
+    return buffer;
+  }
+
+  private TypedProperties getConvertToTypedProperties(Map<String, ?> configs) {
+    TypedProperties typedProperties = new TypedProperties();
+    for (Entry<String, ?> entry : configs.entrySet()) {
+      typedProperties.put(entry.getKey(), entry.getValue());
+    }
+    return typedProperties;
+  }
+}

hudi-utilities/src/main/java/org/apache/hudi/utilities/schema/SchemaRegistryProvider.java

@@ -25,6 +25,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.avro.Schema;
+import org.apache.hudi.utilities.sources.helpers.AvroKafkaSourceHelpers;
 import org.apache.spark.api.java.JavaSparkContext;
 
 import java.io.IOException;
@@ -42,12 +43,19 @@ public class SchemaRegistryProvider extends SchemaProvider {
    * Configs supported.
    */
   public static class Config {
-
     private static final String SRC_SCHEMA_REGISTRY_URL_PROP = "hoodie.deltastreamer.schemaprovider.registry.url";
-    private static final String TARGET_SCHEMA_REGISTRY_URL_PROP =
-        "hoodie.deltastreamer.schemaprovider.registry.targetUrl";
+    private static final String TARGET_SCHEMA_REGISTRY_URL_PROP = "hoodie.deltastreamer.schemaprovider.registry.targetUrl";
+    private static final String CACHE_SCHEMAS = "hoodie.deltastreamer.schemaprovider.registry.cache_enabled";
   }
 
+  private Schema sourceSchema;
+  private Schema targetSchema;
+  private final boolean cacheDisabled;
+  private final boolean injectKafkaFieldSchema;
+  private final String registryUrl;
+  private final String targetRegistryUrl;
+  private final boolean noTargetSchema;
+
   private static String fetchSchemaFromRegistry(String registryUrl) throws IOException {
     URL registry = new URL(registryUrl);
     ObjectMapper mapper = new ObjectMapper();
@@ -58,30 +66,67 @@ private static String fetchSchemaFromRegistry(String registryUrl) throws IOExcep
   public SchemaRegistryProvider(TypedProperties props, JavaSparkContext jssc) {
     super(props, jssc);
     DataSourceUtils.checkRequiredProperties(props, Collections.singletonList(Config.SRC_SCHEMA_REGISTRY_URL_PROP));
+    this.cacheDisabled = !props.getBoolean(Config.CACHE_SCHEMAS, false);
+    this.injectKafkaFieldSchema = props.getBoolean(AvroKafkaSourceHelpers.INJECT_KAFKA_FIELDS, false);
+    this.registryUrl = config.getString(Config.SRC_SCHEMA_REGISTRY_URL_PROP);
+    this.targetRegistryUrl = config.getString(Config.TARGET_SCHEMA_REGISTRY_URL_PROP, registryUrl);
+    this.noTargetSchema = targetRegistryUrl.equals("null");
   }
 
-  private static Schema getSchema(String registryUrl) throws IOException {
-    return new Schema.Parser().parse(fetchSchemaFromRegistry(registryUrl));
+  private static Schema getSchema(String registryUrl, boolean injectKafkaFieldSchema) throws IOException {
+    Schema schema = new Schema.Parser().parse(fetchSchemaFromRegistry(registryUrl));
+    if (injectKafkaFieldSchema) {
+      return AvroKafkaSourceHelpers.addKafkaMetadataFields(schema);
+    }
+    return schema;
   }
 
   @Override
   public Schema getSourceSchema() {
-    String registryUrl = config.getString(Config.SRC_SCHEMA_REGISTRY_URL_PROP);
+    if (cacheDisabled) {
+      return getSourceSchemaFromRegistry();
+    }
+    if (sourceSchema == null) {
+      synchronized (this) {
+        if (sourceSchema == null) {
+          sourceSchema = getSourceSchemaFromRegistry();
+        }
+      }
+    }
+    return sourceSchema;
+  }
+
+  @Override
+  public Schema getTargetSchema() {
+    if (noTargetSchema) {
+      return null;
+    }
+    if (cacheDisabled) {
+      return getTargetSchemaFromRegistry();
+    }
+    if (targetSchema == null) {
+      synchronized (this) {
+        if (targetSchema == null) {
+          targetSchema = getTargetSchemaFromRegistry();
+        }
+      }
+    }
+    return targetSchema;
+  }
+
+  private Schema getSourceSchemaFromRegistry() {
     try {
-      return getSchema(registryUrl);
+      return getSchema(registryUrl, injectKafkaFieldSchema);
     } catch (IOException ioe) {
       throw new HoodieIOException("Error reading source schema from registry :" + registryUrl, ioe);
     }
   }
 
-  @Override
-  public Schema getTargetSchema() {
-    String registryUrl = config.getString(Config.SRC_SCHEMA_REGISTRY_URL_PROP);
-    String targetRegistryUrl = config.getString(Config.TARGET_SCHEMA_REGISTRY_URL_PROP, registryUrl);
+  private Schema getTargetSchemaFromRegistry() {
     try {
-      return getSchema(targetRegistryUrl);
+      return getSchema(targetRegistryUrl, injectKafkaFieldSchema);
     } catch (IOException ioe) {
-      throw new HoodieIOException("Error reading target schema from registry :" + registryUrl, ioe);
+      throw new HoodieIOException("Error reading target schema from registry :" + targetRegistryUrl, ioe);
     }
   }
 }

hudi-utilities/src/main/java/org/apache/hudi/utilities/sources/AvroKafkaSource.java

@@ -20,8 +20,10 @@
 
 import org.apache.hudi.common.config.TypedProperties;
 import org.apache.hudi.common.util.Option;
+import org.apache.hudi.exception.HoodieException;
 import org.apache.hudi.utilities.deltastreamer.HoodieDeltaStreamerMetrics;
 import org.apache.hudi.utilities.schema.SchemaProvider;
+import org.apache.hudi.utilities.sources.helpers.AvroKafkaSourceHelpers;
 import org.apache.hudi.utilities.sources.helpers.KafkaOffsetGen;
 import org.apache.hudi.utilities.sources.helpers.KafkaOffsetGen.CheckpointUtils;
 
@@ -42,18 +44,34 @@
  */
 public class AvroKafkaSource extends AvroSource {
 
+  private static final String KAFKA_AVRO_VALUE_DESERIALIZER = "hoodie.deltastreamer.source.value.deserializer";
   private static final Logger LOG = LogManager.getLogger(AvroKafkaSource.class);
-
   private final KafkaOffsetGen offsetGen;
-
   private final HoodieDeltaStreamerMetrics metrics;
+  private final boolean injectKafkaData;
 
   public AvroKafkaSource(TypedProperties props, JavaSparkContext sparkContext, SparkSession sparkSession,
       SchemaProvider schemaProvider, HoodieDeltaStreamerMetrics metrics) {
     super(props, sparkContext, sparkSession, schemaProvider);
-    this.metrics = metrics;
+
     props.put("key.deserializer", StringDeserializer.class);
-    props.put("value.deserializer", KafkaAvroDeserializer.class);
+    String deserializerClassName = props.getString(KAFKA_AVRO_VALUE_DESERIALIZER, "");
+
+    if (deserializerClassName.isEmpty()) {
+      props.put("value.deserializer", KafkaAvroDeserializer.class);
+    } else {
+      try {
+        props.put("value.deserializer", Class.forName(deserializerClassName));
+      } catch (ClassNotFoundException e) {
+        String error = "Could not load custom avro kafka deserializer: " + deserializerClassName;
+        LOG.error(error);
+        throw new HoodieException(error, e);
+      }
+    }
+
+    this.metrics = metrics;
+    this.injectKafkaData = props.getBoolean(AvroKafkaSourceHelpers.INJECT_KAFKA_FIELDS, false);
+
     offsetGen = new KafkaOffsetGen(props);
   }
 
@@ -70,6 +88,10 @@ protected InputBatch<JavaRDD<GenericRecord>> fetchNewData(Option<String> lastChe
   }
 
   private JavaRDD<GenericRecord> toRDD(OffsetRange[] offsetRanges) {
+    if (injectKafkaData) {
+      return KafkaUtils.createRDD(sparkContext, offsetGen.getKafkaParams(), offsetRanges,
+          LocationStrategies.PreferConsistent()).map(AvroKafkaSourceHelpers::addKafkaFields);
+    }
     return KafkaUtils.createRDD(sparkContext, offsetGen.getKafkaParams(), offsetRanges,
             LocationStrategies.PreferConsistent()).map(obj -> (GenericRecord) obj.value());
   }