Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

UpdateBatchStatement with expressions #1544

Open
ryuichi7 opened this issue Jul 1, 2022 · 0 comments
Open

UpdateBatchStatement with expressions #1544

ryuichi7 opened this issue Jul 1, 2022 · 0 comments

Comments

@ryuichi7
Copy link

ryuichi7 commented Jul 1, 2022

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! 🙏

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant