How to register a StreamingListener with SparkSession?

I have a structured streaming application and I want to register a StreamingListener with it, how do I do that? I only seem to be able to register StreamingQueryListener via sparkSession.streams.addListener().

I know how to add a StreamingListener into StreamingContext, but the problem seems to be in getting StreamingContext, if I create a StreamingContext from SparkSession.sparkContext, I don't think StreamingListener will work if I continue doing streaming from SparkSession instead of that StreamingContext that I created.

Here's my StreamingListener class:

import io.prometheus.client.{CollectorRegistry, Gauge}
import io.prometheus.client.exporter.PushGateway
import org.apache.spark.streaming.scheduler.{StreamingListener, StreamingListenerBatchCompleted}
class MicroBatchStatsListener(val pushGateway: PushGateway, jobName: String) extends StreamingListener {
private val processingTimeGauge = Gauge.build()
.name("processingTimeGauge")
.help("Time it took to process this microbatch")
.register(CollectorRegistry.defaultRegistry)
private val schedulingDelayGauge = Gauge.build()
.name("schedulingDelayGauge")
.help("Scheduling delay of this microbatch")
.register(CollectorRegistry.defaultRegistry)
private val numRecordsGauge = Gauge.build()
.name("numRecordsGauge")
.help("Number of records received in this mircobatch")
.register(CollectorRegistry.defaultRegistry)
override def onBatchCompleted(batchCompleted: StreamingListenerBatchCompleted): Unit = {
val processingTime = if (batchCompleted.batchInfo.processingDelay.isDefined) {
batchCompleted.batchInfo.processingDelay.get / 1000
} else {
0
}
val schedulingDelay = if (batchCompleted.batchInfo.schedulingDelay.isDefined) {
batchCompleted.batchInfo.schedulingDelay.get / 1000
} else {
0
}
val numRecords = batchCompleted.batchInfo.numRecords
processingTimeGauge.set(processingTime)
schedulingDelayGauge.set(schedulingDelay)
numRecordsGauge.set(numRecords)
pushGateway.push(CollectorRegistry.defaultRegistry, jobName)
}
}

if I can't connect it to SparkSession, how do I do the same with StreamingQueryListener? I don't understand how to get these metrics from there:

1. Scheduling delay
2. Processing time

Turns out it is impossible to register a StreamingListener for structured streaming, because StreamingContext is not available if you use structured streaming. This method of getting StreamingContext does not work in this case, because it creates a separate StreamingContext, and if you register listeners with it, they won't work as you use structured streaming.

But all these metrics, can actually be computed with StreamingQueryListener, an example of what metrics it has is the following:

23/06/20 17:59:15 INFO MicroBatchExecution: Streaming query made progress: {
"id" : "4005e0e1-de55-4e9e-a01f-2bd3eec3d77e",
"runId" : "9d847913-72f8-4410-8b5f-0c69b798bc3c",
"name" : null,
"timestamp" : "2023-06-20T14:59:15.184Z",
"batchId" : 3,
"numInputRows" : 10,
"inputRowsPerSecond" : 666.6666666666667,
"processedRowsPerSecond" : 47.16981132075472,
"durationMs" : {
"addBatch" : 31,
"getBatch" : 0,
"getEndOffset" : 0,
"queryPlanning" : 4,
"setOffsetRange" : 0,
"triggerExecution" : 212,
"walCommit" : 112
},
"stateOperators" : [ ],
"sources" : [ {
"description" : "RateStreamV2[rowsPerSecond=10, rampUpTimeSeconds=0, numPartitions=default",
"startOffset" : 2,
"endOffset" : 3,
"numInputRows" : 10,
"inputRowsPerSecond" : 666.6666666666667,
"processedRowsPerSecond" : 47.16981132075472
} ],
"sink" : {
"description" : "org.apache.spark.sql.execution.streaming.ConsoleSinkProvider@3d192c1a"
}
}

We can add up all the values from durationMs and get the total processing time, for example.

Scheduling delays can also be computed from durationMs if you know the order of operations mentioned in that map. So, you'd have to add up all operations that happen before triggerExecution and you will get your scheduling delay.

Number of input records is also there.

All this info is available in the StreamingQueryListener.QueryProgressEvent class of StreamingQueryListener's onQueryProgress method, which is triggered after every microbatch.

UPDATE: found a similar question and provided an answer for it here: https://stackoverflow.com/questions/58959668/structured-spark-streaming-metrics-retrieval/76536197#76536197