set authentication configuration in Pulsar client (#45) (#48)

Co-authored-by: Attila Tóth <[email protected]>

Author: atezs82

README.md

@@ -123,8 +123,8 @@ val df = spark
   .load()
 df.selectExpr("CAST(__key AS STRING)", "CAST(value AS STRING)")
   .as[(String, String)]
-
 ```
+
 > #### Tip
 > For more information on how to use other language bindings for Spark Structured Streaming,
 > see [Structured Streaming Programming Guide](https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html).
@@ -330,6 +330,62 @@ You can use `org.apache.spark.sql.pulsar.JsonUtils.topicOffsets(Map[String, Mess
 
 </table>
 
+#### Authentication
+Should the Pulsar cluster require authentication, credentials can be set in the following way.
+
+The following examples are in Scala.
+```scala
+// Secure connection with authentication, using the same credentials on the
+// Pulsar client and admin interface (if not given explicitly, the client configuration
+// is used for admin as well).
+val df = spark
+  .readStream
+  .format("pulsar")
+  .option("service.url", "pulsar://localhost:6650")
+  .option("admin.url", "http://localhost:8080")
+  .option("pulsar.client.authPluginClassName","org.apache.pulsar.client.impl.auth.AuthenticationToken")
+  .option("pulsar.client.authParams","token:<valid client JWT token>")
+  .option("topicsPattern", "sensitiveTopic")
+  .load()
+df.selectExpr("CAST(__key AS STRING)", "CAST(value AS STRING)")
+  .as[(String, String)]
+// Secure connection with authentication, using different credentials for
+// Pulsar client and admin interfaces.
+val df = spark
+  .readStream
+  .format("pulsar")
+  .option("service.url", "pulsar://localhost:6650")
+  .option("admin.url", "http://localhost:8080")
+  .option("pulsar.admin.authPluginClassName","org.apache.pulsar.client.impl.auth.AuthenticationToken")
+  .option("pulsar.admin.authParams","token:<valid admin JWT token>")
+  .option("pulsar.client.authPluginClassName","org.apache.pulsar.client.impl.auth.AuthenticationToken")
+  .option("pulsar.client.authParams","token:<valid client JWT token>")
+  .option("topicsPattern", "sensitiveTopic")
+  .load()
+df.selectExpr("CAST(__key AS STRING)", "CAST(value AS STRING)")
+  .as[(String, String)]
+
+// Secure connection with client TLS enabled.
+// Note that the certificate file has to be present at the specified
+// path on every machine of the cluster!
+val df = spark
+  .readStream
+  .format("pulsar")
+  .option("service.url", "pulsar+ssl://localhost:6651")
+  .option("admin.url", "http://localhost:8080")
+  .option("pulsar.admin.authPluginClassName","org.apache.pulsar.client.impl.auth.AuthenticationToken")
+  .option("pulsar.admin.authParams","token:<valid admin JWT token>")
+  .option("pulsar.client.authPluginClassName","org.apache.pulsar.client.impl.auth.AuthenticationToken")
+  .option("pulsar.client.authParams","token:<valid client JWT token>")
+  .option("pulsar.client.tlsTrustCertsFilePath","/path/to/tls/cert/cert.pem")
+  .option("pulsar.client.tlsAllowInsecureConnection","false")
+  .option("pulsar.client.tlsHostnameVerificationenable","true")
+  .option("topicsPattern", "sensitiveTopic")
+  .load()
+df.selectExpr("CAST(__key AS STRING)", "CAST(value AS STRING)")
+  .as[(String, String)]
+```
+
 #### Schema of Pulsar source
 * For topics without schema or with primitive schema in Pulsar, messages' payload
 is loaded to a `value` column with the corresponding type with Pulsar schema.
@@ -458,7 +514,15 @@ A possible solution to remove duplicates when reading the written data could be
 
 Client/producer/reader configurations of Pulsar can be set via `DataStreamReader.option`
 with `pulsar.client.`/`pulsar.producer.`/`pulsar.reader.` prefix, e.g,
-`stream.option("pulsar.reader.receiverQueueSize", "1000000")`. For possible Pulsar parameters, check docs at
+`stream.option("pulsar.reader.receiverQueueSize", "1000000")`. 
+Since the connector needs to access the Pulsar Admin interface as well, separate 
+configuration of the admin client can be set via the same method with the
+`pulsar.admin` prefix. For example: `stream.option("pulsar.admin.authParams","token:<token>")`.
+This can be useful if a different authentication plugin or
+token need to be used. If this is not given explicitly, the client
+parameters (with `pulsar.client` prefix) will be used for accessing the admin
+interface as well.
+For possible Pulsar parameters, check docs at
 [Pulsar client libraries](https://pulsar.apache.org/docs/en/client-libraries/).
 
 ## Build Spark Pulsar Connector
@@ -488,11 +552,26 @@ $ cd pulsar-spark
 $ mvn clean install -DskipTests
 ```
 
+If you get the following error during compilation, try running Maven with Java 8:  
+```
+[ERROR] [Error] : Source option 6 is no longer supported. Use 7 or later.
+[ERROR] [Error] : Target option 6 is no longer supported. Use 7 or later.
+```
+
 5. Run the tests.
 
 ```bash
 $ mvn clean install
 ```
+
+Note: by configuring `scalatest-maven-plugin` in the [usual ways](https://www.scalatest.org/user_guide/using_the_scalatest_maven_plugin), individual tests can be executed, if that is needed:
+
+```bash
+mvn -Dsuites=org.apache.spark.sql.pulsar.CachedPulsarClientSuite clean install
+```
+
+This might be handy if test execution is slower, or you get a `java.io.IOException: Too many open files` exception during full suite run.
+
 Once the installation is finished, there is a fat jar generated under both local maven repo and `target` directory.
 
 

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

@@ -22,8 +22,11 @@ import scala.util.control.NonFatal
 import com.google.common.cache._
 import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException}
 
+import org.apache.pulsar.client.api.{ClientBuilder, PulsarClient}
+
 import org.apache.spark.SparkEnv
 import org.apache.spark.internal.Logging
+import org.apache.spark.sql.pulsar.PulsarOptions.{AUTH_PARAMS, AUTH_PLUGIN_CLASS_NAME, TLS_ALLOW_INSECURE_CONNECTION, TLS_HOSTNAME_VERIFICATION_ENABLE, TLS_TRUST_CERTS_FILE_PATH}
 
 private[pulsar] object CachedPulsarClient extends Logging {
 
@@ -62,7 +65,9 @@ private[pulsar] object CachedPulsarClient extends Logging {
       .removalListener(removalListener)
       .build[Seq[(String, Object)], Client](cacheLoader)
 
-  private def createPulsarClient(pulsarConf: ju.Map[String, Object]): Client = {
+  def createPulsarClient(
+                          pulsarConf: ju.Map[String, Object],
+                          pulsarClientBuilder: ClientBuilder = PulsarClient.builder()): Client = {
     val pulsarServiceUrl =
       pulsarConf.get(PulsarOptions.SERVICE_URL_OPTION_KEY).asInstanceOf[String]
     val clientConf = new PulsarConfigUpdater(
@@ -72,11 +77,27 @@ private[pulsar] object CachedPulsarClient extends Logging {
     ).rebuild()
     logInfo(s"Client Conf = ${clientConf}")
     try {
-      val pulsarClient: Client = org.apache.pulsar.client.api.PulsarClient
-        .builder()
+      pulsarClientBuilder
         .serviceUrl(pulsarServiceUrl)
         .loadConf(clientConf)
-        .build();
+      // Set TLS and authentication parameters if they were given
+      if (clientConf.containsKey(AUTH_PLUGIN_CLASS_NAME)) {
+        pulsarClientBuilder.authentication(
+          clientConf.get(AUTH_PLUGIN_CLASS_NAME).toString, clientConf.get(AUTH_PARAMS).toString)
+      }
+      if (clientConf.containsKey(TLS_ALLOW_INSECURE_CONNECTION)) {
+        pulsarClientBuilder.allowTlsInsecureConnection(
+          clientConf.get(TLS_ALLOW_INSECURE_CONNECTION).toString.toBoolean)
+      }
+      if (clientConf.containsKey(TLS_HOSTNAME_VERIFICATION_ENABLE)) {
+        pulsarClientBuilder.enableTlsHostnameVerification(
+          clientConf.get(TLS_HOSTNAME_VERIFICATION_ENABLE).toString.toBoolean)
+      }
+      if (clientConf.containsKey(TLS_TRUST_CERTS_FILE_PATH)) {
+        pulsarClientBuilder.tlsTrustCertsFilePath(
+          clientConf.get(TLS_TRUST_CERTS_FILE_PATH).toString)
+      }
+      val pulsarClient: Client = pulsarClientBuilder.build()
       logDebug(
         s"Created a new instance of PulsarClient for serviceUrl = $pulsarServiceUrl,"
           + s" clientConf = $clientConf.")

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

@@ -25,8 +25,9 @@ import org.apache.spark.internal.Logging
 private[pulsar] case class PulsarConfigUpdater(
     module: String,
     pulsarParams: Map[String, Object],
-    blacklistedKeys: Set[String] = Set())
-    extends Logging {
+    blacklistedKeys: Set[String] = Set(),
+    keysToHideInLog: Set[String] = Set(PulsarOptions.AUTH_PARAMS))
+  extends Logging {
 
   private val map = new ju.HashMap[String, Object](pulsarParams.asJava)
 
@@ -36,10 +37,12 @@ private[pulsar] case class PulsarConfigUpdater(
 
   def set(key: String, value: Object, map: ju.Map[String, Object]): this.type = {
     if (blacklistedKeys.contains(key)) {
-      logInfo(s"$module: Skip $key")
+      logInfo(s"$module: Skip '$key'")
     } else {
       map.put(key, value)
-      logInfo(s"$module: Set $key to $value, earlier value: ${pulsarParams.getOrElse(key, "")}")
+      logInfo(s"$module: Set '$key' to " +
+        s"'${printConfigValue(key, Option(value))}'," +
+        s" earlier value: '${printConfigValue(key, pulsarParams.get(key))}'")
     }
     this
   }
@@ -50,21 +53,17 @@ private[pulsar] case class PulsarConfigUpdater(
 
   def setIfUnset(key: String, value: Object, map: ju.Map[String, Object]): this.type = {
     if (blacklistedKeys.contains(key)) {
-      logInfo(s"$module: Skip $key")
+      logInfo(s"$module: Skip '$key'")
     } else {
       if (!map.containsKey(key)) {
         map.put(key, value)
-        logDebug(s"$module: Set $key to $value")
+        logInfo(s"$module: Set '$key' to " +
+          s"'${printConfigValue(key, pulsarParams.get(key))}'")
       }
     }
     this
   }
 
-  def setAuthenticationConfigIfNeeded(): this.type = {
-    // FIXME: not implemented yet
-    this
-  }
-
   def build(): ju.Map[String, Object] = map
 
   def rebuild(): ju.Map[String, Object] = {
@@ -76,4 +75,21 @@ private[pulsar] case class PulsarConfigUpdater(
     map
   }
 
+  private val HideCompletelyLimit = 6
+  private val ShowFractionOfHiddenValue = 1.0 / 3.0
+  private val CompletelyHiddenMessage = "...<completely hidden>..."
+  private def printConfigValue(key: String,
+                               maybeVal: Option[Object]): String = {
+    val value = maybeVal.map(_.toString).getOrElse("")
+    if (keysToHideInLog.contains(key)) {
+      if (value.length < HideCompletelyLimit) {
+        return CompletelyHiddenMessage
+      } else {
+        return s"${value.take((value.length * ShowFractionOfHiddenValue).toInt)}..."
+      }
+    }
+
+    value
+  }
+
 }

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

@@ -62,6 +62,8 @@ class PulsarContinuousReader(
   val reportDataLoss = reportDataLossFunc(failOnDataLoss)
 
   private var offset: Offset = _
+
+  private var stopped = false
   override def setStartOffset(start: ju.Optional[Offset]): Unit = {
     offset = start.orElse {
       val actualOffsets = SpecificPulsarOffset(
@@ -129,9 +131,12 @@ class PulsarContinuousReader(
     metadataReader.commitCursorToOffset(off)
   }
 
-  override def stop(): Unit = {
-    metadataReader.removeCursor()
-    metadataReader.close()
+  override def stop(): Unit = synchronized {
+    if (!stopped) {
+      metadataReader.removeCursor()
+      metadataReader.close()
+      stopped = true
+    }
   }
 }
 

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

@@ -26,7 +26,7 @@ import org.apache.pulsar.common.naming.TopicName
 import org.apache.pulsar.common.schema.SchemaInfo
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.pulsar.PulsarOptions.TOPIC_OPTION_KEYS
+import org.apache.spark.sql.pulsar.PulsarOptions.{AUTH_PARAMS, AUTH_PLUGIN_CLASS_NAME, TLS_ALLOW_INSECURE_CONNECTION, TLS_HOSTNAME_VERIFICATION_ENABLE, TLS_TRUST_CERTS_FILE_PATH, TOPIC_OPTION_KEYS}
 import org.apache.spark.sql.types.StructType
 
 /**
@@ -38,24 +38,22 @@ private[pulsar] case class PulsarMetadataReader(
     serviceUrl: String,
     adminUrl: String,
     clientConf: ju.Map[String, Object],
+    adminClientConf: ju.Map[String, Object],
     driverGroupIdPrefix: String,
     caseInsensitiveParameters: Map[String, String])
     extends Closeable
     with Logging {
 
   import scala.collection.JavaConverters._
 
-  protected val admin: PulsarAdmin = AdminUtils.buildAdmin(adminUrl, clientConf)
-  protected var client: PulsarClient = null
+  protected val admin: PulsarAdmin = AdminUtils.buildAdmin(adminUrl, adminClientConf)
+  protected var client: PulsarClient = CachedPulsarClient.getOrCreate(clientConf)
 
   private var topics: Seq[String] = _
   private var topicPartitions: Seq[String] = _
 
   override def close(): Unit = {
     admin.close()
-    if (client != null) {
-      client.close()
-    }
   }
 
   def setupCursor(startingPos: PerTopicOffset): Unit = {
@@ -391,9 +389,6 @@ private[pulsar] case class PulsarMetadataReader(
             PulsarSourceUtils.seekableLatestMid(admin.topics().getLastMessageId(tp)))
         } else {
           assert (time > 0, s"time less than 0: $time")
-          if (client == null) {
-            client = PulsarClient.builder().serviceUrl(serviceUrl).build()
-          }
           val reader = client
             .newReader()
             .topic(tp)
@@ -458,9 +453,6 @@ private[pulsar] case class PulsarMetadataReader(
         UserProvidedMessageId(
           PulsarSourceUtils.seekableLatestMid(admin.topics().getLastMessageId(tp)))
       case _ =>
-        if (client == null) {
-          client = PulsarClient.builder().serviceUrl(serviceUrl).build()
-        }
         val reader = client
           .newReader()
           .startMessageId(off)

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

@@ -49,6 +49,7 @@ private[pulsar] class PulsarMicroBatchReader(
 
   private var startTopicOffsets: Map[String, MessageId] = _
   private var endTopicOffsets: Map[String, MessageId] = _
+  private var stopped = false
 
   val reportDataLoss = reportDataLossFunc(failOnDataLoss)
 
@@ -151,9 +152,12 @@ private[pulsar] class PulsarMicroBatchReader(
     }.asJava
   }
 
-  override def stop(): Unit = {
-    metadataReader.removeCursor()
-    metadataReader.close()
+  override def stop(): Unit = synchronized {
+    if (!stopped) {
+      metadataReader.removeCursor()
+      metadataReader.close()
+      stopped = true
+    }
   }
 }
 

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

@@ -19,6 +19,7 @@ import org.apache.pulsar.common.naming.TopicName
 private[pulsar] object PulsarOptions {
 
   // option key prefix for different modules
+  val PULSAR_ADMIN_OPTION_KEY_PREFIX = "pulsar.admin."
   val PULSAR_CLIENT_OPTION_KEY_PREFIX = "pulsar.client."
   val PULSAR_PRODUCER_OPTION_KEY_PREFIX = "pulsar.producer."
   val PULSAR_CONSUMER_OPTION_KEY_PREFIX = "pulsar.consumer."
@@ -47,12 +48,12 @@ private[pulsar] object PulsarOptions {
   val POLL_TIMEOUT_MS = "polltimeoutms"
   val FAIL_ON_DATA_LOSS_OPTION_KEY = "failondataloss"
 
-  val AUTH_PLUGIN_CLASS_NAME = "authpluginclassname"
-  val AUTH_PARAMS = "authparams"
-  val TLS_TRUST_CERTS_FILE_PATH = "tlstrustcertsfilepath"
-  val TLS_ALLOW_INSECURE_CONNECTION = "tlsallowinsecureconnection"
-  val USE_TLS = "usetls"
-  val TLS_HOSTNAME_VERIFICATION_ENABLE = "tlshostnameverificationenable"
+  val AUTH_PLUGIN_CLASS_NAME = "authPluginClassName"
+  val AUTH_PARAMS = "authParams"
+  val TLS_TRUST_CERTS_FILE_PATH = "tlsTrustCertsFilePath"
+  val TLS_ALLOW_INSECURE_CONNECTION = "tlsAllowInsecureConnection"
+  val USE_TLS = "useTls"
+  val TLS_HOSTNAME_VERIFICATION_ENABLE = "tlsHostnameVerificationEnable"
 
 
   val INSTRUCTION_FOR_FAIL_ON_DATA_LOSS_FALSE =