Surgical updates in Slick

SurgicalUpdate.scala

import slick.lifted.TupleShape
import slick.driver.H2Driver.api._
import scala.concurrent.Await
import scala.concurrent.duration.DurationInt
import scala.concurrent.ExecutionContext.Implicits.global

object Example {

  val db = Database.forConfig("db")

  // --- Schema definition
  class UserTable(tag: Tag) extends Table[(Long, String, Int)](tag, "user") {
    def id =   column[Long]  ("user_id", O.PrimaryKey)
    def name = column[String]("name")
    def age =  column[Int]   ("age")
    def * = (id, name, age)
  }

  val users = TableQuery[UserTable]

  // --- The interesting part starts here
  def updateUser(id: Long, maybeNewName: Option[String], maybeNewAge: Option[Int]): DBIO[Int] =
    SurgicalUpdate(users.filter(_.id === id))
      .update(_.name, maybeNewName)
      .update(_.age, maybeNewAge)
      .run()

  updateUser(1, Some("jrf"), None) // UPDATE user SET name = 'jrf' WHERE user.user_id = 1

  updateUser(1, None, Some(30)) // UPDATE user SET age = 30 WHERE user.user_id = 1

  updateUser(1, Some("Julien"), Some(29)) // UPDATE user SET name = 'Julien', age = 29 WHERE user.user_id = 1

  // --- Implementation details
  import scala.language.higherKinds
  case class SurgicalUpdate[E, U, C[_]](query: Query[E, U, C], maybeUpdate: Option[SurgicalUpdate.Update[E]]) {
    def update[F, G, T](projection: E => F, maybeT: Option[T])(implicit shape: Shape[_ <: FlatShapeLevel, F, T, G]): SurgicalUpdate[E, U, C] =
      maybeT match {
        case Some(t) =>
          val u2 = SurgicalUpdate.UpdateImpl(projection, t)
          maybeUpdate match {
            case Some(u1) =>
              implicit val tupleShape: Shape[_ <: FlatShapeLevel, (u1.F, u2.F), (u1.T, u2.T), (u1.G, u2.G)] =
                new TupleShape(u1.shape, u2.shape)
              val tupleUpdate = SurgicalUpdate.UpdateImpl((e: E) => (u1.projection(e), u2.projection(e)), (u1.newValue, u2.newValue))
              SurgicalUpdate(query, Some(tupleUpdate))
            case None =>
              SurgicalUpdate(query, Some(u2))
          }
        case None =>
          this
      }

    def run(): DBIO[Int] =
      maybeUpdate match {
        case Some(update) =>
          import update.shape
          query.map(update.projection).update(update.newValue)
        case None => DBIO.successful(0)
      }
  }

  object SurgicalUpdate {

    def apply[E, U, C[_]](query: Query[E, U, C]): SurgicalUpdate[E, U, C] = SurgicalUpdate(query, None)

    trait Update[E] {
      type F
      type G
      type T
      implicit def shape: Shape[_ <: FlatShapeLevel, F, T, G]
      def projection: E => F
      def newValue: T
    }

    case class UpdateImpl[E, F0, G0, T0](projection: E => F0, newValue: T0)(implicit val shape: Shape[_ <: FlatShapeLevel, F0, T0, G0]) extends Update[E] {
      type F = F0
      type G = G0
      type T = T0
    }
  }

}