EXPOSED-376 batchUpsert does not return database values on conflict

JetBrains · May 14, 2024 · 629b903 · 629b903
1 parent fd519d3
commit 629b903
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 6 deletions.
diff --git a/exposed-core/api/exposed-core.api b/exposed-core/api/exposed-core.api
@@ -2927,6 +2927,7 @@ public class org/jetbrains/exposed/sql/statements/BatchUpsertStatement : org/jet
 	public final fun getKeys ()[Lorg/jetbrains/exposed/sql/Column;
 	public final fun getOnUpdate ()Ljava/util/List;
 	public final fun getOnUpdateExclude ()Ljava/util/List;
+	protected fun isColumnValuePreferredFromResultSet (Lorg/jetbrains/exposed/sql/Column;Ljava/lang/Object;)Z
 	public fun prepareSQL (Lorg/jetbrains/exposed/sql/Transaction;Z)Ljava/lang/String;
 }
 
@@ -2994,6 +2995,7 @@ public class org/jetbrains/exposed/sql/statements/InsertStatement : org/jetbrain
 	public final fun getResultedValues ()Ljava/util/List;
 	public final fun getTable ()Lorg/jetbrains/exposed/sql/Table;
 	protected fun isColumnValuePreferredFromResultSet (Lorg/jetbrains/exposed/sql/Column;Ljava/lang/Object;)Z
+	protected final fun isEntityIdClientSideGeneratedUUID (Lorg/jetbrains/exposed/sql/Column;)Z
 	public final fun isIgnore ()Z
 	public fun prepareSQL (Lorg/jetbrains/exposed/sql/Transaction;Z)Ljava/lang/String;
 	public fun prepared (Lorg/jetbrains/exposed/sql/Transaction;Ljava/lang/String;)Lorg/jetbrains/exposed/sql/statements/api/PreparedStatementApi;

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/BatchUpsertStatement.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/BatchUpsertStatement.kt
@@ -42,4 +42,9 @@ open class BatchUpsertStatement(
         }
         return functionProvider.upsert(table, arguments!!.first(), onUpdate, onUpdateExclude, null, transaction, keys = keys)
     }
+
+    override fun isColumnValuePreferredFromResultSet(column: Column<*>, value: Any?): Boolean {
+        return isEntityIdClientSideGeneratedUUID(column) ||
+            super.isColumnValuePreferredFromResultSet(column, value)
+    }
 }
diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/InsertStatement.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/InsertStatement.kt
@@ -219,4 +219,10 @@ open class InsertStatement<Key : Any>(
             builder.args
         } ?: emptyList()
     }
+
+    protected fun isEntityIdClientSideGeneratedUUID(column: Column<*>) =
+        (column.columnType as? EntityIDColumnType<*>)
+            ?.idColumn
+            ?.takeIf { it.columnType is UUIDColumnType }
+            ?.defaultValueFun != null
 }
diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/UpsertStatement.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/UpsertStatement.kt
@@ -53,10 +53,4 @@ open class UpsertStatement<Key : Any>(
         return isEntityIdClientSideGeneratedUUID(column) ||
             super.isColumnValuePreferredFromResultSet(column, value)
     }
-
-    private fun isEntityIdClientSideGeneratedUUID(column: Column<*>) =
-        (column.columnType as? EntityIDColumnType<*>)
-            ?.idColumn
-            ?.takeIf { it.columnType is UUIDColumnType }
-            ?.defaultValueFun != null
 }
diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/dml/UpsertTests.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/dml/UpsertTests.kt
@@ -655,6 +655,33 @@ class UpsertTests : DatabaseTestsBase() {
         }
     }
 
+    @Test
+    fun testBatchUpsertWithUUIDPrimaryKey() {
+        val tester = object : UUIDTable("batch_upsert_test", "id") {
+            val key = integer("test_key").uniqueIndex()
+            val value = text("test_value")
+        }
+
+        withTables(excludeSettings = TestDB.entries - listOf(TestDB.POSTGRESQL, TestDB.POSTGRESQLNG), tester) {
+            val insertId = tester.insertAndGetId {
+                it[key] = 1
+                it[value] = "one"
+            }
+
+            val upsertId = tester.batchUpsert(
+                data = listOf(1 to "two"),
+                keys = arrayOf(tester.key),
+                onUpdateExclude = listOf(tester.id),
+            ) {
+                this[tester.key] = it.first
+                this[tester.value] = it.second
+            }.first()[tester.id]
+
+            assertEquals(insertId, upsertId)
+            assertEquals("two", tester.selectAll().where { tester.id eq insertId }.first()[tester.value])
+        }
+    }
+
     private object AutoIncTable : Table("auto_inc_table") {
         val id = integer("id").autoIncrement()
         val name = varchar("name", 64)