Adds option for connectionInitSql (zio#639)

sergeda · May 31, 2022 · 9658c65 · 9658c65
1 parent a3d1e27
commit 9658c65
Show file tree

Hide file tree

Showing 4 changed files with 89 additions and 51 deletions.
diff --git a/jdbc-hikaricp/src/main/scala/zio/sql/HikariConnectionPool.scala b/jdbc-hikaricp/src/main/scala/zio/sql/HikariConnectionPool.scala
@@ -1,8 +1,8 @@
 package zio.sql
-import com.zaxxer.hikari.{HikariConfig, HikariDataSource}
-import zio.{Scope, ZIO, ZLayer}
+import com.zaxxer.hikari.{ HikariConfig, HikariDataSource }
+import zio.{ Scope, ZIO, ZLayer }
 
-import java.sql.{Connection, SQLException}
+import java.sql.{ Connection, SQLException }
 
 class HikariConnectionPool private (hikariDataSource: HikariDataSource) extends ConnectionPool {
 
@@ -13,9 +13,10 @@ class HikariConnectionPool private (hikariDataSource: HikariDataSource) extends
    * The managed resource will safely acquire and release the connection, and
    * may be interrupted or timed out if necessary.
    */
-  override def connection: ZIO[Scope, Exception, Connection] = {
-    ZIO.acquireRelease(ZIO.attemptBlocking(hikariDataSource.getConnection).refineToOrDie[SQLException])(con => ZIO.attemptBlocking(hikariDataSource.evictConnection(con)).orDie)
-  }
+  override def connection: ZIO[Scope, Exception, Connection] =
+    ZIO.acquireRelease(ZIO.attemptBlocking(hikariDataSource.getConnection).refineToOrDie[SQLException])(con =>
+      ZIO.attemptBlocking(hikariDataSource.evictConnection(con)).orDie
+    )
 }
 
 object HikariConnectionPool {
@@ -26,9 +27,9 @@ object HikariConnectionPool {
   val live: ZLayer[HikariConnectionPoolConfig, Throwable, HikariConnectionPool] =
     ZLayer.scoped {
       for {
-        config    <- ZIO.service[HikariConnectionPoolConfig]
+        config     <- ZIO.service[HikariConnectionPoolConfig]
         dataSource <- initDataSource(config.toHikariConfig)
-        pool       = new HikariConnectionPool(dataSource)
+        pool        = new HikariConnectionPool(dataSource)
       } yield pool
     }
 }
diff --git a/jdbc-hikaricp/src/main/scala/zio/sql/HikariConnectionPoolConfig.scala b/jdbc-hikaricp/src/main/scala/zio/sql/HikariConnectionPoolConfig.scala
@@ -2,7 +2,6 @@ package zio.sql
 
 import com.zaxxer.hikari.HikariConfig
 
-
 /**
  * Configuration information for the connection pool.
  *
@@ -23,19 +22,24 @@ import com.zaxxer.hikari.HikariConfig
  *                    An in-use connection will never be retired, only when it is idle will it be removed. Should be bigger then 30000
  * @param minimumIdle The property controls the minimum number of idle connections that HikariCP tries to maintain in the pool, including both idle and in-use connections.
  *                    If the idle connections dip below this value, HikariCP will make a best effort to restore them quickly and efficiently.
+ * @param connectionInitSql the SQL to execute on new connections
+ *                    Set the SQL string that will be executed on all new connections when they are
+ *                    created, before they are added to the pool.  If this query fails, it will be
+ *                    treated as a failed connection attempt.
  */
 final case class HikariConnectionPoolConfig(
-                                       url: String,
-                                       userName: String,
-                                       password: String,
-                                       poolSize: Int = 10,
-                                       autoCommit: Boolean = true,
-                                       connectionTimeout: Option[Long] = None,
-                                       idleTimeout: Option[Long] = None,
-                                       initializationFailTimeout: Option[Long] = None,
-                                       maxLifetime: Option[Long] = None,
-                                       minimumIdle: Option[Int] = None
-                                     ) {
+  url: String,
+  userName: String,
+  password: String,
+  poolSize: Int = 10,
+  autoCommit: Boolean = true,
+  connectionTimeout: Option[Long] = None,
+  idleTimeout: Option[Long] = None,
+  initializationFailTimeout: Option[Long] = None,
+  maxLifetime: Option[Long] = None,
+  minimumIdle: Option[Int] = None,
+  connectionInitSql: Option[String] = None
+) {
   private[sql] def toHikariConfig = {
     val hikariConfig = new HikariConfig()
     hikariConfig.setJdbcUrl(this.url)
@@ -48,6 +52,7 @@ final case class HikariConnectionPoolConfig(
     initializationFailTimeout.foreach(hikariConfig.setInitializationFailTimeout)
     maxLifetime.foreach(hikariConfig.setMaxLifetime)
     minimumIdle.foreach(hikariConfig.setMinimumIdle)
+    connectionInitSql.foreach(hikariConfig.setConnectionInitSql)
     hikariConfig
   }
 }
diff --git a/jdbc-hikaricp/src/test/scala/zio/sql/HikariConnectionPoolSpec.scala b/jdbc-hikaricp/src/test/scala/zio/sql/HikariConnectionPoolSpec.scala
@@ -1,16 +1,15 @@
 package zio.sql
 
-import zio.test.TestAspect.{sequential, timeout, withLiveClock}
-import zio.test.{TestEnvironment, _}
-import zio.{ZIO, ZLayer, durationInt}
-
+import zio.test.TestAspect.{ sequential, timeout, withLiveClock }
+import zio.test.{ TestEnvironment, _ }
+import zio.{ durationInt, ZIO, ZLayer }
 
 object HikariConnectionPoolSpec extends ZIOSpecDefault {
 
-
   val mySqlConfigLayer: ZLayer[Any, Throwable, MySqlConfig] =
     ZLayer.scoped {
-      MySqlTestContainer.mysql()
+      MySqlTestContainer
+        .mysql()
         .map(a =>
           MySqlConfig(
             url = a.jdbcUrl,
@@ -20,83 +19,113 @@ object HikariConnectionPoolSpec extends ZIOSpecDefault {
         )
     }
 
-  val hikariPoolConfigLayer: ZLayer[MySqlConfig, Nothing, HikariConnectionPoolConfig] = ZLayer.fromFunction((conf: MySqlConfig) => HikariConnectionPoolConfig(url = conf.url, userName = conf.username, password = conf.password))
-  val poolLayer: ZLayer[HikariConnectionPoolConfig, Nothing, HikariConnectionPool] = HikariConnectionPool.live.orDie
+  val hikariPoolConfigLayer: ZLayer[MySqlConfig, Nothing, HikariConnectionPoolConfig] =
+    ZLayer.fromFunction((conf: MySqlConfig) =>
+      HikariConnectionPoolConfig(url = conf.url, userName = conf.username, password = conf.password)
+    )
+  val poolLayer: ZLayer[HikariConnectionPoolConfig, Nothing, HikariConnectionPool]    = HikariConnectionPool.live.orDie
 
   override def spec: Spec[TestEnvironment, Any] =
     specLayered.provideCustomShared(mySqlConfigLayer.orDie)
 
   def specLayered: Spec[TestEnvironment with MySqlConfig, Any] =
     suite("Hikaricp module")(
       test("Pool size should be configurable") {
-         val poolSize = 20
+        val poolSize = 20
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.getMaximumPoolSize == poolSize)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(poolSize = poolSize))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getMaximumPoolSize == poolSize))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(poolSize = poolSize))) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("Pool size should have 10 connections by default") {
         (for {
           cp <- ZIO.service[HikariConnectionPool]
           _  <- ZIO.replicateZIO(10)(ZIO.scoped(cp.connection))
-        } yield assertTrue(cp.dataSource.getMaximumPoolSize == 10)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getMaximumPoolSize == 10))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer >>> poolLayer)
       } @@ timeout(10.minutes) @@ withLiveClock,
-
       test("It should be possible to acquire connections from the pool") {
         val poolSize = 20
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-          _  <- ZIO.collectAllParDiscard(ZIO.replicate(poolSize)(ZIO.scoped(cp.connection *> ZIO.sleep(500.millisecond))))
-        } yield assert("")(Assertion.anything)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(poolSize = poolSize))) >>> poolLayer)
+          _  <-
+            ZIO.collectAllParDiscard(ZIO.replicate(poolSize)(ZIO.scoped(cp.connection *> ZIO.sleep(500.millisecond))))
+        } yield assert("")(Assertion.anything)).provideSomeLayer[TestEnvironment with MySqlConfig](
+          hikariPoolConfigLayer.map(_.update(_.copy(poolSize = poolSize))) >>> poolLayer
+        )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("Auto commit should be configurable") {
         val autoCommit = false
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.isAutoCommit == autoCommit)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(autoCommit = autoCommit))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.isAutoCommit == autoCommit))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(autoCommit = autoCommit))) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("Auto commit should be true by default") {
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.isAutoCommit)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer >>> poolLayer)
+        } yield assertTrue(cp.dataSource.isAutoCommit))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer >>> poolLayer)
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("Connection timeout should be configurable") {
         val connectionTimeout = 2000L
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.getConnectionTimeout == connectionTimeout)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(connectionTimeout = Some(connectionTimeout)))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getConnectionTimeout == connectionTimeout))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(connectionTimeout = Some(connectionTimeout)))) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("Idle timeout should be configurable") {
         val idleTimeout = 2000L
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.getIdleTimeout == idleTimeout)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(idleTimeout = Some(idleTimeout)))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getIdleTimeout == idleTimeout))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(idleTimeout = Some(idleTimeout)))) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("initialization fail timeout should be configurable") {
         val initializationFailTimeout = 2000L
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.getInitializationFailTimeout == initializationFailTimeout)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(initializationFailTimeout = Some(initializationFailTimeout)))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getInitializationFailTimeout == initializationFailTimeout))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(
+              _.update(_.copy(initializationFailTimeout = Some(initializationFailTimeout)))
+            ) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("max lifetime should be configurable") {
         val maxLifetime = 40000L
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.getMaxLifetime == maxLifetime)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(maxLifetime = Some(maxLifetime)))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getMaxLifetime == maxLifetime))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(maxLifetime = Some(maxLifetime)))) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
       test("minimum idle should be configurable") {
         val minimumIdle = 2
         (for {
           cp <- ZIO.service[HikariConnectionPool]
-        } yield assertTrue(cp.dataSource.getMinimumIdle == minimumIdle)).provideSomeLayer[TestEnvironment with MySqlConfig](hikariPoolConfigLayer.map(_.update(_.copy(minimumIdle = Some(minimumIdle)))) >>> poolLayer)
+        } yield assertTrue(cp.dataSource.getMinimumIdle == minimumIdle))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(minimumIdle = Some(minimumIdle)))) >>> poolLayer
+          )
       } @@ timeout(10.seconds) @@ withLiveClock,
-
+      test("connection init SQL should be configurable") {
+        val initialSql = "SELECT 1 FROM test.test"
+        (for {
+          cp <- ZIO.service[HikariConnectionPool]
+        } yield assertTrue(cp.dataSource.getConnectionInitSql == initialSql))
+          .provideSomeLayer[TestEnvironment with MySqlConfig](
+            hikariPoolConfigLayer.map(_.update(_.copy(connectionInitSql = Some(initialSql)))) >>> poolLayer
+          )
+      } @@ timeout(10.seconds) @@ withLiveClock
     ) @@ sequential
 }
diff --git a/jdbc-hikaricp/src/test/scala/zio/sql/MySqlTestContainer.scala b/jdbc-hikaricp/src/test/scala/zio/sql/MySqlTestContainer.scala
@@ -12,7 +12,10 @@ object MySqlTestContainer {
       ZIO.attemptBlocking {
         val c = new MySQLContainer(
           mysqlImageVersion = Option(imageName).map(DockerImageName.parse)
-        )
+        ).configure { a =>
+          a.withInitScript("test_schema.sql")
+          ()
+        }
         c.start()
         c
       }