LoadYaml.scala

LoadYaml.scala

import java.io.{FileReader, Reader}

import scala.concurrent.Await
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration.Duration

import slick.jdbc.PostgresProfile.api._
import slick.jdbc.PostgresProfile.columnTypes
import slick.jdbc.meta.MPrimaryKey

import io.circe.yaml.parser
import io.circe.{Decoder, DecodingFailure}


object LoadYaml {
  case class Value(repr: String) extends AnyVal
  object Value {
    implicit val decoder: Decoder[Value] = Decoder.instance { cursor =>
      cursor.value.fold(
        jsonNull = Right(Value("NULL")),
        jsonBoolean = b => Right(Value(b.toString.toUpperCase)),
        jsonNumber = n => Right(Value(n.toString)),
        jsonString = s => Right(Value(columnTypes.stringJdbcType.valueToSQLLiteral(s))),
        jsonArray = _ => Left(DecodingFailure("Array values not supported", cursor.history)),
        jsonObject = _ => Left(DecodingFailure("Object values not supported", cursor.history))
      )
    }
  }

  def apply(readers: Reader*): Unit = {
    def rec(xs: Seq[String]) = xs.mkString("(", ", ", ")")

    val res =
      readers.toStream.flatMap { reader =>
        parser.parseDocuments(reader)
          .flatMap(_.flatMap(_.as[Map[String, Seq[Map[String, Value]]]]).toTry.get.toSeq)
      }
    val db = Database.forConfig("db.slick")
    try {
      val action =
        DBIO.sequence(res.map {
          case (tableName, records) =>
            MPrimaryKey.getPrimaryKeys(tableName).flatMap { primaryKeys =>
              DBIO.sequence(records.map { values =>
                val seq = values.mapValues(_.repr).toSeq
                val (pks, plain) = seq.partition(t => primaryKeys.exists(_.column == t._1))
                sql"""
                     |INSERT INTO #$tableName #${rec(seq.map(_._1))}
                     |VALUES #${rec(seq.map(_._2))}
                     |ON CONFLICT #${rec(pks.map(_._1))}
                     |DO UPDATE SET #${plain.map(t => t._1 + " = " + t._2).mkString(", ")}
                     |""".stripMargin.asUpdate
              })
            }
        })
      Await.result(db.run(action), Duration.Inf)
    } finally db.close()
  }

  def main(args: Array[String]): Unit =
    if (args.isEmpty)
      Console.err.println("Expected: one or more names of yaml files to load")
    else
      apply(args.map(new FileReader(_)): _*)
}