Get the result of a JDBC query as a reactive stream with proper resource cleaning

ReactiveResultSet.scala

import java.sql.{PreparedStatement, Statement, Connection, ResultSet}

import akka.stream.scaladsl._
import org.slf4j.Logger

import scala.concurrent.{Future, ExecutionContext}
import scala.util.{Success, Try}
import scala.util.control.NonFatal

// ExecutionContext for blocking ops
case class ECBlocking(value: ExecutionContext) extends AnyVal

// Get the result of a JDBC query as a reactive stream with proper resource cleaning
object ReactiveResultSet {

  def executeQuery[A](query: String, getConnection: => Connection)(mapRow: ResultSet => A)
                     (implicit ec: ECBlocking, logger: Logger): Source[A, Unit] = {
    implicit val ecBlocking = ec.value

    val futStream = Future {
      var conn: Connection = null
      var stmt: PreparedStatement = null
      var rs: ResultSet = null

      try {
        conn = getConnection
        stmt = conn.prepareStatement(query, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)

        // setFetchSize(Integer.MIN_VALUE) is a mysql driver specific way to force streaming results,
        // rather than pulling entire resultset into memory.
        // see http://dev.mysql.com/doc/connector-j/en/connector-j-reference-implementation-notes.html
        if (conn.getMetaData.getURL.matches("jdbc:mysql:.*")) {
          stmt.setFetchSize(Integer.MIN_VALUE)
          logger.debug("statement fetch size set to: " + stmt.getFetchSize + " to force MySQL streaming ")
        }

        rs = stmt.executeQuery()

        val stream = Source.unfoldAsync(()) { _ =>
          Future {
            if (rs.next()) Some(() -> mapRow(rs))
            else None
          }
        }

        stream.alsoTo(Sink.onComplete {
          case _ => Future(cleanup(rs, stmt, conn))
        })
      } catch {
        case NonFatal(failure) =>
          cleanup(rs, stmt, conn)
          Source.failed(failure)
      }
    }

    Source.fromFuture(futStream).flatMapConcat(identity)
  }

  private def cleanup(rs: ResultSet, stmt: Statement, conn: Connection)(implicit logger: Logger): Unit = {
    try {
      if (rs != null) rs.close()
    } catch {
      case NonFatal(exception) => logger.warn(s"Exception closing resultset $exception")
    }
    try {
      if (stmt != null) stmt.close()
    } catch {
      case NonFatal(exception) => logger.warn(s"Exception closing statement $exception")
    }
    try {
      if (conn != null) conn.close()
    } catch {
      case NonFatal(exception) => logger.warn(s"Exception closing connection $exception")
    }
  }

}

TODO improvement: do the sql query ONLY at materialization time

It seems stream.alsoTo(Sink.onComplete closes connection too early.