JetBrains / JetBrains/Exposed

UpdateBatchStatement with expressions

Open
#1,544 0 comments 10 reactions 0 assignees View on GitHub
Dominant language
Kotlin
Stars
9.3k
Forks
798
Avg merge
4d 2h
Merged PRs (30d)
26

Description

Hello, first off, thanks for the great lib!

I have some questions about the existing [BatchUpdateStatement](https://github.com/JetBrains/Exposed/blob/master/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/BatchUpdateStatement.kt).

The override for `update` with expression currently has the following function body:
```kotlin
override fun update(column: Column, value: Expression) = 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:
```kotlin
val data = listOf()
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:
```kotlin
override fun arguments(): Iterable>> = 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](https://github.com/JetBrains/Exposed/blob/master/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Expression.kt#L71) 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):
```kotlin
class BatchUpdateStatementWithExpression(
val table: IdTable<*>,
private val returnKeys: Boolean = false
) : UpdateStatement(table, null) {
val data = ArrayList, Map, Any?>>>()
override val firstDataSet: List, Any?>> get() = data.first().second.toList()

var resultRows: List = 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>> = 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 {
// 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, 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, Any?>) }
}
}

interface BatchUpdateRequest> {
val id: EntityID
}

fun , T : IdTable, E : BatchUpdateRequest> T.batchUpdateWithExpression(
data: Iterable,
returnKeys: Boolean = false,
body: BatchUpdateStatementWithExpression.(E) -> Unit
): List {
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:
```kotlin
val data = listOf()
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! 🙏

Contributor guide

Open the contributing guide

Research direction

Start with exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/BatchUpdateStatement.kt, especially the expression overload of update and the arguments() path. Read QueryBuilder.registerArgument in exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Expression.kt, then verify that batch updates accept expressions, bind their arguments correctly, and preserve the existing batch behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
kotlin, sql
Domain
database
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.