Add clean up for DataSource

debasishg · Apr 5, 2023 · 972d586 · 972d586
1 parent 03f55f5
commit 972d586
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/main/scala/com/softwaremill/realworld/db/Db.scala b/src/main/scala/com/softwaremill/realworld/db/Db.scala
@@ -2,21 +2,30 @@ package com.softwaremill.realworld.db
 
 import com.zaxxer.hikari.{HikariConfig, HikariDataSource}
 import io.getquill.{SnakeCase, SqliteZioJdbcContext}
-import zio.ZLayer
+import zio.{ZIO, ZLayer}
 
+import java.io.Closeable
 import javax.sql.DataSource
 
 object Db:
 
-  private def create(dbConfig: DbConfig): DataSource = {
+  private def create(dbConfig: DbConfig): HikariDataSource = {
     val poolConfig = new HikariConfig()
     poolConfig.setJdbcUrl(dbConfig.jdbcUrl)
     poolConfig.setConnectionInitSql(dbConfig.connectionInitSql)
     new HikariDataSource(poolConfig)
   }
 
   // Used for migration and executing queries.
-  val dataSourceLive: ZLayer[DbConfig, Nothing, DataSource] = ZLayer.fromFunction(create(_))
+  val dataSourceLive: ZLayer[DbConfig, Nothing, DataSource] =
+    ZLayer
+      .service[DbConfig]
+      .flatMap { env =>
+        val dbConfig = env.get
+        ZLayer.scoped {
+          ZIO.fromAutoCloseable(ZIO.succeed(create(dbConfig)))
+        }
+      }
 
   // Quill framework object used for specifying sql queries.
   def quillLive: ZLayer[Any, Nothing, SqliteZioJdbcContext[SnakeCase]] = ZLayer.succeed(new SqliteZioJdbcContext[SnakeCase](SnakeCase))