Support the expression in the update set clause

nakamura-to · nakamura-to · commit 975f86347a0b · 2020-04-29T11:44:28.000+09:00
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/context/Context.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/context/Context.kt
@@ -11,6 +11,11 @@ interface Context {
 sealed class Operand {
     data class Param(val value: InParameter<*>) : Operand()
     data class Prop(val value: PropertyDef<*>) : Operand()
+    sealed class Expr : Operand() {
+        data class Plus(val left: Expr, val right: Expr) : Expr()
+        data class Prop(val value: Operand.Prop) : Expr()
+        data class Param(val value: Operand.Param) : Expr()
+    }
 }
 
 sealed class Criterion {
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/context/UpdateContext.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/context/UpdateContext.kt
@@ -6,7 +6,7 @@ import org.seasar.doma.jdbc.Config
 class UpdateContext(
     val config: Config,
     val entityDef: EntityDef<*>,
-    val set: MutableMap<Operand.Prop, Operand.Param> = mutableMapOf(),
+    val set: MutableMap<Operand.Prop, Operand> = mutableMapOf(),
     val where: MutableList<Criterion> = mutableListOf()
 ) : Context {
 
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/ExpressionDeclaration.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/ExpressionDeclaration.kt
@@ -0,0 +1,23 @@
+package org.seasar.doma.criteria.declaration
+
+import org.seasar.doma.criteria.context.Operand
+import org.seasar.doma.def.PropertyDef
+import org.seasar.doma.jdbc.Config
+
+@Declaration
+class ExpressionDeclaration(private val config: Config) {
+
+    private val support = DeclarationSupport(config)
+
+    infix operator fun <PROPERTY> PropertyDef<PROPERTY>.plus(other: PROPERTY?): Operand.Expr {
+        val prop = support.toProp(this)
+        val param = support.toParam(this, other)
+        return Operand.Expr.Plus(Operand.Expr.Prop(prop), Operand.Expr.Param(param))
+    }
+
+    infix operator fun <PROPERTY> PROPERTY?.plus(other: PropertyDef<PROPERTY>): Operand.Expr {
+        val param = support.toParam(other, this)
+        val prop = support.toProp(other)
+        return Operand.Expr.Plus(Operand.Expr.Param(param), Operand.Expr.Prop(prop))
+    }
+}
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/InsertDeclaration.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/InsertDeclaration.kt
@@ -6,15 +6,17 @@ import org.seasar.doma.def.PropertyDef
 @Declaration
 class InsertDeclaration(private val context: InsertContext) {
     private val support = DeclarationSupport(context.config)
-    private val valuesDeclaration = ValuesDeclaration()
+    private val valuesDeclaration = ValuesDeclaration(context.config)
 
     fun values(block: ValuesDeclaration.(Values) -> Unit) {
-        valuesDeclaration.block(object : Values {
-            override fun <PROPERTY> set(propDef: PropertyDef<PROPERTY>, value: PROPERTY?) {
-                val prop = support.toProp(propDef)
-                val param = support.toParam(propDef, value)
-                context.values[prop] = param
-            }
-        })
+        valuesDeclaration.block(Values())
+    }
+
+    inner class Values {
+        operator fun <PROPERTY> set(propDef: PropertyDef<PROPERTY>, value: PROPERTY?) {
+            val prop = support.toProp(propDef)
+            val param = support.toParam(propDef, value)
+            context.values[prop] = param
+        }
     }
 }
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/UpdateDeclaration.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/UpdateDeclaration.kt
@@ -1,23 +1,31 @@
 package org.seasar.doma.criteria.declaration
 
+import org.seasar.doma.criteria.context.Operand
 import org.seasar.doma.criteria.context.UpdateContext
 import org.seasar.doma.def.PropertyDef
 
 @Declaration
 class UpdateDeclaration(private val context: UpdateContext) {
     private val support = DeclarationSupport(context.config)
-    private val valuesDeclaration = ValuesDeclaration()
+    private val valuesDeclaration = ValuesDeclaration(context.config)
     private val whereDeclaration = WhereDeclaration(context.config) { context.where.add(it) }
 
     fun set(block: ValuesDeclaration.(Values) -> Unit) {
-        valuesDeclaration.block(object : Values {
-            override fun <PROPERTY> set(propDef: PropertyDef<PROPERTY>, value: PROPERTY?) {
-                val prop = support.toProp(propDef)
-                val param = support.toParam(propDef, value)
-                context.set[prop] = param
-            }
-        })
+        valuesDeclaration.block(Values())
     }
 
     fun where(block: WhereDeclaration.() -> Unit) = whereDeclaration.block()
+
+    inner class Values {
+        operator fun <PROPERTY> set(propDef: PropertyDef<PROPERTY>, value: PROPERTY?) {
+            val prop = support.toProp(propDef)
+            val param = support.toParam(propDef, value)
+            context.set[prop] = param
+        }
+
+        operator fun <PROPERTY> set(propDef: PropertyDef<PROPERTY>, expr: Operand.Expr) {
+            val prop = support.toProp(propDef)
+            context.set[prop] = expr
+        }
+    }
 }
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/Values.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/Values.kt
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/ValuesDeclaration.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/declaration/ValuesDeclaration.kt
@@ -1,4 +1,13 @@
 package org.seasar.doma.criteria.declaration
 
+import org.seasar.doma.criteria.context.Operand
+import org.seasar.doma.jdbc.Config
+
 @Declaration
-class ValuesDeclaration
+class ValuesDeclaration(val config: Config) {
+    private val expressionDeclaration = ExpressionDeclaration(config)
+
+    fun expression(block: ExpressionDeclaration.() -> Operand.Expr): Operand.Expr {
+        return expressionDeclaration.block()
+    }
+}
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/query/BuilderSupport.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/query/BuilderSupport.kt
@@ -58,6 +58,20 @@ class BuilderSupport(
         buf.appendParameter(param.value)
     }
 
+    fun expr(expr: Operand.Expr) {
+        when (expr) {
+            is Operand.Expr.Prop -> column(expr.value)
+            is Operand.Expr.Param -> param(expr.value)
+            is Operand.Expr.Plus -> {
+                buf.appendSql("(")
+                expr(expr.left)
+                buf.appendSql(" + ")
+                expr(expr.right)
+                buf.appendSql(")")
+            }
+        }
+    }
+
     fun visitCriterion(index: Int, c: Criterion) {
         when (c) {
             is Criterion.Eq -> equality(c.left, c.right, "=")
@@ -92,6 +106,7 @@ class BuilderSupport(
             return when (operand) {
                 is Operand.Param -> operand.value.wrapper.get() == null
                 is Operand.Prop -> false
+                is Operand.Expr -> false
             }
         }
 
diff --git a/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/query/UpdateBuilder.kt b/doma-criteria/src/main/kotlin/org/seasar/doma/criteria/query/UpdateBuilder.kt
@@ -28,10 +28,14 @@ class UpdateBuilder(
         table(context.entityDef)
         if (context.set.isNotEmpty()) {
             buf.appendSql(" set ")
-            context.set.forEach { (prop, param) ->
-                column(prop)
+            context.set.forEach { (left, right) ->
+                column(left)
                 buf.appendSql(" = ")
-                param(param)
+                when (right) {
+                    is Operand.Prop -> column(right)
+                    is Operand.Param -> param(right)
+                    is Operand.Expr -> expr(right)
+                }
                 buf.appendSql(", ")
             }
             buf.cutBackSql(2)
@@ -59,6 +63,10 @@ class UpdateBuilder(
         support.param(param)
     }
 
+    private fun expr(expr: Operand.Expr) {
+        support.expr(expr)
+    }
+
     private fun visitCriterion(index: Int, c: Criterion) {
         support.visitCriterion(index, c)
     }
diff --git a/doma-criteria/src/test/kotlin/org/seasar/doma/criteria/SqlDslTest.kt b/doma-criteria/src/test/kotlin/org/seasar/doma/criteria/SqlDslTest.kt
@@ -1,5 +1,6 @@
 package org.seasar.doma.criteria
 
+import java.math.BigDecimal
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
 import org.seasar.doma.criteria.entity.Dept_
@@ -234,4 +235,21 @@ internal class SqlDslTest {
         val expected = """update "CATA"."DEPT" t0_ set t0_.NAME = 'hoge' where t0_.ID = 1"""
         assertEquals(expected, sql.formattedSql)
     }
+
+    @Test
+    fun update_expression() {
+        val query = sql {
+            update(::Emp_) { e ->
+                set {
+                    it[e.salary] = expression { e.salary + BigDecimal("1") }
+                }
+                where {
+                    e.id eq 1
+                }
+            }
+        }
+        val sql = query.asSql(config)
+        val expected = """update EMP t0_ set t0_.SALARY = (t0_.SALARY + 1) where t0_.ID = 1"""
+        assertEquals(expected, sql.formattedSql)
+    }
 }
diff --git a/test-doma-criteria/src/test/kotlin/example/SqlDslTest.kt b/test-doma-criteria/src/test/kotlin/example/SqlDslTest.kt
@@ -377,6 +377,22 @@ class SqlDslTest(private val config: Config) {
         assertEquals(1, count)
     }
 
+    @Test
+    fun update_expression() {
+        val query = sql {
+            update(::Department_) { e ->
+                set {
+                    it[e.departmentNo] = expression { e.departmentNo + 100 }
+                }
+                where {
+                    e.departmentId eq 1
+                }
+            }
+        }
+        val count = query.execute(config)
+        assertEquals(1, count)
+    }
+
     @Test
     fun embeddable() {
         val query = sql {
