All Downloads are FREE. Search and download functionalities are using the official Maven repository.

jvmMain.com.river.connector.rdbms.jdbc.JdbcExt.kt Maven / Gradle / Ivy


package com.river.connector.rdbms.jdbc

import com.river.core.ExperimentalRiverApi
import com.river.core.GroupStrategy
import com.river.core.GroupStrategy.TimeWindow
import com.river.core.chunked
import com.river.core.mapAsync
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.invoke
import java.sql.PreparedStatement
import java.sql.ResultSet
import kotlin.reflect.full.primaryConstructor
import kotlin.time.Duration.Companion.milliseconds

typealias Row = Map

/**
 * Executes a single SQL update using a JDBC connection.
 *
 * This function takes an SQL statement [sql] and an optional [prepare] function, which prepares
 * the statement before execution.
 *
 * @param sql The SQL statement to execute.
 * @param prepare A suspend function that prepares the statement before execution.
 * @return A [Flow] of the number of rows affected by the update.
 *
 * Example:
 * ```
 * val jdbc: Jdbc = ...
 * val sql = "UPDATE users SET active = 1 WHERE id = 42"
 *
 * jdbc.singleUpdate(sql)
 *     .collect { rowsAffected ->
 *         println("Rows affected: $rowsAffected")
 *     }
 * ```
 */
@ExperimentalRiverApi
@ExperimentalCoroutinesApi
fun Jdbc.singleUpdate(
    sql: String,
    prepare: suspend PreparedStatement.() -> Unit = {}
): Flow = flow {
    connectionPool
        .borrow { IO { it.prepareStatement(sql).also { prepare(it) }.executeUpdate() } }
        .let { emit(it) }
}

/**
 * Executes a single SQL update using a JDBC connection for each item in the [upstream] flow.
 *
 * This function takes an SQL statement [sql], an [upstream] flow, [concurrency], and an optional
 * [prepare] function, which prepares the statement before execution.
 *
 * @param sql The SQL statement to execute.
 * @param upstream A [Flow] of items to process.
 * @param concurrency The level of concurrency for executing the updates.
 * @param prepare A suspend function that prepares the statement for each item before execution.
 * @return A [Flow] of the number of rows affected by the update for each item.
 *
 * Example:
 * ```
 * val jdbc: Jdbc = ...
 * val sql = "UPDATE users SET active = 1 WHERE id = ?"
 * val userIds = flowOf(1, 2, 3, 4, 5)
 *
 * jdbc.singleUpdate(sql, userIds, prepare = { id ->
 *     setInt(1, id)
 * }).collect { rowsAffected ->
 *     println("Rows affected: $rowsAffected")
 * }
 * ```
 */
@ExperimentalRiverApi
@ExperimentalCoroutinesApi
fun  Jdbc.singleUpdate(
    sql: String,
    upstream: Flow,
    concurrency: Int = 1,
    prepare: suspend PreparedStatement.(T) -> Unit = {}
): Flow =
    upstream
        .mapAsync(concurrency) { item ->
        connectionPool.borrow {
            IO {
                it.prepareStatement(sql)
                    .also { ps -> prepare(ps, item) }
                    .executeUpdate()
            }
        }
    }

/**
 * Executes a batch update for the given SQL statement using the provided [upstream] Flow as input.
 * The batch update is performed in chunks, as specified by the [groupStrategy] parameter, and can be executed concurrently
 * using the specified [concurrency] level.
 *
 * @param sql the SQL statement to be executed
 * @param upstream the Flow of input elements to be used in the batch update
 * @param concurrency the level of concurrency to be used in the batch update (default: 1)
 * @param groupStrategy the chunking strategy to be used for processing the input elements (default: TimeWindow(100, 250.milliseconds))
 * @param prepare a function that prepares the PreparedStatement using the current input element (default: {})
 * @return A [Flow] of the number of rows affected by the update for each batch.
 *
 * Example usage:
 *
 * ```
 *  val data = flowOf(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
 *  val updatedRows = Jdbc.batchUpdate("UPDATE my_table SET value = ? WHERE id = ?", data) { ps, value ->
 *      ps.setInt(1, value)
 *      ps.setInt(2, value)
 *  }
 *  updatedRows.collect { println("Updated $it rows.") }
 * ```
 */
@ExperimentalRiverApi
@ExperimentalCoroutinesApi
fun  Jdbc.batchUpdate(
    sql: String,
    upstream: Flow,
    concurrency: Int = 1,
    groupStrategy: GroupStrategy = TimeWindow(100, 250.milliseconds),
    prepare: suspend PreparedStatement.(T) -> Unit = {}
): Flow =
    upstream
        .chunked(groupStrategy)
        .mapAsync(concurrency) { chunk ->
            connectionPool.borrow {
                IO {
                    logger.debug("Running $sql with ${chunk.size} elements.")

                    it.prepareStatement(sql)
                        .also { ps -> chunk.forEach { prepare(ps, it); ps.addBatch() } }
                        .executeBatch()
                        .size
                        .also { logger.debug("Finished running batch update: $it rows were updated.") }
                }
            }
        }

/**
 * Executes a typed query for the given SQL statement using the provided [fetchSize] as a result set fetch size.
 * The resulting [Flow] emits objects of type [T], which must have a primary constructor and whose parameters are
 * matched with the columns of the result set.
 *
 * @param sql the SQL statement to be executed
 * @param fetchSize the result set fetch size to be used in the query (default: 100)
 * @return a Flow of objects of type [T], whose parameters are matched with the columns of the result set.
 *
 * Example usage:
 *
 * ```
 *  data class Person(val name: String, val age: Int)
 *
 *  val people = Jdbc.typedQuery("SELECT name, age FROM people WHERE age > 18")
 *  people.collect { println(it) }
 * ```
 */
@ExperimentalRiverApi
@ExperimentalCoroutinesApi
inline fun  Jdbc.typedQuery(
    sql: String,
    fetchSize: Int = 100,
): Flow = typedQuery(sql, fetchSize) { }

/**
 * Executes a typed query for the given SQL statement using the provided [fetchSize] as a result set fetch size.
 * The resulting [Flow] emits objects of type [T], which must have a primary constructor and whose parameters are
 * matched with the columns of the result set.
 *
 * @param sql the SQL statement to be executed
 * @param fetchSize the result set fetch size to be used in the query (default: 100)
 * @param prepare a function that prepares the PreparedStatement for the query (default: {})
 * @return a Flow of objects of type [T], whose parameters are matched with the columns of the result set.
 *
 * Example usage:
 * ```
 *  data class Person(val name: String, val age: Int)
 *
 *  val people = Jdbc.typedQuery("SELECT name, age FROM people WHERE age > ?") {
 *      setInt(1, 18)
 *  }
 *
 *  people.collect { println(it) }
 * ```
 */
@ExperimentalRiverApi
@ExperimentalCoroutinesApi
inline fun  Jdbc.typedQuery(
    sql: String,
    fetchSize: Int = 100,
    crossinline prepare: suspend PreparedStatement.() -> Unit,
): Flow =
    query(sql, fetchSize) { prepare() }
        .map { row ->
            val constructor = checkNotNull(T::class.primaryConstructor) {
                "Class ${T::class.simpleName} does not have a primary constructor"
            }

            constructor.callBy(
                constructor.parameters.associateWith { row[it.name] }
            )
        }

/**
 * Executes a query for the given SQL statement using the provided [fetchSize] as a result set fetch size.
 *
 * @param sql the SQL statement to be executed
 * @param fetchSize the result set fetch size to be used in the query (default: 100)
 * @param prepare a function that prepares the PreparedStatement for the query (default: {})
 * @return a Flow of Row objects, representing the result set.
 *
 * Example usage:
 * ```
 *  val result = Jdbc.query("SELECT * FROM my_table WHERE id = ?") {
 *      setInt(1, 10)
 *  }
 *  result.collect { println(it) }
 *```
 */
@ExperimentalRiverApi
@ExperimentalCoroutinesApi
fun Jdbc.query(
    sql: String,
    fetchSize: Int = 100,
    prepare: suspend PreparedStatement.() -> Unit = {}
): Flow = flow {
    connectionPool.borrow { connection ->
        val resultSet: ResultSet = IO {
            val statement = connection.prepareStatement(sql).also { prepare(it) }
            statement.fetchSize = fetchSize
            statement.executeQuery()
        }

        val metaData = resultSet.metaData
        val columns = metaData.columnCount

        while (IO { resultSet.next() }) {
            (1..columns)
                .associate { metaData.getColumnName(it) to resultSet.getObject(it) }
                .let { emit(it) }
        }
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy