Skip to content

Commit

Permalink
feat: EXPOSED-355 Support INSERT...RETURNING statement (#2060)
Browse files Browse the repository at this point in the history
* feat: EXPOSED-355 Support INSERT...RETURNING statement

- Add support for RETURNING clause with insert and upsert statements.
- Implement a ReturningStatement as the basis for update and deletes too.
- Implementation covers PostgreSQL and SQLite only currently.

* feat: EXPOSED-355 Support INSERT...RETURNING statement

Switch ReturningStatement to be iterable instead of auto-executing and returning
a List<ResultRow>.
Add samples to KDocs.
  • Loading branch information
bog-walk committed Apr 28, 2024
1 parent f42315e commit 4f5b037
Show file tree
Hide file tree
Showing 7 changed files with 272 additions and 0 deletions.
18 changes: 18 additions & 0 deletions exposed-core/api/exposed-core.api
Original file line number Diff line number Diff line change
Expand Up @@ -1718,6 +1718,8 @@ public final class org/jetbrains/exposed/sql/QueriesKt {
public static final fun insertIgnore (Lorg/jetbrains/exposed/sql/Table;Lorg/jetbrains/exposed/sql/AbstractQuery;Ljava/util/List;)Ljava/lang/Integer;
public static synthetic fun insertIgnore$default (Lorg/jetbrains/exposed/sql/Table;Lorg/jetbrains/exposed/sql/AbstractQuery;Ljava/util/List;ILjava/lang/Object;)Ljava/lang/Integer;
public static final fun insertIgnoreAndGetId (Lorg/jetbrains/exposed/dao/id/IdTable;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/exposed/dao/id/EntityID;
public static final fun insertReturning (Lorg/jetbrains/exposed/sql/Table;Ljava/util/List;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/exposed/sql/statements/ReturningStatement;
public static synthetic fun insertReturning$default (Lorg/jetbrains/exposed/sql/Table;Ljava/util/List;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lorg/jetbrains/exposed/sql/statements/ReturningStatement;
public static final fun replace (Lorg/jetbrains/exposed/sql/Table;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/exposed/sql/statements/ReplaceStatement;
public static final fun select (Lorg/jetbrains/exposed/sql/FieldSet;Lkotlin/jvm/functions/Function1;)Lorg/jetbrains/exposed/sql/Query;
public static final fun select (Lorg/jetbrains/exposed/sql/FieldSet;Lorg/jetbrains/exposed/sql/Op;)Lorg/jetbrains/exposed/sql/Query;
Expand All @@ -1738,6 +1740,8 @@ public final class org/jetbrains/exposed/sql/QueriesKt {
public static synthetic fun update$default (Lorg/jetbrains/exposed/sql/Table;Lkotlin/jvm/functions/Function1;Ljava/lang/Integer;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)I
public static final fun upsert (Lorg/jetbrains/exposed/sql/Table;[Lorg/jetbrains/exposed/sql/Column;Ljava/util/List;Ljava/util/List;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/exposed/sql/statements/UpsertStatement;
public static synthetic fun upsert$default (Lorg/jetbrains/exposed/sql/Table;[Lorg/jetbrains/exposed/sql/Column;Ljava/util/List;Ljava/util/List;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lorg/jetbrains/exposed/sql/statements/UpsertStatement;
public static final fun upsertReturning (Lorg/jetbrains/exposed/sql/Table;[Lorg/jetbrains/exposed/sql/Column;Ljava/util/List;Ljava/util/List;Ljava/util/List;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/exposed/sql/statements/ReturningStatement;
public static synthetic fun upsertReturning$default (Lorg/jetbrains/exposed/sql/Table;[Lorg/jetbrains/exposed/sql/Column;Ljava/util/List;Ljava/util/List;Ljava/util/List;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lorg/jetbrains/exposed/sql/statements/ReturningStatement;
}

public class org/jetbrains/exposed/sql/Query : org/jetbrains/exposed/sql/AbstractQuery {
Expand Down Expand Up @@ -3000,6 +3004,19 @@ public class org/jetbrains/exposed/sql/statements/ReplaceStatement : org/jetbrai
public fun prepareSQL (Lorg/jetbrains/exposed/sql/Transaction;Z)Ljava/lang/String;
}

public class org/jetbrains/exposed/sql/statements/ReturningStatement : org/jetbrains/exposed/sql/statements/Statement, java/lang/Iterable, kotlin/jvm/internal/markers/KMappedMarker {
public fun <init> (Lorg/jetbrains/exposed/sql/Table;Ljava/util/List;Lorg/jetbrains/exposed/sql/statements/Statement;)V
public fun arguments ()Ljava/lang/Iterable;
public synthetic fun executeInternal (Lorg/jetbrains/exposed/sql/statements/api/PreparedStatementApi;Lorg/jetbrains/exposed/sql/Transaction;)Ljava/lang/Object;
public fun executeInternal (Lorg/jetbrains/exposed/sql/statements/api/PreparedStatementApi;Lorg/jetbrains/exposed/sql/Transaction;)Ljava/sql/ResultSet;
public final fun getMainStatement ()Lorg/jetbrains/exposed/sql/statements/Statement;
public final fun getReturningExpressions ()Ljava/util/List;
public final fun getTable ()Lorg/jetbrains/exposed/sql/Table;
protected final fun getTransaction ()Lorg/jetbrains/exposed/sql/Transaction;
public fun iterator ()Ljava/util/Iterator;
public fun prepareSQL (Lorg/jetbrains/exposed/sql/Transaction;Z)Ljava/lang/String;
}

public class org/jetbrains/exposed/sql/statements/SQLServerBatchInsertStatement : org/jetbrains/exposed/sql/statements/BatchInsertStatement {
public fun <init> (Lorg/jetbrains/exposed/sql/Table;ZZ)V
public synthetic fun <init> (Lorg/jetbrains/exposed/sql/Table;ZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
Expand Down Expand Up @@ -3667,6 +3684,7 @@ public abstract class org/jetbrains/exposed/sql/vendors/FunctionProvider {
public fun regexp (Lorg/jetbrains/exposed/sql/Expression;Lorg/jetbrains/exposed/sql/Expression;ZLorg/jetbrains/exposed/sql/QueryBuilder;)V
public fun replace (Lorg/jetbrains/exposed/sql/Table;Ljava/util/List;Ljava/lang/String;Lorg/jetbrains/exposed/sql/Transaction;Z)Ljava/lang/String;
public static synthetic fun replace$default (Lorg/jetbrains/exposed/sql/vendors/FunctionProvider;Lorg/jetbrains/exposed/sql/Table;Ljava/util/List;Ljava/lang/String;Lorg/jetbrains/exposed/sql/Transaction;ZILjava/lang/Object;)Ljava/lang/String;
public fun returning (Ljava/lang/String;Ljava/util/List;Lorg/jetbrains/exposed/sql/Transaction;)Ljava/lang/String;
public fun second (Lorg/jetbrains/exposed/sql/Expression;Lorg/jetbrains/exposed/sql/QueryBuilder;)V
public fun stdDevPop (Lorg/jetbrains/exposed/sql/Expression;Lorg/jetbrains/exposed/sql/QueryBuilder;)V
public fun stdDevSamp (Lorg/jetbrains/exposed/sql/Expression;Lorg/jetbrains/exposed/sql/QueryBuilder;)V
Expand Down
46 changes: 46 additions & 0 deletions exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Queries.kt
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,23 @@ fun <T : Table> T.insertIgnore(
columns: List<Column<*>> = this.columns.filter { !it.columnType.isAutoInc || it.autoIncColumnType?.nextValExpression != null }
): Int? = InsertSelectStatement(columns, selectQuery, true).execute(TransactionManager.current())

/**
* Represents the SQL statement that inserts new rows into a table and returns specified data from the inserted rows.
*
* @param returning Columns and expressions to include in the returned data. This defaults to all columns in the table.
* @return A [ReturningStatement] that will be executed once iterated over, providing [ResultRow]s containing the specified
* expressions mapped to their resulting data.
* @sample org.jetbrains.exposed.sql.tests.shared.dml.ReturningTests.testInsertReturning
*/
fun <T : Table> T.insertReturning(
returning: List<Expression<*>> = columns,
body: T.(InsertStatement<Number>) -> Unit
): ReturningStatement {
val insert = InsertStatement<Number>(this)
body(insert)
return ReturningStatement(this, returning, insert)
}

/**
* Represents the SQL statement that updates rows of a table.
*
Expand Down Expand Up @@ -434,6 +451,35 @@ fun <T : Table> T.upsert(
execute(TransactionManager.current())
}

/**
* Represents the SQL statement that either inserts a new row into a table, or updates the existing row if insertion would
* violate a unique constraint, and also returns specified data from the modified rows.
*
* @param keys (optional) Columns to include in the condition that determines a unique constraint match. If no columns are
* provided, primary keys will be used. If the table does not have any primary keys, the first unique index will be attempted.
* @param returning Columns and expressions to include in the returned data. This defaults to all columns in the table.
* @param onUpdate List of pairs of specific columns to update and the expressions to update them with.
* If left null, all columns will be updated with the values provided for the insert.
* @param onUpdateExclude List of specific columns to exclude from updating.
* If left null, all columns will be updated with the values provided for the insert.
* @param where Condition that determines which rows to update, if a unique violation is found.
* @return A [ReturningStatement] that will be executed once iterated over, providing [ResultRow]s containing the specified
* expressions mapped to their resulting data.
* @sample org.jetbrains.exposed.sql.tests.shared.dml.ReturningTests.testUpsertReturning
*/
fun <T : Table> T.upsertReturning(
vararg keys: Column<*>,
returning: List<Expression<*>> = columns,
onUpdate: List<Pair<Column<*>, Expression<*>>>? = null,
onUpdateExclude: List<Column<*>>? = null,
where: (SqlExpressionBuilder.() -> Op<Boolean>)? = null,
body: T.(UpsertStatement<Long>) -> Unit
): ReturningStatement {
val update = UpsertStatement<Long>(this, *keys, onUpdate = onUpdate, onUpdateExclude = onUpdateExclude, where = where?.let { SqlExpressionBuilder.it() })
body(update)
return ReturningStatement(this, returning, update)
}

/**
* Represents the SQL statement that either batch inserts new rows into a table, or updates the existing rows if insertions violate unique constraints.
*
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
package org.jetbrains.exposed.sql.statements

import org.jetbrains.exposed.sql.Expression
import org.jetbrains.exposed.sql.IColumnType
import org.jetbrains.exposed.sql.ResultRow
import org.jetbrains.exposed.sql.Table
import org.jetbrains.exposed.sql.Transaction
import org.jetbrains.exposed.sql.statements.api.PreparedStatementApi
import org.jetbrains.exposed.sql.transactions.TransactionManager
import java.sql.ResultSet

/**
* Represents the underlying SQL [mainStatement] that also returns a result set with data from any modified rows.
*
* @param table Table to perform the main statement on and return results from.
* @param returningExpressions Columns or expressions to include in the returned result set.
* @param mainStatement The statement to append the RETURNING clause to. This may be an insert, update, or delete statement.
*/
open class ReturningStatement(
val table: Table,
val returningExpressions: List<Expression<*>>,
val mainStatement: Statement<*>
) : Iterable<ResultRow>, Statement<ResultSet>(mainStatement.type, listOf(table)) {
protected val transaction
get() = TransactionManager.current()

override fun PreparedStatementApi.executeInternal(transaction: Transaction): ResultSet = executeQuery()

override fun arguments(): Iterable<Iterable<Pair<IColumnType<*>, Any?>>> = mainStatement.arguments()

override fun prepareSQL(transaction: Transaction, prepared: Boolean): String {
val mainSql = mainStatement.prepareSQL(transaction, prepared)
return transaction.db.dialect.functionProvider.returning(mainSql, returningExpressions, transaction)
}

override fun iterator(): Iterator<ResultRow> {
val resultIterator = ResultIterator(transaction.exec(this)!!)
return Iterable { resultIterator }.iterator()
}

private inner class ResultIterator(val rs: ResultSet) : Iterator<ResultRow> {
val fieldIndex = returningExpressions.withIndex().associateBy({ it.value }, { it.index })

private var hasNext = false
set(value) {
field = value
if (!field) {
rs.statement?.close()
transaction.openResultSetsCount--
}
}

init {
hasNext = rs.next()
}

override fun hasNext(): Boolean = hasNext

override operator fun next(): ResultRow {
if (!hasNext) throw NoSuchElementException()
val result = ResultRow.create(rs, fieldIndex)
hasNext = rs.next()
return result
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -743,4 +743,23 @@ abstract class FunctionProvider {

/** Appends optional parameters to an EXPLAIN query. */
protected open fun StringBuilder.appendOptionsToExplain(options: String) { append("$options ") }

/**
* Returns the SQL command that performs an insert, update, or delete, and also returns data from any modified rows.
*
* **Note:** This operation is not supported by all vendors, please check the documentation.
*
* @param mainSql SQL string representing the underlying statement before appending a RETURNING clause.
* @param returning Columns and expressions to include in the returned result set.
* @param transaction Transaction where the operation is executed.
*/
open fun returning(
mainSql: String,
returning: List<Expression<*>>,
transaction: Transaction
): String {
transaction.throwUnsupportedException(
"There's no generic SQL for a command with a RETURNING clause. There must be a vendor specific implementation."
)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,18 @@ internal object PostgreSQLFunctionProvider : FunctionProvider() {
}

override fun StringBuilder.appendOptionsToExplain(options: String) { append("($options) ") }

override fun returning(
mainSql: String,
returning: List<Expression<*>>,
transaction: Transaction
): String {
return with(QueryBuilder(true)) {
+"$mainSql RETURNING "
returning.appendTo { +it }
toString()
}
}
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,18 @@ internal object SQLiteFunctionProvider : FunctionProvider() {
val sql = super.explain(false, null, internalStatement, transaction)
return sql.replaceFirst("EXPLAIN ", "EXPLAIN QUERY PLAN ")
}

override fun returning(
mainSql: String,
returning: List<Expression<*>>,
transaction: Transaction
): String {
return with(QueryBuilder(true)) {
+"$mainSql RETURNING "
returning.appendTo { +it }
toString()
}
}
}

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
package org.jetbrains.exposed.sql.tests.shared.dml

import org.jetbrains.exposed.dao.id.IntIdTable
import org.jetbrains.exposed.sql.*
import org.jetbrains.exposed.sql.SqlExpressionBuilder.times
import org.jetbrains.exposed.sql.statements.ReturningStatement
import org.jetbrains.exposed.sql.tests.DatabaseTestsBase
import org.jetbrains.exposed.sql.tests.TestDB
import org.junit.Test
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertIs

class ReturningTests : DatabaseTestsBase() {
private val returningSupportedDb = TestDB.postgreSQLRelatedDB.toSet() + TestDB.SQLITE

object Items : IntIdTable("items") {
val name = varchar("name", 32)
val price = double("price")
}

@Test
fun testInsertReturning() {
withTables(TestDB.enabledDialects() - returningSupportedDb, Items) {
// return all columns by default
val result1 = Items.insertReturning {
it[name] = "A"
it[price] = 99.0
}.single()
assertEquals(1, result1[Items.id].value)
assertEquals("A", result1[Items.name])
assertEquals(99.0, result1[Items.price])

val result2 = Items.insertReturning(listOf(Items.id, Items.name)) {
it[name] = "B"
it[price] = 200.0
}.single()
assertEquals(2, result2[Items.id].value)
assertEquals("B", result2[Items.name])

assertFailsWith<IllegalStateException> { // Items.price not in record set
result2[Items.price]
}

assertEquals(2, Items.selectAll().count())
}
}

@Test
fun testUpsertReturning() {
withTables(TestDB.enabledDialects() - returningSupportedDb, Items) {
// return all columns by default
val result1 = Items.upsertReturning {
it[name] = "A"
it[price] = 99.0
}.single()
assertEquals(1, result1[Items.id].value)
assertEquals("A", result1[Items.name])
assertEquals(99.0, result1[Items.price])

val result2 = Items.upsertReturning(
returning = listOf(Items.name, Items.price),
onUpdate = listOf(Items.price to Items.price.times(10.0))
) {
it[id] = 1
it[name] = "B"
it[price] = 200.0
}.single()
assertEquals("A", result2[Items.name])
assertEquals(990.0, result2[Items.price])

val result3 = Items.upsertReturning(
returning = listOf(Items.name),
onUpdateExclude = listOf(Items.price),
where = { Items.price greater 500.0 }
) {
it[id] = 1
it[name] = "B"
it[price] = 200.0
}.single()
assertEquals("B", result3[Items.name])

assertEquals(1, Items.selectAll().count())
}
}

@Test
fun testReturningWithNoResults() {
withTables(TestDB.enabledDialects() - returningSupportedDb, Items) {
// statement not executed if not iterated over
val stmt = Items.insertReturning {
it[name] = "A"
it[price] = 99.0
}
assertIs<ReturningStatement>(stmt)
assertEquals(0, Items.selectAll().count())
}
}
}

0 comments on commit 4f5b037

Please sign in to comment.