UpdateBatchStatement with expressions

[ryuichi7]

Hello, first off, thanks for the great lib!

I have some questions about the existing BatchUpdateStatement.

The override for update with expression currently has the following function body:

override fun <T, S : T?> update(column: Column<T>, value: Expression<S>) = error("Expressions unsupported in batch update")

I was wondering why that is? Is there a major restriction for that?

I would like to use batch update to do something like the following:

val data = listOf<DataToUpdate>()
BatchUpdateStatement(FooTable).apply {
    data.forEach {
        addBatch(it.id)
        // i.e. SET foo_table.amount = foo_table.amount + ?
        this[FooTable.amount] = FooTable.amount + it.amount
    }
    execute(Transaction.current())
}

But with the current implementation, I get an error because of the set operator not allowing expressions.

What I've done to fix this is write a new version of BatchUpdateStatement that doesn't override update and then overrides the arguments() function with the following:

override fun arguments(): Iterable<Iterable<Pair<IColumnType, Any?>>> = data.map { (id, rowInput) ->
        QueryBuilder(true).run {
            rowInput.forEach {
                registerArgument(it.key, it.value)
            }
            // last input arg should be table id for where clause
            registerArgument(table.id.columnType, id)

            args
        }
    }

By using a QueryBuilder when fetching the arguments for the preparedStatement, I can use registerArgument which will parse the expression correctly.

I would really appreciate any input here with regards to this approach. Is there something I'm missing? Was there a more specific reason why expressions weren't allowed in the current version of BatchUpdateStatement?

I would also be happy to open a PR against the lib if my current approach seems viable.

Here is the full code for my working implementation with returning (assuming using Postgres):

class BatchUpdateStatementWithExpression(
    val table: IdTable<*>,
    private val returnKeys: Boolean = false
) : UpdateStatement(table, null) {
    val data = ArrayList<Pair<EntityID<*>, Map<Column<*>, Any?>>>()
    override val firstDataSet: List<Pair<Column<*>, Any?>> get() = data.first().second.toList()

    var resultRows: List<ResultRow> = emptyList()
        private set

    fun addBatch(id: EntityID<*>) {
        val lastBatch = data.lastOrNull()
        val different by lazy {
            val set1 = firstDataSet.map { it.first }.toSet()
            val set2 = lastBatch!!.second.keys
            (set1 - set2) + (set2 - set1)
        }

        if (data.size > 1 && different.isNotEmpty()) {
            throw BatchDataInconsistentException("Some values missing for batch update. Different columns: $different")
        }

        if (data.isNotEmpty()) {
            data[data.size - 1] = lastBatch!!.copy(second = values.toMap())
            values.clear()
            hasBatchedValues = true
        }
        data.add(id to values)
    }

    override fun prepareSQL(transaction: Transaction): String =
        "${super.prepareSQL(transaction)} WHERE ${transaction.identity(table.id)} = ?"

    override fun PreparedStatementApi.executeInternal(transaction: Transaction): Int {
        val result = if (data.size == 1) executeUpdate() else executeBatch().sum()

        if (returnKeys && resultSet != null) {
            resultRows = processResults(resultSet!!)
        }

        return result
    }

    override fun arguments(): Iterable<Iterable<Pair<IColumnType, Any?>>> = data.map { (id, rowInput) ->
        QueryBuilder(true).run {
            rowInput.forEach {
                registerArgument(it.key, it.value)
            }
            // last input arg should be table id for where clause
            registerArgument(table.id.columnType, id)

            args
        }
    }

    override fun prepared(transaction: Transaction, sql: String): PreparedStatementApi {
        return transaction.connection.prepareStatement(sql, returnKeys)
    }

    private fun processResults(resultSet: ResultSet): List<ResultRow> {
        // map columns to result set index
        val returnedColumns = targetsSet.columns.mapNotNull { col ->
            try {
                resultSet.findColumn(col.name).let { col to it }
            } catch (e: SQLException) {
                null
            }
        }

        val rows = arrayListOf<MutableMap<Column<*>, Any?>>()

        while (resultSet.next()) {
            val columnMappings = returnedColumns.associateTo(mutableMapOf()) {
                it.first to resultSet.getObject(it.second)
            }
            rows.add(columnMappings)
        }

        // cast here because createAndFillValues expects expression instead of column
        @Suppress("UNCHECKED_CAST")
        return rows.map { ResultRow.createAndFillValues(it as Map<Expression<*>, Any?>) }
    }
}

interface BatchUpdateRequest<T : Comparable<T>> {
    val id: EntityID<T>
}

fun <Key : Comparable<Key>, T : IdTable<Key>, E : BatchUpdateRequest<Key>> T.batchUpdateWithExpression(
    data: Iterable<E>,
    returnKeys: Boolean = false,
    body: BatchUpdateStatementWithExpression.(E) -> Unit
): List<ResultRow> {
    if (!data.iterator().hasNext()) return emptyList()

    return BatchUpdateStatementWithExpression(this, returnKeys).apply {
        data.forEach {
            addBatch(it.id)
            body(it)
        }

        execute(TransactionManager.current())
    }.resultRows
}

Which can be used like:

val data = listOf<DataToUpdate>()
FooTable.batchUpdateWithExpression(data, returnKeys = true) {
    // i.e. SET foo_table.amount = foo_table.amount + ?
    this[FooTable.amount] = FooTable.amount + it.amount
}

Any feedback would be greatly appreciated! 🙏