propagate correctly (#162)

streamnative · Sep 11, 2023 · c5b8604 · c5b8604
1 parent ed90ada
commit c5b8604
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 12 deletions.
diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarClientFactory.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarClientFactory.scala
@@ -32,12 +32,13 @@ class DefaultPulsarClientFactory extends PulsarClientFactory {
 
 object PulsarClientFactory {
   val PulsarClientFactoryClassOption = "org.apache.spark.sql.pulsar.PulsarClientFactoryClass"
-  def getOrCreate(sparkConf: SparkConf, params: ju.Map[String, Object]): PulsarClientImpl = {
-    getFactory(sparkConf).getOrCreate(params)
+  def getOrCreate(pulsarClientFactoryClassName: Option[String],
+                  params: ju.Map[String, Object]): PulsarClientImpl = {
+    getFactory(pulsarClientFactoryClassName).getOrCreate(params)
   }
 
-  private def getFactory(sparkConf: SparkConf): PulsarClientFactory = {
-    sparkConf.getOption(PulsarClientFactoryClassOption) match {
+  private def getFactory(pulsarClientFactoryClassName: Option[String]): PulsarClientFactory = {
+    pulsarClientFactoryClassName match {
       case Some(factoryClassName) =>
         Utils.classForName(factoryClassName).getConstructor()
           .newInstance().asInstanceOf[PulsarClientFactory]

diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala
@@ -61,7 +61,8 @@ private[pulsar] case class PulsarHelper(
   import scala.collection.JavaConverters._
 
   protected var client: PulsarClientImpl =
-    PulsarClientFactory.getOrCreate(sparkContext.conf, clientConf)
+    PulsarClientFactory.getOrCreate(
+      sparkContext.conf.getOption(PulsarClientFactory.PulsarClientFactoryClassOption), clientConf)
 
   private var topics: Seq[String] = _
   private var topicPartitions: Seq[String] = _

diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarSinks.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarSinks.scala
@@ -163,7 +163,7 @@ private[pulsar] object PulsarSinks extends Logging {
 
     try {
       PulsarClientFactory
-        .getOrCreate(SparkEnv.get.conf, clientConf)
+        .getOrCreate(None, clientConf)
         .newProducer(schema)
         .topic(topic)
         .loadConf(producerConf)

diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarSource.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarSource.scala
@@ -179,7 +179,9 @@ private[pulsar] class PulsarSource(
       pollTimeoutMs,
       failOnDataLoss,
       subscriptionNamePrefix,
-      jsonOptions)
+      jsonOptions,
+      sqlContext.sparkContext.conf
+        .getOption(PulsarClientFactory.PulsarClientFactoryClassOption))
 
     logInfo(
       "GetBatch generating RDD of offset range: " +

diff --git a/src/main/scala/org/apache/spark/sql/pulsar/PulsarSourceRDD.scala b/src/main/scala/org/apache/spark/sql/pulsar/PulsarSourceRDD.scala
@@ -38,7 +38,8 @@ private[pulsar] abstract class PulsarSourceRDDBase(
     pollTimeoutMs: Int,
     failOnDataLoss: Boolean,
     subscriptionNamePrefix: String,
-    jsonOptions: JSONOptionsInRead)
+    jsonOptions: JSONOptionsInRead,
+    pulsarClientFactoryClassName: Option[String])
     extends RDD[InternalRow](sc, Nil) {
 
   val reportDataLoss = reportDataLossFunc(failOnDataLoss)
@@ -59,7 +60,7 @@ private[pulsar] abstract class PulsarSourceRDDBase(
     val schema: Schema[_] = SchemaUtils.getPSchema(schemaInfo.si)
 
     lazy val reader = PulsarClientFactory
-      .getOrCreate(SparkEnv.get.conf, clientConf)
+      .getOrCreate(pulsarClientFactoryClassName, clientConf)
       .newReader(schema)
       .subscriptionRolePrefix(subscriptionNamePrefix)
       .topic(topic)
@@ -170,7 +171,8 @@ private[pulsar] class PulsarSourceRDD(
     pollTimeoutMs: Int,
     failOnDataLoss: Boolean,
     subscriptionNamePrefix: String,
-    jsonOptions: JSONOptionsInRead)
+    jsonOptions: JSONOptionsInRead,
+    pulsarClientFactoryClassName: Option[String])
     extends PulsarSourceRDDBase(
       sc,
       schemaInfo,
@@ -180,7 +182,8 @@ private[pulsar] class PulsarSourceRDD(
       pollTimeoutMs,
       failOnDataLoss,
       subscriptionNamePrefix,
-      jsonOptions) {
+      jsonOptions,
+      pulsarClientFactoryClassName) {
 
   override def getPreferredLocations(split: Partition): Seq[String] = {
     val part = split.asInstanceOf[PulsarSourceRDDPartition]
@@ -221,7 +224,8 @@ private[pulsar] class PulsarSourceRDD4Batch(
       pollTimeoutMs,
       failOnDataLoss,
       subscriptionNamePrefix,
-      jsonOptions) {
+      jsonOptions,
+      None) {
 
   override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {