如何拦截Spark SQL查询？修改Spark源码实现SQL捕获与动态改写

修改SparkSession核心入口方法 #

直接修改SparkSession.scala中的sql方法，这是spark.sql()调用的最顶层入口，能直接获取原始SQL文本，无需依赖监听器。

步骤1：修改无参数重载的sql方法 #

找到SparkSession类中的sql(String)方法，添加SQL捕获、打印和改写逻辑：

def sql(sqlText: String): DataFrame = {
  // 1. 打印原始SQL到控制台
  println(s"[SQL Capture] Original SQL: $sqlText")
  
  // 2. 执行动态改写逻辑（替换为你的自定义规则）
  val rewrittenSql = rewriteOriginalSql(sqlText)
  println(s"[SQL Capture] Rewritten SQL: $rewrittenSql")
  
  // 3. 使用改写后的SQL执行查询
  Dataset.ofRows(self, sessionState.sqlParser.parsePlan(rewrittenSql))
}

步骤2：实现自定义SQL改写函数 #

在SparkSession类中添加改写逻辑的方法（或单独抽成工具类）：

private def rewriteOriginalSql(originalSql: String): String = {
  // 示例：给所有SELECT查询添加LIMIT 1000（根据你的需求替换逻辑）
  originalSql.trim.toUpperCase match {
    case s if s.startsWith("SELECT") => s"$originalSql LIMIT 1000"
    case _ => originalSql // 非SELECT查询保持原样
  }
}

步骤3：处理带参数的SQL重载方法 #

如果用户使用spark.sql(sqlText, args...)的参数化查询方式，修改对应的重载方法：

def sql(sqlText: String, args: Any*): DataFrame = {
  val interpolatedSql = StringContext.sql(sqlText, args: _*)
  // 直接调用上面修改后的sql方法，自动复用捕获和改写逻辑
  sql(interpolatedSql)
}

为什么监听器方案不适用？ #

你尝试的SparkFirehoseListener、SQLListener等监听器，其触发时机是在SQL解析为逻辑/物理计划之后：

• 监听器拿到的QueryExecution对象中，sql字段仅能从特定类型的逻辑计划（如Command）中提取原始SQL，普通SELECT等查询无法获取完整原始文本。
• 监听器是查询执行的观察者，无法干预查询的执行流程（只能记录，不能改写后重新执行）。

验证效果 #

编译修改后的Spark源码，客户端无需添加任何额外代码，直接调用spark.sql("SELECT * FROM test_table")，控制台会打印原始SQL和改写后的SQL，最终执行的是改写后的查询。

内容的提问来源于stack exchange，提问作者ChoclateFactory