Add unit tests covering streaming query and batch Id logs

Author: BrooksWalls

core/src/test/scala/org/apache/spark/scheduler/PoolSuite.scala

@@ -367,6 +367,43 @@ class PoolSuite extends SparkFunSuite with LocalSparkContext {
     }
   }
 
+  test("Fair Scheduler addTaskSetManager logs include streaming query Id and batch Id") {
+    val xmlPath = getClass.getClassLoader.getResource("fairscheduler.xml").getFile()
+    val conf = new SparkConf().set(SCHEDULER_ALLOCATION_FILE, xmlPath)
+    sc = new SparkContext(LOCAL, APP_NAME, conf)
+    val taskScheduler = new TaskSchedulerImpl(sc)
+
+    val rootPool = new Pool("", FAIR, 0, 0)
+    val schedulableBuilder = new FairSchedulableBuilder(rootPool, sc)
+    schedulableBuilder.buildPools()
+
+    val testQueryId = "test-query-id-5678"
+    val testBatchId = "99"
+    val properties = new Properties()
+    properties.setProperty(schedulableBuilder.FAIR_SCHEDULER_PROPERTIES, "1")
+    properties.setProperty("sql.streaming.queryId", testQueryId)
+    properties.setProperty("streaming.sql.batchId", testBatchId)
+
+    val taskSetManager = createTaskSetManager(0, 1, taskScheduler)
+
+    val logAppender = new LogAppender("pool streaming logs", maxEvents = 1000)
+    val loggerName = classOf[FairSchedulableBuilder].getName
+
+    withLogAppender(logAppender, loggerNames = Seq(loggerName)) {
+      schedulableBuilder.addTaskSetManager(taskSetManager, properties)
+    }
+
+    val logs = logAppender.loggingEvents.map(_.getMessage.getFormattedMessage)
+    val expectedQueryPrefix = s"[queryId = ${testQueryId.take(5)}]"
+    val expectedBatchPrefix = s"[batchId = $testBatchId]"
+    val addedLogs = logs.filter(msg =>
+      msg.contains("Added task set") &&
+        msg.contains(expectedQueryPrefix) && msg.contains(expectedBatchPrefix))
+    assert(addedLogs.nonEmpty,
+      s"Expected 'Added task set' log to contain '$expectedQueryPrefix' " +
+        s"and '$expectedBatchPrefix'.\nCaptured logs:\n${logs.mkString("\n")}")
+  }
+
   private def verifyPool(rootPool: Pool, poolName: String, expectedInitMinShare: Int,
                          expectedInitWeight: Int, expectedSchedulingMode: SchedulingMode): Unit = {
     val selectedPool = rootPool.getSchedulableByName(poolName)

core/src/test/scala/org/apache/spark/scheduler/SchedulableBuilderSuite.scala

@@ -0,0 +1,68 @@
+/*
+ * 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.scheduler
+
+import java.util.Properties
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.scheduler.SchedulableBuilder.{BATCH_ID_KEY, QUERY_ID_KEY}
+
+class SchedulableBuilderSuite extends SparkFunSuite {
+
+  test("schedulingLogStreamingContext - both queryId and batchId present") {
+    val props = new Properties()
+    val queryId = "12345-test-query-id"
+    val batchId = "23"
+
+    props.setProperty(QUERY_ID_KEY, queryId)
+    props.setProperty(BATCH_ID_KEY, batchId)
+
+    val logContext = SchedulableBuilder.schedulingLogStreamingContext(props)
+
+    // log context should include the truncated query Id and the provided batch Id
+    assertResult("[queryId = 12345] [batchId = 23] ")(logContext.message)
+  }
+
+  test("schedulingLogStreamingContext - only queryId and no batchId present") {
+    val props = new Properties()
+    val queryId = "12345-test-query-id"
+
+    props.setProperty(QUERY_ID_KEY, queryId)
+
+    val logContext = SchedulableBuilder.schedulingLogStreamingContext(props)
+
+    // log context should include the truncated query Id but no batch Id
+    assertResult("[queryId = 12345] ")(logContext.message)
+  }
+
+  test("schedulingLogStreamingContext - no queryId or batchId present") {
+    val props = new Properties()
+
+    val logContext = SchedulableBuilder.schedulingLogStreamingContext(props)
+
+    // log context should be an empty string
+    assertResult("")(logContext.message)
+  }
+
+  test("schedulingLogStreamingContext - handles null properties") {
+    val logContext = SchedulableBuilder.schedulingLogStreamingContext(null)
+
+    // log context should be an empty string
+    assertResult("")(logContext.message)
+  }
+}

core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala

@@ -2759,6 +2759,57 @@ class TaskSetManagerSuite
     assert(taskSetManager.taskSetExcludelistHelperOpt.get.isDryRun)
   }
 
+  test("SPARK-56326: Streaming query Id and batch Id are included in scheduling log " +
+    "messages") {
+    sc = new SparkContext("local", "test")
+    sched = new FakeTaskScheduler(sc)
+    val testQueryId = "test-query-id-1234"
+    val testBatchId = "42"
+    // Create a TaskSet with a non-null Properties containing the streaming metadata.
+    val properties = new Properties()
+    properties.setProperty("sql.streaming.queryId", testQueryId)
+    properties.setProperty("streaming.sql.batchId", testBatchId)
+    val taskSet = new TaskSet(Array(new FakeTask(0, 0, Nil)),
+      0, 0, 0, properties, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None)
+
+    val clock = new ManualClock
+    val logAppender = new LogAppender("streaming scheduling logs", maxEvents = 1000)
+    val loggerName = classOf[TaskSetManager].getName
+
+    withLogAppender(logAppender, loggerNames = Seq(loggerName)) {
+      val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock)
+
+      // resourceOffer triggers prepareLaunchingTask which logs "Starting ..."
+      val taskOption = manager.resourceOffer("exec1", "host1", NO_PREF)._1
+      assert(taskOption.isDefined)
+
+      clock.advance(1)
+      // handleSuccessfulTask logs "Finished ..."
+      manager.handleSuccessfulTask(0, createTaskResult(0))
+    }
+
+    val logs = logAppender.loggingEvents.map(_.getMessage.getFormattedMessage)
+
+    val expectedQueryPrefix = s"[queryId = ${testQueryId.take(5)}]"
+    val expectedBatchPrefix = s"[batchId = $testBatchId]"
+
+    // Verify the "Starting" log line includes query Id and batch Id
+    val startingLogs = logs.filter(msg =>
+      msg.contains("Starting") &&
+        msg.contains(expectedQueryPrefix) && msg.contains(expectedBatchPrefix))
+    assert(startingLogs.nonEmpty,
+      s"Expected 'Starting' log to contain '$expectedQueryPrefix' and '$expectedBatchPrefix'." +
+        s"\nCaptured logs:\n${logs.mkString("\n")}")
+
+    // Verify the "Finished" log line includes query Id and batch Id
+    val finishedLogs = logs.filter(msg =>
+      msg.contains("Finished") &&
+        msg.contains(expectedQueryPrefix) && msg.contains(expectedBatchPrefix))
+    assert(finishedLogs.nonEmpty,
+      s"Expected 'Finished' log to contain '$expectedQueryPrefix' and '$expectedBatchPrefix'." +
+        s"\nCaptured logs:\n${logs.mkString("\n")}")
+  }
+
 }
 
 class FakeLongTasks(stageId: Int, partitionId: Int) extends FakeTask(stageId, partitionId) {