PostgreSQLAnalyzer for ScalikeJDBC Users

PostgreSQLAnalyzer.scala

package lib

import java.sql.Timestamp
import org.joda.time._
import scalikejdbc._

/**
 * Usage:
 * scalikekdbc.GlobalSettings.queryCompletionListener = PostgreSQLAnalyzer.queryCompletionListener
 */
object PostgreSQLAnalyzer {

  private val highestCostRegexp = "\\(cost=[\\d\\.]+\\.\\.([\\d\\.]+)".r

  lazy val queryCompletionListener: (String, Seq[Any], Long) => Unit = {
    (statement: String, params: Seq[Any], millis: Long) =>
      {
        DB.readOnly { session =>
          val stmt = session.connection.createStatement()
          val query = sql"explain ${SQLSyntax.createUnsafely(statement, params)}"
          val explainQuery = query.parameters.foldLeft(query.statement) {
            case (query, param) =>
              val value = param match {
                case None => "null"
                case Some(str: String) => s"'${str}'"
                case Some(v) => v.toString
                case str: String => s"'${str}'"
                case datetime: DateTime => s"'${datetime.toString}'"
                case timestamp: Timestamp => s"'${timestamp.toString}'"
                case v => v.toString
              }
              query.replaceFirst("\\?", value)
          }
          var isLowCostQuery = false
          val lines: Traversable[String] = new ResultSetTraversable(stmt.executeQuery(explainQuery)).flatMap { rs =>
            if (!isLowCostQuery) {
              val line = rs.string(1)
              highestCostRegexp.findAllIn(line).matchData.toSeq.headOption.map(_.group(1)) match {
                case Some(cost) if cost.toDouble < 1000 =>
                  isLowCostQuery = true
                  None
                case _ =>
                  Some("  " + line)
              }
            } else {
              None
            }
          }
          if (lines.size > 0) {
            println(lines.mkString("\n"))
            println
          }
        }
      }
  }

}