Make the pulsar client creation factory pluggable (#159)

* fix build

* clean up

* increase timeout

* use client factory in pulsar sink

Author: chaoqin-li1123

.github/workflows/ci.yml

@@ -24,7 +24,7 @@ jobs:
   unittest:
     name: Run unit tests
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 45
     steps:
       - uses: actions/checkout@v3
       - name: Setup Java

src/main/scala/org/apache/spark/sql/pulsar/PulsarClientFactory.scala

@@ -0,0 +1,49 @@
+/*
+ * Licensed 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.spark.sql.pulsar
+
+import java.{util => ju}
+
+import org.apache.pulsar.client.impl.PulsarClientImpl
+
+import org.apache.spark.SparkConf
+import org.apache.spark.util.Utils
+
+trait PulsarClientFactory {
+  def getOrCreate(params: ju.Map[String, Object]): PulsarClientImpl
+}
+
+class DefaultPulsarClientFactory extends PulsarClientFactory {
+  def getOrCreate(params: ju.Map[String, Object]): PulsarClientImpl = {
+    CachedPulsarClient.getOrCreate(params)
+  }
+}
+
+object PulsarClientFactory {
+  val PulsarClientFactoryClassOption = "org.apache.spark.sql.pulsar.PulsarClientFactoryClass"
+  def getOrCreate(sparkConf: SparkConf, params: ju.Map[String, Object]): PulsarClientImpl = {
+    getFactory(sparkConf).getOrCreate(params)
+  }
+
+  private def getFactory(sparkConf: SparkConf): PulsarClientFactory = {
+    sparkConf.getOption(PulsarClientFactoryClassOption) match {
+      case Some(factoryClassName) =>
+        Utils.classForName(factoryClassName).getConstructor()
+          .newInstance().asInstanceOf[PulsarClientFactory]
+      case None => new DefaultPulsarClientFactory()
+    }
+  }
+}
+
+

src/main/scala/org/apache/spark/sql/pulsar/PulsarHelper.scala

@@ -32,6 +32,7 @@ import org.apache.pulsar.common.naming.TopicName
 import org.apache.pulsar.common.schema.SchemaInfo
 import org.apache.pulsar.shade.com.google.common.util.concurrent.Uninterruptibles
 
+import org.apache.spark.SparkContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.connector.read.streaming
 import org.apache.spark.sql.connector.read.streaming.{ReadAllAvailable, ReadLimit}
@@ -52,13 +53,15 @@ private[pulsar] case class PulsarHelper(
     driverGroupIdPrefix: String,
     caseInsensitiveParameters: Map[String, String],
     allowDifferentTopicSchemas: Boolean,
-    predefinedSubscription: Option[String])
+    predefinedSubscription: Option[String],
+    sparkContext: SparkContext)
     extends Closeable
     with Logging {
 
   import scala.collection.JavaConverters._
 
-  protected var client: PulsarClientImpl = CachedPulsarClient.getOrCreate(clientConf)
+  protected var client: PulsarClientImpl =
+    PulsarClientFactory.getOrCreate(sparkContext.conf, clientConf)
 
   private var topics: Seq[String] = _
   private var topicPartitions: Seq[String] = _

src/main/scala/org/apache/spark/sql/pulsar/PulsarProvider.scala

@@ -69,7 +69,8 @@ private[pulsar] class PulsarProvider
         subscriptionNamePrefix,
         caseInsensitiveParams,
         getAllowDifferentTopicSchemas(parameters),
-        getPredefinedSubscription(parameters))) { pulsarHelper =>
+        getPredefinedSubscription(parameters),
+        sqlContext.sparkContext)) { pulsarHelper =>
       pulsarHelper.getAndCheckCompatible(schema)
     }
 
@@ -102,7 +103,8 @@ private[pulsar] class PulsarProvider
       subscriptionNamePrefix,
       caseInsensitiveParams,
       getAllowDifferentTopicSchemas(parameters),
-      getPredefinedSubscription(parameters))
+      getPredefinedSubscription(parameters),
+      sqlContext.sparkContext)
 
     val pSchema = pulsarHelper.getAndCheckCompatible(schema)
     logDebug(s"Schema from Spark: $schema; Schema from Pulsar: ${pSchema}")
@@ -151,7 +153,8 @@ private[pulsar] class PulsarProvider
         subscriptionNamePrefix,
         caseInsensitiveParams,
         getAllowDifferentTopicSchemas(parameters),
-        getPredefinedSubscription(parameters))) { pulsarHelper =>
+        getPredefinedSubscription(parameters),
+        sqlContext.sparkContext)) { pulsarHelper =>
       val perTopicStarts =
         pulsarHelper.offsetForEachTopic(caseInsensitiveParams, EarliestOffset, StartOptionKey)
       val startingOffset = SpecificPulsarOffset(

src/main/scala/org/apache/spark/sql/pulsar/PulsarSinks.scala

@@ -20,6 +20,7 @@ import scala.util.control.NonFatal
 
 import org.apache.pulsar.client.api.{Producer, PulsarClientException, Schema}
 
+import org.apache.spark.SparkEnv
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{AnalysisException, DataFrame, SparkSession, SQLContext}
 import org.apache.spark.sql.catalyst.expressions
@@ -161,8 +162,8 @@ private[pulsar] object PulsarSinks extends Logging {
       schema: Schema[T]): Producer[T] = {
 
     try {
-      CachedPulsarClient
-        .getOrCreate(clientConf)
+      PulsarClientFactory
+        .getOrCreate(SparkEnv.get.conf, clientConf)
         .newProducer(schema)
         .topic(topic)
         .loadConf(producerConf)

src/main/scala/org/apache/spark/sql/pulsar/PulsarSourceRDD.scala

@@ -19,7 +19,7 @@ import java.util.concurrent.TimeUnit
 import org.apache.pulsar.client.api.{Message, MessageId, PulsarClientException, Schema}
 import org.apache.pulsar.client.impl.{BatchMessageIdImpl, MessageIdImpl}
 
-import org.apache.spark.{Partition, SparkContext, TaskContext}
+import org.apache.spark.{Partition, SparkContext, SparkEnv, TaskContext}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.json.JSONOptionsInRead
@@ -58,8 +58,8 @@ private[pulsar] abstract class PulsarSourceRDDBase(
     val deserializer = new PulsarDeserializer(schemaInfo.si, jsonOptions)
     val schema: Schema[_] = SchemaUtils.getPSchema(schemaInfo.si)
 
-    lazy val reader = CachedPulsarClient
-      .getOrCreate(clientConf)
+    lazy val reader = PulsarClientFactory
+      .getOrCreate(SparkEnv.get.conf, clientConf)
       .newReader(schema)
       .subscriptionRolePrefix(subscriptionNamePrefix)
       .topic(topic)

src/test/scala/org/apache/spark/sql/pulsar/PulsarClientFactorySuite.scala

@@ -0,0 +1,68 @@
+package org.apache.spark.sql.pulsar
+
+import org.apache.pulsar.client.impl.PulsarClientImpl
+import org.apache.spark.sql.pulsar.PulsarOptions.{ServiceUrlOptionKey, TopicPattern}
+import org.apache.spark.sql.pulsar.TestPulsarClientFactory.counter
+
+import java.{util => ju}
+
+class TestPulsarClientFactory extends PulsarClientFactory {
+  def getOrCreate(params: ju.Map[String, Object]): PulsarClientImpl = {
+    counter += 1
+    new DefaultPulsarClientFactory().getOrCreate(params)
+  }
+}
+
+object TestPulsarClientFactory {
+  var counter = 0
+}
+
+class PulsarClientFactorySuite extends PulsarSourceTest {
+  test("Set Pulsar client factory class") {
+    sparkContext.conf.set(PulsarClientFactory.PulsarClientFactoryClassOption,
+      "org.apache.spark.sql.pulsar.TestPulsarClientFactory")
+    val topic = newTopic()
+    sendMessages(topic, (101 to 105).map { _.toString }.toArray)
+
+    val reader = spark.readStream
+      .format("pulsar")
+      .option(ServiceUrlOptionKey, serviceUrl)
+      .option(TopicPattern, s"$topic.*")
+
+    val pulsar = reader
+      .load()
+      .selectExpr("CAST(__key AS STRING)", "CAST(value AS STRING)")
+
+    testStream(pulsar)(
+      makeSureGetOffsetCalled,
+      StopStream
+    )
+    // Assert that we are using the test factory.
+    assert(TestPulsarClientFactory.counter > 0)
+  }
+
+  test("Unset Pulsar client factory class") {
+    sparkContext.conf.remove(PulsarClientFactory.PulsarClientFactoryClassOption)
+    val oldCount = TestPulsarClientFactory.counter
+    val topic = newTopic()
+    sendMessages(topic, (101 to 105).map { _.toString }.toArray)
+
+    val reader = spark.readStream
+      .format("pulsar")
+      .option(ServiceUrlOptionKey, serviceUrl)
+      .option(TopicPattern, s"$topic.*")
+
+    val pulsar = reader
+      .load()
+      .selectExpr("CAST(__key AS STRING)", "CAST(value AS STRING)")
+
+    testStream(pulsar)(
+      makeSureGetOffsetCalled,
+      StopStream
+    )
+
+    val newCount = TestPulsarClientFactory.counter
+    // The count doesn't change because we are using the default factory.
+    assert(oldCount == newCount)
+  }
+}