wvlet.airframe.jdbc.PostgreSQLConnectionPool.scala Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of airframe-jdbc_2.13.0-RC2 Show documentation
JDBC connection pool service
The newest version!
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package wvlet.airframe.jdbc

import java.sql.Connection

import com.zaxxer.hikari.{HikariConfig, HikariDataSource}

case class PostgreSQLConfig(
    maxPoolSize: Int = 10,
    autoCommit: Boolean = true, // Enable auto-commit
    // SSL configuration for using RDS
    useSSL: Boolean = true,
    sslFactory: String = "org.postgresql.ssl.NonValidatingFactory"
)

/**
  *
  */
class PostgreSQLConnectionPool(val config: DbConfig, val pgConfig: PostgreSQLConfig = PostgreSQLConfig())
    extends ConnectionPool {
  private val dataSource: HikariDataSource = {
    val connectionPoolConfig = new HikariConfig

    // Set default JDBC parameters
    connectionPoolConfig.setMaximumPoolSize(pgConfig.maxPoolSize) // HikariCP default = 10
    connectionPoolConfig.setAutoCommit(pgConfig.autoCommit)       // Enable auto-commit

    connectionPoolConfig.setDriverClassName("org.postgresql.Driver")
    config.user.foreach(u => connectionPoolConfig.setUsername(u))
    config.password.foreach(p => connectionPoolConfig.setPassword(p))
    if (pgConfig.useSSL) {
      connectionPoolConfig.addDataSourceProperty("ssl", "true")
      connectionPoolConfig.addDataSourceProperty("sslfactory", pgConfig.sslFactory)
    }
    if (config.host.isEmpty) {
      throw new IllegalArgumentException(s"missing postgres jdbc host: ${config}")
    }
    connectionPoolConfig.setJdbcUrl(
      s"jdbc:postgresql://${config.host.get}:${config.port.getOrElse(5432)}/${config.database}")

    info(s"jdbc URL: ${connectionPoolConfig.getJdbcUrl}")
    new HikariDataSource(connectionPoolConfig)
  }

  def withConnection[U](body: Connection => U): U = {
    val conn = dataSource.getConnection
    try {
      body(conn)
    } finally {
      // Return the connection to the pool
      conn.close
    }
  }

  def stop: Unit = {
    info(s"Closing connection pool for ${config}")
    dataSource.close()
  }
}