Adding PulsarAdmin configs to PulsarHelper

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

@@ -47,6 +47,7 @@ import org.apache.spark.sql.types.StructType
 private[pulsar] case class PulsarHelper(
  serviceUrl: String,
  adminUrl: Option[String],
+ adminConf: ju.Map[String, Object],
  clientConf: ju.Map[String, Object],
  driverGroupIdPrefix: String,
  caseInsensitiveParameters: Map[String, String],
@@ -66,7 +67,7 @@ private[pulsar] case class PulsarHelper(
  // will only be called if latestOffset is called and there should
  // be an exception thrown in PulsarProvider if maxBytes is set,
  // and adminUrl is not set
- private lazy val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl.get)
+ private lazy val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl.get, adminConf)
 
  override def close(): Unit = {
  // do nothing
@@ -136,6 +137,9 @@ private[pulsar] case class PulsarHelper(
  try {
  val (subscription, _) = extractSubscription(predefinedSubscription, tp)
  val consumer = CachedConsumer.getOrCreate(tp, subscription, client)
+ // We need to do this because the consumer does not attempt to
+ // reconnect after calling .seek().
+ // TODO: Remove this once we have upgraded to a version so that this is no longer needed
  if (!consumer.isConnected) consumer.getLastMessageId
  consumer.seek(mid)
  } catch {
@@ -517,10 +521,10 @@ private[pulsar] case class PulsarHelper(
  }
 }
 
-class PulsarAdmissionControlHelper(adminUrl: String)
+class PulsarAdmissionControlHelper(adminUrl: String, conf: ju.Map[String, Object])
  extends Logging {
 
- private lazy val pulsarAdmin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).build()
+ private lazy val pulsarAdmin = PulsarAdmin.builder().serviceHttpUrl(adminUrl).loadConf(conf).build()
 
  import scala.collection.JavaConverters._
 

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

@@ -22,6 +22,7 @@ private[pulsar] object PulsarOptions {
 
  // option key prefix for different modules
  val PulsarClientOptionKeyPrefix: String = "pulsar.client."
+ val PulsarAdminOptionKeyPrefix: String = "pulsar.admin."
  val PulsarProducerOptionKeyPrefix: String = "pulsar.producer."
  val PulsarReaderOptionKeyPrefix: String = "pulsar.reader."
 

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

@@ -56,13 +56,14 @@ private[pulsar] class PulsarProvider
  parameters: Map[String, String]): (String, StructType) = {
 
  val caseInsensitiveParams = validateStreamOptions(parameters)
- val (clientConfig, _, serviceUrlConfig, adminUrl) = prepareConfForReader(parameters)
+ val (clientConfig, _, adminConfig, serviceUrlConfig, adminUrl) = prepareConfForReader(parameters)
 
  val subscriptionNamePrefix = s"spark-pulsar-${UUID.randomUUID}"
  val inferredSchema = Utils.tryWithResource(
  PulsarHelper(
  serviceUrlConfig,
  adminUrl,
+ adminConfig,
  clientConfig,
  subscriptionNamePrefix,
  caseInsensitiveParams,
@@ -85,14 +86,15 @@ private[pulsar] class PulsarProvider
  logDebug(s"Creating Pulsar source: $parameters")
 
  val caseInsensitiveParams = validateStreamOptions(parameters)
- val (clientConfig, readerConfig, serviceUrl, adminUrl) = prepareConfForReader(parameters)
+ val (clientConfig, readerConfig, adminConfig, serviceUrl, adminUrl) = prepareConfForReader(parameters)
  logDebug(
  s"Client config: $clientConfig; Reader config: $readerConfig; Service URL: $serviceUrl")
 
  val subscriptionNamePrefix = getSubscriptionPrefix(parameters)
  val pulsarHelper = PulsarHelper(
  serviceUrl,
  adminUrl,
+ adminConfig,
  clientConfig,
  subscriptionNamePrefix,
  caseInsensitiveParams,
@@ -134,11 +136,12 @@ private[pulsar] class PulsarProvider
 
  val subscriptionNamePrefix = getSubscriptionPrefix(parameters, isBatch = true)
 
- val (clientConfig, readerConfig, serviceUrl, adminUrl) = prepareConfForReader(parameters)
+ val (clientConfig, readerConfig, adminConfig, serviceUrl, adminUrl) = prepareConfForReader(parameters)
  val (start, end, schema, pSchema) = Utils.tryWithResource(
  PulsarHelper(
  serviceUrl,
  adminUrl,
+ adminConfig,
  clientConfig,
  subscriptionNamePrefix,
  caseInsensitiveParams,
@@ -258,6 +261,20 @@ private[pulsar] object PulsarProvider extends Logging {
  }
  }
 
+ private def getAdminParams(parameters: Map[String, String]): Map[String, String] = {
+ val lowercaseKeyMap = parameters.keySet
+ .filter(_.startsWith(PulsarAdminOptionKeyPrefix))
+ .map { k =>
+ k.drop(PulsarAdminOptionKeyPrefix.length).toString -> parameters(k)
+ }
+ .toMap
+ lowercaseKeyMap.map { case (k, v) =>
+ clientConfKeys.getOrElse(
+ k,
+ throw new IllegalArgumentException(s"$k not supported by pulsar")) -> v
+ }
+ }
+
  private def getProducerParams(parameters: Map[String, String]): Map[String, String] = {
  getModuleParams(parameters, PulsarProducerOptionKeyPrefix, producerConfKeys)
  }
@@ -507,17 +524,19 @@ private[pulsar] object PulsarProvider extends Logging {
  }
 
  private def prepareConfForReader(parameters: Map[String, String])
- : (ju.Map[String, Object], ju.Map[String, Object], String, Option[String]) = {
+ : (ju.Map[String, Object], ju.Map[String, Object], ju.Map[String, Object], String, Option[String]) = {
 
  val serviceUrl = getServiceUrl(parameters)
  val adminUrl = getAdminUrl(parameters)
  var clientParams = getClientParams(parameters)
  clientParams += (ServiceUrlOptionKey -> serviceUrl)
  val readerParams = getReaderParams(parameters)
+ val adminParams = getAdminParams(parameters)
 
  (
  paramsToPulsarConf("pulsar.client", clientParams),
  paramsToPulsarConf("pulsar.reader", readerParams),
+ paramsToPulsarConf("pulsar.admin", adminParams),
  serviceUrl, adminUrl)
  }
 

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

@@ -1,5 +1,7 @@
 package org.apache.spark.sql.pulsar
 
+import java.{util => ju}
+
 import org.apache.pulsar.client.admin.PulsarAdmin
 import org.apache.pulsar.client.api.MessageId
 import org.apache.pulsar.client.internal.DefaultImplementation
@@ -33,7 +35,7 @@ class PulsarAdmissionControlSuite extends PulsarSourceTest {
  val firstLedger = getLedgerId(firstMid)
  val firstEntry = getEntryId(firstMid)
  require(getLatestOffsets(Set(topic)).size === 1)
- val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl)
+ val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl, new ju.HashMap[String, Object]())
  val offset = admissionControlHelper.latestOffsetForTopicPartition(topic, MessageId.earliest, approxSizeOfInt)
  assert(getLedgerId(offset) == firstLedger && getEntryId(offset) == firstEntry)
  }
@@ -44,7 +46,7 @@ class PulsarAdmissionControlSuite extends PulsarSourceTest {
  val firstMid = messageIds.head._2
  val secondMid = messageIds.apply(1)._2
  require(getLatestOffsets(Set(topic)).size === 1)
- val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl)
+ val admissionControlHelper = new PulsarAdmissionControlHelper(adminUrl, new ju.HashMap[String, Object]())
  val offset = admissionControlHelper.latestOffsetForTopicPartition(topic, firstMid, approxSizeOfInt)
  assert(getLedgerId(offset) == getLedgerId(secondMid) && getEntryId(offset) == getEntryId(secondMid))