[SPARK-50985][SS] Classify Kafka Timestamp Offsets mismatch error instead of assert and throw error for missing server in KafkaTokenProvider

connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json

@@ -30,6 +30,13 @@
       "Specified: <specifiedPartitions> Assigned: <assignedPartitions>"
     ]
   },
+  "KAFKA_TIMESTAMP_OFFSET_DOES_NOT_MATCH_ASSIGNED" : {
+    "message" : [
+      "Partitions specified for Kafka timestamp based <position> offsets don't match what are assigned. Maybe topic partitions are created ",
+      "or deleted while the query is running.",
+      "Specified: <specifiedPartitions> Assigned: <assignedPartitions>"
+    ]
+  },
   "KAFKA_DATA_LOSS" : {
     "message" : [
       "Some data may have been lost because they are not available in Kafka any more;",

connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala

@@ -165,6 +165,18 @@ object KafkaExceptions {
         "specifiedPartitions" -> specifiedPartitions.toString,
         "assignedPartitions" -> assignedPartitions.toString))
   }
+
+  def timestampOffsetDoesNotMatchAssigned(
+      isStartingOffsets: Boolean,
+      specifiedPartitions: Set[TopicPartition],
+      assignedPartitions: Set[TopicPartition]): KafkaIllegalStateException = {
+    new KafkaIllegalStateException(
+      errorClass = "KAFKA_TIMESTAMP_OFFSET_DOES_NOT_MATCH_ASSIGNED",
+      messageParameters = Map(
+        "position" -> (if (isStartingOffsets) "start" else "end"),
+        "specifiedPartitions" -> specifiedPartitions.toString,
+        "assignedPartitions" -> assignedPartitions.toString))
+  }
 }
 
 /**

connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala

@@ -175,9 +175,12 @@ private[kafka010] class KafkaOffsetReaderAdmin(
     : KafkaSourceOffset = {
 
     val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = { partitions =>
-      assert(partitions.asScala == partitionTimestamps.keySet,
-        "If starting/endingOffsetsByTimestamp contains specific offsets, you must specify all " +
-          s"topics. Specified: ${partitionTimestamps.keySet} Assigned: ${partitions.asScala}")
+      val specifiedPartitions = partitionTimestamps.keySet
+      val assignedPartitions = partitions.asScala.toSet
+      if (specifiedPartitions != assignedPartitions) {
+        throw KafkaExceptions.timestampOffsetDoesNotMatchAssigned(
+          isStartingOffsets, specifiedPartitions, assignedPartitions)
+      }
       logDebug(s"Assigned partitions: $partitions. Seeking to $partitionTimestamps")
     }
 

connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala

@@ -213,9 +213,12 @@ private[kafka010] class KafkaOffsetReaderConsumer(
     : KafkaSourceOffset = {
 
     val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = { partitions =>
-      assert(partitions.asScala == partitionTimestamps.keySet,
-        "If starting/endingOffsetsByTimestamp contains specific offsets, you must specify all " +
-          s"topics. Specified: ${partitionTimestamps.keySet} Assigned: ${partitions.asScala}")
+      val specifiedPartitions = partitionTimestamps.keySet
+      val assignedPartitions = partitions.asScala.toSet
+      if (specifiedPartitions != assignedPartitions) {
+        throw KafkaExceptions.timestampOffsetDoesNotMatchAssigned(
+          isStartingOffsets, specifiedPartitions, assignedPartitions)
+      }
       logDebug(s"Partitions assigned to consumer: $partitions. Seeking to $partitionTimestamps")
     }
 

connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderSuite.scala

@@ -25,10 +25,12 @@ import org.apache.kafka.clients.CommonClientConfigs
 import org.apache.kafka.clients.consumer.ConsumerConfig
 import org.apache.kafka.common.{IsolationLevel, TopicPartition}
 
+import org.apache.spark.SparkException
 import org.apache.spark.sql.QueryTest
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.kafka010.KafkaOffsetRangeLimit.{EARLIEST, LATEST}
+import org.apache.spark.sql.kafka010.KafkaSourceProvider.StrategyOnNoMatchStartingOffset
 import org.apache.spark.sql.test.SharedSparkSession
 
 class KafkaOffsetReaderSuite extends QueryTest with SharedSparkSession with KafkaTest {
@@ -203,6 +205,41 @@ class KafkaOffsetReaderSuite extends QueryTest with SharedSparkSession with Kafk
       KafkaOffsetRange(tp2, 0, 3, None)).sortBy(_.topicPartition.toString))
   }
 
+  testWithAllOffsetFetchingSQLConf(
+    "KAFKA_TIMESTAMP_OFFSET_DOES_NOT_MATCH_ASSIGNED error class"
+  ) {
+    val topic = newTopic()
+    testUtils.createTopic(topic, partitions = 3)
+    val reader = createKafkaReader(topic, minPartitions = Some(4))
+
+    // There are three topic partitions, but we only include two in offsets.
+    val tp1 = new TopicPartition(topic, 0)
+    val tp2 = new TopicPartition(topic, 1)
+    val startingOffsets = SpecificTimestampRangeLimit(Map(tp1 -> EARLIEST, tp2 -> EARLIEST),
+      StrategyOnNoMatchStartingOffset.ERROR)
+    val endingOffsets = SpecificTimestampRangeLimit(Map(tp1 -> LATEST, tp2 -> 3),
+      StrategyOnNoMatchStartingOffset.ERROR)
+
+    val ex = if (reader.isInstanceOf[KafkaOffsetReaderConsumer]) {
+      intercept[SparkException] {
+        reader.getOffsetRangesFromUnresolvedOffsets(startingOffsets, endingOffsets)
+      }.getCause.asInstanceOf[KafkaIllegalStateException]
+    } else {
+      intercept[KafkaIllegalStateException] {
+        reader.getOffsetRangesFromUnresolvedOffsets(startingOffsets, endingOffsets)
+      }
+    }
+
+    checkError(
+      exception = ex,
+      condition = "KAFKA_TIMESTAMP_OFFSET_DOES_NOT_MATCH_ASSIGNED",
+      parameters = Map(
+        "position" -> "start",
+        "specifiedPartitions" -> "Set\\(.*,.*\\)",
+        "assignedPartitions" -> "Set\\(.*,.*,.*\\)"),
+      matchPVals = true)
+  }
+
   private def testWithAllOffsetFetchingSQLConf(name: String)(func: => Any): Unit = {
     Seq("true", "false").foreach { useDeprecatedOffsetFetching =>
       val testName = s"$name with useDeprecatedOffsetFetching $useDeprecatedOffsetFetching"

connector/kafka-0-10-token-provider/src/main/resources/error/kafka-token-provider-error-conditions.json

@@ -0,0 +1,8 @@
+{
+  "MISSING_KAFKA_OPTION" : {
+    "message" : [
+      "Kafka option <option> is not set. Please make sure it is set and retry."
+    ],
+    "sqlState" : "42KDF"
+  }
+}

connector/kafka-0-10-token-provider/src/main/scala/org/apache/spark/kafka010/KafkaConfigUpdater.scala

@@ -57,8 +57,14 @@ private[spark] case class KafkaConfigUpdater(module: String, kafkaParams: Map[St
   }
 
   def setAuthenticationConfigIfNeeded(): this.type = {
+    val bootstrapServers = kafkaParams
+      .get(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG)
+      .map(_.asInstanceOf[String])
+      .getOrElse(throw KafkaTokenProviderExceptions.missingKafkaOption(
+        CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG))
+
     val clusterConfig = KafkaTokenUtil.findMatchingTokenClusterConfig(SparkEnv.get.conf,
-      kafkaParams(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG).asInstanceOf[String])
+      bootstrapServers)
     setAuthenticationConfigIfNeeded(clusterConfig)
   }
 

connector/kafka-0-10-token-provider/src/main/scala/org/apache/spark/kafka010/KafkaTokenProviderException.scala

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.kafka010
+
+import org.apache.spark.{ErrorClassesJsonReader, SparkException}
+
+private object ExceptionsHelper {
+  val errorJsonReader: ErrorClassesJsonReader =
+    new ErrorClassesJsonReader(
+      // Note that though we call them "error classes" here, the proper name is "error conditions",
+      // hence why the name of the JSON file is different. We will address this inconsistency as
+      // part of this ticket: https://issues.apache.org/jira/browse/SPARK-47429
+      Seq(getClass.getClassLoader.getResource("error/kafka-token-provider-error-conditions.json")))
+}
+
+object KafkaTokenProviderExceptions {
+  def missingKafkaOption(option: String): SparkException = {
+    val errClass = "MISSING_KAFKA_OPTION"
+    val param = Map("option" -> option)
+
+    new SparkException(
+      message = ExceptionsHelper.errorJsonReader.getErrorMessage(
+        errClass,
+        param),
+      cause = null,
+      errorClass = Some(errClass),
+      messageParameters = param
+    )
+  }
+}

connector/kafka-0-10-token-provider/src/test/scala/org/apache/spark/kafka010/KafkaConfigUpdaterSuite.scala

@@ -25,7 +25,7 @@ import org.apache.kafka.clients.CommonClientConfigs
 import org.apache.kafka.common.config.SaslConfigs
 import org.apache.kafka.common.security.auth.SecurityProtocol.SASL_PLAINTEXT
 
-import org.apache.spark.SparkFunSuite
+import org.apache.spark.{SparkException, SparkFunSuite}
 
 class KafkaConfigUpdaterSuite extends SparkFunSuite with KafkaDelegationTokenTest {
   private val testModule = "testModule"
@@ -161,4 +161,16 @@ class KafkaConfigUpdaterSuite extends SparkFunSuite with KafkaDelegationTokenTes
     assert(updatedParams.size() === 1)
     assert(updatedParams.get(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG) === bootStrapServers)
   }
+
+  test("setAuthenticationConfigIfNeeded handles missing bootstrap servers config") {
+    val ex = intercept[SparkException] {
+      KafkaConfigUpdater(testModule, kafkaParams = Map.empty[String, String])
+        .setAuthenticationConfigIfNeeded()
+    }
+
+    checkError(
+      exception = ex,
+      condition = "MISSING_KAFKA_OPTION",
+      parameters = Map("option" -> CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG))
+  }
 }