jvmMain.kr.jadekim.jext.exposed.dbms.mysql.statement.kt Maven / Gradle / Ivy

Go to download
package kr.jadekim.jext.exposed.dbms.mysql

import org.jetbrains.exposed.sql.*
import org.jetbrains.exposed.sql.Function
import org.jetbrains.exposed.sql.statements.InsertStatement
import org.jetbrains.exposed.sql.statements.UpdateStatement
import org.jetbrains.exposed.sql.transactions.TransactionManager

fun  T.upsert(
    body: T.(InsertStatement) -> Unit,
) = UpsertStatement(this).apply {
    body(this)
    execute(TransactionManager.current())
}

fun  T.upsert(
    body: T.(InsertStatement) -> Unit,
    updateBody: T.(UpdateStatement) -> Unit,
) = UpsertStatement(this, updateStatement = UpdateStatement(this, null).apply { updateBody(this) }).apply {
    body(this)
    execute(TransactionManager.current())
}

class UpsertStatement(
    table: Table,
    isIgnore: Boolean = false,
    private val updateStatement: UpdateStatement? = null,
) : InsertStatement(table, isIgnore) {

    override fun prepareSQL(transaction: Transaction, prepared: Boolean): String {
        val keys = values.keys.filterNot { table.primaryKey?.columns?.contains(it) == true }

        if (keys.isEmpty()) {
            return super.prepareSQL(transaction, prepared)
        }

        val onUpdate = if (updateStatement == null) {
            keys.joinToString { "${transaction.identity(it)} = VALUES(${transaction.identity(it)})" }
        } else {
            val builder = QueryBuilder(true)
            updateStatement.firstDataSet.appendTo(builder) { (col, value) ->
                append("${transaction.identity(col)}=")
                registerArgument(col, value)
            }
            builder.toString()
        }


        return "${super.prepareSQL(transaction, prepared)} ON DUPLICATE KEY UPDATE $onUpdate"
    }

    override fun arguments(): List, Any?>>> {
        val updateArguments = updateStatement?.arguments()?.firstOrNull() ?: return super.arguments()
        val insertArguments = super.arguments().firstOrNull()

        val arguments = mutableListOf, Any?>>().apply {
            if (insertArguments != null) {
                addAll(insertArguments)
            }

            addAll(updateArguments)
        }

        return listOf(arguments)
    }
}

class Values(val column: Column<*>) : Function(VarCharColumnType()) {

    override fun toQueryBuilder(queryBuilder: QueryBuilder) = queryBuilder {
        +"VALUES("
        +column
        +")"
    }
}