Supporting subqueries in insert and update statements (#1328)

* Supporting subqueries in insert and update statements * Fix the added test * Replace expression with wrapWithExpression
JetBrains · Sep 5, 2021 · 2733dc2 · 2733dc2
1 parent 283db83
commit 2733dc2
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/UpdateBuilder.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/statements/UpdateBuilder.kt
@@ -7,8 +7,10 @@ import org.jetbrains.exposed.sql.Column
 import org.jetbrains.exposed.sql.CompositeColumn
 import org.jetbrains.exposed.sql.Expression
 import org.jetbrains.exposed.sql.Op
+import org.jetbrains.exposed.sql.Query
 import org.jetbrains.exposed.sql.SqlExpressionBuilder
 import org.jetbrains.exposed.sql.Table
+import org.jetbrains.exposed.sql.wrapAsExpression
 
 /**
  * @author max
@@ -54,6 +56,8 @@ abstract class UpdateBuilder<out T>(type: StatementType, targets: List<Table>) :
 
     open operator fun <T, S : T, E : Expression<S>> set(column: Column<T>, value: E) = update(column, value)
 
+    open operator fun <S> set(column: Column<S>, value: Query) = update(column, wrapAsExpression(value))
+
     open operator fun <S> set(column: CompositeColumn<S>, value: S) {
         column.getRealColumnsWithValues(value).forEach { (realColumn, itsValue) -> set(realColumn as Column<Any?>, itsValue) }
     }

diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/dml/InsertTests.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/dml/InsertTests.kt
@@ -402,6 +402,35 @@ class InsertTests : DatabaseTestsBase() {
         }
     }
 
+    @Test fun `test subquery in an insert or update statement`() {
+        val tab1 = object : Table("tab1") {
+            val id = varchar("id", 10)
+        }
+        val tab2 = object : Table("tab2") {
+            val id = varchar("id", 10)
+        }
+
+        withTables(tab1, tab2) {
+            // Initial data
+            tab2.insert { it[id] = "foo" }
+            tab2.insert { it[id] = "bar" }
+
+            // Use sub query in an insert
+            tab1.insert { it[id] = tab2.slice(tab2.id).select { tab2.id eq "foo" } }
+
+            // Check inserted data
+            val insertedId = tab1.slice(tab1.id).selectAll().single()[tab1.id]
+            assertEquals("foo", insertedId)
+
+            // Use sub query in an update
+            tab1.update({ tab1.id eq "foo" }) { it[id] = tab2.slice(tab2.id).select { tab2.id eq "bar" } }
+
+            // Check updated data
+            val updatedId = tab1.slice(tab1.id).selectAll().single()[tab1.id]
+            assertEquals("bar", updatedId)
+        }
+    }
+
     /*
     @Test fun testGeneratedKey04() {
         val CharIdTable = object : IdTable<String>("charId") {