Adding EventHub support in Spark jobs

feast-dev · xiaoyongzhu · Apr 25, 2021 · Apr 26, 2021 · May 8, 2021 · May 19, 2021
commit f2cd8be5d003612f9d89dad811b991d96370b4fb
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala
@@ -33,11 +33,12 @@ object BasePipeline {
     val conf = new SparkConf()
 
     jobConfig.store match {
-      case RedisConfig(host, port, ssl) =>
+      case RedisConfig(host, port, auth, ssl) =>
         conf
           .set("spark.redis.host", host)
           .set("spark.redis.port", port.toString)
           .set("spark.redis.ssl", ssl.toString)
+          .set("spark.redis.auth", auth.toString)
       case BigTableConfig(projectId, instanceId) =>
         conf
           .set("spark.bigtable.projectId", projectId)

diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala
@@ -21,6 +21,8 @@ import org.joda.time.{DateTime, DateTimeZone}
 import org.json4s._
 import org.json4s.ext.JavaEnumNameSerializer
 import org.json4s.jackson.JsonMethods.{parse => parseJSON}
+import org.json4s.ext.JavaEnumNameSerializer
+import scala.collection.mutable.ArrayBuffer
 
 object IngestionJob {
   import Modes._
@@ -116,8 +118,21 @@ object IngestionJob {
       .action((x, c) => c.copy(streamingTriggeringSecs = x))
   }
 
+    opt[String](name = "kafka_sasl_auth")
+      .action((x, c) => c.copy(kafkaSASL = Some(x)))
+  }
+
   def main(args: Array[String]): Unit = {
-    parser.parse(args, IngestionJobConfig()) match {
+    println("Debug... Received following argument:")
+    println(args.toList)
+    val args_modified = new Array[String](args.length)
+    for ( i <- 0 to (args_modified.length - 1)) {
+      args_modified(i) = args(i).replace(" }", "}");
+      args_modified(i) = args_modified(i).replace("\\", "\\\"");
+    }
+    println("Remove additional spaces in args:")
+    println(args_modified.toList)
+    parser.parse(args_modified, IngestionJobConfig()) match {
       case Some(config) =>
         println(s"Starting with config $config")
         config.mode match {

diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala
@@ -130,12 +130,13 @@ case class IngestionJobConfig(
     source: Source = null,
     startTime: DateTime = DateTime.now(),
     endTime: DateTime = DateTime.now(),
-    store: StoreConfig = RedisConfig("localhost", 6379, false),
+    store: StoreConfig = RedisConfig("localhost", 6379, "", false),
     metrics: Option[MetricConfig] = None,
     deadLetterPath: Option[String] = None,
     stencilURL: Option[String] = None,
     streamingTriggeringSecs: Int = 0,
     validationConfig: Option[ValidationConfig] = None,
     doNotIngestInvalidRows: Boolean = false,
     checkpointPath: Option[String] = None
+    kafkaSASL: Option[String] = None
 )
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala
@@ -37,6 +37,8 @@ import org.apache.spark.sql.streaming.StreamingQuery
 import org.apache.spark.sql.types.BooleanType
 import org.apache.spark.{SparkEnv, SparkFiles}
 import org.apache.spark.eventhubs._
+import org.apache.kafka.common.security.plain.PlainLoginModule
+import org.apache.kafka.common.security.JaasContext
 
 /**
   * Streaming pipeline (currently in micro-batches mode only, since we need to have multiple sinks: redis & deadletters).
@@ -60,24 +62,36 @@ object StreamingPipeline extends BasePipeline with Serializable {
     val rowValidator  = new RowValidator(featureTable, config.source.eventTimestampColumn)
     val metrics       = new IngestionPipelineMetrics
     val validationUDF = createValidationUDF(sparkSession, config)
-    val connStr = "Endpoint=sb://xiaoyzhufeasttest.servicebus.windows.net/;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=z9obEAyVvD36fZIEvvtNlCRBEDjIrsfNfDAbgDyTbDg=;EntityPath=xiaoyzhufeasttesteh"
-    val ehConf = EventHubsConf(connStr).setStartingPosition(EventPosition.fromStartOfStream)
 
+    val EH_SASL = "org.apache.kafka.common.security.plain.PlainLoginModule required username=\"$ConnectionString\" password=\"Endpoint=sb://xiaoyzhufeasttest.servicebus.windows.net/;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=z9obEAyVvD36fZIEvvtNlCRBEDjIrsfNfDAbgDyTbDg=;EntityPath=driver_trips\";"
 
     val input = config.source match {
       case source: KafkaSource =>
-        sparkSession.readStream
-          .format("kafka")
-          .option("kafka.bootstrap.servers", source.bootstrapServers)
-          .option("subscribe", source.topic)
-          .load()
-      case source: EventHubSource =>
-        sparkSession.readStream
-          .format("eventhubs")
-          .options(ehConf.toMap)
-          .load()
+        if (config.kafkaSASL.nonEmpty)
+        {
+          // if we have authentication enabled
+          sparkSession.readStream
+            .format("kafka")
+            .option("subscribe", source.topic)
+            .option("kafka.bootstrap.servers", source.bootstrapServers)
+            .option("kafka.sasl.mechanism", "PLAIN")
+            .option("kafka.security.protocol", "SASL_SSL")
+            .option("kafka.sasl.jaas.config", config.kafkaSASL.get)
+            .option("kafka.request.timeout.ms", "60000")
+            .option("kafka.session.timeout.ms", "60000")
+            .option("failOnDataLoss", "false")
+            .load()
+        }
+        else
+        {
+          sparkSession.readStream
+            .format("kafka")
+            .option("kafka.bootstrap.servers", source.bootstrapServers)
+            .option("subscribe", source.topic)
+            .load()     
+        }
       case source: MemoryStreamingSource =>
-        source.read
+        source.read        
     }
 
     val parsed = config.source.asInstanceOf[StreamingSource].format match {