A datasource wrapper with connection creation callbacks

DataSourceWithConnectionCallback.kt

//A datasource wrapper with connection creation callbacks


import com.zaxxer.hikari.HikariConfig
import com.zaxxer.hikari.HikariDataSource
import com.zaxxer.hikari.util.DriverDataSource
import org.slf4j.Logger
import org.slf4j.LoggerFactory
import org.springframework.jdbc.core.JdbcTemplate
import org.springframework.jdbc.datasource.AbstractDataSource
import org.springframework.transaction.TransactionStatus
import org.springframework.transaction.support.SimpleTransactionStatus
import org.springframework.transaction.support.TransactionCallback
import org.springframework.transaction.support.TransactionOperations
import java.sql.Connection
import javax.sql.DataSource

/**
 * Usage:
 * 
 * @Configuration
 * class PreWarmedDatasourceConfiguration {
 *     @Bean
 *     @ConfigurationProperties("spring.datasource.hikari")
 *     fun localHikariConfig(): HikariConfig {
 *         return HikariConfig()
 *     }
 * 
 *     @Bean
 *     fun dataSource(): DataSource {
 *         return HikariDataSourceWithConnectionCallbackFactory
 *             .create(localHikariConfig()) { jdbcTemplate: JdbcTemplate ->
 *                 val repo = ServiceRepository(NamedParameterJdbcTemplate(jdbcTemplate))
 *                 repo.findByIds(ns) // warm-up 
 *             }
 *     }
 * }
 */


fun interface ConnectionCreationCallback {
    fun run(connection: Connection)
}

fun interface ConnectionCreationJdbcTemplateCallback {
    fun run(jdbcTemplate: JdbcTemplate)
}

/**
 * Factory functions for Hikari
 */
object HikariDataSourceWithConnectionCallbackFactory {
    fun create(
        config: HikariConfig,
        connectionCreationCallback: ConnectionCreationCallback,
    ): HikariDataSource {
        val configCopy = HikariConfig()
        config.copyStateTo(configCopy)

        configCopy.apply {
            val ds = DriverDataSource(jdbcUrl, driverClassName, dataSourceProperties, username, password)
            this.dataSource = DataSourceWithConnectionCallback(ds, connectionCreationCallback)
            this.jdbcUrl = null
            this.username = null
            this.password = null
        }

        return HikariDataSource(configCopy)
    }

    fun create(
        config: HikariConfig,
        connectionCreationCallback: ConnectionCreationJdbcTemplateCallback,
    ): HikariDataSource {
        val impl = object : AbstractJdbcTemplateConnectionCreationCallback() {
            override fun run(jdbcTemplate: JdbcTemplate) {
                connectionCreationCallback.run(jdbcTemplate)
            }
        }

        return create(config, impl)
    }
}


/**
 * DataSourceWithConnectionCallback implementation 
 */
class DataSourceWithConnectionCallback(
    private val dataSource: DataSource,
    private val connectionCreationCallback: ConnectionCreationCallback
) : DataSource by dataSource {
    override fun getConnection(): Connection {
        return dataSource.getConnection().also { connectionCreationCallback.run(it) }
    }

    override fun getConnection(username: String?, password: String?): Connection? {
        return dataSource.getConnection(username, password).also { connectionCreationCallback.run(it) }
    }
}


/**
 * A bridge from connection to JdbcTemplate
 */
abstract class AbstractJdbcTemplateConnectionCreationCallback() : ConnectionCreationCallback {
    private val logger: Logger = LoggerFactory.getLogger(this.javaClass)

    override fun run(connection: Connection) {
        try {
            val ds = SingleConnectionDataSource(connection)
            run(JdbcTemplate(ds))
        } catch (e: Exception) {
            logger.warn("error found ${e.message}")
        }
    }

    abstract fun run(jdbcTemplate: JdbcTemplate)

    class SingleConnectionDataSource(
        private val con: Connection,
    ) : AbstractDataSource() {
        private val c = NonCloseableConnection(con)
        override fun getConnection(): Connection = c
        override fun getConnection(username: String?, password: String?): Connection = c

        data class NonCloseableConnection(val connection: Connection) : Connection by connection {
            override fun close() {} //do nothing
        }
    }
}


/**
 * A mock transaction manager
 */
class WithoutTransactionOperations() : TransactionOperations {
    override fun <T> execute(action: TransactionCallback<T>): T? {
        return action.doInTransaction(SimpleTransactionStatus(false))
    }

    override fun executeWithoutResult(action: Consumer<TransactionStatus>) {
        action.accept(SimpleTransactionStatus(false))
    }
}