ZIO learning - providing custom service through environments / published by https://github.com/dacr/code-examples-manager #101472db-4a9c-4b7f-99fc-e0a7304cb6fb/8543f302d0f98a9d3bb03374f3435a2e45c462d8

zio-learning-14-service-1.sc

// summary : ZIO learning - providing custom service through environments
// keywords : scala, zio, learning, services, pure-functional, @testable
// publish : gist
// authors : Adam Warski, David Crosson
// license : unknown BUT Machine Learning models training is not allowed by the author
// id : 101472db-4a9c-4b7f-99fc-e0a7304cb6fb
// created-on : 2021-05-01T21:42:05+02:00
// managed-by : https://github.com/dacr/code-examples-manager
// run-with : scala-cli $file

// ---------------------
//> using scala  "3.2.2"
//> using dep "dev.zio::zio:2.0.13"
// ---------------------

import zio.*

// Originally inspired from Adam Warski (but diverged a lot since a while) :
//  - https://softwaremill.com/zio-environment-episode-3/
//  - https://github.com/adamw/zioenv/blob/f373e7d94300456bc1f8d1c188b90c29dc2cf592/core/src/main/scala/zioenv/DB.scala

// ===========================================================================================

case class DBConfig(url: String)

// ===========================================================================================

class ConnectionPool(url: String) {
  def close(): Unit             = ()
  override def toString: String = s"ConnectionPool($url)"
}

// integration with ZIO
object ConnectionPoolIntegration {
  def createConnectionPool(dbConfig: DBConfig): ZIO[Any, Throwable, ConnectionPool] =
    for {
      connectionPool <- ZIO.attempt(new ConnectionPool(dbConfig.url))
      _              <- ZIO.logInfo("connection pool created")
    } yield connectionPool

  def closeConnectionPool: ConnectionPool => ZIO[Any, Nothing, Unit] = (cp: ConnectionPool) =>
    for {
      _ <- ZIO.attempt(cp.close()).catchAll(_ => ZIO.unit)
      _ <- ZIO.logInfo("connection pool closed")
    } yield ()

  val live: ZLayer[DBConfig & Scope, Throwable, ConnectionPool] = ZLayer.fromZIO(
    for {
      dbConfig       <- ZIO.service[DBConfig]
      connectionPool <- ZIO.acquireRelease(createConnectionPool(dbConfig))(closeConnectionPool)
    } yield connectionPool
  )
}

// ===========================================================================================

trait DB {
  def execute(sql: String): Task[Unit]
}

case class DBRelationalLive(cp: ConnectionPool) extends DB {
  override def execute(sql: String): Task[Unit] =
    ZIO.attempt {
      println(s"Running: $sql, on: $cp")
    }
}

object DB {
  val live = ZLayer.fromFunction(DBRelationalLive.apply)

  def execute(sql: String): ZIO[DB, Throwable, Unit] = ZIO.serviceWithZIO(_.execute(sql))
}

// ===========================================================================================

case class User(name: String, email: String)

// -------------------------------------------
object UserModel {
  def insert(u: User): ZIO[DB, Throwable, Unit] = DB.execute(s"INSERT INTO user VALUES ('${u.name}')")
}

// -------------------------------------------
object UserNotifier {
  def notify(u: User, msg: String): Task[Unit] = {
    Clock.currentDateTime.flatMap { dateTime =>
      ZIO.attempt {
        println(s"Sending $msg to ${u.email} @ $dateTime")
      }
    }
  }
}

// -------------------------------------------
object UserRegistration {
  def register(u: User): ZIO[DB, Throwable, User] = {
    for {
      _ <- UserModel.insert(u)
      _ <- UserNotifier.notify(u, "Welcome!")
    } yield u
  }
}
// ===========================================================================================

object Main extends ZIOApp {
  override type Environment = DB & ConnectionPool
  override val environmentTag: Tag[Environment] = Tag[Environment]

  override val bootstrap: ZLayer[ZIOAppArgs, Any, DB & ConnectionPool] =
    ZLayer.make[DB & ConnectionPool](
      ConnectionPoolIntegration.live,
      DB.live,
      ZLayer.succeed(DBConfig("jdbc://localhost")),
      Scope.default
    )

  val logic = for {
    user <- UserRegistration.register(User("adam", "adam@hello.world"))
    _    <- Console.printLine(s"Registered user: $user (layers)")
  } yield ()

  override val run: ZIO[DB, Any, Any] = logic
}

Main.main(args)