TWS session tracking in Scala

%sh
pip install dbldatagen

%python
import dbldatagen as dg
from pyspark.sql.types import IntegerType, StringType, TimestampType
import time
from datetime import datetime, timedelta

# Set table name
table_name = dbutils.widgets.get("table_name")

# Generate session data with user_id, action_type, and timestamp
data_rows = 1000 * 100
base_time = datetime.now()

df_spec = (dg.DataGenerator(spark, name="session_data", rows=data_rows)
           .withColumn("user_id", StringType(), values=['user1', 'user2', 'user3', 'user4', 'user5'])
           .withColumn("action_type", StringType(), values=['login', 'page_view', 'purchase'], 
                      random=True, weights=[1, 5, 1])
           .withColumn("timestamp", TimestampType(), 
                      begin=base_time - timedelta(hours=1),
                      end=base_time,
                      random=True)
           .withColumn("session_value", IntegerType(), minValue=1, maxValue=100))

df = df_spec.build()

# Write to Delta table
df.write.format("delta").mode("overwrite").saveAsTable(table_name)

// Scala code
spark.conf.set(
  "spark.sql.streaming.stateStore.providerClass",
  "com.databricks.sql.streaming.state.RocksDBStateStoreProvider"
)

val tableName = dbutils.widgets.get("table_name")
// Use spark.table() instead of read.format("delta").load()
val df = spark.table(tableName)

package org.apache.spark.sql.streaming

// Case class to hold session state
object SessionState {
  case class SessionState(
    lastTimestamp: Long,
    loginCount: Int,
    pageViewCount: Int,
    purchaseCount: Int,
    sessionValue: Long
  )
}

// Import the RocksDB state store provider
import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
import java.util.UUID
import org.apache.spark.sql.streaming.StatefulProcessor
import org.apache.spark.sql.streaming._

import java.sql.Timestamp
import org.apache.spark.sql.Encoders
import org.apache.spark.sql.streaming.SessionState.SessionState

class SessionTrackingProcessor
  extends StatefulProcessor[String, (String, String, Long, Int), (String, String, Long, Int, Int, Int)] {
  
  @transient protected var _sessionState: ValueState[SessionState] = _
  
  // Session timeout in milliseconds (30 minutes)
  val sessionTimeoutMs: Long = 30 * 60 * 1000
  
  override def init(
      outputMode: OutputMode,
      timeMode: TimeMode): Unit = {
    _sessionState = getHandle.getValueState[SessionState]("sessionState",
      Encoders.product[SessionState], TTLConfig.NONE)
  }
  
  override def handleInputRows(
      key: String,
      inputRows: Iterator[(String, String, Long, Int)],
      timerValues: TimerValues): Iterator[(String, String, Long, Int, Int, Int)] = {
    
    val currentState = Option(_sessionState.get()).getOrElse(
      SessionState(0L, 0, 0, 0, 0L)
    )
    
    val results = inputRows.map { case (userId, actionType, timestamp, value) =>
      // Update counters based on action type
      val newState = actionType match {
        case "login" => currentState.copy(
          lastTimestamp = timestamp,
          loginCount = currentState.loginCount + 1,
          sessionValue = currentState.sessionValue + value
        )
        case "page_view" => currentState.copy(
          lastTimestamp = timestamp,
          pageViewCount = currentState.pageViewCount + 1,
          sessionValue = currentState.sessionValue + value
        )
        case "purchase" => currentState.copy(
          lastTimestamp = timestamp,
          purchaseCount = currentState.purchaseCount + 1,
          sessionValue = currentState.sessionValue + value
        )
        case _ => currentState.copy(lastTimestamp = timestamp)
      }
      
      // Clear old timer if exists
      if (currentState.lastTimestamp > 0) {
        getHandle.deleteTimer(currentState.lastTimestamp + sessionTimeoutMs)
      }
      
      // Register new timer
      val newTimerTime = timestamp + sessionTimeoutMs
      getHandle.registerTimer(newTimerTime)
      
      // Update state
      _sessionState.update(newState)
      
      // Return current session status
      (userId, "ACTIVE", newState.sessionValue, 
       newState.loginCount, newState.pageViewCount, newState.purchaseCount)
    }.toList
    
    results.iterator
  }
  
  override def handleExpiredTimer(
      key: String,
      timerValues: TimerValues,
      expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String, Long, Int, Int, Int)] = {
    
    // Get final state before clearing
    val finalState = Option(_sessionState.get()).getOrElse(
      SessionState(0L, 0, 0, 0, 0L)
    )
    
    // Clear session state
    _sessionState.clear()
    
    // Return final session summary
    Iterator((key, "EXPIRED", finalState.sessionValue, 
             finalState.loginCount, finalState.pageViewCount, finalState.purchaseCount))
  }
}

// Read from Delta table
val inputStream = spark.readStream
  .format("delta")
  .table(tableName)
  .as[(String, String, Long, Int)]
  .map(row => (row._1, row._2, row._3, row._4))

val checkpointLocation = "/tmp/streaming_query/checkpoint_" + UUID.randomUUID().toString
val outputTable = "/tmp/streaming_query/output_table_" + UUID.randomUUID().toString

import spark.implicits._

val sessionStream = inputStream
  .groupByKey(_._1)  // Group by user_id
  .transformWithState(
    new SessionTrackingProcessor(),
    TimeMode.EventTime(),
    OutputMode.Append())

val query = sessionStream.writeStream
  .format("delta")
  .option("checkpointLocation", checkpointLocation)
  .outputMode("append")
  .start(outputTable)

query.processAllAvailable()
query.stop()