Slick Upsert and Select

Scalatra Apr 3, 2019

In our previous post, we wanted to create and update a record from our PostgreSQL database in our Scalatra service to manage inventory data.

We only got as far as using raw SQL to do the query, and this had an added benefit of being an atomic operation. Now we would like to try and implement the same logic using Slick.

Improving the Return Type of our Raw SQL Query

A helpful reference here is http://slick.lightbend.com/doc/3.3.0/sql.html. I would have saved my self some time by sitting down and reading it end to end before I started.

We had the following when we left off last time:

  def create(db: PostgresProfile.backend.DatabaseDef,
               sku: String,
               qty: Int,
               location: String
              ): Future[Seq[(String, Int, String)]] = {
    val query: DBIO[Seq[(String, Int, String)]] =
      sql"""
           INSERT INTO inventory_single (sku, qty, location)
           VALUES ($sku, $qty, $location)
           ON CONFLICT ON CONSTRAINT inventory_single_sku_location_key
              DO UPDATE SET qty = EXCLUDED.qty
           RETURNING sku, qty, location;
        """.as[(String, Int, String)]
    db.run(query)
  }

This code was handling and returning the tuple (String, Int, String) instead of the InventorySingleRecord case class.

Slick provides us with sql, sqlu, and tsql interpolators.

  • sql is used for queries to produce a sequence of tuples, has type DBIO[Seq[<tuples>]] and can use implicit GetResult converters.
  • sqlu is used for queries that produce a row count, has type DBIO[Int]
  • tsql can enforce compile-time type checking, requires access to a configuration that defines the database schema.

We will continue to use the sql interpolator. Let's define our own converter to the InventorySingleRecord class and update this query (you will need to update your test as well).

  implicit val getInventorySingleRecord : GetResult[InventorySingleRecord] =
    GetResult(r => InventorySingleRecord(r.<<, r.<<, r.<<, r.<<))

  def create(db: PostgresProfile.backend.DatabaseDef,
               sku: String,
               qty: Int,
               location: String
              ): Future[Seq[InventorySingleRecord]] = {

    val query: DBIO[Seq[InventorySingleRecord]] =
      sql"""
           INSERT INTO inventory_single (sku, qty, location)
           VALUES ($sku, $qty, $location)
           ON CONFLICT ON CONSTRAINT inventory_single_sku_location_key
              DO UPDATE SET qty = EXCLUDED.qty
           RETURNING id, sku, qty, location;
        """.as[InventorySingleRecord]
    db.run(query)
  }

Now we have the desired type instead of a tuple, but we still have a sequence Seq[InventorySingleRecord], which is not ideal given that this is a create for a single record. This can be addressed with headOption (https://www.garysieling.com/blog/scala-headoption-example) to get the first element and update the return type Option[InventorySingleRecord]

  def create(db: PostgresProfile.backend.DatabaseDef,
               sku: String,
               qty: Int,
               location: String
              ): Future[Option[InventorySingleRecord]] = {

    val query: DBIO[Option[InventorySingleRecord]] =
      sql"""
           INSERT INTO inventory_single (sku, qty, location)
           VALUES ($sku, $qty, $location)
           ON CONFLICT ON CONSTRAINT inventory_single_sku_location_key
              DO UPDATE SET qty = EXCLUDED.qty
           RETURNING id, sku, qty, location;
        """.as[InventorySingleRecord].headOption
    db.run(query)

Our tests then go to validating the option instead of a sequence.

val future = InventorySingleRecordDao.create(database, TEST_SKU, 1, BIN_01)
val result: Option[InventorySingleRecord] = Await.result(future, Duration.Inf)
result should equal(Some(InventorySingleRecord(Some(1), TEST_SKU, 1, BIN_01)))

Implement the raw SQL as a Slick Query

Now let's write the same query using Slick. We will start by using a Scala for comprehension (a nice reference https://medium.com/@scalaisfun/scala-for-comprehension-tricks-9c8b9fe31778).

  def create(db: PostgresProfile.backend.DatabaseDef,
               sku: String,
               qty: Int,
               location: String
              ): Future[Option[InventorySingleRecord]] = {
              
    val upsert = for {
      existing <- {
        this.filter(x => x.location === location && x.sku === sku).forUpdate.result.headOption
      }
    } yield existing
    db.run(upsert.transactionally)
  }

This first step lets us retrieve a record if it already exists, and this is the model we will use to add additional functionality.

We will compose several queries here:

  • One to find the existing record (if it exists) this.filter(x => x.location === location && x.sku === sku).forUpdate.result.headOption,
  • A Query to create TableQuery[InventorySingleRecords] += InventorySingleRecord(Option.empty, sku, qty, location)
  • Finally one to retrieve the value TableQuery[InventorySingleRecords].filter(x => x.location === location && x.sku === sku).result.headOption

I have left the pattern matching unimplemented for the update case just to try and keep things reasonably simple.

  def create(db: PostgresProfile.backend.DatabaseDef,
               sku: String,
               qty: Int,
               location: String
              ): Future[Option[InventorySingleRecord]] = {
    val upsert = for {
      existing <- {
        this.filter(x => x.location === location && x.sku === sku).forUpdate.result.headOption
      }
      _ <- {
        existing match {
          case Some(InventorySingleRecord(_, `sku`, _, `location`)) => 
            // Update
            ???
          case _ => 
            // Create a new record
            TableQuery[InventorySingleRecords] += InventorySingleRecord(Option.empty, sku, qty, location)
        }
      }
      updated <- {
        TableQuery[InventorySingleRecords].filter(x => x.location === location && x.sku === sku).result.headOption
      }
    } yield updated
    db.run(upsert.transactionally)
  }

This should satisfy out tests for a single create, but update is still not implemented.

The logic for doing an update can be implemented as a standard Slick update query.

        existing match {
          case Some(InventorySingleRecord(_, `sku`, _, `location`)) => // Update
            val updateFoo = TableQuery[InventorySingleRecords]
            val q = for {x <- updateFoo if x.location === location && x.sku === sku} yield x.qty
            q.update(qty)

Note that in our update logic we have opted to ignore the existing value for the qty column. You could certainly utilize it if we wanted to log or use it in your update logic, that might look something like this case Some(InventorySingleRecord(_, sku, existingQty, location)) =>

This now should satisfy the tests for both create and update.

Conclusion

We have implemented our create logic in two different ways, one in raw SQL and the other using the Slick functional relational mapper. Hopefully, you found the comparison helpful.

Tags