How To Integrate ZIO Into an Existing Project

ZIO is a powerful framework with a wide ecosystem of libraries, supported by a large developer community. It is advertised as offering more efficient thread processing (fibers), which can reduce resource usage and encourage the use of pure functional programming practices.

But what if you want to introduce it into your existing codebase? Reasons may vary: you may want to optimize a specific algorithm or completely rewrite your legacy project using a modern effects system. In this article, I’ll explore this topic using my previous progress, a REST API example based on Tapir.

View original code here.
View updated code here.

The example includes several repositories, services, and controllers that accept API requests. Their functionality was described in my previous article. Here, we’ll focus on updating the project with the ZIO library.

Updating Repositories and Services Using ZIO Service Pattern

How it was done before (based on the example of UserDao): a class with arguments and functions

class UserDao(service1: Service1, service2: Service2) {
  // ..... your functions
}

How it’s done now: an object that has an internal Service trait with functions. Then, you make the live layer, which fetches required arguments via ZIO.service and creates the trait implementation.

import zio._
// importing services

object UserDao { // name object and ZIO service type shortcut with different name to avoid confusion
  type UserRepository = Service
  
  trait Service {
    // type your functions here, they should return Task[A] (ZIO[Any, Throwable, A]) if they produce errors
    //     or UIO[A] (ZIO[Any, Nothing, A]) if it can't produce any negative results.
    // For example, in services, most functions return ZIO[Any, ErrorInfo, Unit] - it means we put ErrorInfo failures into error type, rather than turning successful result into Either[ErrorInfo, Unit]
  }
  
  val live: ZLayer[Service1 & Service2, Nothing, UserRepository] = ZLayer {
    for {
      service1 <- ZIO.service[Service1]
      service2 <- ZIO.service[Service2]
    } yield {
      new Service {
        // write your implementations of the service using services you fetched using ZIO.service
      }
    }
  }
}

The result layer here is ZLayer[Service1 & Service2, Nothing, dao.UserDao.UserRepository], which translates into:

I need Service1 and Service2 layers and I provide UserRepository layer

Important: provide a type for the live layer; otherwise, you risk triggering an ambiguous layer compile error.

The pattern has various interpretations – you can keep the service inlined, create an implementation class separately, place it as a companion class (as it is done in the project), or put it in a completely separate file.

It’s common practice to place your implementation service as a companion class, but if it’s too big, it’s better to put it in a separate file. Inlined implementation is convenient for making test layers, where most functions are stubbed, and the main one is simplified enough to keep the code readable.

For DAOs, we change the quill library to zio-quill – it is optimized for use with ZIO and returns ZIO monads, which is exactly what we need.

Integrating New ZIO Services Into Controllers

Now that we refactored our service classes into ZIO services, there is a question: how to provide them into controllers and how to actually use them there?

Let’s start with how to use these ZIO services in the TapirSecurity class:

/**
* Configures security endpoint.
*
* @param authentication authentication service.
*/
class TapirSecurity(authentication: ULayer[TapirAuth])

As we see, it accepts ULayer of TapirAuth, which is TapirAuthentication ZIO service.

It has only one endpoint, which is used for secure endpoints:

def tapirSecurityEndpoint(roles: List[RoleType]): PartialServerEndpoint[String, User, Unit, ErrorInfo, Unit, Any, Future] =
  endpoint // base tapir endpoint
    .securityIn(auth.bearer[String]().description("Bearer token from Authorization header")) // defining security input
    .errorOut(
      oneOf[ErrorInfo](
        // returns required http code for different types of ErrorInfo. For secured endpoint you need to define all cases before defining security logic
        oneOfVariant(statusCode(StatusCode.Forbidden).and(jsonBody[Forbidden].description("When user doesn't have role for the endpoint"))),
        oneOfVariant(statusCode(StatusCode.Unauthorized).and(jsonBody[Unauthorized].description("When user doesn't authenticated or token is expired"))),
        oneOfVariant(statusCode(StatusCode.NotFound).and(jsonBody[NotFound].description("When something not found"))),
        oneOfVariant(statusCode(StatusCode.BadRequest).and(jsonBody[BadRequest].description("Bad request"))),
        oneOfVariant(statusCode(StatusCode.InternalServerError).and(jsonBody[InternalServerError].description("For exceptional cases"))),
        // default case below.
        oneOfDefaultVariant(jsonBody[com.example.errors.ErrorMessage].description("Default result").example(com.example.errors.ErrorMessage("Test error message")))
      )
    )
    .serverSecurityLogic(token =>
      ZioUtil.foldRunToFuture(TapirAuthentication.authenticate(token).flatMap { user =>
        // define security logic here. For example, here is authentication, chained with authorization
        isAuthorized(user, roles)
      }.provide(authentication))
    )

As we can see, the util function ZioUtil.foldRunToFuture is used here, which is basically turning ZIO monad into future, used by the endpoint:

/**
 * Exiting zio execution. Use only on edges of zio integration.
 * @tparam T monad result
 * @return future of T
 */
def runToFuture[T](monad: UIO[T]): Future[T] = {
  Unsafe.unsafe { implicit unsafe =>
    zio.Runtime.default.unsafe.runToFuture(monad).future
  }
}

/**
 * Same as runToFuture, but also places both failed and successful results into Either
 * @tparam E error type
 * @tparam T success type
 * @return Future with Either, which contains both failure and success
 */
def foldRunToFuture[E, T](monad: ZIO[Any, E, T]): Future[Either[E, T]] = {
  ZioUtil.runToFuture(monad.fold(error => Left(error), success => Right(success)))
}

This function is used at the edges of ZIO usage, which are located in controller endpoints. There is a way for Tapir to accept ZIO monads, but we assume it only accepts Futures now.

foldRunToFuture is also used to unwrap ZIO into Either, since endpoint requires Either[Fail, Success]

.provide(authentication) function uses ULayer we provided in the arguments to provide means to execute the function, since .authenticate has this implementation:

def authenticate(token: String): ZIO[TapirAuth, ErrorInfo, User] = ZIO.serviceWithZIO[TapirAuth](_.authenticate(token))

We translate it as: I need TapirAuth layer in order to either fail with ErrorInfo or return User, so we provide the layer so the endpoint can work. Without it, the compiler returns an error.

Providing the Layers

So, we know how to make and use these ZIO services, but how can we make and provide them?

Remember these live layers? We use them to wire everything we need:

// ZLayer layers
lazy val config = ZLayer.succeed(ConfigFactory.load())
lazy val postgres: ZLayer[DataSource, Nothing, Quill.Postgres[SnakeCase.type]] = Quill.Postgres.fromNamingStrategy(SnakeCase)
lazy val dataSource: ZLayer[Any, Throwable, DataSource] = Quill.DataSource.fromPrefix("db.default")
lazy val userDao: ZLayer[Quill.Postgres[SnakeCase], Nothing, UserRepository] = UserDao.live
lazy val orderDao: ZLayer[Quill.Postgres[SnakeCase], Nothing, OrderRepository] = OrderDao.live
lazy val productDao: ZLayer[Quill.Postgres[SnakeCase], Nothing, ProductRepository] = ProductDao.live
lazy val orderProductDao: ZLayer[Quill.Postgres[SnakeCase], Nothing, OrderProductRepository] = OrderProductDao.live
lazy val jwtService: ZLayer[Config, Nothing, JwtService] = Jwt.live
lazy val tapirAuth: ZLayer[UserRepository with JwtService, Nothing, TapirAuth] = TapirAuthentication.live
lazy val authService: ZLayer[JwtService with UserRepository, Nothing, Authentication] = AuthService.live
lazy val orderService: ZLayer[OrderProductRepository with ProductRepository with OrderRepository, Nothing, OrdersService] = OrderService.live
lazy val productService: ZLayer[ProductRepository, Nothing, ProductService] = ProductService.live
lazy val adminProductService: ZLayer[ProductRepository, Nothing, AdminProducts] = AdminProductService.live
lazy val adminOrderService: ZLayer[OrderProductRepository with UserRepository with ProductRepository with OrderRepository, Nothing, AdminOrders] = AdminOrderService.live


But these are only ZLayers that depend on other layers to function, while controllers require a ULayer. You could provide every required layer into the controller and use it in the provide function, but that would cause controllers to have a lot of arguments, and it would be inconvenient in general.

Fortunately, we can compact all layers the service needs into one ULayer using ZLayer.make[Service], which automatically wires all required ZLayers into one ULayer, which you can provide into controllers:

// ULayer creation. The reason - to compose layers into single one for convenient transfer to controllers
// Layer wiring is done automatically, you just need to provide all required zlayers.
// orDie forces to throw an error if anything happens in the app's start
lazy val authenticationLayer: ZLayer[Any, Nothing, TapirAuth] = ZLayer.make[TapirAuth](tapirAuth, userDao, dataSource, postgres, jwtService, config).orDie
lazy val authServiceLayer: ZLayer[Any, Nothing, Authentication] = ZLayer.make[Authentication](authService, userDao, dataSource, postgres, jwtService, config).orDie
lazy val orderServiceLayer: ZLayer[Any, Nothing, OrdersService] = ZLayer.make[OrdersService](orderService, orderDao, orderProductDao, productDao, dataSource, postgres).orDie
lazy val productServiceLayer: ZLayer[Any, Nothing, ProductService] = ZLayer.make[ProductService](productService, productDao, dataSource, postgres).orDie
lazy val adminProductServiceLayer: ZLayer[Any, Nothing, AdminProducts] = ZLayer.make[AdminProducts](adminProductService, productDao, dataSource, postgres).orDie
lazy val adminOrderServiceLayer: ZLayer[Any, Nothing, AdminOrders] = ZLayer.make[AdminOrders](adminOrderService, orderDao, userDao, orderProductDao, productDao, dataSource, postgres).orDie

.orDie is used here to get rid of the exception type for the layer, causing an exception if the layer couldn’t be assembled.

Now, with these ULayers, we can provide them into controllers and use them as regular objects.

Testing ZIO Services

For controllers, testing is not much different from usual testing. The only change is that you need to mock ZIO service and create your layer with the service.

Or you could make test layer and fill it with your testing implementation, like this:

val test = ZLayer {
  new Service {
    // implement Service functions with your stubs
  }
}

I personally prefer mocking the services, like here (from OrderControllerUnitTest):

val orderService = mock[OrdersService]
val orderList = List(Order(Util.generateUuid, testUser.id, LocalDateTime.now(), Order.NEW_STATUS, LocalDateTime.now(), "comment"))
when(orderService.findOrdersForUser(testUser.id)).thenReturn(ZIO.succeed(orderList))
val orderController = new OrderController(new TapirSecurity(ZLayer.succeed(authentication)), ZLayer.succeed(orderService))

As for remaking tests for ZIO services themselves, you can use the zio-test library. I have remade one for AuthService. It extends the ZIOSpecDefault trait and uses the "zio.test.sbt.ZTestFramework" framework.

Check out our blog for more technical insights.
If you are looking for a reliable Scala development partner, check out our service offering.

Scala Developer