Quill Context Returning ResultSet/PreparedStatement

JdbcContextBaseStub.scala

package io.getquill

import io.getquill.context.ContextEffect
import io.getquill.context.jdbc.JdbcContextBase
import io.getquill.context.sql.idiom.SqlIdiom

// Need this to be able to defined 'effect' since that's a package-private so this file needs to reside in the 
// `io.getquill` package but it's the only thing that does.
trait JdbcContextBaseStub[Dialect <: SqlIdiom, Naming <: NamingStrategy] extends JdbcContextBase[Dialect, Naming] {
  private[getquill] val effect: ContextEffect[Result] = null
}

PostgresResultSetJdbcContext.scala

package io.test

import io.getquill.context.jdbc.PostgresJdbcContextBase
import io.getquill.{NamingStrategy, PostgresDialect}

// Finally created the needed contexts for the individual dialects
class PostgresResultSetJdbcContext[N <: NamingStrategy](val naming: N)
  extends ResultSetJdbcContext[PostgresDialect, N] with PostgresJdbcContextBase[N] {
}

ResultSetJdbcContext.scala

package io.test
import java.sql.{Connection, PreparedStatement, ResultSet}

import io.getquill.context.jdbc.JdbcContext
import io.getquill.context.sql.idiom.SqlIdiom
import io.getquill.util.ContextLogger
import io.getquill.{JdbcContextBaseStub, NamingStrategy}

import scala.util.{Success, Try}

// This is the meat-and-potatoes of what to do. Use `Result` to define a kind of output that gives you back
// a ResultSet and PreparedStatement that is created. Unfortunately since the more specific types like RunQuerySingleResult
// have already been typed in JdbcContextBase, we need to have a single signature for all of the methods that encapsulate
// every kind of thing that any `execute...` method is doing. Including those that create multiple statements
// etc...
trait ResultSetJdbcContext[Dialect <: SqlIdiom, Naming <: NamingStrategy] extends JdbcContextBaseStub[Dialect, Naming] {

  override type Result[T] = Connection => Iterator[(PreparedStatement, Option[ResultSet])]

  override type PrepareRow = PreparedStatement
  override type ResultRow = ResultSet


  override protected def withConnection[T](f: Connection => Result[T]): Result[T] = ???
  override protected def withConnectionWrapped[T](f: Connection => T): Result[T] = ???
  val log = ContextLogger(classOf[JdbcContext[_, _]])
  override def close(): Unit = ()
  def probe(statement: String): Try[_] = Success(Unit)


  override def executeAction[T](sql: String, prepare: Prepare = identityPrepare): Connection => Iterator[(PreparedStatement, Option[ResultSet])] =
    (conn:Connection) => {
      val (params, ps) = prepare(conn.prepareStatement(sql))
      log.logQuery(sql, params)
      Iterator((ps, None))
    }


  override def executeQuery[T](sql: String, prepare: Prepare = identityPrepare, extractor: Extractor[T] = identityExtractor): Connection => Iterator[(PreparedStatement, Option[ResultSet])] =
    (conn:Connection) => {
      val (params, ps) = prepare(conn.prepareStatement(sql))
      log.logQuery(sql, params)
      val rs = ps.executeQuery()
      Iterator((ps, Some(rs)))
    }

  override def executeQuerySingle[T](sql: String, prepare: Prepare = identityPrepare, extractor: Extractor[T] = identityExtractor): Connection => Iterator[(PreparedStatement, Option[ResultSet])] =
    executeQuery(sql, prepare, extractor)

  override def executeActionReturning[O](sql: String, prepare: Prepare = identityPrepare, extractor: Extractor[O], returningColumn: String): Connection => Iterator[(PreparedStatement, Option[ResultSet])] =
    (conn:Connection) => {
      val (params, ps) = prepare(conn.prepareStatement(sql, Array(returningColumn)))
      log.logQuery(sql, params)
      ps.executeUpdate()
      val rs = ps.getGeneratedKeys
      Iterator((ps, Some(rs)))
    }

  override def executeBatchAction(groups: List[BatchGroup]): Connection => Iterator[(PreparedStatement, Option[ResultSet])] =
    (conn:Connection) => {
      groups.iterator.map {
        case BatchGroup(sql, prepare) =>
          val ps = conn.prepareStatement(sql)
          log.underlying.debug("Batch: {}", sql)
          prepare.foreach { f =>
            val (params, _) = f(ps)
            log.logBatchItem(sql, params)
            ps.addBatch()
          }
          (ps, None)
        // We're not closing PS from our batch groups, this might be the cause of a recent pool-overflow post on gitter
      }
    }

  override def executeBatchActionReturning[T](groups: List[BatchGroupReturning], extractor: Extractor[T]): Connection => Iterator[(PreparedStatement, Option[ResultSet])] =
    (conn:Connection) => {
      groups.iterator.map {
        case BatchGroupReturning(sql, column, prepare) =>
          val ps = conn.prepareStatement(sql, Array(column))
          log.underlying.debug("Batch: {}", sql)
          prepare.foreach { f =>
            val (params, _) = f(ps)
            log.logBatchItem(sql, params)
            ps.addBatch()
          }
          ps.executeBatch()
          val rs = ps.getGeneratedKeys
          (ps, Some(rs))
      }
    }
}