Batched prepared statements

DbBatchedJob.scala

package com.coatue.datascience.util

import java.sql.{Connection, PreparedStatement}

import com.typesafe.scalalogging.Logger
import org.slf4j.LoggerFactory

import slick.driver.PostgresDriver.api._

class DbBatchedJob(db: Database, sql: String, batchSize: Int) extends AutoCloseable {

  private[this] val logger = Logger(LoggerFactory.getLogger(this.getClass))

  private[this] var currentConnection: Connection = _
  private[this] var currentStatement: PreparedStatement = _
  private[this] var count = 0

  def +=[U](f: PreparedStatement => U) = synchronized {
    if (count%batchSize == 0) {
      currentConnection = db.createSession().conn
      currentStatement = currentConnection.prepareStatement(sql)
    }
    count += 1
    f(currentStatement)
    currentStatement.addBatch()
    if (count%batchSize == 0) close()
  }

  override def close() = if (count > 0) {
    val updates = currentStatement.executeBatch()
    logger.info(s"Batch executed ${updates.sum} updates at count=$count")
    currentStatement.close()
    currentConnection.close()
  }
}